Skip to content

Commit 3b3310f

Browse files
authored
fix: reduce repeated history processing and token estimation (#250)
* fix: reduce repeated history processing and token estimation * test: preserve native snapshot cloning boundaries * fix: preserve shared references across semantic snapshot cloning
1 parent bddea42 commit 3b3310f

9 files changed

Lines changed: 444 additions & 164 deletions

File tree

‎packages/agent-modules/context-manager/src/token-estimator.ts‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,8 +199,15 @@ function getLastAssistantUsageInfo(
199199
return undefined;
200200
}
201201

202+
const MAX_CACHED_TEXT_UNITS = 1_048_576;
203+
const MAX_CACHED_TEXT_ENTRIES = 2_048;
204+
202205
export class BpeTokenEstimator implements TokenEstimator {
203206
private readonly countExactTokens: (text: string) => number;
207+
// Content keys remain valid across detached histories, edits and model
208+
// projections. Bound retained UTF-16 text as well as entry count.
209+
private readonly textTokens = new Map<string, number>();
210+
private cachedTextUnits = 0;
204211

205212
/**
206213
* @param encoder Override the default o200k_base tokenizer. Tests inject
@@ -222,12 +229,29 @@ export class BpeTokenEstimator implements TokenEstimator {
222229

223230
estimateTextTokens(text: string): number {
224231
if (!text) return 0;
232+
const cached = this.textTokens.get(text);
233+
if (cached !== undefined) {
234+
this.textTokens.delete(text);
235+
this.textTokens.set(text, cached);
236+
return cached;
237+
}
225238
try {
226239
if (text.length > MAX_EXACT_TOKENIZER_CHARS && hasOversizedAlphanumericRun(text)) {
227240
return estimateTextTokensUpperBound(text);
228241
}
229242
const tokens = this.countExactTokens(text);
230-
return Number.isFinite(tokens) && tokens >= 0 ? tokens : estimateTextTokensUpperBound(text);
243+
if (!Number.isFinite(tokens) || tokens < 0) return estimateTextTokensUpperBound(text);
244+
if (text.length <= MAX_CACHED_TEXT_UNITS) {
245+
while (this.textTokens.size >= MAX_CACHED_TEXT_ENTRIES || this.cachedTextUnits + text.length > MAX_CACHED_TEXT_UNITS) {
246+
const oldest = this.textTokens.keys().next().value;
247+
if (oldest === undefined) break;
248+
this.textTokens.delete(oldest);
249+
this.cachedTextUnits -= oldest.length;
250+
}
251+
this.textTokens.set(text, tokens);
252+
this.cachedTextUnits += text.length;
253+
}
254+
return tokens;
231255
} catch {
232256
return estimateTextTokensUpperBound(text);
233257
}

‎packages/local-runtime-v2/src/service/session-system/messages/history/canonical-history-provider.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -305,7 +305,7 @@ export function createSessionSystemCanonicalHistoryProvider(
305305
`Canonical compaction generation verification failed: ${change.sessionId}`,
306306
);
307307
}
308-
await syncIndex(change.sessionId, paths, true);
308+
await syncIndex(change.sessionId, paths, true, committed);
309309
const result = {
310310
revision: canonicalActiveHistoryRevision(committed),
311311
generation: committedGeneration,
@@ -335,7 +335,7 @@ export function createSessionSystemCanonicalHistoryProvider(
335335
}
336336
options.activity?.notify(sessionId);
337337
}
338-
await syncIndex(sessionId, paths, true);
338+
await syncIndex(sessionId, paths, true, records);
339339
return historySnapshot(records, allowPendingToolCallTail);
340340
}
341341

@@ -372,8 +372,9 @@ export function createSessionSystemCanonicalHistoryProvider(
372372
sessionId: string,
373373
paths: SessionHistoryPaths,
374374
allowMissingParent = false,
375+
verifiedRecords?: readonly CanonicalHistoryEnvelope[],
375376
): Promise<void> {
376-
const active = await files.readActiveStrict(paths.messages);
377+
const active = verifiedRecords ?? (await files.readActiveStrict(paths.messages));
377378
const adapter = indexes.get(sessionId) ?? createCanonicalHistoryIndexAdapter(paths.sessionDir);
378379
indexes.set(sessionId, adapter);
379380
try {

‎packages/local-runtime-v2/src/service/turn-system/agent-host/history/durable-canonical-history-store.ts‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,10 @@ export interface DurableCanonicalHistoryProvider {
2727
read(sessionId: string): Promise<CanonicalHistorySnapshot>;
2828
/** Active post-mutation history; the final tool round may still await results. */
2929
readActive(sessionId: string): Promise<CanonicalHistorySnapshot>;
30-
append(change: CanonicalHistoryChange): Promise<unknown>;
31-
replace(change: CanonicalHistoryChange): Promise<unknown>;
32-
compact?(change: CanonicalHistoryCompactionChange): Promise<unknown>;
30+
/** Return the verified post-write snapshot from the same lane, or request a reread with void. */
31+
append(change: CanonicalHistoryChange): Promise<CanonicalHistorySnapshot | void>;
32+
replace(change: CanonicalHistoryChange): Promise<CanonicalHistorySnapshot | void>;
33+
compact?(change: CanonicalHistoryCompactionChange): Promise<CanonicalHistorySnapshot | void>;
3334
settleTurnTail?(mutation: SettleTurnTailMutation): Promise<TurnHistoryMutationCommit>;
3435
retractTurn?(mutation: RetractTurnMutation): Promise<TurnHistoryMutationCommit>;
3536
}
@@ -102,13 +103,17 @@ export class DurableCanonicalHistoryStore implements CanonicalHistoryStore {
102103
private commit(
103104
kind: 'append' | 'replace',
104105
change: CanonicalHistoryChange,
105-
mutate: (snapshot: CanonicalHistoryChange) => Promise<unknown>,
106+
mutate: (snapshot: CanonicalHistoryChange) => Promise<CanonicalHistorySnapshot | void>,
106107
): Promise<CanonicalHistoryCommit> {
107108
const snapshot = captureSemanticSnapshot(change).value;
108109
validateCanonicalHistoryChange(snapshot, kind);
109110
return this.lane.run(snapshot.sessionId, async () => {
110-
await mutate(snapshot);
111-
return this.readActiveCommitted(snapshot.sessionId);
111+
const committed = await mutate(snapshot);
112+
// Session-owned providers already reread and validate inside their lane.
113+
// Legacy adapters returning void retain the strict post-write read.
114+
return committed === undefined
115+
? this.readActiveCommitted(snapshot.sessionId)
116+
: this.readProviderSnapshot(committed);
112117
});
113118
}
114119

Lines changed: 6 additions & 136 deletions
Original file line numberDiff line numberDiff line change
@@ -1,139 +1,26 @@
1-
const BLOCK_BYTES = 64;
2-
const UTF8_CHUNK_CODE_UNITS = 8 * 1_024;
3-
4-
const INITIAL_STATE = [
5-
0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
6-
] as const;
1+
import { createHash } from 'node:crypto';
72

8-
const ROUND_CONSTANTS = new Uint32Array([
9-
0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
10-
0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
11-
0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
12-
0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
13-
0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
14-
0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
15-
0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
16-
0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
17-
]);
3+
const UTF8_CHUNK_CODE_UNITS = 8 * 1_024;
184

19-
/**
20-
* Dependency-free incremental SHA-256 for AgentHost semantic identity.
21-
* Text is encoded in bounded chunks so hashing one large History string does
22-
* not allocate another payload-sized Buffer. This is intentionally internal;
23-
* persistence digests continue to use their owning adapters.
24-
*/
5+
/** Native SHA-256 with bounded UTF-8 chunks and the existing update boundaries. */
256
export class IncrementalSha256 {
26-
private readonly state = new Uint32Array(INITIAL_STATE);
27-
private readonly block = new Uint8Array(BLOCK_BYTES);
28-
private readonly schedule = new Uint32Array(64);
29-
private readonly textEncoder = new TextEncoder();
30-
private bufferedBytes = 0;
31-
private totalBytes = 0;
7+
private readonly hash = createHash('sha256');
328
private finalized = false;
339

3410
update(value: string): void {
3511
if (this.finalized) throw new Error('SHA-256 digest is already finalized.');
3612
let offset = 0;
3713
while (offset < value.length) {
3814
const end = safeUtf8ChunkEnd(value, offset);
39-
this.updateBytes(this.textEncoder.encode(value.slice(offset, end)));
15+
this.hash.update(value.slice(offset, end), 'utf8');
4016
offset = end;
4117
}
4218
}
4319

4420
digestHex(): string {
4521
if (this.finalized) throw new Error('SHA-256 digest is already finalized.');
4622
this.finalized = true;
47-
const messageBytes = this.totalBytes;
48-
this.block[this.bufferedBytes] = 0x80;
49-
this.bufferedBytes += 1;
50-
if (this.bufferedBytes > 56) {
51-
this.block.fill(0, this.bufferedBytes);
52-
this.compress();
53-
this.bufferedBytes = 0;
54-
}
55-
this.block.fill(0, this.bufferedBytes, 56);
56-
const messageBits = messageBytes * 8;
57-
writeUint32BigEndian(this.block, 56, Math.floor(messageBits / 0x1_0000_0000));
58-
writeUint32BigEndian(this.block, 60, messageBits >>> 0);
59-
this.compress();
60-
return [...this.state].map((word) => word.toString(16).padStart(8, '0')).join('');
61-
}
62-
63-
private updateBytes(bytes: Uint8Array): void {
64-
this.totalBytes += bytes.length;
65-
let offset = 0;
66-
while (offset < bytes.length) {
67-
const copied = Math.min(BLOCK_BYTES - this.bufferedBytes, bytes.length - offset);
68-
this.block.set(bytes.subarray(offset, offset + copied), this.bufferedBytes);
69-
this.bufferedBytes += copied;
70-
offset += copied;
71-
if (this.bufferedBytes !== BLOCK_BYTES) continue;
72-
this.compress();
73-
this.bufferedBytes = 0;
74-
}
75-
}
76-
77-
private compress(): void {
78-
const words = this.schedule;
79-
for (let index = 0; index < 16; index += 1) {
80-
const offset = index * 4;
81-
words[index] =
82-
(arrayValue(this.block, offset) << 24) |
83-
(arrayValue(this.block, offset + 1) << 16) |
84-
(arrayValue(this.block, offset + 2) << 8) |
85-
arrayValue(this.block, offset + 3);
86-
}
87-
for (let index = 16; index < 64; index += 1) {
88-
const prior15 = arrayValue(words, index - 15);
89-
const prior2 = arrayValue(words, index - 2);
90-
const sigma0 = rotateRight(prior15, 7) ^ rotateRight(prior15, 18) ^ (prior15 >>> 3);
91-
const sigma1 = rotateRight(prior2, 17) ^ rotateRight(prior2, 19) ^ (prior2 >>> 10);
92-
words[index] =
93-
(arrayValue(words, index - 16) + sigma0 + arrayValue(words, index - 7) + sigma1) >>> 0;
94-
}
95-
this.compressRounds(words);
96-
}
97-
98-
private compressRounds(words: Uint32Array): void {
99-
let a = arrayValue(this.state, 0);
100-
let b = arrayValue(this.state, 1);
101-
let c = arrayValue(this.state, 2);
102-
let d = arrayValue(this.state, 3);
103-
let e = arrayValue(this.state, 4);
104-
let f = arrayValue(this.state, 5);
105-
let g = arrayValue(this.state, 6);
106-
let h = arrayValue(this.state, 7);
107-
for (let index = 0; index < 64; index += 1) {
108-
const upperSigma1 = rotateRight(e, 6) ^ rotateRight(e, 11) ^ rotateRight(e, 25);
109-
const choose = (e & f) ^ (~e & g);
110-
const temporary1 =
111-
(h +
112-
upperSigma1 +
113-
choose +
114-
arrayValue(ROUND_CONSTANTS, index) +
115-
arrayValue(words, index)) >>>
116-
0;
117-
const upperSigma0 = rotateRight(a, 2) ^ rotateRight(a, 13) ^ rotateRight(a, 22);
118-
const majority = (a & b) ^ (a & c) ^ (b & c);
119-
const temporary2 = (upperSigma0 + majority) >>> 0;
120-
h = g;
121-
g = f;
122-
f = e;
123-
e = (d + temporary1) >>> 0;
124-
d = c;
125-
c = b;
126-
b = a;
127-
a = (temporary1 + temporary2) >>> 0;
128-
}
129-
this.state[0] = (arrayValue(this.state, 0) + a) >>> 0;
130-
this.state[1] = (arrayValue(this.state, 1) + b) >>> 0;
131-
this.state[2] = (arrayValue(this.state, 2) + c) >>> 0;
132-
this.state[3] = (arrayValue(this.state, 3) + d) >>> 0;
133-
this.state[4] = (arrayValue(this.state, 4) + e) >>> 0;
134-
this.state[5] = (arrayValue(this.state, 5) + f) >>> 0;
135-
this.state[6] = (arrayValue(this.state, 6) + g) >>> 0;
136-
this.state[7] = (arrayValue(this.state, 7) + h) >>> 0;
23+
return this.hash.digest('hex');
13724
}
13825
}
13926

@@ -152,20 +39,3 @@ function isHighSurrogate(value: number): boolean {
15239
function isLowSurrogate(value: number): boolean {
15340
return value >= 0xdc00 && value <= 0xdfff;
15441
}
155-
156-
function rotateRight(value: number, bits: number): number {
157-
return (value >>> bits) | (value << (32 - bits));
158-
}
159-
160-
function writeUint32BigEndian(target: Uint8Array, offset: number, value: number): void {
161-
target[offset] = (value >>> 24) & 0xff;
162-
target[offset + 1] = (value >>> 16) & 0xff;
163-
target[offset + 2] = (value >>> 8) & 0xff;
164-
target[offset + 3] = value & 0xff;
165-
}
166-
167-
function arrayValue(values: Uint8Array | Uint32Array, index: number): number {
168-
const value = values[index];
169-
if (value === undefined) throw new RangeError(`SHA-256 array index is out of bounds: ${index}.`);
170-
return value;
171-
}

‎packages/local-runtime-v2/src/service/turn-system/agent-host/history/semantic-identity.ts‎

Lines changed: 47 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,39 @@ export interface SemanticSnapshot<T> {
55
readonly fingerprint: string;
66
}
77

8+
const ownedValues = new WeakSet<object>();
9+
const fingerprints = new WeakMap<object, string>();
10+
811
/**
912
* Detach callback-owned History/event data before it becomes an in-run
1013
* identity or deferred delivery payload. Unsupported or cyclic values fail
1114
* closed. The digest is streamed so identity memory does not scale with the
1215
* encoded payload size.
1316
*/
1417
export function captureSemanticSnapshot<T>(value: T): SemanticSnapshot<T> {
15-
const snapshot = structuredClone(value);
16-
const encoder = new SemanticIdentityEncoder(new IncrementalSha256());
17-
encodeValue(snapshot, new WeakSet<object>(), encoder);
18+
// Clone an external graph in one operation: independent subtree clones can
19+
// change aliasing and getter order. Reuse only an already-owned root.
20+
const snapshot =
21+
typeof value === 'object' && value !== null && ownedValues.has(value)
22+
? value
23+
: freezeSemanticValue(structuredClone(value), new WeakSet<object>());
24+
let fingerprint: string | undefined;
1825
return {
19-
value: freezeSemanticValue(snapshot, new WeakSet<object>()),
20-
fingerprint: encoder.digest(),
26+
value: snapshot,
27+
// Value-only consumers still validate and detach eagerly, but never encode
28+
// or hash the history. Only values frozen by this module may be reused.
29+
get fingerprint() {
30+
if (fingerprint !== undefined) return fingerprint;
31+
const object = typeof snapshot === 'object' && snapshot !== null ? snapshot : undefined;
32+
fingerprint = object ? fingerprints.get(object) : undefined;
33+
if (fingerprint === undefined) {
34+
const encoder = new SemanticIdentityEncoder(new IncrementalSha256());
35+
encodeValue(snapshot, new WeakSet<object>(), encoder);
36+
fingerprint = encoder.digest();
37+
if (object) fingerprints.set(object, fingerprint);
38+
}
39+
return fingerprint;
40+
},
2141
};
2242
}
2343

@@ -54,13 +74,28 @@ class SemanticIdentityEncoder {
5474
}
5575
}
5676

57-
function freezeSemanticValue<T>(value: T, seen: WeakSet<object>): T {
58-
if (typeof value !== 'object' || value === null || seen.has(value)) return value;
59-
seen.add(value);
60-
Reflect.ownKeys(value).forEach((key) => {
61-
freezeSemanticValue(Reflect.get(value, key), seen);
62-
});
63-
return Object.freeze(value);
77+
function freezeSemanticValue<T>(value: T, ancestors: WeakSet<object>): T {
78+
if (typeof value !== 'object' || value === null) {
79+
if (!['undefined', 'string', 'boolean', 'number'].includes(typeof value) && value !== null) {
80+
throw new TypeError(`Unsupported semantic identity value: ${typeof value}.`);
81+
}
82+
return value;
83+
}
84+
if (ownedValues.has(value)) return value;
85+
if (ancestors.has(value)) throw new TypeError('Cyclic semantic identity values are unsupported.');
86+
const prototype = Object.getPrototypeOf(value);
87+
if (!Array.isArray(value) && prototype !== Object.prototype && prototype !== null) {
88+
throw new TypeError('Semantic identity values must contain only plain objects and arrays.');
89+
}
90+
if (Object.getOwnPropertySymbols(value).length > 0) {
91+
throw new TypeError('Symbol-keyed semantic identity values are unsupported.');
92+
}
93+
ancestors.add(value);
94+
for (const key of Object.keys(value)) freezeSemanticValue(Reflect.get(value, key), ancestors);
95+
ancestors.delete(value);
96+
Object.freeze(value);
97+
ownedValues.add(value);
98+
return value;
6499
}
65100

66101
function encodeValue(

0 commit comments

Comments
 (0)