File size: 4,681 Bytes
7a1ad33
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
/**
 * @license
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { randomUUID } from 'node:crypto';
import { ContextManager } from '../contextManager.js';
import { AgentChatHistory } from '../../core/agentChatHistory.js';
import type { Content } from '@google/genai';
import type { ContextProfile } from '../config/profiles.js';
import { ContextEnvironmentImpl } from '../pipeline/environmentImpl.js';
import { ContextTracer } from '../tracer.js';
import { ContextEventBus } from '../eventBus.js';
import { PipelineOrchestrator } from '../pipeline/orchestrator.js';
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
import { StaticTokenCalculator } from '../utils/contextTokenCalculator.js';
import { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
import { registerBuiltInBehaviors } from '../graph/builtinBehaviors.js';

export interface TurnSummary {
  turnIndex: number;
  tokensBeforeBackground: number;
  tokensAfterBackground: number;
}

export class SimulationHarness {
  readonly chatHistory: AgentChatHistory;
  contextManager!: ContextManager;
  env!: ContextEnvironmentImpl;
  orchestrator!: PipelineOrchestrator;
  readonly eventBus: ContextEventBus;
  config!: ContextProfile;
  private tracer!: ContextTracer;
  private currentTurnIndex = 0;
  private tokenTrajectory: TurnSummary[] = [];

  static async create(
    config: ContextProfile,
    mockLlmClient: BaseLlmClient,
    mockTempDir = '/tmp/sim',
  ): Promise<SimulationHarness> {
    const harness = new SimulationHarness();
    await harness.init(config, mockLlmClient, mockTempDir);
    return harness;
  }

  private constructor() {
    this.chatHistory = new AgentChatHistory();
    this.eventBus = new ContextEventBus();
  }

  private async init(
    config: ContextProfile,
    mockLlmClient: BaseLlmClient,
    mockTempDir: string,
  ) {
    this.config = config;

    this.tracer = new ContextTracer({
      targetDir: mockTempDir,
      sessionId: 'sim-session',
    });

    const behaviorRegistry = new NodeBehaviorRegistry();
    registerBuiltInBehaviors(behaviorRegistry);
    const calculator = new StaticTokenCalculator(1, behaviorRegistry);

    this.env = new ContextEnvironmentImpl(
      () => mockLlmClient,
      'sim-prompt',
      'sim-session',
      mockTempDir,
      mockTempDir,
      this.tracer,
      1, // 1 char per token average for estimation (but estimator uses 0.33)
      this.eventBus,
      calculator,
      behaviorRegistry,
    );

    this.orchestrator = new PipelineOrchestrator(
      config.buildPipelines(this.env),
      config.buildAsyncPipelines(this.env),
      this.env,
      this.tracer,
    );
    this.contextManager = new ContextManager(
      config,
      this.env,
      this.tracer,
      this.orchestrator,
      this.chatHistory,
      calculator,
    );
  }

  async simulateTurn(messages: Content[]) {
    // In the new turn-based flow, we simulate the 'next' prompt or turn
    // by calling renderHistory on the pending content.

    // For the purpose of the simulation, we'll treat the first message as the 'pending' one
    // if it hasn't been added to history yet.
    const pendingContent = messages[messages.length - 1];

    // 1. Render to trigger sync and management
    const { processedNodes } = await this.contextManager.renderHistory({
      id: randomUUID(),
      content: pendingContent,
    });

    const tokensBefore =
      this.env.tokenCalculator.calculateConcreteListTokens(processedNodes);

    // 2. Append the new messages to durable history
    const currentHistory = this.chatHistory.get();
    const turns = messages.map((m) => ({ id: randomUUID(), content: m }));
    this.chatHistory.set([...currentHistory, ...turns]);

    // 3. Wait for any async pipelines triggered by the sync
    await this.contextManager.waitForPipelines();

    // 4. Measure tokens after background processors (requires another render or sync check)
    // In the new model, we'd need to re-render to see the effect of async processors
    // that might have finished.
    const { processedNodes: nodesAfter } =
      await this.contextManager.renderHistory();

    const tokensAfter =
      this.env.tokenCalculator.calculateConcreteListTokens(nodesAfter);

    this.tokenTrajectory.push({
      turnIndex: this.currentTurnIndex++,
      tokensBeforeBackground: tokensBefore,
      tokensAfterBackground: tokensAfter,
    });
  }

  async getGoldenState() {
    const { history: finalProjection, baseUnits } =
      await this.contextManager.renderHistory();
    return {
      tokenTrajectory: this.tokenTrajectory,
      finalProjection,
      baseUnits,
    };
  }
}