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

2026-07-11 02:20:48 -07:00 · 2026-05-01 15:04:39 -07:00
parent 408afd3c5a
commit de8fdcfa16
52 changed files with 2133 additions and 1364 deletions
@@ -3,21 +3,11 @@
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
-import type { Content, Part } from '@google/genai';
-import type { ConcreteNode } from './types.js';
-
-export interface NodeSerializationWriter {
-  appendContent(content: Content): void;
-  appendModelPart(part: Part): void;
-  appendUserPart(part: Part): void;
-  flushModelParts(): void;
-}
+import type { Part } from '@google/genai';
+import type { ConcreteNode, NodeType } from './types.js';

 export interface NodeBehavior<T extends ConcreteNode = ConcreteNode> {
-  readonly type: T['type'];
-
-  /** Serializes the node into the Gemini Content structure. */
-  serialize(node: T, writer: NodeSerializationWriter): void;
+  readonly type: NodeType;

  /**
   * Generates a structural representation of the node for the purpose
@@ -27,13 +17,13 @@ export interface NodeBehavior<T extends ConcreteNode = ConcreteNode> {
 }

 export class NodeBehaviorRegistry {
-  private readonly behaviors = new Map<string, NodeBehavior<ConcreteNode>>();
+  private readonly behaviors = new Map<NodeType, NodeBehavior<ConcreteNode>>();

  register<T extends ConcreteNode>(behavior: NodeBehavior<T>) {
    this.behaviors.set(behavior.type, behavior);
  }

-  get(type: string): NodeBehavior<ConcreteNode> {
+  get(type: NodeType): NodeBehavior<ConcreteNode> {
    const behavior = this.behaviors.get(type);
    if (!behavior) {
      throw new Error(`Unregistered Node type: ${type}`);
@@ -3,160 +3,72 @@
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
-import type { Part } from '@google/genai';
 import type { NodeBehavior, NodeBehaviorRegistry } from './behaviorRegistry.js';
-import type {
-  UserPrompt,
-  AgentThought,
-  ToolExecution,
-  MaskedTool,
-  AgentYield,
-  Snapshot,
-  RollingSummary,
-  SystemEvent,
+import {
+  type UserPrompt,
+  type AgentThought,
+  type ToolExecution,
+  type MaskedTool,
+  type AgentYield,
+  type Snapshot,
+  type RollingSummary,
+  type SystemEvent,
+  NodeType,
 } from './types.js';

 export const UserPromptBehavior: NodeBehavior<UserPrompt> = {
-  type: 'USER_PROMPT',
-  getEstimatableParts(prompt) {
-    const parts: Part[] = [];
-    for (const sp of prompt.semanticParts) {
-      switch (sp.type) {
-        case 'text':
-          parts.push({ text: sp.text });
-          break;
-        case 'inline_data':
-          parts.push({ inlineData: { mimeType: sp.mimeType, data: sp.data } });
-          break;
-        case 'file_data':
-          parts.push({
-            fileData: { mimeType: sp.mimeType, fileUri: sp.fileUri },
-          });
-          break;
-        case 'raw_part':
-          parts.push(sp.part);
-          break;
-        default:
-          break;
-      }
-    }
-    return parts;
-  },
-  serialize(prompt, writer) {
-    const parts = this.getEstimatableParts(prompt);
-    if (parts.length > 0) {
-      writer.flushModelParts();
-      writer.appendContent({ role: 'user', parts });
-    }
+  type: NodeType.USER_PROMPT,
+  getEstimatableParts(node) {
+    return [node.payload];
  },
 };

 export const AgentThoughtBehavior: NodeBehavior<AgentThought> = {
-  type: 'AGENT_THOUGHT',
-  getEstimatableParts(thought) {
-    return [{ text: thought.text }];
-  },
-  serialize(thought, writer) {
-    writer.appendModelPart({ text: thought.text });
+  type: NodeType.AGENT_THOUGHT,
+  getEstimatableParts(node) {
+    return [node.payload];
  },
 };

 export const ToolExecutionBehavior: NodeBehavior<ToolExecution> = {
-  type: 'TOOL_EXECUTION',
-  getEstimatableParts(tool) {
-    return [
-      { functionCall: { id: tool.id, name: tool.toolName, args: tool.intent } },
-      {
-        functionResponse: {
-          id: tool.id,
-          name: tool.toolName,
-          response:
-            typeof tool.observation === 'string'
-              ? { message: tool.observation }
-              : tool.observation,
-        },
-      },
-    ];
-  },
-  serialize(tool, writer) {
-    const parts = this.getEstimatableParts(tool);
-    writer.appendModelPart(parts[0]);
-    writer.flushModelParts();
-    writer.appendUserPart(parts[1]);
+  type: NodeType.TOOL_EXECUTION,
+  getEstimatableParts(node) {
+    return [node.payload];
  },
 };

 export const MaskedToolBehavior: NodeBehavior<MaskedTool> = {
-  type: 'MASKED_TOOL',
-  getEstimatableParts(tool) {
-    return [
-      {
-        functionCall: {
-          id: tool.id,
-          name: tool.toolName,
-          args: tool.intent ?? {},
-        },
-      },
-      {
-        functionResponse: {
-          id: tool.id,
-          name: tool.toolName,
-          response:
-            typeof tool.observation === 'string'
-              ? { message: tool.observation }
-              : (tool.observation ?? {}),
-        },
-      },
-    ];
-  },
-  serialize(tool, writer) {
-    const parts = this.getEstimatableParts(tool);
-    writer.appendModelPart(parts[0]);
-    writer.flushModelParts();
-    writer.appendUserPart(parts[1]);
+  type: NodeType.MASKED_TOOL,
+  getEstimatableParts(node) {
+    return [node.payload];
  },
 };

 export const AgentYieldBehavior: NodeBehavior<AgentYield> = {
-  type: 'AGENT_YIELD',
-  getEstimatableParts(yieldNode) {
-    return [{ text: yieldNode.text }];
-  },
-  serialize() {
-    // AGENT_YIELD is a synthetic marker node used for internal graph tracking.
-    // We intentionally do NOT serialize it to the LLM to prevent prompt corruption.
+  type: NodeType.AGENT_YIELD,
+  getEstimatableParts() {
+    return [];
  },
 };

 export const SystemEventBehavior: NodeBehavior<SystemEvent> = {
-  type: 'SYSTEM_EVENT',
-  getEstimatableParts() {
-    return [];
-  },
-  serialize(node, writer) {
-    writer.flushModelParts();
+  type: NodeType.SYSTEM_EVENT,
+  getEstimatableParts(node) {
+    return [node.payload];
  },
 };

 export const SnapshotBehavior: NodeBehavior<Snapshot> = {
-  type: 'SNAPSHOT',
+  type: NodeType.SNAPSHOT,
  getEstimatableParts(node) {
-    return [{ text: node.text }];
-  },
-  serialize(node, writer) {
-    writer.flushModelParts();
-    writer.appendUserPart({ text: node.text });
+    return [node.payload];
  },
 };

 export const RollingSummaryBehavior: NodeBehavior<RollingSummary> = {
-  type: 'ROLLING_SUMMARY',
+  type: NodeType.ROLLING_SUMMARY,
  getEstimatableParts(node) {
-    return [{ text: node.text }];
-  },
-  serialize(node, writer) {
-    writer.flushModelParts();
-    writer.appendUserPart({ text: node.text });
+    return [node.payload];
  },
 };

@@ -3,52 +3,53 @@
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */
-import type { Content, Part } from '@google/genai';
+
+import type { Content } from '@google/genai';
 import type { ConcreteNode } from './types.js';
-import type {
-  NodeSerializationWriter,
-  NodeBehaviorRegistry,
-} from './behaviorRegistry.js';
+import { debugLogger } from '../../utils/debugLogger.js';

-class NodeSerializer implements NodeSerializationWriter {
-  private history: Content[] = [];
-  private currentModelParts: Part[] = [];
+/**
+ * Reconstructs a valid Gemini Chat History from a list of Concrete Nodes.
+ * This process is "role-alternation-aware" and uses turnId to
+ * preserve original turn boundaries even if multiple turns have the same role.
+ */
+export function fromGraph(nodes: readonly ConcreteNode[]): Content[] {
+  debugLogger.log(
+    `[fromGraph] Reconstructing history from ${nodes.length} nodes`,
+  );

-  appendContent(content: Content) {
-    this.flushModelParts();
-    this.history.push(content);
-  }
+  const history: Content[] = [];
+  let currentTurn: (Content & { _turnId?: string }) | null = null;

-  appendModelPart(part: Part) {
-    this.currentModelParts.push(part);
-  }
+  for (const node of nodes) {
+    const turnId = node.turnId;

-  appendUserPart(part: Part) {
-    this.flushModelParts();
-    this.history.push({ role: 'user', parts: [part] });
-  }
-
-  flushModelParts() {
-    if (this.currentModelParts.length > 0) {
-      this.history.push({ role: 'model', parts: [...this.currentModelParts] });
-      this.currentModelParts = [];
+    // We start a new turn if:
+    // 1. We don't have a current turn.
+    // 2. The role changes (Standard alternation).
+    // 3. The turnId changes (Preserving distinct turns of the same role).
+    if (
+      !currentTurn ||
+      currentTurn.role !== node.role ||
+      currentTurn._turnId !== turnId
+    ) {
+      currentTurn = {
+        role: node.role,
+        parts: [node.payload],
+        _turnId: turnId,
+      };
+      history.push(currentTurn);
+    } else {
+      currentTurn.parts = [...(currentTurn.parts || []), node.payload];
    }
  }

-  getContents(): Content[] {
-    this.flushModelParts();
-    return this.history;
+  // Final cleanup: remove our internal tracking field
+  for (const turn of history) {
+    const t = turn as Content & { _turnId?: string };
+    delete t._turnId;
  }
-}

-export function fromGraph(
-  nodes: readonly ConcreteNode[],
-  registry: NodeBehaviorRegistry,
-): Content[] {
-  const writer = new NodeSerializer();
-  for (const node of nodes) {
-    const behavior = registry.get(node.type);
-    behavior.serialize(node, writer);
-  }
-  return writer.getContents();
+  debugLogger.log(`[fromGraph] Reconstructed ${history.length} turns`);
+  return history;
 }
@@ -8,41 +8,20 @@ import { ContextGraphBuilder } from './toGraph.js';
 import type { Content } from '@google/genai';
 import type { HistoryEvent } from '../../core/agentChatHistory.js';
 import { fromGraph } from './fromGraph.js';
-import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
-import type { NodeBehaviorRegistry } from './behaviorRegistry.js';

 export class ContextGraphMapper {
  private readonly nodeIdentityMap = new WeakMap<object, string>();
+  private readonly builder: ContextGraphBuilder;

-  constructor(private readonly registry: NodeBehaviorRegistry) {}
+  constructor() {
+    this.builder = new ContextGraphBuilder(this.nodeIdentityMap);
+  }

-  private builder?: ContextGraphBuilder;
-
-  applyEvent(
-    event: HistoryEvent,
-    tokenCalculator: ContextTokenCalculator,
-  ): ConcreteNode[] {
-    if (!this.builder) {
-      this.builder = new ContextGraphBuilder(
-        tokenCalculator,
-        this.nodeIdentityMap,
-      );
-    }
-
-    if (event.type === 'CLEAR') {
-      this.builder.clear();
-      return [];
-    }
-
-    if (event.type === 'SYNC_FULL') {
-      this.builder.clear();
-    }
-
-    this.builder.processHistory(event.payload);
-    return this.builder.getNodes();
+  applyEvent(event: HistoryEvent): ConcreteNode[] {
+    return this.builder.processHistory(event.payload);
  }

  fromGraph(nodes: readonly ConcreteNode[]): Content[] {
-    return fromGraph(nodes, this.registry);
+    return fromGraph(nodes);
  }
 }
@@ -6,17 +6,14 @@

 import type { Content } from '@google/genai';
 import type { ConcreteNode } from './types.js';
-import { debugLogger } from '../../utils/debugLogger.js';
-import type {
-  ContextEnvironment,
-  ContextTracer,
-} from '../pipeline/environment.js';
-import type { PipelineOrchestrator } from '../pipeline/orchestrator.js';
+import type { ContextTracer } from '../tracer.js';
 import type { ContextProfile } from '../config/profiles.js';
+import type { PipelineOrchestrator } from '../pipeline/orchestrator.js';
+import type { ContextEnvironment } from '../pipeline/environment.js';

 /**
- * Orchestrates the final render: takes a working buffer view (The Nodes),
- * applies the Immediate Sanitization pipeline, and enforces token boundaries.
+ * Maps the Episodic Context Graph back into a raw Gemini Content[] array for transmission.
+ * It applies synchronous context management (GC backstop) if the budget is exceeded.
 */
 export async function render(
  nodes: readonly ConcreteNode[],
@@ -24,28 +21,40 @@ export async function render(
  sidecar: ContextProfile,
  tracer: ContextTracer,
  env: ContextEnvironment,
-  protectedIds: Set<string>,
-): Promise<Content[]> {
+  protectionReasons: Map<string, string> = new Map(),
+  headerTokens: number = 0,
+): Promise<{ history: Content[]; didApplyManagement: boolean }> {
  if (!sidecar.config.budget) {
    const contents = env.graphMapper.fromGraph(nodes);
    tracer.logEvent('Render', 'Render Context to LLM (No Budget)', {
      renderedContext: contents,
    });
-    return contents;
+    return { history: contents, didApplyManagement: false };
  }

  const maxTokens = sidecar.config.budget.maxTokens;
-  const currentTokens = env.tokenCalculator.calculateConcreteListTokens(nodes);
+  const graphTokens = env.tokenCalculator.calculateConcreteListTokens(nodes);
+  const currentTokens = graphTokens + headerTokens;

-  // V0: Always protect the first node (System Prompt) and the last turn
-  if (nodes.length > 0) {
-    protectedIds.add(nodes[0].id);
-    if (nodes[0].logicalParentId) protectedIds.add(nodes[0].logicalParentId);
+  const protectedIds = new Set(protectionReasons.keys());

-    const lastNode = nodes[nodes.length - 1];
-    protectedIds.add(lastNode.id);
-    if (lastNode.logicalParentId) protectedIds.add(lastNode.logicalParentId);
-  }
+  tracer.logEvent('Render', 'Budget Audit', {
+    maxTokens,
+    retainedTokens: sidecar.config.budget.retainedTokens,
+    graphTokens,
+    headerTokens,
+    currentTokens,
+    pressure: (currentTokens / maxTokens).toFixed(2),
+    isOverBudget: currentTokens > maxTokens,
+  });
+
+  tracer.logEvent('Render', 'Estimation Calibration', {
+    breakdown: env.tokenCalculator.calculateTokenBreakdown(nodes),
+  });
+
+  tracer.logEvent('Render', 'Protection Audit', {
+    reasons: Object.fromEntries(protectionReasons),
+  });

  if (currentTokens <= maxTokens) {
    tracer.logEvent(
@@ -56,15 +65,14 @@ export async function render(
    tracer.logEvent('Render', 'Render Context for LLM', {
      renderedContext: contents,
    });
-    return contents;
+    return { history: contents, didApplyManagement: false };
  }

+  const targetDelta = currentTokens - sidecar.config.budget.retainedTokens;
  tracer.logEvent(
    'Render',
    `View exceeds maxTokens (${currentTokens} > ${maxTokens}). Hitting Synchronous Pressure Barrier.`,
-  );
-  debugLogger.log(
-    `Context Manager Synchronous Barrier triggered: View at ${currentTokens} tokens (limit: ${maxTokens}).`,
+    { targetDelta },
  );

  // Calculate exactly which nodes aged out of the retainedTokens budget to form our target delta
@@ -87,16 +95,6 @@ export async function render(
    protectedIds,
  );

-  const finalTokens =
-    env.tokenCalculator.calculateConcreteListTokens(processedNodes);
-  tracer.logEvent(
-    'Render',
-    `Finished rendering. Final token count: ${finalTokens}.`,
-  );
-  debugLogger.log(
-    `Context Manager finished. Final actual token count: ${finalTokens}.`,
-  );
-
  // Apply skipList logic to abstract over summarized nodes
  const skipList = new Set<string>();
  for (const node of processedNodes) {
@@ -111,5 +109,5 @@ export async function render(
  tracer.logEvent('Render', 'Render Sanitized Context for LLM', {
    renderedContextSanitized: contents,
  });
-  return contents;
+  return { history: contents, didApplyManagement: true };
 }
@@ -5,294 +5,227 @@
 */

 import type { Content, Part } from '@google/genai';
-import type {
-  ConcreteNode,
-  Episode,
-  SemanticPart,
-  ToolExecution,
-  AgentThought,
-  AgentYield,
-  UserPrompt,
-} from './types.js';
-import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
-import { randomUUID } from 'node:crypto';
-import { isRecord } from '../../utils/markdownUtils.js';
+import { type ConcreteNode, NodeType } from './types.js';
+import { randomUUID, createHash } from 'node:crypto';
+import { debugLogger } from '../../utils/debugLogger.js';

-// We remove the global nodeIdentityMap and instead rely on one passed from ContextGraphMapper
-export function getStableId(
-  obj: object,
-  nodeIdentityMap: WeakMap<object, string>,
-): string {
-  let id = nodeIdentityMap.get(obj);
-  if (!id) {
-    id = randomUUID();
-    nodeIdentityMap.set(obj, id);
-  }
-  return id;
+interface PartWithSynthId extends Part {
+  _synthId?: string;
 }

-function isCompleteEpisode(ep: Partial<Episode>): ep is Episode {
+// Global WeakMap to cache hashes for Part objects.
+// This optimizes getStableId by avoiding redundant stringify/hash operations
+// on the same object instances across multiple management passes.
+const PART_HASH_CACHE = new WeakMap<object, string>();
+
+function isTextPart(part: Part): part is Part & { text: string } {
+  return typeof part.text === 'string';
+}
+
+function isInlineDataPart(
+  part: Part,
+): part is Part & { inlineData: { data: string } } {
  return (
-    typeof ep.id === 'string' &&
-    Array.isArray(ep.concreteNodes) &&
-    ep.concreteNodes.length > 0
+    typeof part.inlineData === 'object' &&
+    part.inlineData !== null &&
+    typeof part.inlineData.data === 'string'
  );
 }

-export class ContextGraphBuilder {
-  private episodes: Episode[] = [];
-  private currentEpisode: Partial<Episode> | null = null;
-  private pendingCallParts: Map<string, Part> = new Map();
-  private pendingCallPartsWithoutId: Part[] = [];
+function isFileDataPart(
+  part: Part,
+): part is Part & { fileData: { fileUri: string } } {
+  return (
+    typeof part.fileData === 'object' &&
+    part.fileData !== null &&
+    typeof part.fileData.fileUri === 'string'
+  );
+}

+function isFunctionCallPart(
+  part: Part,
+): part is Part & { functionCall: { id: string; name: string } } {
+  return (
+    typeof part.functionCall === 'object' &&
+    part.functionCall !== null &&
+    typeof part.functionCall.name === 'string'
+  );
+}
+
+function isFunctionResponsePart(
+  part: Part,
+): part is Part & { functionResponse: { id: string; name: string } } {
+  return (
+    typeof part.functionResponse === 'object' &&
+    part.functionResponse !== null &&
+    typeof part.functionResponse.name === 'string'
+  );
+}
+
+/**
+ * Generates a stable ID for an object reference using a WeakMap.
+ * Falls back to content-based hashing for Part-like objects to ensure
+ * stability across object re-creations (e.g. during history mapping).
+ */
+export function getStableId(
+  obj: object,
+  nodeIdentityMap: WeakMap<object, string>,
+  turnSalt: string = '',
+  partIdx: number = 0,
+): string {
+  let id = nodeIdentityMap.get(obj);
+  if (id) return id;
+
+  const cachedHash = PART_HASH_CACHE.get(obj);
+  if (cachedHash) {
+    id = `${cachedHash}_${turnSalt}_${partIdx}`;
+    nodeIdentityMap.set(obj, id);
+    return id;
+  }
+
+  const part = obj as PartWithSynthId;
+  let contentHash: string | undefined;
+
+  // If the object already has a synthetic ID property, use it.
+  if (typeof part._synthId === 'string') {
+    id = part._synthId;
+  } else if (isTextPart(part)) {
+    contentHash = createHash('sha256').update(part.text).digest('hex');
+    id = `text_${contentHash}_${turnSalt}_${partIdx}`;
+  } else if (isInlineDataPart(part)) {
+    contentHash = createHash('sha256')
+      .update(part.inlineData.data)
+      .digest('hex');
+    id = `media_${contentHash}_${turnSalt}_${partIdx}`;
+  } else if (isFileDataPart(part)) {
+    contentHash = createHash('sha256')
+      .update(part.fileData.fileUri)
+      .digest('hex');
+    id = `file_${contentHash}_${turnSalt}_${partIdx}`;
+  } else if (isFunctionCallPart(part)) {
+    contentHash = createHash('sha256')
+      .update(
+        `call:${part.functionCall.name}:${JSON.stringify(part.functionCall.args)}`,
+      )
+      .digest('hex');
+    id = `call_h_${contentHash}_${turnSalt}_${partIdx}`;
+  } else if (isFunctionResponsePart(part)) {
+    contentHash = createHash('sha256')
+      .update(
+        `resp:${part.functionResponse.name}:${JSON.stringify(part.functionResponse.response)}`,
+      )
+      .digest('hex');
+    id = `resp_h_${contentHash}_${turnSalt}_${partIdx}`;
+  }
+
+  if (contentHash) {
+    PART_HASH_CACHE.set(obj, contentHash);
+  }
+
+  if (!id) {
+    id = randomUUID();
+  }
+
+  nodeIdentityMap.set(obj, id);
+  return id;
+}
+
+/**
+ * Builds a 1:1 Mirror Graph from Chat History.
+ * Every Part in history is mapped to exactly one ConcreteNode.
+ */
+export class ContextGraphBuilder {
  constructor(
-    private readonly tokenCalculator: ContextTokenCalculator,
    private readonly nodeIdentityMap: WeakMap<object, string> = new WeakMap(),
  ) {}

-  clear() {
-    this.episodes = [];
-    this.currentEpisode = null;
-    this.pendingCallParts.clear();
-    this.pendingCallPartsWithoutId = [];
-  }
+  processHistory(history: readonly Content[]): ConcreteNode[] {
+    const nodes: ConcreteNode[] = [];

-  processHistory(history: readonly Content[]) {
-    const finalizeEpisode = () => {
-      if (this.currentEpisode && isCompleteEpisode(this.currentEpisode)) {
-        this.episodes.push(this.currentEpisode);
-      }
-      this.currentEpisode = null;
-    };
+    // Tracks occurrences of identical turn content to ensure unique stable IDs
+    const seenHashes = new Map<string, number>();

-    for (const msg of history) {
+    for (let turnIdx = 0; turnIdx < history.length; turnIdx++) {
+      const msg = history[turnIdx];
      if (!msg.parts) continue;

-      if (msg.role === 'user') {
-        const hasToolResponses = msg.parts.some((p) => !!p.functionResponse);
-        const hasUserParts = msg.parts.some(
-          (p) => !!p.text || !!p.inlineData || !!p.fileData,
-        );
-
-        if (hasToolResponses) {
-          this.currentEpisode = parseToolResponses(
-            msg,
-            this.currentEpisode,
-            this.pendingCallParts,
-            this.pendingCallPartsWithoutId,
-            this.tokenCalculator,
-            this.nodeIdentityMap,
+      // Defensive: Skip legacy environment header if it's the first turn.
+      // We now manage this as an orthogonal late-addition header.
+      if (turnIdx === 0 && 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.')
+        ) {
+          debugLogger.log(
+            '[ContextGraphBuilder] Skipping legacy environment header turn from graph.',
          );
+          continue;
        }
+      }

-        if (hasUserParts) {
-          finalizeEpisode();
-          this.currentEpisode = parseUserParts(msg, this.nodeIdentityMap);
+      // Generate a stable salt for this turn based on its role and content
+      const turnContent = JSON.stringify(msg.parts);
+      const h = createHash('md5')
+        .update(`${msg.role}:${turnContent}`)
+        .digest('hex');
+      const occurrence = (seenHashes.get(h) || 0) + 1;
+      seenHashes.set(h, occurrence);
+      const turnSalt = `${h}_${occurrence}`;
+      const turnId = getStableId(msg, this.nodeIdentityMap, turnSalt, -1);
+
+      if (msg.role === 'user') {
+        for (let partIdx = 0; partIdx < msg.parts.length; partIdx++) {
+          const part = msg.parts[partIdx];
+          const apiId =
+            isFunctionResponsePart(part) &&
+            typeof part.functionResponse.id === 'string'
+              ? `resp_${part.functionResponse.id}_${turnSalt}_${partIdx}`
+              : isFunctionCallPart(part) &&
+                  typeof part.functionCall.id === 'string'
+                ? `call_${part.functionCall.id}_${turnSalt}_${partIdx}`
+                : undefined;
+          const id =
+            apiId || getStableId(part, this.nodeIdentityMap, turnSalt, partIdx);
+          const node: ConcreteNode = {
+            id,
+            timestamp: Date.now(),
+            type: isFunctionResponsePart(part)
+              ? NodeType.TOOL_EXECUTION
+              : NodeType.USER_PROMPT,
+            role: 'user',
+            payload: part,
+            turnId,
+          };
+          nodes.push(node);
        }
      } else if (msg.role === 'model') {
-        this.currentEpisode = parseModelParts(
-          msg,
-          this.currentEpisode,
-          this.pendingCallParts,
-          this.pendingCallPartsWithoutId,
-          this.nodeIdentityMap,
-        );
-      }
-    }
-  }
-
-  getNodes(): ConcreteNode[] {
-    const copy = [...this.episodes];
-    if (this.currentEpisode) {
-      const activeEp = {
-        ...this.currentEpisode,
-        concreteNodes: [...(this.currentEpisode.concreteNodes || [])],
-      };
-      finalizeYield(activeEp);
-      if (isCompleteEpisode(activeEp)) {
-        copy.push(activeEp);
-      }
-    }
-
-    const nodes: ConcreteNode[] = [];
-    for (const ep of copy) {
-      if (ep.concreteNodes) {
-        for (const child of ep.concreteNodes) {
-          nodes.push(child);
+        for (let partIdx = 0; partIdx < msg.parts.length; partIdx++) {
+          const part = msg.parts[partIdx];
+          const apiId =
+            isFunctionCallPart(part) && typeof part.functionCall.id === 'string'
+              ? `call_${part.functionCall.id}_${turnSalt}_${partIdx}`
+              : undefined;
+          const id =
+            apiId || getStableId(part, this.nodeIdentityMap, turnSalt, partIdx);
+          const node: ConcreteNode = {
+            id,
+            timestamp: Date.now(),
+            type: isFunctionCallPart(part)
+              ? NodeType.TOOL_EXECUTION
+              : NodeType.AGENT_THOUGHT,
+            role: 'model',
+            payload: part,
+            turnId,
+          };
+          nodes.push(node);
        }
      }
    }
+
+    debugLogger.log(
+      `[ContextGraphBuilder] Mirror Graph built with ${nodes.length} nodes.`,
+    );
    return nodes;
  }
 }
-
-function parseToolResponses(
-  msg: Content,
-  currentEpisode: Partial<Episode> | null,
-  pendingCallParts: Map<string, Part>,
-  pendingCallPartsWithoutId: Part[],
-  tokenCalculator: ContextTokenCalculator,
-  nodeIdentityMap: WeakMap<object, string>,
-): Partial<Episode> {
-  if (!currentEpisode) {
-    currentEpisode = {
-      id: getStableId(msg, nodeIdentityMap),
-
-      concreteNodes: [],
-    };
-  }
-
-  const parts = msg.parts || [];
-  for (const part of parts) {
-    if (part.functionResponse) {
-      const callId = part.functionResponse.id || '';
-      let matchingCall = pendingCallParts.get(callId);
-
-      if (!matchingCall && pendingCallPartsWithoutId.length > 0) {
-        const idx = pendingCallPartsWithoutId.findIndex(
-          (p) => p.functionCall?.name === part.functionResponse!.name,
-        );
-        if (idx !== -1) {
-          matchingCall = pendingCallPartsWithoutId[idx];
-          pendingCallPartsWithoutId.splice(idx, 1);
-        } else {
-          matchingCall = pendingCallPartsWithoutId.shift();
-        }
-      }
-
-      const intentTokens = matchingCall
-        ? tokenCalculator.estimateTokensForParts([matchingCall])
-        : 0;
-      const obsTokens = tokenCalculator.estimateTokensForParts([part]);
-
-      const step: ToolExecution = {
-        id: getStableId(part, nodeIdentityMap),
-        timestamp: Date.now(),
-        type: 'TOOL_EXECUTION',
-        toolName: part.functionResponse.name || 'unknown',
-        intent: isRecord(matchingCall?.functionCall?.args)
-          ? matchingCall.functionCall.args
-          : {},
-        observation: isRecord(part.functionResponse.response)
-          ? part.functionResponse.response
-          : {},
-        tokens: {
-          intent: intentTokens,
-          observation: obsTokens,
-        },
-      };
-
-      currentEpisode.concreteNodes = [
-        ...(currentEpisode.concreteNodes || []),
-        step,
-      ];
-      if (callId) pendingCallParts.delete(callId);
-    }
-  }
-  return currentEpisode;
-}
-
-function parseUserParts(
-  msg: Content,
-  nodeIdentityMap: WeakMap<object, string>,
-): Partial<Episode> {
-  const semanticParts: SemanticPart[] = [];
-  const parts = msg.parts || [];
-  for (const p of parts) {
-    if (p.text !== undefined)
-      semanticParts.push({ type: 'text', text: p.text });
-    else if (p.inlineData)
-      semanticParts.push({
-        type: 'inline_data',
-        mimeType: p.inlineData.mimeType || '',
-        data: p.inlineData.data || '',
-      });
-    else if (p.fileData)
-      semanticParts.push({
-        type: 'file_data',
-        mimeType: p.fileData.mimeType || '',
-        fileUri: p.fileData.fileUri || '',
-      });
-    else if (!p.functionResponse)
-      semanticParts.push({ type: 'raw_part', part: p }); // Preserve unknowns
-  }
-
-  const baseObj = parts.length > 0 ? parts[0] : msg;
-  const trigger: UserPrompt = {
-    id: getStableId(baseObj, nodeIdentityMap),
-    timestamp: Date.now(),
-    type: 'USER_PROMPT',
-    semanticParts,
-  };
-  return {
-    id: getStableId(msg, nodeIdentityMap),
-
-    concreteNodes: [trigger],
-  };
-}
-
-function parseModelParts(
-  msg: Content,
-  currentEpisode: Partial<Episode> | null,
-  pendingCallParts: Map<string, Part>,
-  pendingCallPartsWithoutId: Part[],
-  nodeIdentityMap: WeakMap<object, string>,
-): Partial<Episode> {
-  if (!currentEpisode) {
-    currentEpisode = {
-      id: getStableId(msg, nodeIdentityMap),
-
-      concreteNodes: [],
-    };
-  }
-
-  const parts = msg.parts || [];
-  for (const part of parts) {
-    if (part.functionCall) {
-      const callId = part.functionCall.id || '';
-      if (callId) {
-        pendingCallParts.set(callId, part);
-      } else {
-        const lastIdx = pendingCallPartsWithoutId.length - 1;
-        const lastPart = pendingCallPartsWithoutId[lastIdx];
-
-        if (
-          lastPart &&
-          lastPart.functionCall &&
-          lastPart.functionCall.name === part.functionCall.name
-        ) {
-          // Replace the previous chunk with the more complete one
-          pendingCallPartsWithoutId[lastIdx] = part;
-        } else {
-          pendingCallPartsWithoutId.push(part);
-        }
-      }
-    } else if (part.text) {
-      const thought: AgentThought = {
-        id: getStableId(part, nodeIdentityMap),
-        timestamp: Date.now(),
-        type: 'AGENT_THOUGHT',
-        text: part.text,
-      };
-
-      currentEpisode.concreteNodes = [
-        ...(currentEpisode.concreteNodes || []),
-        thought,
-      ];
-    }
-  }
-  return currentEpisode;
-}
-
-function finalizeYield(currentEpisode: Partial<Episode>) {
-  if (currentEpisode.concreteNodes && currentEpisode.concreteNodes.length > 0) {
-    const yieldNode: AgentYield = {
-      id: randomUUID(),
-      timestamp: Date.now(),
-      type: 'AGENT_YIELD',
-      text: 'Yield', // Synthesized yield since we don't have the original concrete node
-    };
-    const existingNodes = currentEpisode.concreteNodes || [];
-    currentEpisode.concreteNodes = [...existingNodes, yieldNode];
-  }
-}
@@ -6,24 +6,22 @@

 import type { Part } from '@google/genai';

-export type NodeType =
-  // Organic Concrete Nodes
-  | 'USER_PROMPT'
-  | 'SYSTEM_EVENT'
-  | 'AGENT_THOUGHT'
-  | 'TOOL_EXECUTION'
-  | 'AGENT_YIELD'
+/**
+ * Basic Node Interface
+ * Every element in the Context Graph is a Node.
+ */

-  // Synthetic Concrete Nodes
-  | 'SNAPSHOT'
-  | 'ROLLING_SUMMARY'
-  | 'MASKED_TOOL'
+export enum NodeType {
+  USER_PROMPT = 'USER_PROMPT',
+  SYSTEM_EVENT = 'SYSTEM_EVENT',
+  AGENT_THOUGHT = 'AGENT_THOUGHT',
+  TOOL_EXECUTION = 'TOOL_EXECUTION',
+  MASKED_TOOL = 'MASKED_TOOL',
+  AGENT_YIELD = 'AGENT_YIELD',
+  SNAPSHOT = 'SNAPSHOT',
+  ROLLING_SUMMARY = 'ROLLING_SUMMARY',
+}

-  // Logical Nodes
-  | 'TASK'
-  | 'EPISODE';
-
-/** Base interface for all nodes in the Episodic Context Graph */
 export interface Node {
  readonly id: string;
  readonly type: NodeType;
@@ -32,11 +30,20 @@ export interface Node {
 /**
 * Concrete Nodes: The atomic, renderable pieces of data.
 * These are the actual "planks" of the Nodes of Theseus.
+ *
+ * Each ConcreteNode is now a 1:1 wrapper around a Gemini Part,
+ * ensuring 100% fidelity during reconstruction.
 */
 export interface BaseConcreteNode extends Node {
+  readonly type: NodeType;
  readonly timestamp: number;
-  /** The ID of the Logical Node (e.g., Episode) that structurally owns this node */
-  readonly logicalParentId?: string;
+  /** The role of the turn this part belongs to */
+  readonly role: 'user' | 'model';
+  /** The original, high-fidelity Part object from the API */
+  readonly payload: Part;
+
+  /** The ID of the specific turn in history this node belongs to. Unique per turn. */
+  readonly turnId: string;

  /** If this node replaced a single node 1:1 (e.g., masking), this points to the original */
  readonly replacesId?: string;
@@ -45,50 +52,19 @@ export interface BaseConcreteNode extends Node {
  readonly abstractsIds?: readonly string[];
 }

-/**
- * Semantic Parts for User Prompts
- */
-export interface SemanticTextPart {
-  readonly type: 'text';
-  readonly text: string;
-}
-
-export interface SemanticInlineDataPart {
-  readonly type: 'inline_data';
-  readonly mimeType: string;
-  readonly data: string;
-}
-
-export interface SemanticFileDataPart {
-  readonly type: 'file_data';
-  readonly mimeType: string;
-  readonly fileUri: string;
-}
-
-export interface SemanticRawPart {
-  readonly type: 'raw_part';
-  readonly part: Part;
-}
-
-export type SemanticPart =
-  | SemanticTextPart
-  | SemanticInlineDataPart
-  | SemanticFileDataPart
-  | SemanticRawPart;
-
 /**
 * Trigger Nodes
 * Events that wake the agent up and initiate an Episode.
 */
 export interface UserPrompt extends BaseConcreteNode {
-  readonly type: 'USER_PROMPT';
-  readonly semanticParts: readonly SemanticPart[];
+  readonly type: NodeType.USER_PROMPT;
+  readonly role: 'user';
 }

 export interface SystemEvent extends BaseConcreteNode {
-  readonly type: 'SYSTEM_EVENT';
+  readonly type: NodeType.SYSTEM_EVENT;
  readonly name: string;
-  readonly payload: Record<string, unknown>;
+  readonly payload: Part; // System events are usually injected as user text parts
 }

 export type EpisodeTrigger = UserPrompt | SystemEvent;
@@ -98,30 +74,16 @@ export type EpisodeTrigger = UserPrompt | SystemEvent;
 * The internal autonomous actions taken by the agent during its loop.
 */
 export interface AgentThought extends BaseConcreteNode {
-  readonly type: 'AGENT_THOUGHT';
-  readonly text: string;
+  readonly type: NodeType.AGENT_THOUGHT;
+  readonly role: 'model';
 }

 export interface ToolExecution extends BaseConcreteNode {
-  readonly type: 'TOOL_EXECUTION';
-  readonly toolName: string;
-  readonly intent: Record<string, unknown>;
-  readonly observation: string | Record<string, unknown>;
-  readonly tokens: {
-    readonly intent: number;
-    readonly observation: number;
-  };
+  readonly type: NodeType.TOOL_EXECUTION;
 }

 export interface MaskedTool extends BaseConcreteNode {
-  readonly type: 'MASKED_TOOL';
-  readonly toolName: string;
-  readonly intent?: Record<string, unknown>;
-  readonly observation?: string | Record<string, unknown>;
-  readonly tokens: {
-    readonly intent: number;
-    readonly observation: number;
-  };
+  readonly type: NodeType.MASKED_TOOL;
 }

 export type EpisodeStep = AgentThought | ToolExecution | MaskedTool;
@@ -131,8 +93,8 @@ export type EpisodeStep = AgentThought | ToolExecution | MaskedTool;
 * The final message where the agent yields control back to the user.
 */
 export interface AgentYield extends BaseConcreteNode {
-  readonly type: 'AGENT_YIELD';
-  readonly text: string;
+  readonly type: NodeType.AGENT_YIELD;
+  readonly role: 'model';
 }

 /**
@@ -140,13 +102,11 @@ export interface AgentYield extends BaseConcreteNode {
 * Processors that generate summaries emit explicit synthetic nodes.
 */
 export interface Snapshot extends BaseConcreteNode {
-  readonly type: 'SNAPSHOT';
-  readonly text: string;
+  readonly type: NodeType.SNAPSHOT;
 }

 export interface RollingSummary extends BaseConcreteNode {
-  readonly type: 'ROLLING_SUMMARY';
-  readonly text: string;
+  readonly type: NodeType.ROLLING_SUMMARY;
 }

 export type SyntheticLeaf = Snapshot | RollingSummary;
@@ -161,62 +121,34 @@ export type ConcreteNode =
  | Snapshot
  | RollingSummary;

-/**
- * Logical Nodes
- * These define hierarchy and grouping. They do not directly render to Gemini.
- */
-export interface Episode extends Node {
-  readonly type: 'EPISODE';
-  /** References to the Concrete Node IDs that conceptually belong to this Episode. */
-  concreteNodes: readonly ConcreteNode[];
-}
-
-export interface Task extends Node {
-  readonly type: 'TASK';
-  readonly goal: string;
-  readonly status: 'active' | 'completed' | 'failed';
-  /** References to the Episode IDs that belong to this task */
-  readonly episodeIds: readonly string[];
-}
-
-export type LogicalNode = Task | Episode;
-
-export function isEpisode(node: Node): node is Episode {
-  return node.type === 'EPISODE';
-}
-
-export function isTask(node: Node): node is Task {
-  return node.type === 'TASK';
-}
-
 export function isAgentThought(node: Node): node is AgentThought {
-  return node.type === 'AGENT_THOUGHT';
+  return node.type === NodeType.AGENT_THOUGHT;
 }

 export function isAgentYield(node: Node): node is AgentYield {
-  return node.type === 'AGENT_YIELD';
+  return node.type === NodeType.AGENT_YIELD;
 }

 export function isToolExecution(node: Node): node is ToolExecution {
-  return node.type === 'TOOL_EXECUTION';
+  return node.type === NodeType.TOOL_EXECUTION;
 }

 export function isMaskedTool(node: Node): node is MaskedTool {
-  return node.type === 'MASKED_TOOL';
+  return node.type === NodeType.MASKED_TOOL;
 }

 export function isUserPrompt(node: Node): node is UserPrompt {
-  return node.type === 'USER_PROMPT';
+  return node.type === NodeType.USER_PROMPT;
 }

 export function isSystemEvent(node: Node): node is SystemEvent {
-  return node.type === 'SYSTEM_EVENT';
+  return node.type === NodeType.SYSTEM_EVENT;
 }

 export function isSnapshot(node: Node): node is Snapshot {
-  return node.type === 'SNAPSHOT';
+  return node.type === NodeType.SNAPSHOT;
 }

 export function isRollingSummary(node: Node): node is RollingSummary {
-  return node.type === 'ROLLING_SUMMARY';
+  return node.type === NodeType.ROLLING_SUMMARY;
 }