diff --git a/packages/core/src/context/graph/fromGraph.test.ts b/packages/core/src/context/graph/fromGraph.test.ts new file mode 100644 index 0000000000..0017fd7465 --- /dev/null +++ b/packages/core/src/context/graph/fromGraph.test.ts @@ -0,0 +1,202 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { fromGraph } from './fromGraph.js'; +import { NodeType, type ConcreteNode } from './types.js'; +import { NodeIdService } from './nodeIdService.js'; + +describe('fromGraph', () => { + it('should reconstruct an empty history from empty nodes', () => { + expect(fromGraph([])).toEqual([]); + }); + + it('should reconstruct a single turn from a single node', () => { + const nodes: ConcreteNode[] = [ + { + id: 'node_1', + turnId: 'turn_durable_1', + role: 'user', + type: NodeType.USER_PROMPT, + payload: { text: 'hello' }, + timestamp: 100, + }, + ]; + + const history = fromGraph(nodes); + expect(history).toEqual([ + { + id: 'durable_1', + content: { + role: 'user', + parts: [{ text: 'hello' }], + }, + }, + ]); + }); + + it('should coalesce adjacent nodes with the same turnId into a single turn', () => { + const nodes: ConcreteNode[] = [ + { + id: 'node_1', + turnId: 'turn_durable_1', + role: 'user', + type: NodeType.USER_PROMPT, + payload: { text: 'hello' }, + timestamp: 100, + }, + { + id: 'node_2', + turnId: 'turn_durable_1', + role: 'user', + type: NodeType.USER_PROMPT, + payload: { text: 'world' }, + timestamp: 101, + }, + ]; + + const history = fromGraph(nodes); + expect(history).toEqual([ + { + id: 'durable_1', + content: { + role: 'user', + parts: [{ text: 'hello' }, { text: 'world' }], + }, + }, + ]); + }); + + it('should split turns when the role changes', () => { + const nodes: ConcreteNode[] = [ + { + id: 'node_1', + turnId: 'turn_durable_1', + role: 'user', + type: NodeType.USER_PROMPT, + payload: { text: 'hello' }, + timestamp: 100, + }, + { + id: 'node_2', + turnId: 'turn_durable_2', + role: 'model', + type: NodeType.AGENT_THOUGHT, + payload: { text: 'hi' }, + timestamp: 101, + }, + ]; + + const history = fromGraph(nodes); + expect(history).toEqual([ + { + id: 'durable_1', + content: { + role: 'user', + parts: [{ text: 'hello' }], + }, + }, + { + id: 'durable_2', + content: { + role: 'model', + parts: [{ text: 'hi' }], + }, + }, + ]); + }); + + it('should split turns when the turnId changes, even if role is the same', () => { + const nodes: ConcreteNode[] = [ + { + id: 'node_1', + turnId: 'turn_durable_1', + role: 'user', + type: NodeType.USER_PROMPT, + payload: { text: 'hello' }, + timestamp: 100, + }, + { + id: 'node_2', + turnId: 'turn_durable_2', + role: 'user', + type: NodeType.USER_PROMPT, + payload: { text: 'world' }, + timestamp: 101, + }, + ]; + + const history = fromGraph(nodes); + expect(history).toEqual([ + { + id: 'durable_1', + content: { + role: 'user', + parts: [{ text: 'hello' }], + }, + }, + { + id: 'durable_2', + content: { + role: 'user', + parts: [{ text: 'world' }], + }, + }, + ]); + }); + + it('should correctly strip the turn_ prefix from turnId', () => { + const nodes: ConcreteNode[] = [ + { + id: 'node_1', + turnId: 'turn_my_stable_id_123', + role: 'user', + type: NodeType.USER_PROMPT, + payload: { text: 'hello' }, + timestamp: 100, + }, + ]; + + const history = fromGraph(nodes); + expect(history[0].id).toBe('my_stable_id_123'); + }); + + it('should handle orphan nodes gracefully', () => { + const nodes: ConcreteNode[] = [ + { + id: 'node_1', + role: 'user', + type: NodeType.USER_PROMPT, + payload: { text: 'orphan part' }, + timestamp: 100, + } as unknown as ConcreteNode, + ]; + + const history = fromGraph(nodes); + expect(history[0].id).toBe('orphan'); + expect(history[0].content.parts).toEqual([{ text: 'orphan part' }]); + }); + + it('should register identities with the NodeIdService if provided', () => { + const idService = new NodeIdService(); + const payload = { text: 'hello' }; + const nodes: ConcreteNode[] = [ + { + id: 'node_1', + turnId: 'turn_1', + role: 'user', + type: NodeType.USER_PROMPT, + payload, + timestamp: 100, + }, + ]; + + fromGraph(nodes, idService); + + // The payload object reference should map to the node ID + expect(idService.get(payload)).toBe('node_1'); + }); +}); diff --git a/packages/core/src/utils/historyHardening.test.ts b/packages/core/src/utils/historyHardening.test.ts new file mode 100644 index 0000000000..cb4949481e --- /dev/null +++ b/packages/core/src/utils/historyHardening.test.ts @@ -0,0 +1,223 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { + hardenHistory, + SYNTHETIC_THOUGHT_SIGNATURE, +} from './historyHardening.js'; +import type { HistoryTurn } from '../core/agentChatHistory.js'; +import { deriveStableId } from './cryptoUtils.js'; +import type { Part } from '@google/genai'; + +describe('hardenHistory', () => { + it('should return an empty array if input is empty', () => { + expect(hardenHistory([])).toEqual([]); + }); + + it('should coalesce adjacent turns of the same role', () => { + const history: HistoryTurn[] = [ + { id: '1', content: { role: 'user', parts: [{ text: 'hello' }] } }, + { id: '2', content: { role: 'user', parts: [{ text: 'world' }] } }, + ]; + const hardened = hardenHistory(history); + expect(hardened.length).toBe(1); + expect(hardened[0].content.parts).toEqual([ + { text: 'hello' }, + { text: 'world' }, + ]); + expect(hardened[0].id).toBe('1'); // Inherits ID of the first turn in the sequence + }); + + it('should inject thoughtSignature into the first functionCall of a model turn if missing', () => { + const history: HistoryTurn[] = [ + { id: '1', content: { role: 'user', parts: [{ text: 'do it' }] } }, + { + id: '2', + content: { + role: 'model', + parts: [{ functionCall: { name: 'myTool', args: {} } }], + }, + }, + { + id: '3', + content: { + role: 'user', + parts: [ + { + functionResponse: { + name: 'myTool', + response: { ok: true }, + }, + }, + ], + }, + }, + ]; + + const hardened = hardenHistory(history); + const modelPart = hardened[1].content.parts![0]; + expect(modelPart).toHaveProperty( + 'thoughtSignature', + SYNTHETIC_THOUGHT_SIGNATURE, + ); + }); + + it('should inject a sentinel user turn if history ends with a model turn', () => { + const history: HistoryTurn[] = [ + { id: '1', content: { role: 'user', parts: [{ text: 'hello' }] } }, + { id: '2', content: { role: 'model', parts: [{ text: 'hi' }] } }, + ]; + + const hardened = hardenHistory(history); + expect(hardened.length).toBe(3); + expect(hardened[2].content.role).toBe('user'); + expect(hardened[2].content.parts![0]).toEqual({ text: 'Please continue.' }); + expect(hardened[2].id).toBe(deriveStableId(['2', 'sentinel_end'])); + }); + + it('should inject a sentinel user turn if history starts with a model turn', () => { + const history: HistoryTurn[] = [ + { id: '1', content: { role: 'model', parts: [{ text: 'hi' }] } }, + { id: '2', content: { role: 'user', parts: [{ text: 'hello' }] } }, + ]; + + const hardened = hardenHistory(history, { + sentinels: { continuation: 'Custom start' }, + }); + expect(hardened.length).toBe(3); + expect(hardened[0].content.role).toBe('user'); + expect(hardened[0].content.parts![0]).toEqual({ text: 'Custom start' }); + expect(hardened[0].id).toBe(deriveStableId(['1', 'sentinel_start'])); + }); + + it('should inject sentinel responses for missing functionResponses', () => { + const history: HistoryTurn[] = [ + { id: '1', content: { role: 'user', parts: [{ text: 'do it' }] } }, + { + id: '2', + content: { + role: 'model', + parts: [ + { + functionCall: { id: 'call_1', name: 'toolA', args: {} }, + thoughtSignature: 'sig', + }, + { functionCall: { id: 'call_2', name: 'toolB', args: {} } }, + ], + }, + }, + // Note: Turn 3 is missing, so toolA and toolB have no responses + ]; + + const hardened = hardenHistory(history, { + sentinels: { lostToolResponse: 'Lost.' }, + }); + + // The history should now be: User -> Model -> User (sentinel responses) -> User (sentinel end) + // Wait, the sentinel responses turn will satisfy the "ends with user" rule. + expect(hardened.length).toBe(3); + expect(hardened[2].content.role).toBe('user'); + expect(hardened[2].content.parts).toHaveLength(2); + + const resp1 = hardened[2].content.parts![0].functionResponse; + expect(resp1?.id).toBe('call_1'); + expect(resp1?.response).toEqual({ error: 'Lost.' }); + + const resp2 = hardened[2].content.parts![1].functionResponse; + expect(resp2?.id).toBe('call_2'); + expect(resp2?.response).toEqual({ error: 'Lost.' }); + + expect(hardened[2].id).toBe(deriveStableId(['2', 'sentinel_resp'])); + }); + + it('should drop orphaned functionResponses', () => { + const history: HistoryTurn[] = [ + { id: '1', content: { role: 'user', parts: [{ text: 'hello' }] } }, + { id: '2', content: { role: 'model', parts: [{ text: 'hi' }] } }, + { + id: '3', + content: { + role: 'user', + parts: [ + { text: 'text is kept' }, + { + functionResponse: { id: 'orphan_1', name: 'toolA', response: {} }, + }, + ], + }, + }, + ]; + + const hardened = hardenHistory(history); + expect(hardened.length).toBe(3); + expect(hardened[2].content.parts).toHaveLength(1); + expect(hardened[2].content.parts![0]).toEqual({ text: 'text is kept' }); + }); + + it('should hoist and re-order tool responses to match functionCall order', () => { + const history: HistoryTurn[] = [ + { id: '1', content: { role: 'user', parts: [{ text: 'do it' }] } }, + { + id: '2', + content: { + role: 'model', + parts: [ + { + functionCall: { id: 'c1', name: 'toolA', args: {} }, + thoughtSignature: 'sig', + }, + { functionCall: { id: 'c2', name: 'toolB', args: {} } }, + ], + }, + }, + { + id: '3', + content: { + role: 'user', + parts: [ + { text: 'some text' }, + { functionResponse: { id: 'c2', name: 'toolB', response: {} } }, + { functionResponse: { id: 'c1', name: 'toolA', response: {} } }, + ], + }, + }, + ]; + + const hardened = hardenHistory(history); + expect(hardened[2].content.parts).toHaveLength(3); + + // Order should be: resp(c1) -> resp(c2) -> text + const p0 = hardened[2].content.parts![0]; + const p1 = hardened[2].content.parts![1]; + const p2 = hardened[2].content.parts![2]; + + expect(p0.functionResponse?.id).toBe('c1'); + expect(p1.functionResponse?.id).toBe('c2'); + expect(p2.text).toBe('some text'); + }); + + it('should scrub non-standard properties from parts', () => { + const history: HistoryTurn[] = [ + { + id: '1', + content: { + role: 'user', + parts: [ + { + text: 'hello', + extraProp: 'should be removed', + } as unknown as Part, + ], + }, + }, + ]; + + const hardened = hardenHistory(history); + expect(hardened[0].content.parts![0]).not.toHaveProperty('extraProp'); + expect(hardened[0].content.parts![0]).toHaveProperty('text', 'hello'); + }); +});