Fix bulk of remaining issues with generalist profile (#26073)

packages/core/src/context/contextManager.ts

@@ -6,7 +6,7 @@
 
 import type { Content } from '@google/genai';
 import type { AgentChatHistory } from '../core/agentChatHistory.js';
-import type { ConcreteNode } from './graph/types.js';
+import { isToolExecution, type ConcreteNode } from './graph/types.js';
 import type { ContextEventBus } from './eventBus.js';
 import type { ContextTracer } from './tracer.js';
 import type { ContextEnvironment } from './pipeline/environment.js';
@@ -15,6 +15,9 @@ import type { PipelineOrchestrator } from './pipeline/orchestrator.js';
 import { HistoryObserver } from './historyObserver.js';
 import { render } from './graph/render.js';
 import { ContextWorkingBufferImpl } from './pipeline/contextWorkingBuffer.js';
+import { debugLogger } from '../utils/debugLogger.js';
+import { hardenHistory } from '../utils/historyHardening.js';
+import { checkContextInvariants } from './utils/invariantChecker.js';
 
 export class ContextManager {
   // The master state containing the pristine graph and current active graph.
@@ -27,21 +30,30 @@ export class ContextManager {
   private readonly orchestrator: PipelineOrchestrator;
   private readonly historyObserver: HistoryObserver;
 
+  // Cache for Anomaly 3 (Redundant Renders)
+  private lastRenderCache?: {
+    nodesHash: string;
+    result: { history: Content[]; didApplyManagement: boolean };
+  };
+
   constructor(
     private readonly sidecar: ContextProfile,
     private readonly env: ContextEnvironment,
     private readonly tracer: ContextTracer,
     orchestrator: PipelineOrchestrator,
     chatHistory: AgentChatHistory,
+    private readonly headerProvider?: () => Promise<Content | undefined>,
   ) {
     this.eventBus = env.eventBus;
     this.orchestrator = orchestrator;
 
+    // Provide the orchestrator with a way to fetch the latest nodes from the live buffer
+    this.orchestrator.setNodeProvider(() => this.buffer.nodes);
+
     this.historyObserver = new HistoryObserver(
       chatHistory,
       this.env.eventBus,
       this.tracer,
-      this.env.tokenCalculator,
       this.env.graphMapper,
     );
 
@@ -69,6 +81,13 @@ export class ContextManager {
     this.historyObserver.start();
   }
 
+  /**
+   * Returns a promise that resolves when all currently executing async pipelines have finished.
+   */
+  async waitForPipelines(): Promise<void> {
+    return this.orchestrator.waitForPipelines();
+  }
+
   /**
    * Safely stops background async pipelines and clears event listeners.
    */
@@ -98,6 +117,15 @@ export class ContextManager {
     if (currentTokens > this.sidecar.config.budget.retainedTokens) {
       const agedOutNodes = new Set<string>();
       let rollingTokens = 0;
+
+      // Identify active tool calls that must NEVER be truncated
+      const protectedIds = this.getProtectedNodeIds(this.buffer.nodes);
+      if (protectedIds.size > 0) {
+        debugLogger.log(
+          `[ContextManager] Pinning ${protectedIds.size} active tool call nodes to prevent truncation.`,
+        );
+      }
+
       // Walk backwards finding nodes that fall out of the retained budget
       for (let i = this.buffer.nodes.length - 1; i >= 0; i--) {
         const node = this.buffer.nodes[i];
@@ -105,7 +133,10 @@ export class ContextManager {
           node,
         ]);
         if (rollingTokens > this.sidecar.config.budget.retainedTokens) {
-          agedOutNodes.add(node.id);
+          // Only age out if not protected
+          if (!protectedIds.has(node.id)) {
+            agedOutNodes.add(node.id);
+          }
         }
       }
 
@@ -123,6 +154,54 @@ export class ContextManager {
     }
   }
 
+  /**
+   * Identifies 'pinned' nodes that should not be truncated.
+   * This includes:
+   * 1. The entire last turn (Recent context).
+   * 2. Active tool calls (calls without responses in the graph).
+   */
+  private getProtectedNodeIds(
+    nodes: readonly ConcreteNode[],
+    extraProtectedIds: Set<string> = new Set(),
+  ): Map<string, string> {
+    const protectionMap = new Map<string, string>();
+    if (nodes.length === 0) return protectionMap;
+
+    // 1. Identify all nodes belonging to the last turn (Recent context)
+    const lastNode = nodes[nodes.length - 1];
+    const lastTurnId = lastNode.turnId;
+
+    for (const node of nodes) {
+      if (node.turnId === lastTurnId) {
+        protectionMap.set(node.id, 'recent_turn');
+      }
+    }
+
+    // 2. Identify active tool calls that must NEVER be truncated
+    const calls = nodes.filter((n) => isToolExecution(n) && n.role === 'model');
+    const responses = new Set(
+      nodes
+        .filter((n) => isToolExecution(n) && n.role === 'user')
+        .map((n) => n.payload.functionResponse?.id)
+        .filter((id): id is string => !!id),
+    );
+
+    for (const call of calls) {
+      const id = call.payload.functionCall?.id;
+      // If we have a call but no response in the current graph, it's 'in flight'
+      if (id && !responses.has(id)) {
+        protectionMap.set(call.id, 'in_flight_tool_call');
+      }
+    }
+
+    // 3. Any externally requested protections
+    for (const id of extraProtectedIds) {
+      protectionMap.set(id, 'external_active_task');
+    }
+
+    return protectionMap;
+  }
+
   /**
    * Retrieves the raw, uncompressed Episodic Context Graph graph.
    * Useful for internal tool rendering (like the trace viewer).
@@ -157,22 +236,78 @@ export class ContextManager {
    * This is the primary method called by the agent framework before sending a request.
    */
   async renderHistory(
+    pendingRequest?: Content,
     activeTaskIds: Set<string> = new Set(),
-  ): Promise<Content[]> {
+  ): Promise<{ history: Content[]; didApplyManagement: boolean }> {
     this.tracer.logEvent('ContextManager', 'Starting rendering of LLM context');
 
+    // 1. Synchronous Pressure Barrier: Wait for background management pipelines to finish.
+    // This ensures that the render sees the results of recent pushes (Anomaly 2).
+    await this.orchestrator.waitForPipelines();
+
+    let nodes = this.buffer.nodes;
+
+    // If we have a pending request, we need to build a 'preview' graph for this render.
+    if (pendingRequest) {
+      const previewNodes = this.env.graphMapper.applyEvent({
+        type: 'PUSH',
+        payload: [pendingRequest],
+      });
+      nodes = [...nodes, ...previewNodes];
+    }
+
+    // 2. Fetch Header and calculate tokens
+    const header = this.headerProvider
+      ? await this.headerProvider()
+      : undefined;
+    const headerTokens = header
+      ? this.env.tokenCalculator.calculateContentTokens(header)
+      : 0;
+
+    // 3. Cache Check (Anomaly 3): If nodes haven't changed, return previous result.
+    // We combine the graph hash with a hash of the header to ensure total freshness.
+    const graphHash = nodes.map((n) => n.id).join('|');
+    const headerHash = header ? JSON.stringify(header.parts) : 'no-header';
+    const totalHash = `${graphHash}::${headerHash}`;
+
+    if (this.lastRenderCache?.nodesHash === totalHash) {
+      debugLogger.log(
+        '[ContextManager] Render cache hit. Skipping redundant render.',
+      );
+      return this.lastRenderCache.result;
+    }
+
+    const protectionReasons = this.getProtectedNodeIds(nodes, activeTaskIds);
+
     // Apply final GC Backstop pressure barrier synchronously before mapping
-    const finalHistory = await render(
-      this.buffer.nodes,
+    const { history: renderedHistory, didApplyManagement } = await render(
+      nodes,
       this.orchestrator,
       this.sidecar,
       this.tracer,
       this.env,
-      activeTaskIds,
+      protectionReasons,
+      headerTokens,
     );
 
+    // Structural validation in debug mode
+    checkContextInvariants(this.buffer.nodes, 'RenderHistory');
+
     this.tracer.logEvent('ContextManager', 'Finished rendering');
 
-    return finalHistory;
+    const combinedHistory = header
+      ? [header, ...renderedHistory]
+      : renderedHistory;
+
+    const result = {
+      history: hardenHistory(combinedHistory, {
+        sentinels: this.sidecar.sentinels,
+      }),
+      didApplyManagement,
+    };
+
+    // Update cache
+    this.lastRenderCache = { nodesHash: totalHash, result };
+    return result;
   }
 }