395 lines
13 KiB
TypeScript
395 lines
13 KiB
TypeScript
import { describe, expect, it, vi } from 'vitest';
|
|
import {
|
|
createAssistantMessageEventStream,
|
|
type AssistantMessage,
|
|
type Context,
|
|
} from '@earendil-works/pi-ai';
|
|
import { convertToLlm, type AgentMessage, type StreamFn } from '@earendil-works/pi-agent-core';
|
|
import type { AgentConfig } from '@/lib/orchestration/registry/types';
|
|
import type { AgentTurnSummary } from '@/lib/orchestration/types';
|
|
import type { StatelessChatRequest, StatelessEvent } from '@/lib/types/chat';
|
|
import { buildAgent } from '@/lib/agent/runtime/build-agent';
|
|
import { createDirectorCompactionRuntime } from '@/lib/chat/pi/director-compaction';
|
|
import { buildCallAgentTool } from '@/lib/chat/pi/tools/call-agent';
|
|
import {
|
|
buildReadSceneTool,
|
|
type DirectorSceneEvidenceMetadata,
|
|
type DirectorSceneEvidencePacket,
|
|
} from '@/lib/chat/pi/tools/read-scene';
|
|
|
|
const mocks = vi.hoisted(() => ({
|
|
childContexts: [] as string[],
|
|
}));
|
|
|
|
vi.mock('@/lib/agent/runtime/stream-fn', async (importOriginal) => {
|
|
const actual = await importOriginal<typeof import('@/lib/agent/runtime/stream-fn')>();
|
|
const { createAssistantMessageEventStream: createStream } = await import('@earendil-works/pi-ai');
|
|
return {
|
|
...actual,
|
|
createCallLlmStreamFn: vi.fn(
|
|
() =>
|
|
((_model, context) => {
|
|
mocks.childContexts.push(contextText(context));
|
|
const stream = createStream();
|
|
const message = assistantText(SCENE_TRUTH);
|
|
queueMicrotask(() => stream.push({ type: 'done', reason: 'stop', message }));
|
|
return stream;
|
|
}) as StreamFn,
|
|
),
|
|
};
|
|
});
|
|
const EMPTY_USAGE = {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
};
|
|
|
|
const SUMMARY_MARKER = 'COMPACTED_MEMORY_MARKER';
|
|
const TARGET_SCENE_ID = 'scene-light';
|
|
const SCENE_TRUTH = 'Light energy drives ATP production.';
|
|
const STALE_PACKET_MARKER = 'STALE_PRECOMPACTION_PACKET';
|
|
|
|
function assistantText(text: string): AssistantMessage {
|
|
return {
|
|
role: 'assistant',
|
|
content: [{ type: 'text', text }],
|
|
api: 'test',
|
|
provider: 'test',
|
|
model: 'deterministic-test',
|
|
usage: EMPTY_USAGE,
|
|
stopReason: 'stop',
|
|
timestamp: Date.now(),
|
|
};
|
|
}
|
|
|
|
function assistantToolCall(
|
|
id: string,
|
|
name: string,
|
|
args: Record<string, unknown>,
|
|
): AssistantMessage {
|
|
return {
|
|
...assistantText(''),
|
|
content: [{ type: 'toolCall', id, name, arguments: args }],
|
|
stopReason: 'toolUse',
|
|
};
|
|
}
|
|
|
|
function messageText(message: Context['messages'][number]): string {
|
|
if (message.role === 'user') {
|
|
return typeof message.content === 'string'
|
|
? message.content
|
|
: message.content.map((part) => (part.type === 'text' ? part.text : '[image]')).join('\n');
|
|
}
|
|
if (message.role === 'assistant') {
|
|
return message.content
|
|
.map((part) =>
|
|
part.type === 'text'
|
|
? part.text
|
|
: part.type === 'toolCall'
|
|
? `${part.name}(${JSON.stringify(part.arguments)})`
|
|
: '',
|
|
)
|
|
.join('\n');
|
|
}
|
|
return `${message.toolName}: ${message.content
|
|
.map((part) => (part.type === 'text' ? part.text : '[image]'))
|
|
.join('\n')}`;
|
|
}
|
|
|
|
function contextText(context: Context): string {
|
|
return context.messages.map(messageText).join('\n');
|
|
}
|
|
|
|
function deterministicSummaryStream(): StreamFn {
|
|
return ((_model, _context, _options) => {
|
|
const stream = createAssistantMessageEventStream();
|
|
const message = assistantText(
|
|
`${SUMMARY_MARKER}: preserve the early classroom goal and scene identity ${TARGET_SCENE_ID}; ` +
|
|
'the previously attached scene evidence was already consumed and must be fetched again.',
|
|
);
|
|
queueMicrotask(() => stream.push({ type: 'done', reason: 'stop', message }));
|
|
return stream;
|
|
}) as StreamFn;
|
|
}
|
|
|
|
function longHistory(): AgentMessage[] {
|
|
return Array.from({ length: 10 }, (_, index) => ({
|
|
role: 'user' as const,
|
|
content:
|
|
index === 0
|
|
? `Early goal: explain the light-reaction mechanism using scene identity ${TARGET_SCENE_ID}. ${'classroom context '.repeat(45)}`
|
|
: `history-${index}: ${'important classroom evidence '.repeat(45)}`,
|
|
timestamp: index,
|
|
}));
|
|
}
|
|
|
|
function makeBody(): StatelessChatRequest {
|
|
return {
|
|
messages: [],
|
|
storeState: {
|
|
stage: {
|
|
id: 'stage-1',
|
|
name: 'Photosynthesis',
|
|
createdAt: 1,
|
|
updatedAt: 50,
|
|
},
|
|
outlines: [
|
|
{
|
|
id: 'outline-light',
|
|
type: 'slide',
|
|
title: 'Light reactions',
|
|
description: 'How light energy becomes chemical energy.',
|
|
keyPoints: ['chlorophyll absorbs light', 'ATP and NADPH are produced'],
|
|
order: 2,
|
|
},
|
|
],
|
|
scenes: [
|
|
{
|
|
id: TARGET_SCENE_ID,
|
|
outlineId: 'outline-light',
|
|
stageId: 'stage-1',
|
|
title: 'Light reactions',
|
|
order: 2,
|
|
type: 'slide',
|
|
updatedAt: 42,
|
|
content: {
|
|
type: 'slide',
|
|
canvas: {
|
|
elements: [
|
|
{
|
|
id: 'equation',
|
|
type: 'text',
|
|
content: SCENE_TRUTH,
|
|
left: 40,
|
|
top: 60,
|
|
width: 400,
|
|
height: 80,
|
|
},
|
|
],
|
|
} as never,
|
|
},
|
|
},
|
|
],
|
|
currentSceneId: TARGET_SCENE_ID,
|
|
mode: 'autonomous',
|
|
whiteboardOpen: false,
|
|
},
|
|
config: { agentIds: ['teacher-1'] },
|
|
apiKey: '',
|
|
} as StatelessChatRequest;
|
|
}
|
|
|
|
const teacher: AgentConfig = {
|
|
id: 'teacher-1',
|
|
name: 'AI Teacher',
|
|
role: 'teacher',
|
|
persona: 'Teach only from supplied evidence.',
|
|
avatar: '',
|
|
color: '#3366ff',
|
|
allowedActions: [],
|
|
priority: 10,
|
|
createdAt: new Date('2026-01-01T00:00:00Z'),
|
|
updatedAt: new Date('2026-01-01T00:00:00Z'),
|
|
isDefault: true,
|
|
};
|
|
|
|
describe('Pi Director post-compaction refetch', () => {
|
|
it('COMPACTION-REFETCH-01 preserves memory, refetches, and binds only fresh scene evidence', async () => {
|
|
mocks.childContexts.length = 0;
|
|
const body = makeBody();
|
|
const pendingSceneEvidence = new Map<string, DirectorSceneEvidencePacket>();
|
|
const installedPackets: DirectorSceneEvidencePacket[] = [];
|
|
const consumedPackets: DirectorSceneEvidencePacket[][] = [];
|
|
let installCount = 0;
|
|
|
|
const readScene = buildReadSceneTool({
|
|
body,
|
|
onEvidence: (evidence) => {
|
|
installCount += 1;
|
|
const installed =
|
|
installCount === 1
|
|
? { ...evidence, content: `${evidence.content}\n${STALE_PACKET_MARKER}` }
|
|
: evidence;
|
|
installedPackets.push(installed);
|
|
pendingSceneEvidence.set(installed.details.sceneId, installed);
|
|
},
|
|
});
|
|
const takeSceneEvidence = () => {
|
|
if (pendingSceneEvidence.size === 0) return undefined;
|
|
const packets = [...pendingSceneEvidence.values()];
|
|
pendingSceneEvidence.clear();
|
|
consumedPackets.push(packets);
|
|
return {
|
|
content: packets.map((packet) => packet.content).join('\n\n'),
|
|
metadata: packets.map(
|
|
(packet): DirectorSceneEvidenceMetadata => ({
|
|
sceneId: packet.details.sceneId,
|
|
title: packet.details.title,
|
|
sceneType: packet.details.sceneType,
|
|
order: packet.details.order,
|
|
revision: packet.details.revision,
|
|
source: packet.details.source,
|
|
}),
|
|
),
|
|
};
|
|
};
|
|
|
|
// Establish and consume a pre-compaction packet through the same request-scoped
|
|
// evidence store. This packet must never be the one attached after compaction.
|
|
await readScene.execute('pre-compaction-read', { sceneId: TARGET_SCENE_ID });
|
|
const preCompactionPacket = takeSceneEvidence()?.content;
|
|
expect(preCompactionPacket).toContain(STALE_PACKET_MARKER);
|
|
expect(pendingSceneEvidence.size).toBe(0);
|
|
|
|
const compactionRuntime = createDirectorCompactionRuntime({
|
|
streamFn: deterministicSummaryStream(),
|
|
contextWindow: 600,
|
|
maxOutputTokens: 200,
|
|
settings: { enabled: true, reserveTokens: 160, keepRecentTokens: 120 },
|
|
});
|
|
const events: StatelessEvent[] = [];
|
|
const childTurns: AgentTurnSummary[] = [];
|
|
const callAgent = buildCallAgentTool({
|
|
body,
|
|
agentConfigs: [teacher],
|
|
send: async (event) => {
|
|
events.push(event);
|
|
},
|
|
languageModel: {} as never,
|
|
onAgentDone: (summary) => childTurns.push(summary),
|
|
onActionDone: () => {},
|
|
thinkingConfig: { mode: 'disabled', enabled: false },
|
|
abortSignal: new AbortController().signal,
|
|
maxAgentTurns: 2,
|
|
getAgentTurnCount: () => childTurns.length,
|
|
getAgentResponses: () => childTurns,
|
|
getWhiteboardLedger: () => [],
|
|
maxActionsPerAgent: 0,
|
|
enableWhiteboardTools: false,
|
|
takeSceneEvidence,
|
|
});
|
|
|
|
const directorContexts: string[] = [];
|
|
const directorStream = ((_model, context) => {
|
|
const text = contextText(context);
|
|
directorContexts.push(text);
|
|
const stream = createAssistantMessageEventStream();
|
|
const hasReadResult = context.messages.some(
|
|
(message) => message.role === 'toolResult' && message.toolName === 'read_scene',
|
|
);
|
|
const hasCallAgentResult = context.messages.some(
|
|
(message) => message.role === 'toolResult' && message.toolName === 'call_agent',
|
|
);
|
|
let message: AssistantMessage;
|
|
if (!text.includes(SUMMARY_MARKER)) {
|
|
message = assistantText('MISSING_COMPACTION_SUMMARY');
|
|
} else if (!hasReadResult) {
|
|
message = assistantToolCall('post-read', 'read_scene', { sceneId: TARGET_SCENE_ID });
|
|
} else if (!hasCallAgentResult) {
|
|
message = assistantToolCall('post-delegate', 'call_agent', {
|
|
agentId: teacher.id,
|
|
instruction: 'Use the freshly read scene evidence to explain the mechanism.',
|
|
});
|
|
} else {
|
|
message = assistantText('Post-compaction evidence-based delegation completed.');
|
|
}
|
|
const reason = message.stopReason === 'toolUse' ? 'toolUse' : 'stop';
|
|
queueMicrotask(() => stream.push({ type: 'done', reason, message }));
|
|
return stream;
|
|
}) as StreamFn;
|
|
|
|
const toolTrace: Array<{
|
|
toolName: string;
|
|
args: unknown;
|
|
isError: boolean;
|
|
details: unknown;
|
|
}> = [];
|
|
const director = buildAgent({
|
|
streamFn: directorStream,
|
|
systemPrompt: 'Use read_scene before delegating a scene-dependent answer.',
|
|
tools: [readScene, callAgent],
|
|
allowedToolNames: new Set(['read_scene', 'call_agent']),
|
|
history: longHistory(),
|
|
transformContext: compactionRuntime.transformContext,
|
|
convertToLlm,
|
|
afterToolCall: (context) => {
|
|
toolTrace.push({
|
|
toolName: context.toolCall.name,
|
|
args: context.args,
|
|
isError: context.isError,
|
|
details: context.result.details,
|
|
});
|
|
return undefined;
|
|
},
|
|
});
|
|
|
|
try {
|
|
await director.prompt(
|
|
'Continue the early classroom goal. Answer the mechanism question whose scene identity was established earlier.',
|
|
);
|
|
await director.waitForIdle();
|
|
|
|
const compactionTrace = compactionRuntime.getTrace();
|
|
const postRead = toolTrace.find((entry) => entry.toolName === 'read_scene');
|
|
const postDelegation = toolTrace.find((entry) => entry.toolName === 'call_agent');
|
|
const visibleChildAnswer = childTurns.map((turn) => turn.contentPreview).join('\n');
|
|
|
|
expect(compactionTrace.triggerCount).toBeGreaterThanOrEqual(1);
|
|
expect(compactionTrace.failures).toEqual([]);
|
|
expect(directorContexts[0]).toContain(SUMMARY_MARKER);
|
|
expect(postRead).toMatchObject({
|
|
args: { sceneId: TARGET_SCENE_ID },
|
|
isError: false,
|
|
details: {
|
|
status: 'ok',
|
|
sceneId: TARGET_SCENE_ID,
|
|
revision: '42',
|
|
source: 'request_start_snapshot',
|
|
},
|
|
});
|
|
expect(postDelegation).toMatchObject({
|
|
isError: false,
|
|
details: {
|
|
agentId: teacher.id,
|
|
sceneEvidence: [
|
|
{
|
|
sceneId: TARGET_SCENE_ID,
|
|
revision: '42',
|
|
source: 'request_start_snapshot',
|
|
},
|
|
],
|
|
},
|
|
});
|
|
expect(installedPackets.length).toBeGreaterThanOrEqual(2);
|
|
expect(consumedPackets).toHaveLength(2);
|
|
expect(consumedPackets[1]?.[0]).not.toBe(consumedPackets[0]?.[0]);
|
|
expect(mocks.childContexts[0]).toContain(`sceneId=${TARGET_SCENE_ID}`);
|
|
expect(mocks.childContexts[0]).toContain('revision=42');
|
|
expect(mocks.childContexts[0]).toContain('source=request_start_snapshot');
|
|
expect(mocks.childContexts[0]).not.toContain(STALE_PACKET_MARKER);
|
|
expect(mocks.childContexts[0]).toContain(SCENE_TRUTH);
|
|
expect(visibleChildAnswer).toContain(SCENE_TRUTH);
|
|
expect(events.some((event) => event.type === 'agent_start')).toBe(true);
|
|
console.info(
|
|
'[COMPACTION-REFETCH-01]',
|
|
JSON.stringify({
|
|
compactionTriggerCount: compactionTrace.triggerCount,
|
|
postCompactionReadSceneCount: toolTrace.filter(
|
|
(entry) => entry.toolName === 'read_scene' && !entry.isError,
|
|
).length,
|
|
callAgentSceneEvidence: (
|
|
postDelegation?.details as { sceneEvidence?: DirectorSceneEvidenceMetadata[] }
|
|
)?.sceneEvidence,
|
|
stalePacketReused: mocks.childContexts[0]?.includes(STALE_PACKET_MARKER) ?? false,
|
|
finalChildAnswer: visibleChildAnswer,
|
|
}),
|
|
);
|
|
} finally {
|
|
compactionRuntime.dispose();
|
|
}
|
|
});
|
|
});
|