attocode 0.1.0

package/dist/src/tracing/trace-collector.js

@@ -0,0 +1,727 @@
+/**
+ * Lesson 26: Trace Collector
+ *
+ * Central hub that coordinates tracing, metrics, and cache analysis.
+ * Subscribes to agent events and produces comprehensive JSONL output.
+ *
+ * Key responsibilities:
+ * 1. Capture full LLM request/response details
+ * 2. Track tool execution with timing
+ * 3. Analyze cache efficiency per request
+ * 4. Export structured JSONL for analysis
+ *
+ * @example
+ * ```typescript
+ * const collector = new TraceCollector({
+ *   outputDir: '.traces',
+ *   captureMessageContent: true,
+ * });
+ *
+ * collector.startSession('task-id', 'Write a fizzbuzz function', 'claude-3-sonnet');
+ *
+ * // During agent execution...
+ * collector.recordLLMRequest(requestData);
+ * collector.recordLLMResponse(responseData);
+ * collector.recordToolExecution(toolData);
+ *
+ * // When done
+ * const trace = await collector.endSession({ success: true });
+ * ```
+ */
+import { mkdir, appendFile } from 'fs/promises';
+import { join } from 'path';
+import { createTracer } from '../observability/tracer.js';
+import { createCacheBoundaryTracker } from './cache-boundary-tracker.js';
+import { DEFAULT_TRACE_CONFIG } from './types.js';
+// =============================================================================
+// TRACE COLLECTOR
+// =============================================================================
+/**
+ * Central trace collection hub.
+ */
+export class TraceCollector {
+    config;
+    tracer;
+    cacheTracker;
+    // Session state
+    sessionId = null;
+    traceId = null;
+    task = '';
+    model = '';
+    startTime = 0;
+    // Current iteration state
+    currentIteration = 0;
+    iterationSpan = null;
+    iterationStartTime = 0;
+    currentIterationTrace = null;
+    // Accumulated data
+    iterations = [];
+    pendingRequests = new Map();
+    pendingTools = new Map();
+    // JSONL output
+    outputPath = null;
+    writeQueue = Promise.resolve();
+    constructor(config = {}) {
+        this.config = { ...DEFAULT_TRACE_CONFIG, ...config };
+        this.tracer = createTracer('trace-collector');
+        this.cacheTracker = createCacheBoundaryTracker();
+    }
+    // ===========================================================================
+    // SESSION LIFECYCLE
+    // ===========================================================================
+    /**
+     * Start a new tracing session.
+     */
+    async startSession(sessionId, task, model, metadata = {}) {
+        if (this.sessionId) {
+            throw new Error('Session already in progress. Call endSession() first.');
+        }
+        this.sessionId = sessionId;
+        this.task = task;
+        this.model = model;
+        this.startTime = Date.now();
+        this.currentIteration = 0;
+        this.iterations = [];
+        this.cacheTracker.reset();
+        // Start root span
+        const rootSpan = this.tracer.startSpan('session', {
+            kind: 'internal',
+            attributes: {
+                'session.id': sessionId,
+                'session.task': task,
+                'session.model': model,
+            },
+        });
+        this.traceId = rootSpan.traceId;
+        // Set up JSONL output
+        if (this.config.enabled) {
+            await this.initializeOutput();
+            await this.writeEntry({
+                _type: 'session.start',
+                _ts: new Date().toISOString(),
+                traceId: this.traceId,
+                sessionId,
+                task,
+                model,
+                metadata,
+            });
+        }
+    }
+    /**
+     * End the current session and return the complete trace.
+     */
+    async endSession(result) {
+        if (!this.sessionId || !this.traceId) {
+            throw new Error('No session in progress');
+        }
+        const endTime = Date.now();
+        const durationMs = endTime - this.startTime;
+        // End any pending iteration
+        if (this.iterationSpan) {
+            this.tracer.endSpan(this.iterationSpan);
+            this.iterationSpan = null;
+        }
+        // Calculate aggregated metrics
+        const metrics = this.calculateAggregatedMetrics();
+        // Build session trace
+        const sessionTrace = {
+            sessionId: this.sessionId,
+            traceId: this.traceId,
+            task: this.task,
+            model: this.model,
+            startTime: this.startTime,
+            endTime,
+            durationMs,
+            status: result.success ? 'completed' : 'failed',
+            iterations: this.iterations,
+            metrics,
+            result,
+            metadata: {},
+        };
+        // Write session end entry
+        if (this.config.enabled) {
+            await this.writeEntry({
+                _type: 'session.end',
+                _ts: new Date().toISOString(),
+                traceId: this.traceId,
+                sessionId: this.sessionId,
+                status: sessionTrace.status,
+                durationMs,
+                metrics: {
+                    inputTokens: metrics.inputTokens,
+                    outputTokens: metrics.outputTokens,
+                    cacheReadTokens: metrics.cacheReadTokens,
+                    cacheWriteTokens: metrics.cacheWriteTokens,
+                    estimatedCost: metrics.estimatedCost,
+                    toolCalls: metrics.toolCalls,
+                    llmCalls: metrics.llmCalls,
+                    duration: durationMs,
+                    errors: metrics.errors,
+                    retries: metrics.retries,
+                },
+            });
+            await this.flush();
+        }
+        // Reset state
+        this.sessionId = null;
+        this.traceId = null;
+        this.task = '';
+        this.model = '';
+        return sessionTrace;
+    }
+    // ===========================================================================
+    // EVENT RECORDING
+    // ===========================================================================
+    /**
+     * Record an event.
+     */
+    async record(event) {
+        if (!this.config.enabled || !this.sessionId)
+            return;
+        switch (event.type) {
+            case 'llm.request':
+                await this.recordLLMRequest(event.data);
+                break;
+            case 'llm.response':
+                await this.recordLLMResponse(event.data);
+                break;
+            case 'tool.start':
+                await this.recordToolStart(event.data);
+                break;
+            case 'tool.end':
+                await this.recordToolEnd(event.data);
+                break;
+            case 'iteration.start':
+                this.startIteration(event.data.iterationNumber);
+                break;
+            case 'iteration.end':
+                await this.endIteration();
+                break;
+            case 'error':
+                await this.recordError(event.data);
+                break;
+        }
+    }
+    /**
+     * Record an LLM request.
+     */
+    async recordLLMRequest(data) {
+        const span = this.tracer.startSpan('llm.request', {
+            kind: 'client',
+            attributes: {
+                'llm.model': data.model,
+                'llm.provider': data.provider,
+                'llm.message_count': data.messages.length,
+                'llm.tool_count': data.tools?.length ?? 0,
+            },
+        });
+        // Track cache boundaries
+        if (this.config.analyzeCacheBoundaries) {
+            this.cacheTracker.recordRequest({
+                systemPrompt: data.systemPrompt ?? data.messages.find(m => m.role === 'system')?.content ?? '',
+                messages: data.messages,
+                toolDefinitions: data.tools,
+            });
+        }
+        // Store pending request
+        this.pendingRequests.set(data.requestId, {
+            span,
+            data,
+            startTime: Date.now(),
+        });
+        // Estimate input tokens
+        const estimatedInputTokens = this.estimateInputTokens(data);
+        // Write JSONL entry
+        await this.writeEntry({
+            _type: 'llm.request',
+            _ts: new Date().toISOString(),
+            traceId: this.traceId,
+            requestId: data.requestId,
+            model: data.model,
+            messageCount: data.messages.length,
+            toolCount: data.tools?.length ?? 0,
+            estimatedInputTokens,
+        });
+    }
+    /**
+     * Record an LLM response.
+     */
+    async recordLLMResponse(data) {
+        const pending = this.pendingRequests.get(data.requestId);
+        if (!pending) {
+            console.warn(`No pending request found for ${data.requestId}`);
+            return;
+        }
+        const { span, data: requestData, startTime } = pending;
+        this.pendingRequests.delete(data.requestId);
+        // Update cache tracker with actual API response
+        let cacheBreakdown;
+        if (this.config.analyzeCacheBoundaries && data.usage.cacheReadTokens !== undefined) {
+            cacheBreakdown = this.cacheTracker.recordResponse({
+                cacheReadTokens: data.usage.cacheReadTokens,
+                cacheWriteTokens: data.usage.cacheWriteTokens ?? 0,
+                totalInputTokens: data.usage.inputTokens,
+                outputTokens: data.usage.outputTokens,
+            });
+        }
+        else {
+            cacheBreakdown = {
+                cacheReadTokens: data.usage.cacheReadTokens ?? 0,
+                cacheWriteTokens: data.usage.cacheWriteTokens ?? 0,
+                freshTokens: data.usage.inputTokens - (data.usage.cacheReadTokens ?? 0),
+                hitRate: data.usage.cacheReadTokens
+                    ? data.usage.cacheReadTokens / data.usage.inputTokens
+                    : 0,
+                estimatedSavings: 0,
+                breakpoints: [],
+            };
+        }
+        // Calculate token breakdown
+        const tokens = {
+            input: data.usage.inputTokens,
+            output: data.usage.outputTokens,
+            total: data.usage.inputTokens + data.usage.outputTokens,
+            breakdown: {
+                systemPrompt: this.estimateTokens(requestData.systemPrompt ?? ''),
+                messages: data.usage.inputTokens - this.estimateTokens(requestData.systemPrompt ?? ''),
+                toolDefinitions: requestData.tools ? this.estimateToolTokens(requestData.tools) : 0,
+                toolResults: 0,
+            },
+        };
+        // Build LLM request trace
+        const llmTrace = {
+            requestId: data.requestId,
+            traceId: this.traceId,
+            spanId: span.spanId,
+            timestamp: startTime,
+            durationMs: data.durationMs,
+            model: requestData.model,
+            provider: requestData.provider,
+            request: {
+                messages: this.config.captureMessageContent
+                    ? this.traceMessages(requestData.messages)
+                    : [],
+                tools: requestData.tools?.map(t => ({
+                    name: t.name,
+                    description: t.description,
+                    parametersSchema: t.parametersSchema,
+                    estimatedTokens: this.estimateTokens(JSON.stringify(t)),
+                })),
+                parameters: requestData.parameters,
+            },
+            response: {
+                content: this.config.captureMessageContent ? data.content : '[content not captured]',
+                toolCalls: data.toolCalls,
+                stopReason: data.stopReason,
+            },
+            tokens,
+            cache: cacheBreakdown,
+            // Store actual cost from provider (e.g., OpenRouter) when available
+            actualCost: data.usage.cost,
+        };
+        // Store in current iteration
+        if (this.currentIterationTrace) {
+            this.currentIterationTrace.llmRequest = llmTrace;
+        }
+        // Update span
+        this.tracer.setAttributes(span, {
+            'llm.input_tokens': data.usage.inputTokens,
+            'llm.output_tokens': data.usage.outputTokens,
+            'llm.cache_read_tokens': data.usage.cacheReadTokens ?? 0,
+            'llm.cache_hit_rate': cacheBreakdown.hitRate,
+            'llm.stop_reason': data.stopReason,
+            'llm.tool_calls': data.toolCalls?.length ?? 0,
+        });
+        this.tracer.endSpan(span);
+        // Write JSONL entry
+        await this.writeEntry({
+            _type: 'llm.response',
+            _ts: new Date().toISOString(),
+            traceId: this.traceId,
+            requestId: data.requestId,
+            durationMs: data.durationMs,
+            tokens,
+            cache: cacheBreakdown,
+            stopReason: data.stopReason,
+            toolCallCount: data.toolCalls?.length ?? 0,
+        });
+    }
+    /**
+     * Record tool execution start.
+     */
+    async recordToolStart(data) {
+        const span = this.tracer.startSpan(`tool.${data.toolName}`, {
+            kind: 'internal',
+            attributes: {
+                'tool.name': data.toolName,
+                'tool.execution_id': data.executionId,
+            },
+        });
+        this.pendingTools.set(data.executionId, {
+            span,
+            data,
+            startTime: Date.now(),
+        });
+    }
+    /**
+     * Record tool execution end.
+     */
+    async recordToolEnd(data) {
+        const pending = this.pendingTools.get(data.executionId);
+        if (!pending) {
+            console.warn(`No pending tool execution found for ${data.executionId}`);
+            return;
+        }
+        const { span, data: startData, startTime } = pending;
+        this.pendingTools.delete(data.executionId);
+        // Build tool trace
+        const toolTrace = {
+            executionId: data.executionId,
+            traceId: this.traceId,
+            spanId: span.spanId,
+            toolName: startData.toolName,
+            arguments: startData.arguments,
+            startTime,
+            durationMs: data.durationMs,
+            status: data.status,
+        };
+        // Add result or error
+        if (data.status === 'success' && data.result !== undefined) {
+            const resultStr = typeof data.result === 'string'
+                ? data.result
+                : JSON.stringify(data.result);
+            toolTrace.result = {
+                type: typeof data.result === 'string' ? 'string' : 'object',
+                value: this.config.captureToolResults
+                    ? this.truncateResult(data.result)
+                    : '[result not captured]',
+                truncated: resultStr.length > this.config.maxResultSize,
+                originalSize: resultStr.length,
+            };
+        }
+        if (data.error) {
+            toolTrace.error = {
+                name: data.error.name,
+                message: data.error.message,
+                stack: data.error.stack,
+            };
+        }
+        // Store in current iteration
+        if (this.currentIterationTrace) {
+            if (!this.currentIterationTrace.toolExecutions) {
+                this.currentIterationTrace.toolExecutions = [];
+            }
+            this.currentIterationTrace.toolExecutions.push(toolTrace);
+        }
+        // Update span
+        this.tracer.setAttributes(span, {
+            'tool.status': data.status,
+            'tool.duration_ms': data.durationMs,
+        });
+        if (data.status === 'error') {
+            this.tracer.setError(span, data.error ?? new Error('Unknown error'));
+        }
+        this.tracer.endSpan(span, {
+            code: data.status === 'success' ? 'ok' : 'error',
+            message: data.status === 'error' ? data.error?.message : undefined,
+        });
+        // Write JSONL entry
+        await this.writeEntry({
+            _type: 'tool.execution',
+            _ts: new Date().toISOString(),
+            traceId: this.traceId,
+            executionId: data.executionId,
+            toolName: startData.toolName,
+            durationMs: data.durationMs,
+            status: data.status,
+            resultSize: toolTrace.result?.originalSize,
+        });
+    }
+    /**
+     * Record an error.
+     */
+    async recordError(data) {
+        await this.writeEntry({
+            _type: 'error',
+            _ts: new Date().toISOString(),
+            traceId: this.traceId,
+            errorCode: data.code,
+            errorMessage: data.message,
+            context: data.context,
+            recoverable: data.recoverable,
+        });
+    }
+    // ===========================================================================
+    // ITERATION MANAGEMENT
+    // ===========================================================================
+    /**
+     * Start a new iteration.
+     */
+    startIteration(iterationNumber) {
+        this.currentIteration = iterationNumber;
+        this.iterationStartTime = Date.now();
+        this.iterationSpan = this.tracer.startSpan(`iteration.${iterationNumber}`, {
+            kind: 'internal',
+            attributes: {
+                'iteration.number': iterationNumber,
+            },
+        });
+        this.currentIterationTrace = {
+            iterationNumber,
+            traceId: this.traceId,
+            spanId: this.iterationSpan.spanId,
+            startTime: this.iterationStartTime,
+            toolExecutions: [],
+        };
+    }
+    /**
+     * End the current iteration.
+     */
+    async endIteration() {
+        if (!this.iterationSpan || !this.currentIterationTrace)
+            return;
+        const endTime = Date.now();
+        const durationMs = endTime - this.iterationStartTime;
+        // Calculate iteration metrics
+        const llmRequest = this.currentIterationTrace.llmRequest;
+        // Use actual cost from provider if available, otherwise calculate
+        const actualCost = llmRequest?.actualCost;
+        const calculatedCost = this.calculateCost(llmRequest?.tokens.input ?? 0, llmRequest?.tokens.output ?? 0, llmRequest?.cache.cacheReadTokens ?? 0);
+        const metrics = {
+            inputTokens: llmRequest?.tokens.input ?? 0,
+            outputTokens: llmRequest?.tokens.output ?? 0,
+            cacheHitRate: llmRequest?.cache.hitRate ?? 0,
+            toolCallCount: this.currentIterationTrace.toolExecutions?.length ?? 0,
+            totalCost: actualCost ?? calculatedCost,
+        };
+        // Complete iteration trace
+        const iterationTrace = {
+            ...this.currentIterationTrace,
+            durationMs,
+            metrics,
+        };
+        this.iterations.push(iterationTrace);
+        // End span
+        this.tracer.setAttributes(this.iterationSpan, {
+            'iteration.duration_ms': durationMs,
+            'iteration.input_tokens': metrics.inputTokens,
+            'iteration.output_tokens': metrics.outputTokens,
+            'iteration.cache_hit_rate': metrics.cacheHitRate,
+            'iteration.tool_calls': metrics.toolCallCount,
+        });
+        this.tracer.endSpan(this.iterationSpan);
+        this.iterationSpan = null;
+        this.currentIterationTrace = null;
+    }
+    // ===========================================================================
+    // JSONL OUTPUT
+    // ===========================================================================
+    /**
+     * Initialize output file.
+     */
+    async initializeOutput() {
+        if (!this.config.outputDir)
+            return;
+        await mkdir(this.config.outputDir, { recursive: true });
+        const filename = this.config.filePattern
+            .replace('{sessionId}', this.sessionId ?? 'unknown')
+            .replace('{timestamp}', Date.now().toString());
+        this.outputPath = join(this.config.outputDir, filename);
+    }
+    /**
+     * Write a JSONL entry.
+     */
+    async writeEntry(entry) {
+        if (!this.outputPath)
+            return;
+        // Queue writes to avoid race conditions
+        this.writeQueue = this.writeQueue.then(async () => {
+            try {
+                await appendFile(this.outputPath, JSON.stringify(entry) + '\n');
+            }
+            catch (err) {
+                console.error('Failed to write trace entry:', err);
+            }
+        });
+    }
+    /**
+     * Flush pending writes.
+     */
+    async flush() {
+        await this.writeQueue;
+    }
+    // ===========================================================================
+    // HELPER METHODS
+    // ===========================================================================
+    /**
+     * Convert messages to traced format.
+     */
+    traceMessages(messages) {
+        return messages.map(msg => ({
+            role: msg.role,
+            content: msg.content,
+            estimatedTokens: this.estimateTokens(msg.content),
+            toolCallId: msg.toolCallId,
+            toolCalls: msg.toolCalls,
+            contentHash: this.hashContent(msg.content),
+        }));
+    }
+    /**
+     * Estimate tokens from text.
+     */
+    estimateTokens(text) {
+        return Math.ceil(text.length / 4);
+    }
+    /**
+     * Estimate input tokens for a request.
+     */
+    estimateInputTokens(data) {
+        let tokens = 0;
+        for (const msg of data.messages) {
+            tokens += this.estimateTokens(msg.content);
+        }
+        if (data.tools) {
+            tokens += this.estimateToolTokens(data.tools);
+        }
+        return tokens;
+    }
+    /**
+     * Estimate tokens for tool definitions.
+     */
+    estimateToolTokens(tools) {
+        let tokens = 0;
+        for (const tool of tools) {
+            tokens += this.estimateTokens(tool.name);
+            tokens += this.estimateTokens(tool.description);
+            tokens += this.estimateTokens(JSON.stringify(tool.parametersSchema));
+        }
+        return tokens;
+    }
+    /**
+     * Truncate result if too large.
+     */
+    truncateResult(result) {
+        const str = typeof result === 'string' ? result : JSON.stringify(result);
+        if (str.length <= this.config.maxResultSize) {
+            return result;
+        }
+        return str.substring(0, this.config.maxResultSize) + '... [truncated]';
+    }
+    /**
+     * Calculate cost for tokens.
+     */
+    calculateCost(inputTokens, outputTokens, cachedTokens) {
+        // Claude 3 Sonnet pricing
+        const inputCostPer1k = 0.003;
+        const outputCostPer1k = 0.015;
+        const cachedCostPer1k = 0.0003; // ~10x cheaper
+        const uncachedInput = inputTokens - cachedTokens;
+        return ((uncachedInput / 1000) * inputCostPer1k +
+            (cachedTokens / 1000) * cachedCostPer1k +
+            (outputTokens / 1000) * outputCostPer1k);
+    }
+    /**
+     * Simple content hash.
+     */
+    hashContent(content) {
+        let hash = 0;
+        for (let i = 0; i < content.length; i++) {
+            const char = content.charCodeAt(i);
+            hash = ((hash << 5) - hash) + char;
+            hash = hash & hash;
+        }
+        return hash.toString(16);
+    }
+    /**
+     * Calculate aggregated metrics across all iterations.
+     */
+    calculateAggregatedMetrics() {
+        let inputTokens = 0;
+        let outputTokens = 0;
+        let cacheReadTokens = 0;
+        let cacheWriteTokens = 0;
+        let estimatedCost = 0;
+        let toolCalls = 0;
+        let totalCacheHitRate = 0;
+        let tokensSavedByCache = 0;
+        for (const iteration of this.iterations) {
+            inputTokens += iteration.metrics.inputTokens;
+            outputTokens += iteration.metrics.outputTokens;
+            cacheReadTokens += iteration.llmRequest?.cache.cacheReadTokens ?? 0;
+            cacheWriteTokens += iteration.llmRequest?.cache.cacheWriteTokens ?? 0;
+            estimatedCost += iteration.metrics.totalCost;
+            toolCalls += iteration.metrics.toolCallCount;
+            totalCacheHitRate += iteration.metrics.cacheHitRate;
+        }
+        const llmCalls = this.iterations.length;
+        const avgCacheHitRate = llmCalls > 0 ? totalCacheHitRate / llmCalls : 0;
+        // Cost saved by caching
+        const costWithoutCache = this.calculateCost(inputTokens, outputTokens, 0);
+        const costSavedByCache = costWithoutCache - estimatedCost;
+        tokensSavedByCache = cacheReadTokens; // Tokens that would have been re-sent
+        return {
+            inputTokens,
+            outputTokens,
+            cacheReadTokens,
+            cacheWriteTokens,
+            estimatedCost,
+            toolCalls,
+            llmCalls,
+            duration: Date.now() - this.startTime,
+            errors: 0, // TODO: track errors
+            retries: 0, // TODO: track retries
+            avgCacheHitRate,
+            tokensSavedByCache,
+            costSavedByCache,
+        };
+    }
+    // ===========================================================================
+    // PUBLIC ACCESSORS
+    // ===========================================================================
+    /**
+     * Get the underlying tracer.
+     */
+    getTracer() {
+        return this.tracer;
+    }
+    /**
+     * Get the cache tracker.
+     */
+    getCacheTracker() {
+        return this.cacheTracker;
+    }
+    /**
+     * Get current session ID.
+     */
+    getSessionId() {
+        return this.sessionId;
+    }
+    /**
+     * Get current trace ID.
+     */
+    getTraceId() {
+        return this.traceId;
+    }
+    /**
+     * Check if a session is active.
+     */
+    isSessionActive() {
+        return this.sessionId !== null;
+    }
+    /**
+     * Get iteration count.
+     */
+    getIterationCount() {
+        return this.iterations.length;
+    }
+}
+// =============================================================================
+// FACTORY FUNCTION
+// =============================================================================
+/**
+ * Create a new trace collector.
+ */
+export function createTraceCollector(config = {}) {
+    return new TraceCollector(config);
+}
+//# sourceMappingURL=trace-collector.js.map