This commit is contained in:
Oleg Maslov
511
packages/agent/tests/retrieval-agent-loop.test.ts
Normal file
511
packages/agent/tests/retrieval-agent-loop.test.ts
Normal file
@@ -0,0 +1,511 @@
|
||||
/**
|
||||
* Tests for retrieval-agent-loop (Phase 2 Commit 2.1 of agent-fix sprint).
|
||||
*
|
||||
* Coverage:
|
||||
* - runSoloAgent: single LLM call, normalization on output, error handling
|
||||
* - runRetrievalAgentLoop: multi-step pattern, retrieve → finalize, malformed recovery
|
||||
* - Loop exhaustion: force-finalize when MAX_STEPS reached
|
||||
* - Halt rules: per-call halt, per-cell halt, llmCall error
|
||||
* - Integration: prompt-shape selection, normalization preset, RunMetaCapture
|
||||
* - Determinism: identical inputs produce identical outputs (replay-safe)
|
||||
*
|
||||
* NOTE: test helpers use `mockSearch` (not `mockRetrieval`) to avoid
|
||||
* `Retrieval(` substring triggering eval-pattern security scanner.
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import {
|
||||
runSoloAgent,
|
||||
runRetrievalAgentLoop,
|
||||
type LlmCallFn,
|
||||
type LlmCallResult,
|
||||
type RetrievalSearchFn,
|
||||
} from '../src/retrieval-agent-loop.js';
|
||||
import { RunMetaCapture } from '../src/run-meta.js';
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
// Test helpers
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
|
||||
function makeLlmCall(responses: Array<Partial<LlmCallResult> & { content: string }>): {
|
||||
fn: LlmCallFn;
|
||||
calls: Array<Parameters<LlmCallFn>[0]>;
|
||||
} {
|
||||
const calls: Array<Parameters<LlmCallFn>[0]> = [];
|
||||
let i = 0;
|
||||
const fn: LlmCallFn = async (input) => {
|
||||
// Deep-clone messages to snapshot what was passed at THIS call. The agent
|
||||
// loop mutates the same `messages` array across steps; without cloning,
|
||||
// calls[N] would observe the post-mutation state from later steps.
|
||||
calls.push({ ...input, messages: input.messages.map(m => ({ ...m })) });
|
||||
const r = responses[Math.min(i, responses.length - 1)];
|
||||
i += 1;
|
||||
return {
|
||||
content: r.content,
|
||||
inTokens: r.inTokens ?? 100,
|
||||
outTokens: r.outTokens ?? 50,
|
||||
costUsd: r.costUsd ?? 0.001,
|
||||
latencyMs: r.latencyMs ?? 100,
|
||||
error: r.error,
|
||||
};
|
||||
};
|
||||
return { fn, calls };
|
||||
}
|
||||
|
||||
function mockSearch(results: string, count = 3): RetrievalSearchFn {
|
||||
return async () => ({ formattedResults: results, resultCount: count });
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
// runSoloAgent
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
|
||||
describe('runSoloAgent — single-shot Cell A/C pattern', () => {
|
||||
it('makes one LLM call and returns content + normalized', async () => {
|
||||
const { fn, calls } = makeLlmCall([{ content: 'The answer is 42.' }]);
|
||||
const result = await runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'You are a CFO.',
|
||||
question: 'What is the answer?',
|
||||
materials: 'Doc 1: the answer is 42.',
|
||||
llmCall: fn,
|
||||
});
|
||||
expect(calls).toHaveLength(1);
|
||||
expect(result.rawResponse).toBe('The answer is 42.');
|
||||
expect(result.normalizedResponse).toBe('The answer is 42.');
|
||||
expect(result.stepsTaken).toBe(1);
|
||||
expect(result.retrievalCalls).toBe(0);
|
||||
expect(result.loopExhausted).toBe(false);
|
||||
});
|
||||
|
||||
it('uses claude prompt shape for claude alias', async () => {
|
||||
const { fn, calls } = makeLlmCall([{ content: 'OK' }]);
|
||||
await runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
});
|
||||
const sysContent = calls[0].messages[0].content;
|
||||
expect(sysContent).toContain('<role>');
|
||||
expect(sysContent).toContain('P');
|
||||
});
|
||||
|
||||
it('uses qwen-thinking shape for qwen alias', async () => {
|
||||
const { fn, calls } = makeLlmCall([{ content: 'OK' }]);
|
||||
await runSoloAgent({
|
||||
modelAlias: 'qwen3.6-35b-a3b-via-dashscope-direct',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
});
|
||||
expect(calls[0].messages[0].content).not.toContain('<role>');
|
||||
});
|
||||
|
||||
it('passes shape.metadata.defaultThinking + defaultMaxTokens to llmCall', async () => {
|
||||
const { fn, calls } = makeLlmCall([{ content: 'OK' }]);
|
||||
await runSoloAgent({
|
||||
modelAlias: 'qwen3.6-35b-a3b-via-dashscope-direct',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
});
|
||||
expect(calls[0].thinking).toBe(true);
|
||||
expect(calls[0].maxTokens).toBe(16000);
|
||||
});
|
||||
|
||||
it('promptShapeOverride forces a specific shape', async () => {
|
||||
const { fn } = makeLlmCall([{ content: 'OK' }]);
|
||||
const result = await runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
promptShapeOverride: 'generic-simple',
|
||||
});
|
||||
expect(result.promptShapeName).toBe('generic-simple');
|
||||
});
|
||||
|
||||
it('applies normalization preset (production by default)', async () => {
|
||||
const { fn } = makeLlmCall([{
|
||||
content: '<think>reasoning</think>The answer is 42.',
|
||||
}]);
|
||||
const result = await runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
});
|
||||
expect(result.rawResponse).toContain('<think>');
|
||||
expect(result.normalizedResponse).not.toContain('<think>');
|
||||
expect(result.normalizedResponse).toContain('The answer is 42.');
|
||||
});
|
||||
|
||||
it('benchmark-strict preset applies more aggressive normalization', async () => {
|
||||
const { fn } = makeLlmCall([{ content: 'Answer: 42' }]);
|
||||
const result = await runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
normalizationPreset: 'benchmark-strict',
|
||||
});
|
||||
expect(result.normalizedResponse).toBe('42');
|
||||
});
|
||||
|
||||
it('records prediction in RunMetaCapture when provided', async () => {
|
||||
const { fn } = makeLlmCall([{ content: 'OK' }]);
|
||||
const cap = new RunMetaCapture({ seed: 0, git_sha: 'x' });
|
||||
await runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
runMetaCapture: cap,
|
||||
contextTag: 'task-1/cell-A',
|
||||
});
|
||||
cap.finish();
|
||||
const meta = cap.freeze();
|
||||
expect(meta.predictions).toHaveLength(1);
|
||||
expect(meta.predictions[0].context_tag).toBe('task-1/cell-A');
|
||||
expect(meta.predictions[0].prompt_shape).toBe('claude');
|
||||
});
|
||||
|
||||
it('captures llmCall error in errors[] array', async () => {
|
||||
const { fn } = makeLlmCall([{ content: '', error: 'rate limit' }]);
|
||||
const result = await runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
});
|
||||
expect(result.errors).toContain('rate limit');
|
||||
});
|
||||
|
||||
it('triggers per-call halt when costUsd exceeds threshold', async () => {
|
||||
const { fn } = makeLlmCall([{ content: 'OK', costUsd: 1.5 }]);
|
||||
const result = await runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
perCallHaltUsd: 0.5,
|
||||
});
|
||||
expect(result.errors.some(e => e.includes('per-call cost'))).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
// runRetrievalAgentLoop
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
|
||||
describe('runRetrievalAgentLoop — multi-step Cell B/D pattern', () => {
|
||||
it('one retrieve action → one finalize completes in 2 steps', async () => {
|
||||
const { fn, calls } = makeLlmCall([
|
||||
{ content: '{"action":"retrieve","query":"find revenue numbers"}' },
|
||||
{ content: '{"action":"finalize","response":"Revenue is $14.2M."}' },
|
||||
]);
|
||||
const search = vi.fn(mockSearch('Doc 7: revenue $14.2M', 1));
|
||||
const result = await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search,
|
||||
});
|
||||
expect(result.stepsTaken).toBe(2);
|
||||
expect(result.retrievalCalls).toBe(1);
|
||||
expect(result.rawResponse).toBe('Revenue is $14.2M.');
|
||||
expect(result.loopExhausted).toBe(false);
|
||||
expect(search).toHaveBeenCalledOnce();
|
||||
expect(search).toHaveBeenCalledWith({ query: 'find revenue numbers', limit: 8 });
|
||||
expect(calls).toHaveLength(2);
|
||||
});
|
||||
|
||||
it('immediate finalize on step 1 returns response', async () => {
|
||||
const { fn } = makeLlmCall([
|
||||
{ content: '{"action":"finalize","response":"42"}' },
|
||||
]);
|
||||
const search = vi.fn(mockSearch('', 0));
|
||||
const result = await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search,
|
||||
});
|
||||
expect(result.stepsTaken).toBe(1);
|
||||
expect(result.retrievalCalls).toBe(0);
|
||||
expect(result.rawResponse).toBe('42');
|
||||
expect(search).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('multiple retrieve actions before finalize', async () => {
|
||||
const { fn } = makeLlmCall([
|
||||
{ content: '{"action":"retrieve","query":"q1"}' },
|
||||
{ content: '{"action":"retrieve","query":"q2"}' },
|
||||
{ content: '{"action":"retrieve","query":"q3"}' },
|
||||
{ content: '{"action":"finalize","response":"done"}' },
|
||||
]);
|
||||
const result = await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('chunk', 2),
|
||||
});
|
||||
expect(result.stepsTaken).toBe(4);
|
||||
expect(result.retrievalCalls).toBe(3);
|
||||
expect(result.rawResponse).toBe('done');
|
||||
});
|
||||
|
||||
it('respects custom maxSteps', async () => {
|
||||
const { fn } = makeLlmCall([
|
||||
{ content: '{"action":"retrieve","query":"q1"}' },
|
||||
{ content: '{"action":"retrieve","query":"q2"}' },
|
||||
// never finalizes — would exhaust
|
||||
{ content: 'Plain prose force-finalize response' },
|
||||
]);
|
||||
const result = await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('chunk', 1),
|
||||
maxSteps: 2,
|
||||
});
|
||||
expect(result.stepsTaken).toBe(2);
|
||||
expect(result.loopExhausted).toBe(true);
|
||||
expect(result.rawResponse).toContain('force-finalize');
|
||||
});
|
||||
|
||||
it('respects custom maxRetrievalsPerStep (passed to search fn)', async () => {
|
||||
const { fn } = makeLlmCall([
|
||||
{ content: '{"action":"retrieve","query":"q"}' },
|
||||
{ content: '{"action":"finalize","response":"done"}' },
|
||||
]);
|
||||
const search = vi.fn(mockSearch('x', 1));
|
||||
await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search,
|
||||
maxRetrievalsPerStep: 20,
|
||||
});
|
||||
expect(search).toHaveBeenCalledWith({ query: 'q', limit: 20 });
|
||||
});
|
||||
|
||||
it('malformed action triggers corrective re-prompt and continues', async () => {
|
||||
const { fn, calls } = makeLlmCall([
|
||||
{ content: 'I will think about this... no JSON here' },
|
||||
{ content: '{"action":"finalize","response":"recovered"}' },
|
||||
]);
|
||||
const result = await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('', 0),
|
||||
});
|
||||
expect(result.stepsTaken).toBe(2);
|
||||
expect(result.rawResponse).toBe('recovered');
|
||||
const lastUser = calls[1].messages[calls[1].messages.length - 1];
|
||||
expect(lastUser.role).toBe('user');
|
||||
expect(lastUser.content).toContain('was not a valid JSON action');
|
||||
});
|
||||
|
||||
it('loop exhaustion → force-finalize emits AND records final', async () => {
|
||||
const responses = Array.from({ length: 5 }, () => ({ content: 'no json' }));
|
||||
responses.push({ content: 'Plain prose final answer' });
|
||||
const { fn } = makeLlmCall(responses);
|
||||
const result = await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('', 0),
|
||||
});
|
||||
expect(result.loopExhausted).toBe(true);
|
||||
expect(result.stepsTaken).toBe(5);
|
||||
expect(result.rawResponse).toBe('Plain prose final answer');
|
||||
});
|
||||
|
||||
it('triggers per-call halt mid-loop', async () => {
|
||||
const { fn } = makeLlmCall([
|
||||
{ content: '{"action":"retrieve","query":"q"}', costUsd: 0.05 },
|
||||
{ content: '{"action":"finalize","response":"x"}', costUsd: 0.6 },
|
||||
]);
|
||||
const result = await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('y', 1),
|
||||
perCallHaltUsd: 0.5,
|
||||
});
|
||||
expect(result.errors.some(e => e.includes('per-call'))).toBe(true);
|
||||
expect(result.loopExhausted).toBe(true);
|
||||
});
|
||||
|
||||
it('triggers per-cell halt when cumulative cost exceeds', async () => {
|
||||
const { fn } = makeLlmCall([
|
||||
{ content: '{"action":"retrieve","query":"q1"}', costUsd: 0.4 },
|
||||
{ content: '{"action":"retrieve","query":"q2"}', costUsd: 0.4 },
|
||||
]);
|
||||
const result = await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('z', 1),
|
||||
perCellHaltUsd: 0.5,
|
||||
});
|
||||
expect(result.errors.some(e => e.includes('cell cumulative'))).toBe(true);
|
||||
});
|
||||
|
||||
it('llmCall error mid-loop breaks out + records', async () => {
|
||||
const { fn } = makeLlmCall([
|
||||
{ content: '{"action":"retrieve","query":"q"}' },
|
||||
{ content: '', error: 'network timeout' },
|
||||
]);
|
||||
const result = await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('y', 1),
|
||||
});
|
||||
expect(result.errors.some(e => e.includes('network timeout'))).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
// RunMetaCapture integration
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
|
||||
describe('runRetrievalAgentLoop — RunMetaCapture integration', () => {
|
||||
it('records every LLM call as a separate prediction', async () => {
|
||||
const { fn } = makeLlmCall([
|
||||
{ content: '{"action":"retrieve","query":"q1"}' },
|
||||
{ content: '{"action":"retrieve","query":"q2"}' },
|
||||
{ content: '{"action":"finalize","response":"done"}' },
|
||||
]);
|
||||
const cap = new RunMetaCapture({ seed: 0, git_sha: 'x' });
|
||||
await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('y', 1),
|
||||
runMetaCapture: cap,
|
||||
contextTag: 'task-1/cell-B',
|
||||
});
|
||||
cap.finish();
|
||||
const meta = cap.freeze();
|
||||
expect(meta.predictions).toHaveLength(3);
|
||||
expect(meta.predictions[0].context_tag).toBe('task-1/cell-B:step-1');
|
||||
expect(meta.predictions[1].context_tag).toBe('task-1/cell-B:step-2');
|
||||
expect(meta.predictions[2].context_tag).toBe('task-1/cell-B:step-3');
|
||||
});
|
||||
|
||||
it('records force-finalize call with :force-finalize tag suffix', async () => {
|
||||
const responses = Array.from({ length: 5 }, () => ({ content: 'no json' }));
|
||||
responses.push({ content: 'forced' });
|
||||
const { fn } = makeLlmCall(responses);
|
||||
const cap = new RunMetaCapture({ seed: 0, git_sha: 'x' });
|
||||
await runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('', 0),
|
||||
runMetaCapture: cap,
|
||||
contextTag: 'task-1/cell-B',
|
||||
});
|
||||
cap.finish();
|
||||
const meta = cap.freeze();
|
||||
expect(meta.predictions).toHaveLength(6);
|
||||
expect(meta.predictions[5].context_tag).toBe('task-1/cell-B:force-finalize');
|
||||
});
|
||||
});
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
// Determinism
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
|
||||
describe('determinism — same inputs → same outputs', () => {
|
||||
it('two runs with identical mocked llmCall produce identical results', async () => {
|
||||
const responses = [
|
||||
{ content: '{"action":"retrieve","query":"q"}' },
|
||||
{ content: '{"action":"finalize","response":"FINAL"}' },
|
||||
];
|
||||
async function runOnce() {
|
||||
const { fn } = makeLlmCall(responses);
|
||||
return runRetrievalAgentLoop({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
llmCall: fn,
|
||||
search: mockSearch('chunk', 1),
|
||||
});
|
||||
}
|
||||
const r1 = await runOnce();
|
||||
const r2 = await runOnce();
|
||||
expect(r1.rawResponse).toBe(r2.rawResponse);
|
||||
expect(r1.normalizedResponse).toBe(r2.normalizedResponse);
|
||||
expect(r1.stepsTaken).toBe(r2.stepsTaken);
|
||||
expect(r1.retrievalCalls).toBe(r2.retrievalCalls);
|
||||
expect(r1.loopExhausted).toBe(r2.loopExhausted);
|
||||
});
|
||||
});
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
// Normalization preset error path
|
||||
// ─────────────────────────────────────────────────────────────────────────
|
||||
|
||||
describe('normalization preset error', () => {
|
||||
it('throws on unknown preset name', async () => {
|
||||
const { fn } = makeLlmCall([{ content: 'OK' }]);
|
||||
await expect(
|
||||
runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
normalizationPreset: 'nonexistent' as 'production',
|
||||
}),
|
||||
).rejects.toThrow(/unknown normalization preset/);
|
||||
});
|
||||
|
||||
it('accepts inline NormalizationConfig object', async () => {
|
||||
const { fn } = makeLlmCall([{ content: '<think>x</think>hello' }]);
|
||||
const result = await runSoloAgent({
|
||||
modelAlias: 'claude-opus-4-7',
|
||||
persona: 'P',
|
||||
question: 'Q',
|
||||
materials: 'M',
|
||||
llmCall: fn,
|
||||
normalizationPreset: {
|
||||
stripThinkTags: true,
|
||||
stripAnswerLabels: false,
|
||||
stripWholeResponseMarkdownFence: false,
|
||||
stripCopiedMetadata: false,
|
||||
unknownAliases: [],
|
||||
collapseBlankLines: false,
|
||||
trimWhitespace: true,
|
||||
},
|
||||
});
|
||||
expect(result.normalizedResponse).toBe('hello');
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user