Files
waggle-os/packages/agent/tests/prompt-assembler.test.ts
Oleg Maslov b20b138fe4 moving
2026-09-02 10:14:22 +02:00

677 lines
25 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import { describe, it, expect } from 'vitest';
import type { MemoryFrame, Importance, FrameType, FrameSource } from '@waggle/core';
import type { AgentPersona } from '../src/personas.js';
import type { TaskShape } from '../src/task-shape.js';
import type { ContextFrames } from '../src/orchestrator.js';
import {
PromptAssembler,
type AssembleInput,
type RecalledMemory,
} from '../src/prompt-assembler.js';
// ── Fixtures ─────────────────────────────────────────────────────────
let nextFrameId = 1;
function frame(
content: string,
opts: { type?: FrameType; importance?: Importance; source?: FrameSource } = {},
): MemoryFrame {
return {
id: nextFrameId++,
frame_type: opts.type ?? 'I',
gop_id: 'gop-1',
t: 0,
base_frame_id: null,
content,
importance: opts.importance ?? 'normal',
source: opts.source ?? 'user_stated',
access_count: 0,
created_at: '2026-04-17T10:00:00Z',
last_accessed: '2026-04-17T10:00:00Z',
};
}
function persona(overrides: Partial<AgentPersona> = {}): AgentPersona {
return {
id: 'researcher',
name: 'Researcher',
description: 'You are a researcher focused on evidence-based analysis.',
icon: 'search',
systemPrompt: '',
modelPreference: 'claude-opus-4-7',
tools: ['search_memory', 'save_memory'],
workspaceAffinity: [],
suggestedCommands: [],
defaultWorkflow: null,
tagline: 'Evidence-first thinker.',
...overrides,
};
}
function emptyContext(): ContextFrames {
return {
stateFrames: [],
recentChanges: [],
activeWork: [],
keyEntities: [],
personalPreferences: [],
};
}
function emptyRecalled(): RecalledMemory {
return { workspace: [], personal: [], scanSafe: true };
}
function shape(type: TaskShape['type'], confidence = 0.8): TaskShape {
return {
type,
confidence,
signals: [],
complexity: 'moderate',
};
}
function baseInput(overrides: Partial<AssembleInput> = {}): AssembleInput {
return {
corePrompt: 'You are Waggle, a memory-first AI colleague.',
persona: persona(),
context: emptyContext(),
recalled: emptyRecalled(),
query: 'test query',
tier: 'mid',
...overrides,
};
}
function defaultScaffold(body: string): string {
return `If the user specifies a response format, follow it exactly. Otherwise: ${body}`;
}
// ── Tests ────────────────────────────────────────────────────────────
describe('PromptAssembler.assemble', () => {
const assembler = new PromptAssembler();
it('includes Identity and Persona in every output', () => {
const out = assembler.assemble(baseInput());
expect(out.system).toContain('# Identity');
expect(out.system).toContain('## Persona: Researcher');
expect(out.debug.sectionsIncluded).toContain('Identity');
expect(out.debug.sectionsIncluded).toContain('Persona');
});
it('packages the persona operating instructions exactly once', () => {
const marker = 'PERSONA_OPERATING_RAIL_UNIQUE';
const out = assembler.assemble(baseInput({
persona: persona({ systemPrompt: `${marker}\nAlways ground claims in evidence.` }),
}));
expect(out.system).toContain(marker);
expect(out.system.match(new RegExp(marker, 'g'))).toHaveLength(1);
});
it('small tier caps State frames at 3', () => {
const frames: MemoryFrame[] = [];
for (let i = 0; i < 10; i++) {
frames.push(frame(`State frame ${i}`, { type: 'I', importance: 'important' }));
}
const out = assembler.assemble(
baseInput({
tier: 'small',
context: { ...emptyContext(), stateFrames: frames },
}),
);
const stateBlock = out.system.split('# State\n')[1]?.split('\n\n')[0] ?? '';
const stateLines = stateBlock.split('\n').filter(l => l.startsWith('-'));
expect(stateLines.length).toBeLessThanOrEqual(3);
});
it('frontier tier emits no Response-format section', () => {
const out = assembler.assemble(
baseInput({ tier: 'frontier', taskShape: shape('plan-execute', 0.9) }),
);
expect(out.system).not.toContain('# Response format');
expect(out.responseScaffold).toBeNull();
expect(out.debug.scaffoldApplied).toBe(false);
});
it('mid + plan-execute yields "State plan. Execute. Report." scaffold', () => {
const out = assembler.assemble(
baseInput({ tier: 'mid', taskShape: shape('plan-execute', 0.8) }),
);
expect(out.responseScaffold).toBe(defaultScaffold('State plan. Execute. Report.'));
expect(out.system).toContain('# Response format');
expect(out.debug.scaffoldApplied).toBe(true);
});
it('small + compare yields the long assumption/trade-offs/recommendation scaffold', () => {
const out = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.7) }),
);
expect(out.responseScaffold).toBe(
defaultScaffold('State the assumption. List the trade-offs. Give the recommendation.'),
);
});
it('makes every generic review scaffold explicitly subordinate to the user response format', () => {
for (const tier of ['small', 'mid'] as const) {
for (const scaffoldStyle of ['compression', 'expansion'] as const) {
const out = assembler.assemble(
baseInput({
query: 'Review this release decision and explain the issues.',
tier,
taskShape: shape('review', 0.9),
}),
{ scaffoldStyle },
);
expect(out.responseScaffold).toMatch(
/^If the user specifies a response format, follow it exactly\. Otherwise:/,
);
expect(out.debug.scaffoldApplied).toBe(true);
expect(out.debug.exclusiveResponseContract).toBe(false);
expect(out.debug.scaffoldSuppressed).toBe(false);
}
}
});
it.each([
'Return JSON only.',
'Return only the code.',
'Reply with exactly "PASS" and nothing else.',
'Only use JSON.parse and explain the result.',
'Never answer with only JSON; include a narrative.',
])('does not infer free-form language and keeps the scaffold safely conditional: %s', (query) => {
const out = assembler.assemble(
baseInput({ query, tier: 'small', taskShape: shape('review', 0.9) }),
);
expect(out.responseScaffold).toMatch(
/^If the user specifies a response format, follow it exactly\. Otherwise:/,
);
expect(out.debug.exclusiveResponseContract).toBe(false);
expect(out.debug.scaffoldSuppressed).toBe(false);
});
it('lets code-owned contracts explicitly suppress a scaffold without magic prompt wording', () => {
const out = assembler.assemble(
baseInput({
query: 'Review the release evidence.',
tier: 'small',
taskShape: shape('review', 0.9),
}),
{ exclusiveResponseContract: true },
);
expect(out.responseScaffold).toBeNull();
expect(out.debug.exclusiveResponseContract).toBe(true);
expect(out.debug.scaffoldSuppressed).toBe(true);
});
it('treats an explicitly bounded rewrite as closed-world and suppresses outside context', () => {
const out = assembler.assemble(
baseInput({
query: 'Rewrite this into a crisp executive memo. Preserve the facts and add no new claims: API tests pass.',
tier: 'mid',
taskShape: shape('decide', 0.9),
context: {
...emptyContext(),
stateFrames: [frame('State says shipping now is safe.')],
recentChanges: [frame('Recent changes say all gaps are closed.', { type: 'P' })],
activeWork: [{ category: 'task', content: 'Ship immediately.', priority: 1 }],
},
recalled: {
workspace: [],
personal: [],
scanSafe: true,
renderedText: '# Recalled Memories\n- Shipping now is safe.',
},
}),
);
expect(out.responseScaffold).toBeNull();
expect(out.debug.closedWorldRewrite).toBe(true);
expect(out.debug.scaffoldSuppressed).toBe(true);
expect(out.debug.sectionsIncluded).not.toContain('State');
expect(out.debug.sectionsIncluded).not.toContain('Recent changes');
expect(out.debug.sectionsIncluded).not.toContain('Active work');
expect(out.debug.sectionsIncluded).not.toContain('Recalled memory');
expect(out.system).not.toContain('State says shipping now is safe.');
expect(out.system).not.toContain('Recent changes say all gaps are closed.');
expect(out.system).not.toContain('Ship immediately.');
expect(out.system).not.toContain('Shipping now is safe.');
expect(out.system).toContain('# Closed-world rewrite');
expect(out.system).toContain('Do not add implications, explanations, rationale, risks');
expect(out.debug.sectionsIncluded.at(-1)).toBe('Closed-world rewrite');
});
it('does not infer a closed-world boundary from an ordinary rewrite request', () => {
const out = assembler.assemble(
baseInput({
query: 'Rewrite this product launch note to sound clearer.',
tier: 'mid',
taskShape: shape('draft', 0.9),
}),
);
expect(out.debug.closedWorldRewrite).toBe(false);
expect(out.system).not.toContain('# Closed-world rewrite');
});
it('recognizes a boundary-first closed-world rewrite directive', () => {
const out = assembler.assemble(
baseInput({
query: 'Using only the provided text, condense this into three bullets.',
tier: 'mid',
taskShape: shape('decide', 0.9),
}),
);
expect(out.debug.closedWorldRewrite).toBe(true);
expect(out.responseScaffold).toBeNull();
expect(out.system).toContain('# Closed-world rewrite');
});
it('does not mistake a quoted transform phrase for a rewrite directive', () => {
const out = assembler.assemble(
baseInput({
query: 'Explain what “rewrite this” means without adding new facts.',
tier: 'mid',
taskShape: shape('review', 0.9),
}),
);
expect(out.debug.closedWorldRewrite).toBe(false);
expect(out.responseScaffold).toBe(defaultScaffold('Briefly state assumption, then recommendation.'));
expect(out.system).not.toContain('# Closed-world rewrite');
});
it('draft shape emits no scaffold at any tier', () => {
for (const tier of ['small', 'mid', 'frontier'] as const) {
const out = assembler.assemble(
baseInput({ tier, taskShape: shape('draft', 0.9) }),
);
expect(out.responseScaffold).toBeNull();
expect(out.debug.scaffoldApplied).toBe(false);
}
});
it('mixed shape emits no scaffold at any tier', () => {
for (const tier of ['small', 'mid', 'frontier'] as const) {
const out = assembler.assemble(
baseInput({ tier, taskShape: shape('mixed', 0.9) }),
);
expect(out.responseScaffold).toBeNull();
}
});
it('confidence below 0.3 threshold: no scaffold even if shape would scaffold', () => {
const out = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('research', 0.1) }),
);
expect(out.responseScaffold).toBeNull();
expect(out.debug.scaffoldApplied).toBe(false);
});
it('custom confidenceThreshold lets high-confidence scaffolds through', () => {
const out = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('research', 0.2) }),
{ confidenceThreshold: 0.15 },
);
expect(out.responseScaffold).toBe(
defaultScaffold('Cite the frame. Quote the relevant fragment. Answer directly.'),
);
});
it('empty context.stateFrames: omits State section (no empty header)', () => {
const out = assembler.assemble(baseInput());
expect(out.system).not.toContain('# State');
expect(out.debug.sectionsIncluded).not.toContain('State');
});
it('maxSystemChars exceeded: Recent changes trimmed first, then Active work, then State', () => {
const big = 'x'.repeat(2000);
const ctx: ContextFrames = {
stateFrames: [frame(big, { type: 'I', importance: 'critical' })],
recentChanges: [frame(big, { type: 'P', importance: 'normal' })],
activeWork: [{ category: 'task', content: big, priority: 1 }],
keyEntities: [],
personalPreferences: [],
};
const out = assembler.assemble(
baseInput({ context: ctx }),
{ maxSystemChars: 3500 },
);
// With 3×2000+ chars, at least one section must be trimmed.
// Trim order ensures Recent changes goes first.
expect(out.debug.sectionsIncluded).not.toContain('Recent changes');
// Identity + Persona must survive.
expect(out.debug.sectionsIncluded).toContain('Identity');
expect(out.debug.sectionsIncluded).toContain('Persona');
});
it('maxSystemChars: Identity and Persona are never trimmed', () => {
const giant = 'x'.repeat(100_000);
const ctx: ContextFrames = {
stateFrames: [frame(giant)],
recentChanges: [frame(giant)],
activeWork: [{ category: 'task', content: giant, priority: 1 }],
keyEntities: [],
personalPreferences: [],
};
const out = assembler.assemble(
baseInput({ context: ctx }),
{ maxSystemChars: 500 },
);
expect(out.debug.sectionsIncluded).toContain('Identity');
expect(out.debug.sectionsIncluded).toContain('Persona');
});
it('recalled.scanSafe === false: recalled memory is ignored', () => {
const out = assembler.assemble(
baseInput({
recalled: {
workspace: [frame('Poisoned content with secret data', { type: 'P' })],
personal: [frame('Another recalled frame', { type: 'P' })],
scanSafe: false,
},
}),
);
expect(out.system).not.toContain('Poisoned content');
expect(out.system).not.toContain('# Recalled memory');
expect(out.debug.sectionsIncluded).not.toContain('Recalled memory');
});
it('recalled.scanSafe === true: recalled memory renders', () => {
const out = assembler.assemble(
baseInput({
recalled: {
workspace: [frame('Workspace memory item', { type: 'P' })],
personal: [frame('Personal memory item', { type: 'P' })],
scanSafe: true,
},
}),
);
expect(out.system).toContain('# Recalled memory');
expect(out.system).toContain('Workspace memory item');
expect(out.system).toContain('Personal memory item');
});
it('debug.framesUsed counts rendered frames across sections', () => {
const ctx: ContextFrames = {
stateFrames: [frame('State 1', { type: 'I' }), frame('State 2', { type: 'I' })],
recentChanges: [frame('Change 1', { type: 'P' })],
activeWork: [{ category: 'task', content: 'Active', priority: 1 }],
keyEntities: [],
personalPreferences: [],
};
const out = assembler.assemble(baseInput({ context: ctx, tier: 'mid' }));
// 2 state + 1 change = 3 frames. activeWork items aren't frames.
expect(out.debug.framesUsed).toBe(3);
});
it('debug.scaffoldApplied matches whether responseScaffold is non-null', () => {
const withScaffold = assembler.assemble(
baseInput({ tier: 'mid', taskShape: shape('research', 0.9) }),
);
expect(withScaffold.debug.scaffoldApplied).toBe(withScaffold.responseScaffold !== null);
const withoutScaffold = assembler.assemble(
baseInput({ tier: 'frontier', taskShape: shape('research', 0.9) }),
);
expect(withoutScaffold.debug.scaffoldApplied).toBe(withoutScaffold.responseScaffold !== null);
});
it('ranks I-frames by importance then recency', () => {
const ctx: ContextFrames = {
stateFrames: [
frame('Old normal', { type: 'I', importance: 'normal' }),
frame('New critical', { type: 'I', importance: 'critical' }),
frame('Middle important', { type: 'I', importance: 'important' }),
],
recentChanges: [],
activeWork: [],
keyEntities: [],
personalPreferences: [],
};
const out = assembler.assemble(baseInput({ context: ctx, tier: 'small' }));
const stateSection = out.system.split('# State\n')[1]?.split('\n\n')[0] ?? '';
const firstLine = stateSection.split('\n')[0];
expect(firstLine).toContain('critical');
expect(firstLine).toContain('New critical');
});
it('dedupes frames by id', () => {
const f = frame('Duplicate', { type: 'I' });
const ctx: ContextFrames = {
stateFrames: [f, f, f],
recentChanges: [],
activeWork: [],
keyEntities: [],
personalPreferences: [],
};
const out = assembler.assemble(baseInput({ context: ctx, tier: 'mid' }));
expect(out.debug.framesUsed).toBe(1);
});
it('debug.totalChars matches system.length exactly', () => {
const out = assembler.assemble(baseInput());
expect(out.debug.totalChars).toBe(out.system.length);
});
it('userPrefix is empty string in v1', () => {
const out = assembler.assemble(baseInput());
expect(out.userPrefix).toBe('');
});
it('tierOverride takes precedence over input.tier', () => {
const out = assembler.assemble(
baseInput({ tier: 'frontier', taskShape: shape('research', 0.9) }),
{ tierOverride: 'small' },
);
expect(out.debug.tier).toBe('small');
// Small + research + high confidence → scaffold applies
expect(out.responseScaffold).not.toBeNull();
});
});
// ── v5: scaffoldStyle (compression vs expansion) ─────────────────────
//
// See PromptAssembler v5 brief §7.3, §12.2.
//
// v5 adds an EXPANSION variant alongside v4's COMPRESSION scaffolds.
// The hypothesis: dense instruction-tuned models (Gemma family) benefit
// from expansion scaffolds that give them more structure to fill, rather
// than compression scaffolds that tell them to say less.
describe('PromptAssembler.assemble — v5 scaffoldStyle', () => {
const assembler = new PromptAssembler();
it('default style equals explicit compression (v4 parity)', () => {
const noStyle = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
);
const explicit = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
{ scaffoldStyle: 'compression' },
);
expect(noStyle.responseScaffold).toBe(explicit.responseScaffold);
expect(noStyle.system).toBe(explicit.system);
});
it('compression + small + compare preserves the v4 body after the safety qualifier', () => {
const out = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
{ scaffoldStyle: 'compression' },
);
expect(out.responseScaffold).toBe(
defaultScaffold('State the assumption. List the trade-offs. Give the recommendation.'),
);
});
it('compression + mid + plan-execute preserves the v4 body after the safety qualifier', () => {
const out = assembler.assemble(
baseInput({ tier: 'mid', taskShape: shape('plan-execute', 0.8) }),
{ scaffoldStyle: 'compression' },
);
expect(out.responseScaffold).toBe(defaultScaffold('State plan. Execute. Report.'));
});
it('compression + small + research preserves the v4 body after the safety qualifier', () => {
const out = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('research', 0.8) }),
{ scaffoldStyle: 'compression' },
);
expect(out.responseScaffold).toBe(
defaultScaffold('Cite the frame. Quote the relevant fragment. Answer directly.'),
);
});
it('expansion + small + compare includes Factors and Confidence keywords', () => {
const out = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
{ scaffoldStyle: 'expansion' },
);
expect(out.responseScaffold).not.toBeNull();
expect(out.responseScaffold).toContain('Factors');
expect(out.responseScaffold).toContain('confidence');
expect(out.responseScaffold).toContain('Trade-offs');
expect(out.responseScaffold).toContain('Recommendation');
expect(out.responseScaffold).toContain('Assumption');
});
it('expansion + small + decide produces same template as compare (shared analysis block)', () => {
const compareOut = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
{ scaffoldStyle: 'expansion' },
);
const decideOut = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('decide', 0.8) }),
{ scaffoldStyle: 'expansion' },
);
const reviewOut = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('review', 0.8) }),
{ scaffoldStyle: 'expansion' },
);
expect(decideOut.responseScaffold).toBe(compareOut.responseScaffold);
expect(reviewOut.responseScaffold).toBe(compareOut.responseScaffold);
});
it('expansion + small + plan-execute uses named-phases template', () => {
const out = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('plan-execute', 0.8) }),
{ scaffoldStyle: 'expansion' },
);
expect(out.responseScaffold).not.toBeNull();
expect(out.responseScaffold).toContain('named phases');
expect(out.responseScaffold).toContain('Goal');
expect(out.responseScaffold).toContain('Steps');
expect(out.responseScaffold).toContain('Dependencies');
expect(out.responseScaffold).toContain('Blockers');
expect(out.responseScaffold).toContain('timeline');
expect(out.responseScaffold).toContain('critical-path');
});
it('expansion + small + research uses three-part template', () => {
const out = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('research', 0.8) }),
{ scaffoldStyle: 'expansion' },
);
expect(out.responseScaffold).not.toBeNull();
expect(out.responseScaffold).toContain('Direct answer');
expect(out.responseScaffold).toContain('Source');
expect(out.responseScaffold).toContain('Context');
expect(out.responseScaffold).toContain('three parts');
});
it('expansion + mid produces shorter variant than small for same shape', () => {
const small = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
{ scaffoldStyle: 'expansion' },
);
const mid = assembler.assemble(
baseInput({ tier: 'mid', taskShape: shape('compare', 0.8) }),
{ scaffoldStyle: 'expansion' },
);
expect(mid.responseScaffold).not.toBeNull();
expect(mid.responseScaffold!.length).toBeLessThan(small.responseScaffold!.length);
});
it('expansion + draft shape → null at every tier (creative never scaffolded)', () => {
for (const tier of ['small', 'mid', 'frontier'] as const) {
const out = assembler.assemble(
baseInput({ tier, taskShape: shape('draft', 0.9) }),
{ scaffoldStyle: 'expansion' },
);
expect(out.responseScaffold).toBeNull();
expect(out.debug.scaffoldApplied).toBe(false);
}
});
it('expansion + mixed shape → null at every tier', () => {
for (const tier of ['small', 'mid', 'frontier'] as const) {
const out = assembler.assemble(
baseInput({ tier, taskShape: shape('mixed', 0.9) }),
{ scaffoldStyle: 'expansion' },
);
expect(out.responseScaffold).toBeNull();
}
});
it('expansion + frontier → null (expansion forced to compression at frontier)', () => {
// Per §7.3: "Frontier tier receives compression scaffold regardless of
// scaffoldStyle request." Since compression[*][frontier] = null in v4,
// expansion at frontier also yields null. Don't risk v4's F > E
// replication by applying expansion to frontier models.
for (const shapeType of ['compare', 'decide', 'review', 'plan-execute', 'research'] as const) {
const out = assembler.assemble(
baseInput({ tier: 'frontier', taskShape: shape(shapeType, 0.9) }),
{ scaffoldStyle: 'expansion' },
);
expect(out.responseScaffold).toBeNull();
}
});
it('expansion respects confidence threshold (same gate as compression)', () => {
const out = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.1) }),
{ scaffoldStyle: 'expansion' },
);
expect(out.responseScaffold).toBeNull();
});
it('debug.scaffoldStyle reflects the requested style', () => {
const defaultOut = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
);
expect(defaultOut.debug.scaffoldStyle).toBe('compression');
const compressionOut = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
{ scaffoldStyle: 'compression' },
);
expect(compressionOut.debug.scaffoldStyle).toBe('compression');
const expansionOut = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('compare', 0.8) }),
{ scaffoldStyle: 'expansion' },
);
expect(expansionOut.debug.scaffoldStyle).toBe('expansion');
});
it('debug.scaffoldStyle populated even when no scaffold is emitted (frontier / draft / low-conf)', () => {
const frontierOut = assembler.assemble(
baseInput({ tier: 'frontier', taskShape: shape('compare', 0.9) }),
{ scaffoldStyle: 'expansion' },
);
expect(frontierOut.debug.scaffoldStyle).toBe('expansion');
expect(frontierOut.debug.scaffoldApplied).toBe(false);
const draftOut = assembler.assemble(
baseInput({ tier: 'small', taskShape: shape('draft', 0.9) }),
{ scaffoldStyle: 'expansion' },
);
expect(draftOut.debug.scaffoldStyle).toBe('expansion');
expect(draftOut.debug.scaffoldApplied).toBe(false);
});
});