This commit is contained in:
Oleg Maslov
290
packages/agent/tests/trace-recorder.test.ts
Normal file
290
packages/agent/tests/trace-recorder.test.ts
Normal file
@@ -0,0 +1,290 @@
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { MindDB, ExecutionTraceStore } from '@waggle/core';
|
||||
import {
|
||||
TraceRecorder,
|
||||
truncate as truncateTraceText,
|
||||
scrubSecrets,
|
||||
} from '../src/trace-recorder.js';
|
||||
|
||||
describe('truncateTraceText', () => {
|
||||
it('leaves short strings untouched', () => {
|
||||
expect(truncateTraceText('hi', 10)).toBe('hi');
|
||||
});
|
||||
|
||||
it('appends truncation notice when exceeding cap', () => {
|
||||
const out = truncateTraceText('a'.repeat(120), 50);
|
||||
expect(out).toMatch(/^a{50}\n\[\.\.\. 70ch truncated\]$/);
|
||||
});
|
||||
|
||||
it('boundary: equal length is not truncated', () => {
|
||||
expect(truncateTraceText('abcd', 4)).toBe('abcd');
|
||||
});
|
||||
});
|
||||
|
||||
describe('scrubSecrets', () => {
|
||||
it('redacts obvious secret keys at top level', () => {
|
||||
const out = scrubSecrets({ password: 'hunter2', username: 'marko' });
|
||||
expect(out).toEqual({ password: '[REDACTED]', username: 'marko' });
|
||||
});
|
||||
|
||||
it('redacts nested objects', () => {
|
||||
const out = scrubSecrets({
|
||||
auth: { apiKey: 'sk-abc', region: 'eu' },
|
||||
other: 'x',
|
||||
});
|
||||
expect(out).toEqual({
|
||||
auth: { apiKey: '[REDACTED]', region: 'eu' },
|
||||
other: 'x',
|
||||
});
|
||||
});
|
||||
|
||||
it('leaves arrays alone (not scanned recursively)', () => {
|
||||
const out = scrubSecrets({ items: [1, 2, 3] });
|
||||
expect(out).toEqual({ items: [1, 2, 3] });
|
||||
});
|
||||
|
||||
it('handles common alternate naming (snake_case + camelCase)', () => {
|
||||
const out = scrubSecrets({
|
||||
api_key: 'x', apiKey: 'y', access_token: 'z', refreshToken: 'w', bearer: 'b',
|
||||
});
|
||||
expect(out).toEqual({
|
||||
api_key: '[REDACTED]', apiKey: '[REDACTED]',
|
||||
access_token: '[REDACTED]', refreshToken: '[REDACTED]',
|
||||
bearer: '[REDACTED]',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('TraceRecorder', () => {
|
||||
let db: MindDB;
|
||||
let store: ExecutionTraceStore;
|
||||
let recorder: TraceRecorder;
|
||||
|
||||
beforeEach(() => {
|
||||
db = new MindDB(':memory:');
|
||||
store = new ExecutionTraceStore(db);
|
||||
recorder = new TraceRecorder(store, { maxToolResultChars: 1000 });
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
db.close();
|
||||
});
|
||||
|
||||
describe('start', () => {
|
||||
it('creates a trace and returns handle with id', () => {
|
||||
const handle = recorder.start({ input: 'Hello world' });
|
||||
expect(handle.id).toBeGreaterThan(0);
|
||||
expect(handle.startedAt).toBeGreaterThan(0);
|
||||
|
||||
const row = store.get(handle.id);
|
||||
expect(row?.outcome).toBe('pending');
|
||||
});
|
||||
});
|
||||
|
||||
describe('recordReasoning', () => {
|
||||
it('buffers reasoning steps in memory before flush', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
recorder.recordReasoning(handle, 'Thinking step 1');
|
||||
|
||||
expect(recorder.peekReasoning(handle)).toHaveLength(1);
|
||||
// Not yet persisted
|
||||
expect(store.getParsed(handle.id)?.payload.reasoning).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('truncates long reasoning content', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
recorder.recordReasoning(handle, 'x'.repeat(5000));
|
||||
|
||||
const step = recorder.peekReasoning(handle)[0];
|
||||
expect(step.content.length).toBeLessThanOrEqual(1050); // 1000 + trailer
|
||||
expect(step.content).toContain('truncated');
|
||||
});
|
||||
});
|
||||
|
||||
describe('tool call lifecycle', () => {
|
||||
it('start + complete records duration and scrubs args', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
const callId = recorder.startToolCall(handle, 'http_get', {
|
||||
url: 'https://example.com',
|
||||
apiKey: 'sk-secret',
|
||||
});
|
||||
// Give the event loop a tick so durationMs is non-zero (optional)
|
||||
recorder.completeToolCall(handle, callId, 'response body');
|
||||
|
||||
const buffered = recorder.peekToolCalls(handle);
|
||||
expect(buffered).toHaveLength(1);
|
||||
expect(buffered[0].tool).toBe('http_get');
|
||||
expect(buffered[0].args.apiKey).toBe('[REDACTED]');
|
||||
expect(buffered[0].args.url).toBe('https://example.com');
|
||||
expect(buffered[0].result).toBe('response body');
|
||||
expect(buffered[0].ok).toBe(true);
|
||||
});
|
||||
|
||||
it('completeToolCall is no-op for unknown call id', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
expect(() => recorder.completeToolCall(handle, 'bogus', 'x')).not.toThrow();
|
||||
expect(recorder.peekToolCalls(handle)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('pendingToolCount tracks in-flight calls', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
const a = recorder.startToolCall(handle, 'a', {});
|
||||
recorder.startToolCall(handle, 'b', {});
|
||||
expect(recorder.pendingToolCount(handle)).toBe(2);
|
||||
recorder.completeToolCall(handle, a, 'done');
|
||||
expect(recorder.pendingToolCount(handle)).toBe(1);
|
||||
});
|
||||
|
||||
it('recordToolCall captures one-shot calls with given duration', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
recorder.recordToolCall(handle, {
|
||||
tool: 'cached_op',
|
||||
args: { token: 'abc' },
|
||||
result: 'ok',
|
||||
ok: true,
|
||||
durationMs: 42,
|
||||
timestamp: '2026-04-14T12:00:00Z',
|
||||
});
|
||||
const buffered = recorder.peekToolCalls(handle);
|
||||
expect(buffered[0].durationMs).toBe(42);
|
||||
expect(buffered[0].args.token).toBe('[REDACTED]');
|
||||
});
|
||||
|
||||
it('truncates long tool results', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
const callId = recorder.startToolCall(handle, 'big_read', {});
|
||||
recorder.completeToolCall(handle, callId, 'z'.repeat(5000));
|
||||
|
||||
const result = recorder.peekToolCalls(handle)[0].result;
|
||||
expect(result.length).toBeLessThanOrEqual(1050);
|
||||
expect(result).toContain('truncated');
|
||||
});
|
||||
});
|
||||
|
||||
describe('recordArtifact', () => {
|
||||
it('deduplicates artifacts', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
recorder.recordArtifact(handle, '/a.ts');
|
||||
recorder.recordArtifact(handle, '/a.ts');
|
||||
recorder.recordArtifact(handle, '/b.ts');
|
||||
expect(recorder.peekArtifacts(handle)).toEqual(['/a.ts', '/b.ts']);
|
||||
});
|
||||
});
|
||||
|
||||
describe('flush', () => {
|
||||
it('moves buffered events to the store and clears buffers', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
recorder.recordReasoning(handle, 'r1');
|
||||
recorder.recordArtifact(handle, '/f.ts');
|
||||
const callId = recorder.startToolCall(handle, 't', {});
|
||||
recorder.completeToolCall(handle, callId, 'res');
|
||||
|
||||
recorder.flush(handle);
|
||||
|
||||
expect(recorder.peekReasoning(handle)).toHaveLength(0);
|
||||
expect(recorder.peekToolCalls(handle)).toHaveLength(0);
|
||||
expect(recorder.peekArtifacts(handle)).toHaveLength(0);
|
||||
|
||||
const persisted = store.getParsed(handle.id);
|
||||
expect(persisted?.payload.reasoning).toHaveLength(1);
|
||||
expect(persisted?.payload.toolCalls).toHaveLength(1);
|
||||
expect(persisted?.payload.artifacts).toEqual(['/f.ts']);
|
||||
});
|
||||
|
||||
it('is a cheap no-op when nothing is buffered', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
expect(() => recorder.flush(handle)).not.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe('finalize', () => {
|
||||
it('flushes + sets outcome + output', () => {
|
||||
const handle = recorder.start({ input: 'Hello' });
|
||||
recorder.recordReasoning(handle, 'hmm');
|
||||
const row = recorder.finalize(handle, { outcome: 'success', output: 'World' });
|
||||
|
||||
expect(row?.outcome).toBe('success');
|
||||
const parsed = store.getParsed(handle.id);
|
||||
expect(parsed?.payload.output).toBe('World');
|
||||
expect(parsed?.payload.reasoning).toHaveLength(1);
|
||||
expect(parsed?.finalized_at).not.toBeNull();
|
||||
});
|
||||
|
||||
it('records cost and tokens', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
recorder.finalize(handle, {
|
||||
outcome: 'verified',
|
||||
output: 'done',
|
||||
costUsd: 0.05,
|
||||
tokens: { input: 100, output: 50 },
|
||||
});
|
||||
|
||||
const parsed = store.getParsed(handle.id);
|
||||
expect(parsed?.cost_usd).toBeCloseTo(0.05);
|
||||
expect(parsed?.payload.tokens).toEqual({ input: 100, output: 50 });
|
||||
});
|
||||
|
||||
it('clears pending tool calls after finalize', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
recorder.startToolCall(handle, 'a', {});
|
||||
expect(recorder.pendingToolCount(handle)).toBe(1);
|
||||
recorder.finalize(handle, { outcome: 'abandoned', output: '' });
|
||||
expect(recorder.pendingToolCount(handle)).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('markCorrected', () => {
|
||||
it('delegates to the store', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
recorder.finalize(handle, { outcome: 'success', output: 'v1' });
|
||||
recorder.markCorrected(handle, 'too terse');
|
||||
|
||||
const parsed = store.getParsed(handle.id);
|
||||
expect(parsed?.outcome).toBe('corrected');
|
||||
expect(parsed?.payload.correctionFeedback).toBe('too terse');
|
||||
});
|
||||
});
|
||||
|
||||
describe('wireAgentLoopCallbacks', () => {
|
||||
it('correlates onToolUse / onToolResult pairs', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
const { onToolUse, onToolResult } = recorder.wireAgentLoopCallbacks(handle);
|
||||
|
||||
onToolUse('read_file', { path: '/a' });
|
||||
onToolResult('read_file', { path: '/a' }, 'file contents');
|
||||
|
||||
const buffered = recorder.peekToolCalls(handle);
|
||||
expect(buffered).toHaveLength(1);
|
||||
expect(buffered[0].tool).toBe('read_file');
|
||||
expect(buffered[0].result).toBe('file contents');
|
||||
});
|
||||
|
||||
it('matches concurrent calls to the same tool by input (FIFO stack per key)', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
const { onToolUse, onToolResult } = recorder.wireAgentLoopCallbacks(handle);
|
||||
|
||||
onToolUse('search', { q: 'a' });
|
||||
onToolUse('search', { q: 'b' });
|
||||
onToolResult('search', { q: 'a' }, 'A');
|
||||
onToolResult('search', { q: 'b' }, 'B');
|
||||
|
||||
const buffered = recorder.peekToolCalls(handle);
|
||||
expect(buffered).toHaveLength(2);
|
||||
const byTool = buffered.map(c => [c.args.q, c.result]);
|
||||
expect(byTool).toContainEqual(['a', 'A']);
|
||||
expect(byTool).toContainEqual(['b', 'B']);
|
||||
});
|
||||
|
||||
it('falls back to one-shot recording if onToolResult arrives without onToolUse', () => {
|
||||
const handle = recorder.start({ input: 'x' });
|
||||
const { onToolResult } = recorder.wireAgentLoopCallbacks(handle);
|
||||
|
||||
onToolResult('ghost', { x: 1 }, 'mystery');
|
||||
|
||||
const buffered = recorder.peekToolCalls(handle);
|
||||
expect(buffered).toHaveLength(1);
|
||||
expect(buffered[0].tool).toBe('ghost');
|
||||
expect(buffered[0].durationMs).toBe(0);
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user