diff --git a/packages/core/src/context/contextManager.golden.test.ts b/packages/core/src/context/contextManager.golden.test.ts index e2966d205a..e0772d367c 100644 --- a/packages/core/src/context/contextManager.golden.test.ts +++ b/packages/core/src/context/contextManager.golden.test.ts @@ -148,9 +148,14 @@ describe('ContextManager Golden Tests', () => { it('should process history and match golden snapshot', async () => { const history = createLargeHistory(); - ( - contextManager as unknown as { pristineEpisodes: Episode[] } - ).pristineEpisodes = (contextManager as unknown as import("../pipeline.js").ContextWorkingBuffer).env.irMapper.toIr(history, (contextManager as unknown as import("../pipeline.js").ContextWorkingBuffer).env.tokenCalculator); + // Use the actual public methods or carefully type the internal state for testing + // To seed the manager purely for testing without invoking generateContent, we bypass the pipeline: + const managerAsAny = contextManager as unknown as { + pristineEpisodes: Episode[]; + env: { irMapper: { toIr(h: unknown, t: unknown): Episode[] }, tokenCalculator: unknown } + }; + managerAsAny.pristineEpisodes = managerAsAny.env.irMapper.toIr(history, managerAsAny.env.tokenCalculator); + const result = await contextManager.projectCompressedHistory(); expect(result).toMatchSnapshot(); }); diff --git a/packages/core/src/context/processors/blobDegradationProcessor.test.ts b/packages/core/src/context/processors/blobDegradationProcessor.test.ts index 4ca041101d..cfb84e6336 100644 --- a/packages/core/src/context/processors/blobDegradationProcessor.test.ts +++ b/packages/core/src/context/processors/blobDegradationProcessor.test.ts @@ -6,6 +6,7 @@ import { describe, it, expect } from 'vitest'; import { BlobDegradationProcessor } from './blobDegradationProcessor.js'; import { + createMockProcessArgs, createMockEnvironment, createDummyNode, } from '../testing/contextTestUtils.js'; @@ -34,11 +35,7 @@ describe('BlobDegradationProcessor', () => { const targets = [prompt]; - const result = await processor.process({ - buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - targets, - inbox: undefined as unknown as import('../pipeline.js').InboxSnapshot, - }); + const result = await processor.process(createMockProcessArgs(targets)); expect(result.length).toBe(1); const modifiedPrompt = result[0] as UserPrompt; @@ -51,7 +48,7 @@ describe('BlobDegradationProcessor', () => { expect(modifiedPrompt.semanticParts[2]).toEqual(parts[2]); // The inline_data part should be replaced with text - const degradedPart = modifiedPrompt.semanticParts[1] as import('../ir/types.js').TextPart; + const degradedPart = modifiedPrompt.semanticParts[1] as unknown as { type: string, text: string }; expect(degradedPart.type).toBe('text'); expect(degradedPart.text).toContain('[Multi-Modal Blob (image/png, 0.00MB) degraded to text'); }); @@ -76,11 +73,7 @@ describe('BlobDegradationProcessor', () => { const targets = [prompt]; - const result = await processor.process({ - buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - targets, - inbox: undefined as unknown as import('../pipeline.js').InboxSnapshot, - }); + const result = await processor.process(createMockProcessArgs(targets)); const modifiedPrompt = result[0] as UserPrompt; expect(modifiedPrompt.semanticParts.length).toBe(2); @@ -96,11 +89,7 @@ describe('BlobDegradationProcessor', () => { const processor = BlobDegradationProcessor.create(env, {}); const targets: Array = []; - const result = await processor.process({ - buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - targets, - inbox: undefined as unknown as import('../pipeline.js').InboxSnapshot, - }); + const result = await processor.process(createMockProcessArgs(targets)); expect(result).toBe(targets); }); diff --git a/packages/core/src/context/processors/nodeDistillationProcessor.test.ts b/packages/core/src/context/processors/nodeDistillationProcessor.test.ts index 7a4606f1b0..bc095697f3 100644 --- a/packages/core/src/context/processors/nodeDistillationProcessor.test.ts +++ b/packages/core/src/context/processors/nodeDistillationProcessor.test.ts @@ -6,50 +6,49 @@ import { describe, it, expect, vi } from 'vitest'; import { NodeDistillationProcessor } from './nodeDistillationProcessor.js'; import { + createMockProcessArgs, createMockEnvironment, createDummyNode, createDummyToolNode, createMockGenerateContentResponse } from '../testing/contextTestUtils.js'; import type { UserPrompt, AgentThought, ToolExecution } from '../ir/types.js'; +import type { BaseLlmClient } from '../../core/baseLlmClient.js'; describe('NodeDistillationProcessor', () => { it('should trigger summarization via LLM for long text parts', async () => { const mockLlmClient = { generateContent: vi.fn().mockResolvedValue(createMockGenerateContentResponse('Mocked Summary!')), // length = 15 - }; + } as unknown as BaseLlmClient; + // Use charsPerToken=1 naturally. const env = createMockEnvironment({ - llmClient: mockLlmClient as unknown as import("../pipeline.js").ContextWorkingBuffer, + llmClient: mockLlmClient, }); const processor = NodeDistillationProcessor.create(env, { nodeThresholdTokens: 10, }); + const longText = 'A'.repeat(50); // 50 chars - const prompt = createDummyNode('ep1', 'USER_PROMPT', 3800, { + const prompt = createDummyNode('ep1', 'USER_PROMPT', 50, { semanticParts: [ - { type: 'text', text: 'This text is way longer than 10 characters and needs compression' } + { type: 'text', text: longText } ], }, 'prompt-id') as UserPrompt; - const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 1500, { - text: 'The model is thinking something incredibly long and verbose that exceeds 10 chars', + const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 50, { + text: longText, }, 'thought-id') as AgentThought; - const tool = createDummyToolNode('ep1', 50, 1000, { - observation: { result: 'Massive tool JSON observation payload' }, - tokens: { intent: 50, observation: 1000 } + const tool = createDummyToolNode('ep1', 5, 500, { + observation: { result: 'A'.repeat(500) }, }, 'tool-id'); const targets = [prompt, thought, tool]; - const result = await processor.process({ - buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - targets, - inbox: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - }); + const result = await processor.process(createMockProcessArgs(targets)); expect(result.length).toBe(3); @@ -57,15 +56,15 @@ describe('NodeDistillationProcessor', () => { const compressedPrompt = result[0] as UserPrompt; expect(compressedPrompt.id).not.toBe(prompt.id); expect(compressedPrompt.semanticParts[0].type).toBe('text'); - expect((compressedPrompt.semanticParts[0] as unknown as import("../pipeline.js").ContextWorkingBuffer).text).toBe('Mocked Summary!'); + expect((compressedPrompt.semanticParts[0] as unknown as {text: string}).text).toBe('Mocked Summary!'); + // 2. Agent Thought const compressedThought = result[1] as AgentThought; - expect(compressedThought.id).toMatch(/^mock-uuid-/); expect(compressedThought.id).not.toBe(thought.id); expect(compressedThought.text).toBe('Mocked Summary!'); + // 3. Tool Execution const compressedTool = result[2] as ToolExecution; - expect(compressedTool.id).toMatch(/^mock-uuid-/); expect(compressedTool.id).not.toBe(tool.id); expect(compressedTool.observation).toEqual({ summary: 'Mocked Summary!' }); @@ -75,34 +74,31 @@ describe('NodeDistillationProcessor', () => { it('should ignore nodes that are below the threshold', async () => { const mockLlmClient = { generateContent: vi.fn().mockResolvedValue(createMockGenerateContentResponse('S')), // length = 1 - }; + } as unknown as BaseLlmClient; const env = createMockEnvironment({ - llmClient: mockLlmClient as unknown as import("../pipeline.js").ContextWorkingBuffer, + llmClient: mockLlmClient, }); const processor = NodeDistillationProcessor.create(env, { nodeThresholdTokens: 100, // Very high threshold }); + const shortText = 'Short text'; // 10 chars - const prompt = createDummyNode('ep1', 'USER_PROMPT', 3800, { + const prompt = createDummyNode('ep1', 'USER_PROMPT', 10, { semanticParts: [ - { type: 'text', text: 'Short text' } // Below threshold + { type: 'text', text: shortText } ], }, 'prompt-id') as UserPrompt; - const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 1500, { - text: 'Short thought', // Below threshold + const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 13, { + text: 'Short thought', }, 'thought-id') as AgentThought; const targets = [prompt, thought]; - const result = await processor.process({ - buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - targets, - inbox: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - }); + const result = await processor.process(createMockProcessArgs(targets)); expect(result.length).toBe(2); diff --git a/packages/core/src/context/processors/nodeTruncationProcessor.test.ts b/packages/core/src/context/processors/nodeTruncationProcessor.test.ts index d2f98bb64e..325163b54d 100644 --- a/packages/core/src/context/processors/nodeTruncationProcessor.test.ts +++ b/packages/core/src/context/processors/nodeTruncationProcessor.test.ts @@ -3,118 +3,88 @@ * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi } from 'vitest'; +import { describe, it, expect } from 'vitest'; import { NodeTruncationProcessor } from './nodeTruncationProcessor.js'; import { + createMockProcessArgs, createMockEnvironment, createDummyNode, } from '../testing/contextTestUtils.js'; import type { UserPrompt, AgentThought, AgentYield } from '../ir/types.js'; -import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; describe('NodeTruncationProcessor', () => { it('should truncate nodes that exceed maxTokensPerNode', async () => { + // env.tokenCalculator uses charsPerToken=1 natively. const env = createMockEnvironment(); - const mockTokenCalculator = new ContextTokenCalculator(1, env.behaviorRegistry) as unknown as import("../pipeline.js").ContextWorkingBuffer; - mockTokenCalculator.tokensToChars = vi.fn().mockReturnValue(10); // Limit is 10 chars - - mockTokenCalculator.estimateTokensForString = vi.fn((text: string) => { - if (text.includes('OMITTED')) return 1; // Summary size - return 100; // Original size - }); - mockTokenCalculator.estimateTokensForParts = vi.fn(() => 1); - - (env as unknown as import("../pipeline.js").ContextWorkingBuffer).tokenCalculator = mockTokenCalculator; const processor = NodeTruncationProcessor.create(env, { - maxTokensPerNode: 1, // Will equal 10 chars limit + maxTokensPerNode: 10, // 10 chars limit }); + const longText = 'A'.repeat(50); // 50 tokens - const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, { - semanticParts: [ - { type: 'text', text: 'This text is way longer than 10 characters and needs truncation' } - ], + const prompt = createDummyNode('ep1', 'USER_PROMPT', 50, { + semanticParts: [{ type: 'text', text: longText }] }, 'prompt-id') as UserPrompt; - const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 100, { - text: 'The model is thinking something incredibly long and verbose that exceeds 10 chars', + const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 50, { + text: longText, }, 'thought-id') as AgentThought; - const yieldNode = createDummyNode('ep1', 'AGENT_YIELD', 100, { - text: 'Final output yield that is also extremely long', + const yieldNode = createDummyNode('ep1', 'AGENT_YIELD', 50, { + text: longText, }, 'yield-id') as AgentYield; const targets = [prompt, thought, yieldNode]; - const result = await processor.process({ - buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - targets, - inbox: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - }); + const result = await processor.process(createMockProcessArgs(targets)); expect(result.length).toBe(3); // 1. User Prompt const squashedPrompt = result[0] as UserPrompt; - expect(squashedPrompt.id).toBe('mock-uuid-1'); expect(squashedPrompt.id).not.toBe(prompt.id); expect(squashedPrompt.semanticParts[0].type).toBe('text'); - expect((squashedPrompt.semanticParts[0] as unknown as import("../pipeline.js").ContextWorkingBuffer).text).toContain('[... OMITTED'); + expect((squashedPrompt.semanticParts[0] as unknown as { text: string }).text).toContain('[... OMITTED'); // 2. Agent Thought const squashedThought = result[1] as AgentThought; - expect(squashedThought.id).toBe('mock-uuid-2'); expect(squashedThought.id).not.toBe(thought.id); expect(squashedThought.text).toContain('[... OMITTED'); // 3. Agent Yield const squashedYield = result[2] as AgentYield; - expect(squashedYield.id).toBe('mock-uuid-3'); expect(squashedYield.id).not.toBe(yieldNode.id); expect(squashedYield.text).toContain('[... OMITTED'); }); it('should ignore nodes that are below maxTokensPerNode', async () => { const env = createMockEnvironment(); - const mockTokenCalculator = new ContextTokenCalculator(1, env.behaviorRegistry) as unknown as import("../pipeline.js").ContextWorkingBuffer; - mockTokenCalculator.tokensToChars = vi.fn().mockReturnValue(100); - - mockTokenCalculator.estimateTokensForString = vi.fn((text: string) => text.length); - mockTokenCalculator.estimateTokensForParts = vi.fn(() => 5); - mockTokenCalculator.getTokenCost = vi.fn(() => 5); - - (env as unknown as import("../pipeline.js").ContextWorkingBuffer).tokenCalculator = mockTokenCalculator; const processor = NodeTruncationProcessor.create(env, { - maxTokensPerNode: 100, + maxTokensPerNode: 100, // 100 chars limit }); + const shortText = 'Short text'; // 10 chars - const prompt = createDummyNode('ep1', 'USER_PROMPT', 5, { - semanticParts: [ - { type: 'text', text: 'Short text' } // 10 chars - ], + const prompt = createDummyNode('ep1', 'USER_PROMPT', 10, { + semanticParts: [{ type: 'text', text: shortText }] }, 'prompt-id') as UserPrompt; - const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 5, { + const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 13, { text: 'Short thought', // 13 chars }, 'thought-id') as AgentThought; const targets = [prompt, thought]; - const result = await processor.process({ - buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - targets, - inbox: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - }); + const result = await processor.process(createMockProcessArgs(targets)); expect(result.length).toBe(2); // 1. User Prompt (untouched) const squashedPrompt = result[0] as UserPrompt; expect(squashedPrompt.id).toBe(prompt.id); - expect((squashedPrompt.semanticParts[0] as unknown as import("../pipeline.js").ContextWorkingBuffer).text).not.toContain('[... OMITTED'); + expect((squashedPrompt.semanticParts[0] as unknown as { text: string }).text).not.toContain('[... OMITTED'); // 2. Agent Thought (untouched) const untouchedThought = result[1] as AgentThought; diff --git a/packages/core/src/context/processors/rollingSummaryProcessor.test.ts b/packages/core/src/context/processors/rollingSummaryProcessor.test.ts index d8aae65b0e..528aab97d9 100644 --- a/packages/core/src/context/processors/rollingSummaryProcessor.test.ts +++ b/packages/core/src/context/processors/rollingSummaryProcessor.test.ts @@ -5,7 +5,8 @@ */ import { describe, it, expect } from 'vitest'; import { RollingSummaryProcessor } from './rollingSummaryProcessor.js'; -import { createMockEnvironment, createDummyNode } from '../testing/contextTestUtils.js'; +import { createMockProcessArgs, + createMockEnvironment, createDummyNode } from '../testing/contextTestUtils.js'; describe('RollingSummaryProcessor', () => { it('should initialize with correct default options', () => { @@ -32,7 +33,7 @@ describe('RollingSummaryProcessor', () => { createDummyNode('ep1', 'AGENT_YIELD', 50, { text: text50 }, 'id3'), ]; - const result = await processor.process({ targets, buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, inbox: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer }); + const result = await processor.process(createMockProcessArgs(targets)); // 3 nodes at 50 cost each. // The first node (id1) is the initial USER_PROMPT and is always skipped by RollingSummaryProcessor. @@ -58,7 +59,7 @@ describe('RollingSummaryProcessor', () => { createDummyNode('ep1', 'AGENT_THOUGHT', 10, { text: text10 }, 'id2'), ]; - const result = await processor.process({ targets, buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, inbox: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer }); + const result = await processor.process(createMockProcessArgs(targets)); // Deficit accumulator reaches 10. This is < 100 limit, and total summarizable nodes < 2 anyway. expect(result.length).toBe(2); diff --git a/packages/core/src/context/processors/toolMaskingProcessor.test.ts b/packages/core/src/context/processors/toolMaskingProcessor.test.ts index 6c892afa00..7c85f1bca5 100644 --- a/packages/core/src/context/processors/toolMaskingProcessor.test.ts +++ b/packages/core/src/context/processors/toolMaskingProcessor.test.ts @@ -3,51 +3,46 @@ * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi } from 'vitest'; +import { describe, it, expect } from 'vitest'; import { ToolMaskingProcessor } from './toolMaskingProcessor.js'; import { + createMockProcessArgs, createMockEnvironment, createDummyToolNode, } from '../testing/contextTestUtils.js'; +import type { ToolExecution } from '../ir/types.js'; describe('ToolMaskingProcessor', () => { it('should write large strings to disk and replace them with a masked pointer', async () => { const env = createMockEnvironment(); - // 1 token = 1 char for simplicity - // Fake token calculator says new tokens are 5 - env.tokenCalculator.estimateTokensForParts = vi.fn().mockReturnValue(5); - env.tokenCalculator.getTokenCost = vi.fn().mockReturnValue(150); + // env uses charsPerToken=1 natively. + // original string lengths > stringLengthThresholdTokens (which is 10) will be masked const processor = ToolMaskingProcessor.create(env, { stringLengthThresholdTokens: 10, }); + const longString = 'A'.repeat(500); // 500 chars - const toolStep = createDummyToolNode('ep1', 50, 100, { + const toolStep = createDummyToolNode('ep1', 50, 500, { observation: { - result: 'this is a really long string that should get masked out because it exceeds 10 chars', - metadata: 'short', + result: longString, + metadata: 'short', // 5 chars, will not be masked }, }); - const result = await processor.process({ - buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - targets: [toolStep], - inbox: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - }); + const result = await processor.process(createMockProcessArgs([toolStep])); expect(result.length).toBe(1); - const masked = result[0]; + const masked = result[0] as ToolExecution; // It should have generated a new ID because it modified it expect(masked.id).not.toBe(toolStep.id); // It should have masked the observation - const obs = (masked as unknown as import("../pipeline.js").ContextWorkingBuffer).observation; + const obs = masked.observation as { result: string, metadata: string }; expect(obs.result).toContain(''); expect(obs.metadata).toBe('short'); // Untouched - - // Transformation logged }); it('should skip unmaskable tools', async () => { @@ -57,7 +52,6 @@ describe('ToolMaskingProcessor', () => { stringLengthThresholdTokens: 10, }); - const toolStep = createDummyToolNode('ep1', 10, 10, { toolName: 'activate_skill', observation: { @@ -65,11 +59,7 @@ describe('ToolMaskingProcessor', () => { } }); - const result = await processor.process({ - buffer: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - targets: [toolStep], - inbox: undefined as unknown as import('../pipeline.js').ContextWorkingBuffer, - }); + const result = await processor.process(createMockProcessArgs([toolStep])); // Returned the exact same object reference expect(result[0]).toBe(toolStep); diff --git a/packages/core/src/context/sidecar/inbox.ts b/packages/core/src/context/sidecar/inbox.ts index cdd40ee17f..1b96dc835b 100644 --- a/packages/core/src/context/sidecar/inbox.ts +++ b/packages/core/src/context/sidecar/inbox.ts @@ -36,6 +36,16 @@ export class InboxSnapshotImpl implements InboxSnapshot { getMessages(topic: string): ReadonlyArray> { const raw = this.messages.filter((m) => m.topic === topic); + /* + * Architectural Justification for Unchecked Cast: + * The Inbox is a heterogeneous event bus designed to support arbitrary, declarative + * routing via configuration files (where topics are just strings). Because TypeScript + * completely erases generic type information () at runtime, the central array + * can only hold `unknown` payloads. To enforce strict type safety without a central + * registry (which would break decoupling) or heavy runtime validation (Zod schemas), + * we must assert the type boundary here. The contract relies on the Worker and Processor + * agreeing on the payload structure associated with the configured topic string. + */ // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion return raw as ReadonlyArray>; } diff --git a/packages/core/src/context/sidecar/registry.test.ts b/packages/core/src/context/sidecar/registry.test.ts index a1d98094af..b51523b7d0 100644 --- a/packages/core/src/context/sidecar/registry.test.ts +++ b/packages/core/src/context/sidecar/registry.test.ts @@ -26,7 +26,7 @@ describe('SidecarRegistry', () => { const workerDef: ContextWorkerDef = { id: 'TestWorker', schema: { type: 'object' }, - create: () => ({} as unknown as import('../pipeline.js').ContextProcessor), + create: () => ({} as unknown as import('../pipeline.js').ContextWorker), }; registry.registerWorker(workerDef); @@ -54,10 +54,10 @@ describe('SidecarRegistry', () => { registry.registerWorker({ id: 'TestWorker', schema: { title: 'workerSchema' }, - create: () => ({} as unknown as import('../pipeline.js').ContextProcessor), + create: () => ({} as unknown as import('../pipeline.js').ContextWorker), }); - const schemas = registry.getSchemas() as unknown as Array>; + const schemas = registry.getSchemas() as unknown as Array<{title?: string}>; expect(schemas.length).toBe(2); expect(schemas.find(s => s.title === 'processorSchema')).toBeDefined(); expect(schemas.find(s => s.title === 'workerSchema')).toBeDefined(); @@ -66,7 +66,7 @@ describe('SidecarRegistry', () => { it('should safely clear the registry', () => { const registry = new SidecarRegistry(); registry.registerProcessor({ id: 'TestProcessor', schema: {}, create: () => ({} as unknown as import('../pipeline.js').ContextProcessor) }); - registry.registerWorker({ id: 'TestWorker', schema: {}, create: () => ({} as unknown as import('../pipeline.js').ContextProcessor) }); + registry.registerWorker({ id: 'TestWorker', schema: {}, create: () => ({} as unknown as import('../pipeline.js').ContextWorker) }); registry.clear(); diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts index 413b1168f8..84163f8aa0 100644 --- a/packages/core/src/context/testing/contextTestUtils.ts +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -127,6 +127,44 @@ export function createMockEnvironment( * Creates a block of synthetic conversation history designed to consume a specific number of tokens. * Assumes roughly 4 characters per token for standard English text. */ +import { InboxSnapshotImpl } from '../sidecar/inbox.js'; +import type { ContextWorkingBuffer, InboxMessage, ProcessArgs } from '../pipeline.js'; + +export class FakeContextWorkingBuffer implements ContextWorkingBuffer { + readonly nodes: readonly ConcreteNode[]; + private readonly nodesById = new Map(); + private readonly nodesByEpisode = new Map(); + + constructor(nodes: readonly ConcreteNode[]) { + this.nodes = nodes; + for (const node of nodes) { + this.nodesById.set(node.id, node); + const parentId = node.logicalParentId || 'orphan'; + const epNodes = this.nodesByEpisode.get(parentId) || []; + epNodes.push(node); + this.nodesByEpisode.set(parentId, epNodes); + } + } + + getPristineNode(id: string): ConcreteNode | undefined { + return this.nodesById.get(id); + } + + getLineage(id: string): readonly ConcreteNode[] { + const node = this.nodesById.get(id); + if (!node) return []; + return this.nodesByEpisode.get(node.logicalParentId || 'orphan') || []; + } +} + +export function createMockProcessArgs(targets: ConcreteNode[], bufferNodes: ConcreteNode[] = [], inboxMessages: InboxMessage[] = []): ProcessArgs { + return { + targets, + buffer: new FakeContextWorkingBuffer(bufferNodes.length ? bufferNodes : targets), + inbox: new InboxSnapshotImpl(inboxMessages), + }; +} + export function createSyntheticHistory( numTurns: number, tokensPerTurn: number,