291 lines
10 KiB
TypeScript
291 lines
10 KiB
TypeScript
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);
|
|
});
|
|
});
|
|
});
|