MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-07-08 17:16:48 -07:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
joshualitt
2026-05-01 15:04:39 -07:00
committed by GitHub
parent 408afd3c5a
commit de8fdcfa16
52 changed files with 2133 additions and 1364 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/context/processors/blobDegradationProcessor.test.ts
+35 -43
View File
@@ -4,7 +4,6 @@
* SPDX-License-Identifier: Apache-2.0
*/
import assert from 'node:assert';
import { describe, it, expect } from 'vitest';
import { createBlobDegradationProcessor } from './blobDegradationProcessor.js';
import {
@@ -12,7 +11,7 @@ import {
createMockEnvironment,
createDummyNode,
} from '../testing/contextTestUtils.js';
import type { UserPrompt, SemanticPart, ConcreteNode } from '../graph/types.js';
import { type ConcreteNode, NodeType } from '../graph/types.js';
describe('BlobDegradationProcessor', () => {
it('should ignore text parts and only target inline_data and file_data', async () => {
@@ -28,35 +27,31 @@ describe('BlobDegradationProcessor', () => {
env,
);
const parts: SemanticPart[] = [
{ type: 'text', text: 'Hello' },
{ type: 'inline_data', mimeType: 'image/png', data: fakeData },
{ type: 'text', text: 'World' },
];
const node1 = createDummyNode('ep1', NodeType.USER_PROMPT, 10, {
payload: { text: 'Hello' },
});
const node2 = createDummyNode('ep1', NodeType.USER_PROMPT, 100, {
payload: { inlineData: { mimeType: 'image/png', data: fakeData } },
});
const node3 = createDummyNode('ep1', NodeType.USER_PROMPT, 10, {
payload: { text: 'World' },
});
const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, {
semanticParts: parts,
}) as UserPrompt;
const targets = [prompt];
const targets = [node1, node2, node3];
const result = await processor.process(createMockProcessArgs(targets));
expect(result.length).toBe(1);
const modifiedPrompt = result[0] as UserPrompt;
expect(result.length).toBe(3);
expect(modifiedPrompt.id).not.toBe(prompt.id);
expect(modifiedPrompt.semanticParts.length).toBe(3);
// Text nodes should be untouched
expect(result[0]).toBe(node1);
expect(result[2]).toBe(node3);
// Text parts should be untouched
expect(modifiedPrompt.semanticParts[0]).toEqual(parts[0]);
expect(modifiedPrompt.semanticParts[2]).toEqual(parts[2]);
// The inline_data part should be replaced with text
const degradedPart = modifiedPrompt.semanticParts[1];
expect(degradedPart.type).toBe('text');
assert(degradedPart.type === 'text');
expect(degradedPart.text).toContain(
// The inline_data node should be replaced with text
const degradedNode = result[1];
expect(degradedNode.id).not.toBe(node2.id);
expect(degradedNode.replacesId).toBe(node2.id);
expect(degradedNode.payload.text).toContain(
'[Multi-Modal Blob (image/png, 0.00MB) degraded to text',
);
});
@@ -69,29 +64,26 @@ describe('BlobDegradationProcessor', () => {
env,
);
// Tokens for fileData = 258.
// Degraded text = "[File Reference (video/mp4) degraded to text to preserve context window. Original URI: gs://test1]"
// Degraded text length ~100 characters.
// Since charsPerToken=1, degraded text = 100 tokens.
// Tokens saved = 258 - 100 = 158. This is > 0, so it WILL degrade it!
const node1 = createDummyNode('ep1', NodeType.USER_PROMPT, 100, {
payload: {
fileData: { mimeType: 'video/mp4', fileUri: 'gs://test1' },
},
});
const node2 = createDummyNode('ep1', NodeType.USER_PROMPT, 100, {
payload: {
fileData: { mimeType: 'video/mp4', fileUri: 'gs://test2' },
},
});
const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, {
semanticParts: [
{ type: 'file_data', mimeType: 'video/mp4', fileUri: 'gs://test1' },
{ type: 'file_data', mimeType: 'video/mp4', fileUri: 'gs://test2' },
],
}) as UserPrompt;
const targets = [prompt];
const targets = [node1, node2];
const result = await processor.process(createMockProcessArgs(targets));
const modifiedPrompt = result[0] as UserPrompt;
expect(modifiedPrompt.semanticParts.length).toBe(2);
expect(result.length).toBe(2);
// Both parts should be degraded
expect(modifiedPrompt.semanticParts[0].type).toBe('text');
expect(modifiedPrompt.semanticParts[1].type).toBe('text');
// Both nodes should be degraded
expect(result[0].payload.text).toContain('degraded to text');
expect(result[1].payload.text).toContain('degraded to text');
});
it('should return exactly the targets array if targets are empty', async () => {
packages/core/src/context/processors/blobDegradationProcessor.ts
+40 -85
View File
@@ -8,7 +8,7 @@ import type { JSONSchemaType } from 'ajv';
import type { ProcessArgs, ContextProcessor } from '../pipeline.js';
import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import type { ConcreteNode, UserPrompt } from '../graph/types.js';
import type { ConcreteNode } from '../graph/types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
@@ -55,95 +55,50 @@ export function createBlobDegradationProcessor(
// Forward scan, looking for bloated non-text parts to degrade
for (const node of targets) {
switch (node.type) {
case 'USER_PROMPT': {
let modified = false;
const newParts = [...node.semanticParts];
const payload = node.payload;
let newText = '';
let tokensSaved = 0;
for (let j = 0; j < node.semanticParts.length; j++) {
const part = node.semanticParts[j];
if (part.type === 'text') continue;
if (payload.inlineData?.data && payload.inlineData?.mimeType) {
await ensureDir();
const ext = payload.inlineData.mimeType.split('/')[1] || 'bin';
const fileName = `blob_${Date.now()}_${randomUUID()}.${ext}`;
const filePath = path.join(blobOutputsDir, fileName);
let newText = '';
let tokensSaved = 0;
const buffer = Buffer.from(payload.inlineData.data, 'base64');
await fs.writeFile(filePath, buffer);
switch (part.type) {
case 'inline_data': {
await ensureDir();
const ext = part.mimeType.split('/')[1] || 'bin';
const fileName = `blob_${Date.now()}_${randomUUID()}.${ext}`;
const filePath = path.join(blobOutputsDir, fileName);
const mb = (buffer.byteLength / 1024 / 1024).toFixed(2);
newText = `[Multi-Modal Blob (${payload.inlineData.mimeType}, ${mb}MB) degraded to text to preserve context window. Saved to: ${filePath}]`;
const buffer = Buffer.from(part.data, 'base64');
await fs.writeFile(filePath, buffer);
const oldTokens = env.tokenCalculator.estimateTokensForParts([
payload,
]);
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
} else if (payload.fileData?.mimeType && payload.fileData?.fileUri) {
newText = `[File Reference (${payload.fileData.mimeType}) degraded to text to preserve context window. Original URI: ${payload.fileData.fileUri}]`;
const oldTokens = env.tokenCalculator.estimateTokensForParts([
payload,
]);
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
}
const mb = (buffer.byteLength / 1024 / 1024).toFixed(2);
newText = `[Multi-Modal Blob (${part.mimeType}, ${mb}MB) degraded to text to preserve context window. Saved to: ${filePath}]`;
const oldTokens = env.tokenCalculator.estimateTokensForParts([
{
inlineData: { mimeType: part.mimeType, data: part.data },
},
]);
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
break;
}
case 'file_data': {
newText = `[File Reference (${part.mimeType}) degraded to text to preserve context window. Original URI: ${part.fileUri}]`;
const oldTokens = env.tokenCalculator.estimateTokensForParts([
{
fileData: {
mimeType: part.mimeType,
fileUri: part.fileUri,
},
},
]);
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
break;
}
case 'raw_part': {
newText = `[Unknown Part degraded to text to preserve context window.]`;
const oldTokens = env.tokenCalculator.estimateTokensForParts([
part.part,
]);
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
break;
}
default:
break;
}
if (newText && tokensSaved > 0) {
newParts[j] = { type: 'text', text: newText };
modified = true;
}
}
if (modified) {
const degradedNode: UserPrompt = {
...node,
id: randomUUID(),
semanticParts: newParts,
replacesId: node.id,
};
returnedNodes.push(degradedNode);
} else {
returnedNodes.push(node);
}
break;
}
default:
returnedNodes.push(node);
break;
if (newText && tokensSaved > 0) {
returnedNodes.push({
...node,
id: randomUUID(),
payload: { text: newText },
replacesId: node.id,
turnId: node.turnId,
});
} else {
returnedNodes.push(node);
}
}
packages/core/src/context/processors/nodeDistillationProcessor.test.ts
+21 -15
View File
@@ -4,7 +4,6 @@
* SPDX-License-Identifier: Apache-2.0
*/
import assert from 'node:assert';
import { describe, it, expect } from 'vitest';
import { createNodeDistillationProcessor } from './nodeDistillationProcessor.js';
import {
@@ -14,6 +13,7 @@ import {
createDummyToolNode,
createMockLlmClient,
} from '../testing/contextTestUtils.js';
import { NodeType } from '../graph/types.js';
import type {
UserPrompt,
AgentThought,
@@ -41,20 +41,20 @@ describe('NodeDistillationProcessor', () => {
const prompt = createDummyNode(
'ep1',
'USER_PROMPT',
NodeType.USER_PROMPT,
50,
{
semanticParts: [{ type: 'text', text: longText }],
payload: { text: longText },
},
'prompt-id',
) as UserPrompt;
const thought = createDummyNode(
'ep1',
'AGENT_THOUGHT',
NodeType.AGENT_THOUGHT,
50,
{
text: longText,
payload: { text: longText },
},
'thought-id',
) as AgentThought;
@@ -64,7 +64,13 @@ describe('NodeDistillationProcessor', () => {
5,
500,
{
observation: { result: 'A'.repeat(500) },
role: 'user',
payload: {
functionResponse: {
name: 'dummy_tool',
response: { result: 'A'.repeat(500) },
},
},
},
'tool-id',
);
@@ -78,19 +84,19 @@ describe('NodeDistillationProcessor', () => {
// 1. User Prompt
const compressedPrompt = result[0] as UserPrompt;
expect(compressedPrompt.id).not.toBe(prompt.id);
expect(compressedPrompt.semanticParts[0].type).toBe('text');
assert(compressedPrompt.semanticParts[0].type === 'text');
expect(compressedPrompt.semanticParts[0].text).toBe('Mocked Summary!');
expect(compressedPrompt.payload.text).toBe('Mocked Summary!');
// 2. Agent Thought
const compressedThought = result[1] as AgentThought;
expect(compressedThought.id).not.toBe(thought.id);
expect(compressedThought.text).toBe('Mocked Summary!');
expect(compressedThought.payload.text).toBe('Mocked Summary!');
// 3. Tool Execution
const compressedTool = result[2] as ToolExecution;
expect(compressedTool.id).not.toBe(tool.id);
expect(compressedTool.observation).toEqual({ summary: 'Mocked Summary!' });
expect(compressedTool.payload.functionResponse?.response).toEqual({
summary: 'Mocked Summary!',
});
expect(mockLlmClient.generateContent).toHaveBeenCalledTimes(3);
});
@@ -114,20 +120,20 @@ describe('NodeDistillationProcessor', () => {
const prompt = createDummyNode(
'ep1',
'USER_PROMPT',
NodeType.USER_PROMPT,
10,
{
semanticParts: [{ type: 'text', text: shortText }],
payload: { text: shortText },
},
'prompt-id',
) as UserPrompt;
const thought = createDummyNode(
'ep1',
'AGENT_THOUGHT',
NodeType.AGENT_THOUGHT,
13,
{
text: 'Short thought',
payload: { text: 'Short thought' },
},
'thought-id',
) as AgentThought;
packages/core/src/context/processors/nodeDistillationProcessor.ts
+71 -88
View File
@@ -6,10 +6,14 @@
import { randomUUID } from 'node:crypto';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import type { ConcreteNode } from '../graph/types.js';
import { type ConcreteNode, NodeType } from '../graph/types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { getResponseText } from '../../utils/partUtils.js';
import {
getResponseText,
updatePart,
cloneFunctionResponse,
} from '../../utils/partUtils.js';
import { LlmRole } from '../../telemetry/llmRole.js';
export interface NodeDistillationProcessorOptions {
@@ -56,7 +60,7 @@ export function createNodeDistillationProcessor(
},
});
return getResponseText(response) || text;
} catch (e) {
} catch (e: unknown) {
debugLogger.warn(
`NodeDistillationProcessor failed to summarize ${contextInfo}`,
e,
@@ -77,58 +81,31 @@ export function createNodeDistillationProcessor(
// Scan the target working buffer and unconditionally apply the configured hyperparameter threshold
for (const node of targets) {
const payload = node.payload;
switch (node.type) {
case 'USER_PROMPT': {
let modified = false;
const newParts = [...node.semanticParts];
for (let j = 0; j < node.semanticParts.length; j++) {
const part = node.semanticParts[j];
if (part.type !== 'text') continue;
if (part.text.length > thresholdChars) {
const summary = await generateSummary(part.text, 'User Prompt');
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: summary },
]);
const oldTokens = env.tokenCalculator.estimateTokensForParts([
{ text: part.text },
]);
if (newTokens < oldTokens) {
newParts[j] = { type: 'text', text: summary };
modified = true;
}
}
}
if (modified) {
returnedNodes.push({
...node,
id: randomUUID(),
semanticParts: newParts,
replacesId: node.id,
});
} else {
returnedNodes.push(node);
}
break;
}
case 'AGENT_THOUGHT': {
if (node.text.length > thresholdChars) {
const summary = await generateSummary(node.text, 'Agent Thought');
case NodeType.USER_PROMPT:
case NodeType.AGENT_THOUGHT: {
const text = payload.text;
if (text && text.length > thresholdChars) {
const summary = await generateSummary(text, node.type);
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: summary },
]);
const oldTokens = env.tokenCalculator.getTokenCost(node);
const oldTokens = env.tokenCalculator.estimateTokensForParts([
{ text },
]);
if (newTokens < oldTokens) {
const distilledPayload = updatePart(payload, { text: summary });
returnedNodes.push({
...node,
id: randomUUID(),
text: summary,
payload: distilledPayload,
replacesId: node.id,
timestamp: node.timestamp,
turnId: node.turnId,
});
break;
}
@@ -137,54 +114,60 @@ export function createNodeDistillationProcessor(
break;
}
case 'TOOL_EXECUTION': {
const rawObs = node.observation;
let stringifiedObs = '';
if (typeof rawObs === 'string') {
stringifiedObs = rawObs;
} else {
try {
stringifiedObs = JSON.stringify(rawObs);
} catch {
stringifiedObs = String(rawObs);
case NodeType.TOOL_EXECUTION: {
if (payload.functionResponse) {
const rawObs = payload.functionResponse.response;
let stringifiedObs = '';
if (typeof rawObs === 'string') {
stringifiedObs = rawObs;
} else {
try {
stringifiedObs = JSON.stringify(rawObs);
} catch {
stringifiedObs = String(rawObs);
}
}
}
if (stringifiedObs.length > thresholdChars) {
const summary = await generateSummary(
stringifiedObs,
node.toolName || 'unknown',
);
const newObsObject = { summary };
if (stringifiedObs.length > thresholdChars) {
const summary = await generateSummary(
stringifiedObs,
payload.functionResponse.name || 'unknown',
);
const newObsObject = { summary };
const newObsTokens = env.tokenCalculator.estimateTokensForParts([
{
functionResponse: {
name: node.toolName || 'unknown',
response: newObsObject,
id: node.id,
},
},
]);
const newFR = cloneFunctionResponse(payload.functionResponse);
newFR.response = newObsObject;
const oldObsTokens =
node.tokens?.observation ??
env.tokenCalculator.getTokenCost(node);
const intentTokens = node.tokens?.intent ?? 0;
const newObsTokens = env.tokenCalculator.estimateTokensForParts(
[
{
functionResponse: newFR,
},
],
);
if (newObsTokens < oldObsTokens) {
returnedNodes.push({
...node,
id: randomUUID(),
observation: newObsObject,
tokens: {
intent: intentTokens,
observation: newObsTokens,
},
replacesId: node.id,
});
break;
const oldObsTokens = env.tokenCalculator.estimateTokensForParts(
[payload],
);
if (newObsTokens < oldObsTokens) {
const newFR = cloneFunctionResponse(payload.functionResponse);
newFR.response = newObsObject;
const distilledPayload = updatePart(payload, {
functionResponse: newFR,
});
returnedNodes.push({
...node,
id: randomUUID(),
payload: distilledPayload,
replacesId: node.id,
timestamp: node.timestamp,
turnId: node.turnId,
});
break;
}
}
}
returnedNodes.push(node);
packages/core/src/context/processors/nodeTruncationProcessor.test.ts
+21 -20
View File
@@ -4,7 +4,6 @@
* SPDX-License-Identifier: Apache-2.0
*/
import assert from 'node:assert';
import { describe, it, expect } from 'vitest';
import { createNodeTruncationProcessor } from './nodeTruncationProcessor.js';
import {
@@ -12,7 +11,12 @@ import {
createMockEnvironment,
createDummyNode,
} from '../testing/contextTestUtils.js';
import type { UserPrompt, AgentThought, AgentYield } from '../graph/types.js';
import {
NodeType,
type UserPrompt,
type AgentThought,
type AgentYield,
} from '../graph/types.js';
describe('NodeTruncationProcessor', () => {
it('should truncate nodes that exceed maxTokensPerNode', async () => {
@@ -31,30 +35,30 @@ describe('NodeTruncationProcessor', () => {
const prompt = createDummyNode(
'ep1',
'USER_PROMPT',
NodeType.USER_PROMPT,
50,
{
semanticParts: [{ type: 'text', text: longText }],
payload: { text: longText },
},
'prompt-id',
) as UserPrompt;
const thought = createDummyNode(
'ep1',
'AGENT_THOUGHT',
NodeType.AGENT_THOUGHT,
50,
{
text: longText,
payload: { text: longText },
},
'thought-id',
) as AgentThought;
const yieldNode = createDummyNode(
'ep1',
'AGENT_YIELD',
NodeType.AGENT_YIELD,
50,
{
text: longText,
payload: { text: longText },
},
'yield-id',
) as AgentYield;
@@ -68,19 +72,17 @@ describe('NodeTruncationProcessor', () => {
// 1. User Prompt
const squashedPrompt = result[0] as UserPrompt;
expect(squashedPrompt.id).not.toBe(prompt.id);
expect(squashedPrompt.semanticParts[0].type).toBe('text');
assert(squashedPrompt.semanticParts[0].type === 'text');
expect(squashedPrompt.semanticParts[0].text).toContain('[... OMITTED');
expect(squashedPrompt.payload.text).toContain('[... OMITTED');
// 2. Agent Thought
const squashedThought = result[1] as AgentThought;
expect(squashedThought.id).not.toBe(thought.id);
expect(squashedThought.text).toContain('[... OMITTED');
expect(squashedThought.payload.text).toContain('[... OMITTED');
// 3. Agent Yield
const squashedYield = result[2] as AgentYield;
expect(squashedYield.id).not.toBe(yieldNode.id);
expect(squashedYield.text).toContain('[... OMITTED');
expect(squashedYield.payload.text).toContain('[... OMITTED');
});
it('should ignore nodes that are below maxTokensPerNode', async () => {
@@ -98,20 +100,20 @@ describe('NodeTruncationProcessor', () => {
const prompt = createDummyNode(
'ep1',
'USER_PROMPT',
NodeType.USER_PROMPT,
10,
{
semanticParts: [{ type: 'text', text: shortText }],
payload: { text: shortText },
},
'prompt-id',
) as UserPrompt;
const thought = createDummyNode(
'ep1',
'AGENT_THOUGHT',
NodeType.AGENT_THOUGHT,
13,
{
text: 'Short thought', // 13 chars
payload: { text: 'Short thought' }, // 13 chars
},
'thought-id',
) as AgentThought;
@@ -125,12 +127,11 @@ describe('NodeTruncationProcessor', () => {
// 1. User Prompt (untouched)
const squashedPrompt = result[0] as UserPrompt;
expect(squashedPrompt.id).toBe(prompt.id);
assert(squashedPrompt.semanticParts[0].type === 'text');
expect(squashedPrompt.semanticParts[0].text).not.toContain('[... OMITTED');
expect(squashedPrompt.payload.text).not.toContain('[... OMITTED');
// 2. Agent Thought (untouched)
const untouchedThought = result[1] as AgentThought;
expect(untouchedThought.id).toBe(thought.id);
expect(untouchedThought.text).not.toContain('[... OMITTED');
expect(untouchedThought.payload.text).not.toContain('[... OMITTED');
});
});
packages/core/src/context/processors/nodeTruncationProcessor.ts
+15 -60
View File
@@ -73,69 +73,24 @@ export function createNodeTruncationProcessor(
const returnedNodes: ConcreteNode[] = [];
for (const node of targets) {
switch (node.type) {
case 'USER_PROMPT': {
let modified = false;
const newParts = [...node.semanticParts];
const payload = node.payload;
const text = payload.text;
for (let j = 0; j < node.semanticParts.length; j++) {
const part = node.semanticParts[j];
if (part.type === 'text') {
const squashResult = tryApplySquash(part.text, limitChars);
if (squashResult) {
newParts[j] = { type: 'text', text: squashResult.text };
modified = true;
}
}
}
if (modified) {
returnedNodes.push({
...node,
id: randomUUID(),
semanticParts: newParts,
replacesId: node.id,
});
} else {
returnedNodes.push(node);
}
break;
if (text) {
const squashResult = tryApplySquash(text, limitChars);
if (squashResult) {
returnedNodes.push({
...node,
id: randomUUID(),
payload: { ...payload, text: squashResult.text },
replacesId: node.id,
turnId: node.turnId,
});
continue;
}
case 'AGENT_THOUGHT': {
const squashResult = tryApplySquash(node.text, limitChars);
if (squashResult) {
returnedNodes.push({
...node,
id: randomUUID(),
text: squashResult.text,
replacesId: node.id,
});
} else {
returnedNodes.push(node);
}
break;
}
case 'AGENT_YIELD': {
const squashResult = tryApplySquash(node.text, limitChars);
if (squashResult) {
returnedNodes.push({
...node,
id: randomUUID(),
text: squashResult.text,
replacesId: node.id,
});
} else {
returnedNodes.push(node);
}
break;
}
default:
returnedNodes.push(node);
break;
}
returnedNodes.push(node);
}
return returnedNodes;
packages/core/src/context/processors/rollingSummaryProcessor.test.ts
+30 -11
View File
@@ -10,6 +10,7 @@ import {
createMockEnvironment,
createDummyNode,
} from '../testing/contextTestUtils.js';
import { NodeType } from '../graph/types.js';
describe('RollingSummaryProcessor', () => {
it('should initialize with correct default options', () => {
@@ -43,13 +44,25 @@ describe('RollingSummaryProcessor', () => {
const targets = [
createDummyNode(
'ep1',
'USER_PROMPT',
NodeType.USER_PROMPT,
50,
{ semanticParts: [{ type: 'text', text: text50 }] },
{ payload: { text: text50 } },
'id1',
),
createDummyNode('ep1', 'AGENT_THOUGHT', 50, { text: text50 }, 'id2'),
createDummyNode('ep1', 'AGENT_YIELD', 50, { text: text50 }, 'id3'),
createDummyNode(
'ep1',
NodeType.AGENT_THOUGHT,
50,
{ payload: { text: text50 } },
'id2',
),
createDummyNode(
'ep1',
NodeType.AGENT_YIELD,
50,
{ payload: { text: text50 } },
'id3',
),
];
const result = await processor.process(createMockProcessArgs(targets));
@@ -59,8 +72,8 @@ describe('RollingSummaryProcessor', () => {
// Node id2 adds 50 deficit. Node id3 adds 50 deficit. Total = 100 deficit, which hits the target break point.
// Thus, id2 and id3 are summarized into a new ROLLING_SUMMARY node.
expect(result.length).toBe(2);
expect(result[0].type).toBe('USER_PROMPT');
expect(result[1].type).toBe('ROLLING_SUMMARY');
expect(result[0].type).toBe(NodeType.USER_PROMPT);
expect(result[1].type).toBe(NodeType.ROLLING_SUMMARY);
});
it('should preserve targets if deficit does not trigger summary', async () => {
@@ -80,19 +93,25 @@ describe('RollingSummaryProcessor', () => {
const targets = [
createDummyNode(
'ep1',
'USER_PROMPT',
NodeType.USER_PROMPT,
10,
{ semanticParts: [{ type: 'text', text: text10 }] },
{ payload: { text: text10 } },
'id1',
),
createDummyNode('ep1', 'AGENT_THOUGHT', 10, { text: text10 }, 'id2'),
createDummyNode(
'ep1',
NodeType.AGENT_THOUGHT,
10,
{ payload: { text: text10 } },
'id2',
),
];
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);
expect(result[0].type).toBe('USER_PROMPT');
expect(result[1].type).toBe('AGENT_THOUGHT');
expect(result[0].type).toBe(NodeType.USER_PROMPT);
expect(result[1].type).toBe(NodeType.AGENT_THOUGHT);
});
});
packages/core/src/context/processors/rollingSummaryProcessor.ts
+17 -14
View File
@@ -11,7 +11,11 @@ import type {
BackstopTargetOptions,
} from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import type { ConcreteNode, RollingSummary } from '../graph/types.js';
import {
type ConcreteNode,
type RollingSummary,
NodeType,
} from '../graph/types.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { LlmRole } from '../../telemetry/llmRole.js';
@@ -45,16 +49,14 @@ export function createRollingSummaryProcessor(
): Promise<string> => {
let transcript = '';
for (const node of nodes) {
const payload = node.payload;
let nodeContent = '';
if ('text' in node && typeof node.text === 'string') {
nodeContent = node.text;
} else if ('semanticParts' in node) {
nodeContent = JSON.stringify(node.semanticParts);
} else if ('observation' in node) {
nodeContent =
typeof node.observation === 'string'
? node.observation
: JSON.stringify(node.observation);
if (payload.text) {
nodeContent = payload.text;
} else if (payload.functionCall) {
nodeContent = `CALL: ${payload.functionCall.name}(${JSON.stringify(payload.functionCall.args)})`;
} else if (payload.functionResponse) {
nodeContent = `RESPONSE: ${JSON.stringify(payload.functionResponse.response)}`;
}
transcript += `[${node.type}]: ${nodeContent}\n`;
}
@@ -125,10 +127,11 @@ export function createRollingSummaryProcessor(
const summaryNode: RollingSummary = {
id: newId,
logicalParentId: newId,
type: 'ROLLING_SUMMARY',
timestamp: Date.now(),
text: snapshotText,
turnId: newId,
type: NodeType.ROLLING_SUMMARY,
timestamp: nodesToSummarize[nodesToSummarize.length - 1].timestamp,
role: 'user',
payload: { text: snapshotText },
abstractsIds: nodesToSummarize.map((n) => n.id),
};
packages/core/src/context/processors/stateSnapshotAsyncProcessor.test.ts
+22 -3
View File
@@ -10,6 +10,7 @@ import {
createDummyNode,
createMockProcessArgs,
} from '../testing/contextTestUtils.js';
import { NodeType } from '../graph/types.js';
import type { InboxMessage } from '../pipeline.js';
import type { InboxSnapshotImpl } from '../pipeline/inbox.js';
@@ -25,8 +26,20 @@ describe('StateSnapshotAsyncProcessor', () => {
{ type: 'point-in-time' },
);
const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A');
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
const nodeA = createDummyNode(
'ep1',
NodeType.USER_PROMPT,
50,
{},
'node-A',
);
const nodeB = createDummyNode(
'ep1',
NodeType.AGENT_THOUGHT,
60,
{},
'node-B',
);
const targets = [nodeA, nodeB];
await worker.process(createMockProcessArgs(targets, targets, []));
@@ -56,7 +69,13 @@ describe('StateSnapshotAsyncProcessor', () => {
{ type: 'accumulate' },
);
const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C');
const nodeC = createDummyNode(
'ep2',
NodeType.USER_PROMPT,
50,
{},
'node-C',
);
const targets = [nodeC];
const inboxMessages: InboxMessage[] = [
packages/core/src/context/processors/stateSnapshotAsyncProcessor.ts
+8 -6
View File
@@ -7,7 +7,7 @@ import { randomUUID } from 'node:crypto';
import type { JSONSchemaType } from 'ajv';
import type { AsyncContextProcessor, ProcessArgs } from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import type { ConcreteNode } from '../graph/types.js';
import { type ConcreteNode, NodeType } from '../graph/types.js';
import { SnapshotGenerator } from '../utils/snapshotGenerator.js';
import { debugLogger } from '../../utils/debugLogger.js';
@@ -73,13 +73,14 @@ export function createStateSnapshotAsyncProcessor(
previousConsumedIds = latest.payload.consumedIds;
// Prepend a synthetic node representing the previous rolling state
const snapshotId = randomUUID();
const previousStateNode: ConcreteNode = {
id: randomUUID(),
logicalParentId: '',
type: 'SNAPSHOT',
id: snapshotId,
turnId: snapshotId,
type: NodeType.SNAPSHOT,
timestamp: latest.timestamp,
text: latest.payload.newText,
role: 'user',
payload: { text: latest.payload.newText },
};
nodesToSummarize = [previousStateNode, ...targets];
@@ -101,6 +102,7 @@ export function createStateSnapshotAsyncProcessor(
newText: snapshotText,
consumedIds: newConsumedIds,
type: processorType,
timestamp: targets[targets.length - 1].timestamp,
});
} catch (e) {
debugLogger.error(
packages/core/src/context/processors/stateSnapshotProcessor.test.ts
+54 -11
View File
@@ -10,6 +10,7 @@ import {
createDummyNode,
createMockProcessArgs,
} from '../testing/contextTestUtils.js';
import { NodeType } from '../graph/types.js';
import type { InboxSnapshotImpl } from '../pipeline/inbox.js';
describe('StateSnapshotProcessor', () => {
@@ -22,7 +23,7 @@ describe('StateSnapshotProcessor', () => {
target: 'incremental',
},
);
const targets = [createDummyNode('ep1', 'USER_PROMPT')];
const targets = [createDummyNode('ep1', NodeType.USER_PROMPT)];
const result = await processor.process(createMockProcessArgs(targets));
expect(result).toBe(targets); // Strict equality
});
@@ -37,9 +38,27 @@ describe('StateSnapshotProcessor', () => {
},
);
const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A');
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C');
const nodeA = createDummyNode(
'ep1',
NodeType.USER_PROMPT,
50,
{},
'node-A',
);
const nodeB = createDummyNode(
'ep1',
NodeType.AGENT_THOUGHT,
60,
{},
'node-B',
);
const nodeC = createDummyNode(
'ep2',
NodeType.USER_PROMPT,
50,
{},
'node-C',
);
const targets = [nodeA, nodeB, nodeC];
@@ -62,7 +81,7 @@ describe('StateSnapshotProcessor', () => {
// Should remove A and B, insert Snapshot, keep C
expect(result.length).toBe(2);
expect(result[0].type).toBe('SNAPSHOT');
expect(result[0].type).toBe(NodeType.SNAPSHOT);
expect(result[1].id).toBe('node-C');
// Should consume the message
@@ -83,7 +102,13 @@ describe('StateSnapshotProcessor', () => {
// Make deficit 0 so we don't fall through to the sync backstop and fail the test that way
// node-A is MISSING (user deleted it)
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
const nodeB = createDummyNode(
'ep1',
NodeType.AGENT_THOUGHT,
60,
{},
'node-B',
);
const targets = [nodeB];
const messages = [
@@ -117,15 +142,33 @@ describe('StateSnapshotProcessor', () => {
{ target: 'max' },
); // Summarize all
const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A');
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C');
const nodeA = createDummyNode(
'ep1',
NodeType.USER_PROMPT,
50,
{},
'node-A',
);
const nodeB = createDummyNode(
'ep1',
NodeType.AGENT_THOUGHT,
60,
{},
'node-B',
);
const nodeC = createDummyNode(
'ep2',
NodeType.USER_PROMPT,
50,
{},
'node-C',
);
const targets = [nodeA, nodeB, nodeC];
const result = await processor.process(createMockProcessArgs(targets));
// Should synthesize a new snapshot synchronously
expect(env.llmClient.generateContent).toHaveBeenCalled();
expect(result.length).toBe(2); // nodeA is skipped as "system prompt", snapshot + nodeA
expect(result[1].type).toBe('SNAPSHOT');
expect(result.length).toBe(1); // nodeA is no longer protected, so everything is snapshotted
expect(result[0].type).toBe(NodeType.SNAPSHOT);
});
});
packages/core/src/context/processors/stateSnapshotProcessor.ts
+13 -16
View File
@@ -11,7 +11,7 @@ import type {
BackstopTargetOptions,
} from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import type { ConcreteNode, Snapshot } from '../graph/types.js';
import { type ConcreteNode, type Snapshot, NodeType } from '../graph/types.js';
import { SnapshotGenerator } from '../utils/snapshotGenerator.js';
import { debugLogger } from '../../utils/debugLogger.js';
@@ -61,6 +61,7 @@ export function createStateSnapshotProcessor(
newText: string;
consumedIds: string[];
type: string;
timestamp: number;
}>('PROPOSED_SNAPSHOT');
if (proposedSnapshots.length > 0) {
@@ -75,7 +76,7 @@ export function createStateSnapshotProcessor(
);
for (const proposed of sorted) {
const { consumedIds, newText } = proposed.payload;
const { consumedIds, newText, timestamp } = proposed.payload;
// Verify all consumed IDs still exist sequentially in targets
const targetIds = new Set(targets.map((t) => t.id));
@@ -87,10 +88,11 @@ export function createStateSnapshotProcessor(
const snapshotNode: Snapshot = {
id: newId,
logicalParentId: newId,
type: 'SNAPSHOT',
timestamp: Date.now(),
text: newText,
turnId: newId,
type: NodeType.SNAPSHOT,
timestamp: timestamp ?? Date.now(),
role: 'user',
payload: { text: newText },
abstractsIds: consumedIds,
};
@@ -131,12 +133,6 @@ export function createStateSnapshotProcessor(
// Scan oldest to newest
for (const node of targets) {
if (node.id === targets[0].id && node.type === 'USER_PROMPT') {
// Keep system prompt if it's the very first node
// In a real system, system prompt is protected, but we double check
continue;
}
nodesToSummarize.push(node);
deficitAccumulator += env.tokenCalculator.getTokenCost(node);
@@ -153,10 +149,11 @@ export function createStateSnapshotProcessor(
const newId = randomUUID();
const snapshotNode: Snapshot = {
id: newId,
logicalParentId: newId,
type: 'SNAPSHOT',
timestamp: Date.now(),
text: snapshotText,
turnId: newId,
type: NodeType.SNAPSHOT,
timestamp: nodesToSummarize[nodesToSummarize.length - 1].timestamp,
role: 'user',
payload: { text: snapshotText },
abstractsIds: nodesToSummarize.map((n) => n.id),
};
packages/core/src/context/processors/toolMaskingProcessor.test.ts
+60 -19
View File
@@ -25,9 +25,16 @@ describe('ToolMaskingProcessor', () => {
const longString = 'A'.repeat(500); // 500 chars
const toolStep = createDummyToolNode('ep1', 50, 500, {
observation: {
result: longString,
metadata: 'short', // 5 chars, will not be masked
role: 'model',
payload: {
functionResponse: {
name: 'dummy_tool',
id: 'dummy_id',
response: {
result: longString,
metadata: 'short', // 5 chars, will not be masked
},
},
},
});
@@ -40,7 +47,10 @@ describe('ToolMaskingProcessor', () => {
expect(masked.id).not.toBe(toolStep.id);
// It should have masked the observation
const obs = masked.observation as { result: string; metadata: string };
const obs = masked.payload.functionResponse?.response as {
result: string;
metadata: string;
};
expect(obs.result).toContain('<tool_output_masked>');
expect(obs.metadata).toBe('short'); // Untouched
});
@@ -53,10 +63,15 @@ describe('ToolMaskingProcessor', () => {
});
const toolStep = createDummyToolNode('ep1', 10, 10, {
toolName: 'activate_skill',
observation: {
result:
'this is a really long string that normally would get masked but wont because of the tool name',
payload: {
functionCall: {
name: 'activate_skill',
id: 'dummy_id',
args: {
result:
'this is a really long string that normally would get masked but wont because of the tool name',
},
},
},
});
@@ -76,23 +91,49 @@ describe('ToolMaskingProcessor', () => {
const longString = 'A'.repeat(500);
const toolStep = createDummyToolNode('ep1', 50, 500, {
intent: originalIntent,
observation: {
result: longString,
payload: {
functionCall: {
name: 'ls',
id: 'call_123',
args: originalIntent,
},
},
});
const result = await processor.process(createMockProcessArgs([toolStep]));
// We also need a response node if we want to test "observation is masked"
// Wait, the test says "strictly preserve the original intent args when only the observation is masked"
// But ToolMaskingProcessor processes nodes individually now.
// If we have a ToolExecution node with a functionCall, it masks the args.
// If we have a ToolExecution node with a functionResponse, it masks the response.
expect(result.length).toBe(1);
const masked = result[0] as ToolExecution;
const responseStep = createDummyToolNode('ep1', 50, 500, {
payload: {
functionResponse: {
name: 'ls',
id: 'call_123',
response: {
result: longString,
},
},
},
});
expect(masked.id).not.toBe(toolStep.id);
const result = await processor.process(
createMockProcessArgs([toolStep, responseStep]),
);
const obs = masked.observation as { result: string };
expect(result.length).toBe(2);
const maskedCall = result[0] as ToolExecution;
const maskedObs = result[1] as ToolExecution;
// Intent was short, so it should be the same node (or at least same content)
expect(maskedCall.payload.functionCall?.args).toEqual(originalIntent);
// Observation was long, so it should be masked
expect(maskedObs.id).not.toBe(responseStep.id);
const obs = maskedObs.payload.functionResponse?.response as {
result: string;
};
expect(obs.result).toContain('<tool_output_masked>');
// The intent MUST be perfectly preserved and not fall back to {} or undefined incorrectly
expect(masked.intent).toEqual(originalIntent);
});
});
packages/core/src/context/processors/toolMaskingProcessor.ts
+120 -144
View File
@@ -8,7 +8,7 @@ import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import type { ConcreteNode, ToolExecution } from '../graph/types.js';
import type { ConcreteNode } from '../graph/types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
import {
@@ -18,7 +18,11 @@ import {
ENTER_PLAN_MODE_TOOL_NAME,
EXIT_PLAN_MODE_TOOL_NAME,
} from '../../tools/tool-names.js';
import type { Part } from '@google/genai';
import {
updatePart,
cloneFunctionCall,
cloneFunctionResponse,
} from '../../utils/partUtils.js';
export interface ToolMaskingProcessorOptions {
stringLengthThresholdTokens: number;
@@ -138,149 +142,121 @@ export function createToolMaskingProcessor(
const returnedNodes: ConcreteNode[] = [];
for (const node of targets) {
switch (node.type) {
case 'TOOL_EXECUTION': {
const toolName = node.toolName;
if (toolName && UNMASKABLE_TOOLS.has(toolName)) {
returnedNodes.push(node);
break;
}
const callId = node.id || Date.now().toString();
const maskAsync = async (
obj: MaskableValue,
nodeType: string,
): Promise<{ masked: MaskableValue; changed: boolean }> => {
if (typeof obj === 'string') {
if (obj.length > limitChars && !isAlreadyMasked(obj)) {
const newString = await handleMasking(
obj,
toolName || 'unknown',
callId,
nodeType,
);
return { masked: newString, changed: true };
}
return { masked: obj, changed: false };
}
if (Array.isArray(obj)) {
let changed = false;
const masked: MaskableValue[] = [];
for (const item of obj) {
const res = await maskAsync(item, nodeType);
if (res.changed) changed = true;
masked.push(res.masked);
}
return { masked, changed };
}
if (typeof obj === 'object' && obj !== null) {
let changed = false;
const masked: Record<string, MaskableValue> = {};
for (const [key, value] of Object.entries(obj)) {
const res = await maskAsync(value, nodeType);
if (res.changed) changed = true;
masked[key] = res.masked;
}
return { masked, changed };
}
return { masked: obj, changed: false };
};
const rawIntent = node.intent;
const rawObs = node.observation;
if (!isMaskableRecord(rawIntent) || !isMaskableValue(rawObs)) {
returnedNodes.push(node);
break;
}
const intentRes = await maskAsync(rawIntent, 'intent');
const obsRes = await maskAsync(rawObs, 'observation');
if (intentRes.changed || obsRes.changed) {
const maskedIntent = isMaskableRecord(intentRes.masked)
? (intentRes.masked as Record<string, unknown>)
: undefined;
// Ensure we strictly preserve the original intent if it was unchanged and is a record
const finalIntent = intentRes.changed
? maskedIntent
: isMaskableRecord(rawIntent)
? (rawIntent as Record<string, unknown>)
: undefined;
// Handle observation explicitly as string vs object
const maskedObs =
typeof obsRes.masked === 'string'
? ({ message: obsRes.masked } as Record<string, unknown>)
: isMaskableRecord(obsRes.masked)
? (obsRes.masked as Record<string, unknown>)
: undefined;
// Ensure we strictly preserve the original observation if it was unchanged
const finalObs = obsRes.changed
? maskedObs
: typeof rawObs === 'string'
? ({ message: rawObs } as Record<string, unknown>)
: isMaskableRecord(rawObs)
? (rawObs as Record<string, unknown>)
: undefined;
const newIntentTokens =
env.tokenCalculator.estimateTokensForParts([
{
functionCall: {
name: toolName || 'unknown',
args: finalIntent,
id: callId,
},
},
]);
let obsPart: Record<string, unknown> = {};
if (maskedObs) {
obsPart = {
functionResponse: {
name: toolName || 'unknown',
response: finalObs,
id: callId,
},
};
}
const newObsTokens = env.tokenCalculator.estimateTokensForParts([
obsPart as Part,
]);
const tokensSaved =
env.tokenCalculator.getTokenCost(node) -
(newIntentTokens + newObsTokens);
if (tokensSaved > 0) {
const maskedNode: ToolExecution = {
...node,
id: randomUUID(), // Modified, so generate new ID
intent: finalIntent ?? node.intent,
observation: finalObs ?? node.observation,
tokens: {
intent: newIntentTokens,
observation: newObsTokens,
},
replacesId: node.id,
};
returnedNodes.push(maskedNode);
} else {
returnedNodes.push(node);
}
} else {
returnedNodes.push(node);
}
break;
}
default:
returnedNodes.push(node);
break;
if (node.type !== 'TOOL_EXECUTION') {
returnedNodes.push(node);
continue;
}
const payload = node.payload;
const toolName =
payload.functionCall?.name || payload.functionResponse?.name;
if (toolName && UNMASKABLE_TOOLS.has(toolName)) {
returnedNodes.push(node);
continue;
}
const callId =
payload.functionCall?.id || payload.functionResponse?.id || 'unknown';
const maskAsync = async (
obj: MaskableValue,
nodeType: string,
): Promise<{ masked: MaskableValue; changed: boolean }> => {
if (typeof obj === 'string') {
if (obj.length > limitChars && !isAlreadyMasked(obj)) {
const newString = await handleMasking(
obj,
toolName || 'unknown',
callId,
nodeType,
);
return { masked: newString, changed: true };
}
return { masked: obj, changed: false };
}
if (Array.isArray(obj)) {
let changed = false;
const masked: MaskableValue[] = [];
for (const item of obj) {
const res = await maskAsync(item, nodeType);
if (res.changed) changed = true;
masked.push(res.masked);
}
return { masked, changed };
}
if (typeof obj === 'object' && obj !== null) {
let changed = false;
const masked: Record<string, MaskableValue> = {};
for (const [key, value] of Object.entries(obj)) {
const res = await maskAsync(value, nodeType);
if (res.changed) changed = true;
masked[key] = res.masked;
}
return { masked, changed };
}
return { masked: obj, changed: false };
};
if (payload.functionCall) {
const rawIntent = payload.functionCall.args;
if (isMaskableRecord(rawIntent)) {
const res = await maskAsync(rawIntent, 'intent');
if (res.changed) {
const newFC = cloneFunctionCall(payload.functionCall);
let maskedRecord: Record<string, unknown>;
if (isMaskableRecord(res.masked)) {
maskedRecord = res.masked;
} else {
maskedRecord = { message: String(res.masked) };
}
newFC.args = maskedRecord;
const maskedPart = updatePart(payload, {
functionCall: newFC,
});
returnedNodes.push({
...node,
id: randomUUID(),
payload: maskedPart,
replacesId: node.id,
turnId: node.turnId,
});
continue;
}
}
} else if (payload.functionResponse) {
const rawObs = payload.functionResponse.response;
if (isMaskableValue(rawObs)) {
const res = await maskAsync(rawObs, 'observation');
if (res.changed) {
const newFR = cloneFunctionResponse(payload.functionResponse);
let maskedRecord: Record<string, unknown>;
if (isMaskableRecord(res.masked)) {
maskedRecord = res.masked;
} else {
maskedRecord = { message: String(res.masked) };
}
newFR.response = maskedRecord;
const maskedPart = updatePart(payload, {
functionResponse: newFR,
});
returnedNodes.push({
...node,
id: randomUUID(),
payload: maskedPart,
replacesId: node.id,
turnId: node.turnId,
});
continue;
}
}
}
returnedNodes.push(node);
}
return returnedNodes;