fix(core): Fix chat corruption bug in context manager. (#26534)

packages/core/src/context/graph/render.test.ts

@@ -0,0 +1,64 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi } from 'vitest';
+import { render } from './render.js';
+import type { ConcreteNode } from './types.js';
+import { NodeType } from './types.js';
+import type { ContextEnvironment } from '../pipeline/environment.js';
+import type { ContextTracer } from '../tracer.js';
+import type { ContextProfile } from '../config/profiles.js';
+import type { PipelineOrchestrator } from '../pipeline/orchestrator.js';
+import type { Part } from '@google/genai';
+
+describe('render', () => {
+  it('should filter out previewNodeIds', async () => {
+    const mockNodes: ConcreteNode[] = [
+      {
+        id: '1',
+        type: NodeType.USER_PROMPT,
+        payload: {} as Part,
+      } as unknown as ConcreteNode,
+      {
+        id: '2',
+        type: NodeType.AGENT_THOUGHT,
+        payload: {} as Part,
+      } as unknown as ConcreteNode,
+      {
+        id: 'preview-1',
+        type: NodeType.USER_PROMPT,
+        payload: {} as Part,
+      } as unknown as ConcreteNode,
+    ];
+    const previewNodeIds = new Set(['preview-1']);
+
+    const orchestrator = {} as PipelineOrchestrator;
+    const sidecar = { config: {} } as ContextProfile; // No budget
+    const env = {
+      graphMapper: {
+        fromGraph: vi.fn((nodes: readonly ConcreteNode[]) =>
+          nodes.map((n) => ({ text: n.id })),
+        ),
+      },
+    } as unknown as ContextEnvironment;
+    const tracer = {
+      logEvent: vi.fn(),
+    } as unknown as ContextTracer;
+
+    const result = await render(
+      mockNodes,
+      orchestrator,
+      sidecar,
+      tracer,
+      env,
+      new Map(),
+      0,
+      previewNodeIds,
+    );
+
+    expect(result.history).toEqual([{ text: '1' }, { text: '2' }]);
+  });
+});

packages/core/src/context/graph/render.ts

@@ -23,9 +23,11 @@ export async function render(
   env: ContextEnvironment,
   protectionReasons: Map<string, string> = new Map(),
   headerTokens: number = 0,
+  previewNodeIds: ReadonlySet<string> = new Set(),
 ): Promise<{ history: Content[]; didApplyManagement: boolean }> {
   if (!sidecar.config.budget) {
-    const contents = env.graphMapper.fromGraph(nodes);
+    const visibleNodes = nodes.filter((n) => !previewNodeIds.has(n.id));
+    const contents = env.graphMapper.fromGraph(visibleNodes);
     tracer.logEvent('Render', 'Render Context to LLM (No Budget)', {
       renderedContext: contents,
     });
@@ -61,13 +63,13 @@ export async function render(
       'Render',
       `View is within maxTokens (${currentTokens} <= ${maxTokens}). Returning view.`,
     );
-    const contents = env.graphMapper.fromGraph(nodes);
+    const visibleNodes = nodes.filter((n) => !previewNodeIds.has(n.id));
+    const contents = env.graphMapper.fromGraph(visibleNodes);
     tracer.logEvent('Render', 'Render Context for LLM', {
       renderedContext: contents,
     });
     return { history: contents, didApplyManagement: false };
   }
-
   const targetDelta = currentTokens - sidecar.config.budget.retainedTokens;
   tracer.logEvent(
     'Render',
@@ -103,7 +105,9 @@ export async function render(
     }
   }
 
-  const visibleNodes = processedNodes.filter((n) => !skipList.has(n.id));
+  const visibleNodes = processedNodes.filter(
+    (n) => !skipList.has(n.id) && !previewNodeIds.has(n.id),
+  );
 
   const contents = env.graphMapper.fromGraph(visibleNodes);
   tracer.logEvent('Render', 'Render Sanitized Context for LLM', {

packages/core/src/context/graph/toGraph.test.ts

@@ -0,0 +1,40 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect } from 'vitest';
+import { ContextGraphBuilder } from './toGraph.js';
+import type { Content } from '@google/genai';
+import type { BaseConcreteNode } from './types.js';
+
+describe('ContextGraphBuilder', () => {
+  describe('toGraph', () => {
+    it('should skip legacy <session_context> headers even if they appear later in the history', () => {
+      const history: Content[] = [
+        { role: 'user', parts: [{ text: 'Message 1' }] },
+        { role: 'model', parts: [{ text: 'Reply 1' }] },
+        {
+          role: 'user',
+          parts: [
+            {
+              text: '<session_context>\nThis is the Gemini CLI\nSome context...',
+            },
+          ],
+        },
+        { role: 'user', parts: [{ text: 'Message 2' }] },
+      ];
+
+      const builder = new ContextGraphBuilder();
+      const nodes = builder.processHistory(history);
+
+      // We expect the first two messages and the last one to be present
+      // The session context message should be filtered out
+      expect(nodes.length).toBe(3);
+      expect((nodes[0] as BaseConcreteNode).payload.text).toBe('Message 1');
+      expect((nodes[1] as BaseConcreteNode).payload.text).toBe('Reply 1');
+      expect((nodes[2] as BaseConcreteNode).payload.text).toBe('Message 2');
+    });
+  });
+});

packages/core/src/context/graph/toGraph.ts

@@ -149,13 +149,13 @@ export class ContextGraphBuilder {
       const msg = history[turnIdx];
       if (!msg.parts) continue;
 
-      // Defensive: Skip legacy environment header if it's the first turn.
+      // Defensive: Skip legacy environment header regardless of where it appears.
       // We now manage this as an orthogonal late-addition header.
-      if (turnIdx === 0 && msg.role === 'user' && msg.parts.length === 1) {
+      if (msg.role === 'user' && msg.parts.length === 1) {
         const text = msg.parts[0].text;
         if (
           text?.startsWith('<session_context>') &&
-          text?.includes('This is the Gemini CLI.')
+          text?.includes('This is the Gemini CLI')
         ) {
           debugLogger.log(
             '[ContextGraphBuilder] Skipping legacy environment header turn from graph.',