feat(tracing): add support for Rust Tracing (#30)

* feature(traces): support rust tracing with spans

* Make the extension hyper optimized for Rust Tracing

* Add support JSON logs

* Responsive Span Visualizer

* Change click behavior

* chore(package.json): bumps version to 0.5.0

---------

Co-authored-by: Priyank Chodisetti <priyank.ch@gmail.com>

Author: arthurgousset

package-lock.json

@@ -2,7 +2,7 @@
   "name": "traceback",
   "displayName": "TraceBack",
   "description": "A VS Code extension that brings telemetry data (traces, logs, and metrics) into your code.",
-  "version": "0.4.14",
+  "version": "0.5.0",
   "publisher": "hyperdrive-eng",
   "repository": {
     "type": "git",
@@ -35,11 +35,11 @@
     "configuration": {
       "title": "TraceBack",
       "properties": {
-          "traceback.claudeApiKey": {
-              "type": "string",
-              "default": "",
-              "description": "API key for Claude AI service"
-          }
+        "traceback.claudeApiKey": {
+          "type": "string",
+          "default": "",
+          "description": "API key for Claude AI service"
+        }
       }
     },
     "commands": [
@@ -111,6 +111,23 @@
         "command": "traceback.inspectVariableFromContext",
         "title": "Inspect Value",
         "icon": "$(eye)"
+      },
+      {
+        "command": "traceback.showSpanVisualizer",
+        "title": "Show Span Visualizer",
+        "category": "TraceBack"
+      },
+      {
+        "command": "traceback.importLogs",
+        "title": "Import Logs from File",
+        "category": "Traceback",
+        "icon": "$(file-add)"
+      },
+      {
+        "command": "traceback.pasteLogs",
+        "title": "Import Logs from Clipboard",
+        "category": "Traceback",
+        "icon": "$(clippy)"
       }
     ],
     "viewsContainers": {
@@ -160,6 +177,21 @@
           "command": "traceback.openSettings",
           "when": "view == logExplorer",
           "group": "navigation@3"
+        },
+        {
+          "command": "traceback.showSpanVisualizer",
+          "when": "view == logExplorer",
+          "group": "navigation@4"
+        },
+        {
+          "command": "traceback.importLogs",
+          "when": "view == tracebackLogs",
+          "group": "navigation"
+        },
+        {
+          "command": "traceback.pasteLogs",
+          "when": "view == tracebackLogs",
+          "group": "navigation"
         }
       ],
       "view/item/context": [
@@ -181,6 +213,10 @@
       {
         "command": "traceback.openSettings",
         "title": "TraceBack: Open Settings"
+      },
+      {
+        "command": "traceback.showSpanVisualizer",
+        "title": "TraceBack: Show Span Visualizer"
       }
     ]
   },
@@ -189,12 +225,24 @@
     "compile": "webpack --mode development",
     "watch": "webpack --mode development --watch",
     "pretest": "npm run compile",
-    "test": "node ./out/test/runTest.js",
+    "test": "jest",
     "package": "vsce package",
     "lint": "eslint src --ext ts"
   },
+  "jest": {
+    "preset": "ts-jest",
+    "testEnvironment": "node",
+    "testMatch": [
+      "<rootDir>/src/**/*.test.ts"
+    ],
+    "moduleFileExtensions": [
+      "ts",
+      "js"
+    ]
+  },
   "devDependencies": {
     "@types/glob": "^7.2.0",
+    "@types/jest": "^29.5.14",
     "@types/node": "^16.18.36",
     "@types/node-fetch": "^2.6.12",
     "@types/vscode": "^1.74.0",
@@ -203,6 +251,8 @@
     "@vscode/vsce": "^3.3.2",
     "eslint": "^8.42.0",
     "glob": "^8.1.0",
+    "jest": "^29.7.0",
+    "ts-jest": "^29.3.2",
     "ts-loader": "^9.5.2",
     "typescript": "^5.1.3",
     "webpack": "^5.99.3",

package.json

@@ -1,5 +1,5 @@
 import * as vscode from 'vscode';
-import { LogEntry, Span } from './logExplorer';
+import { RustLogEntry } from './logExplorer';
 import { CallerAnalysis } from './claudeService';
 import { ClaudeService } from './claudeService';
 import * as path from 'path';
@@ -141,7 +141,7 @@ export class CallStackExplorerProvider implements vscode.TreeDataProvider<CallSt
     new vscode.EventEmitter<CallStackTreeItem | undefined | null | void>();
   readonly onDidChangeTreeData: vscode.Event<CallStackTreeItem | undefined | null | void> = this._onDidChangeTreeData.event;
 
-  private currentLogEntry: LogEntry | undefined;
+  private currentLogEntry: RustLogEntry | undefined;
   private callerAnalysis: CallerNode[] = [];
   private claudeService: ClaudeService = ClaudeService.getInstance();
   private isAnalyzing: boolean = false;
@@ -152,7 +152,7 @@ export class CallStackExplorerProvider implements vscode.TreeDataProvider<CallSt
   /**
    * Set the spans for the current log entry and refresh the view
    */
-  public setLogEntry(log: LogEntry | undefined, isAnalyzing: boolean = false): void {
+  public setLogEntry(log: RustLogEntry | undefined, isAnalyzing: boolean = false): void {
     this.currentLogEntry = log;
     this.callerAnalysis = [];
     this.isAnalyzing = isAnalyzing;
@@ -414,19 +414,15 @@ export class CallStackExplorerProvider implements vscode.TreeDataProvider<CallSt
   async analyzeCallers(
     currentLogLine: string,
     staticSearchString: string,
-    allLogs: LogEntry[],
+    allLogs: RustLogEntry[],
     potentialCallers: Array<{ filePath: string; lineNumber: number; code: string; functionName: string; functionRange?: vscode.Range }>
   ): Promise<void> {
     try {
       this.isAnalyzing = true;
       this._onDidChangeTreeData.fire();
 
       vscode.window.showInformationMessage('Computing call stack analysis...');
-      const allLogLines = allLogs.map(log =>
-        log.message ||
-        log.jsonPayload?.fields?.message ||
-        ''
-      ).filter(msg => msg);
+      const allLogLines = allLogs.map(log => log.message || '').filter(msg => msg);
 
       const analysis = await this.claudeService.analyzeCallers(
         currentLogLine,

src/callStackExplorer.ts

@@ -28,7 +28,9 @@ export interface CallerAnalysis {
 export interface RegexPattern {
     pattern: string;
     description: string;
-    extractionMap: Record<string, string>;
+    extractionMap: {
+        [key: string]: string;  // Maps regex capture group names to RustLogEntry fields
+    };
 }
 
 export class ClaudeService {
@@ -59,6 +61,16 @@ export class ClaudeService {
         await vscode.workspace.getConfiguration('traceback').update('claudeApiKey', key, true);
     }
 
+    /**
+     * Analyze a log line to extract structured information
+     * 
+     * @param logLine The log line to analyze
+     * @param language The programming language of the codebase
+     * @returns {Promise<LLMLogAnalysis>} Analysis results including:
+     * 1. A static search string for finding similar logs
+     * 2. Extracted variables and their values
+     * 3. Include an "extractionMap" that maps regex capture group names to RustLogEntry field names
+     */
     public async analyzeLog(logMessage: string, language: string): Promise<LLMLogAnalysis> {
         if (!this.apiKey) {
             throw new Error('Claude API key not set. Please set your API key first.');
@@ -143,7 +155,7 @@ export class ClaudeService {
         const request = {
             messages: [{
                 role: 'user',
-                content: `Analyze this log message and extract:
+                content: `Analyze this log message from Rust's tracing library and extract:
 1. Think and infer a possibly longest static substring that can be searched in the code base.
 2. Key-value pairs of any variables or dynamic values in the log
 
@@ -153,17 +165,26 @@ Rules for static search string:
 - Predicted staticSearchString should be exact substring of logMessage. 
 - logMessage.substring(staticSearchString) should be true.
 - No regular expressions allowed.
+- For Rust tracing library logs, focus on the span name and event message.
 
 Rules for variables:
-- Extract all key-value pairs and dynamic values
+- Extract all key-value pairs and dynamic values from tracing fields
 - Preserve variable names as they appear in the log
+- Look for structured fields in the format key="value" or key=value
 
 Examples:
-Input: "[PlaceOrder] user_id=\"3790d414-165b-11f0-8ee4-96dac6adf53a\" user_currency=\"USD\""
-Static: "PlaceOrder"  (Note: brackets and log level removed)
+Input: "[TRACE traceback::span] span_name=\"PlaceOrder\" user_id=\"3790d414-165b-11f0-8ee4-96dac6adf53a\" user_currency=\"USD\""
+Static: "PlaceOrder"  (Note: brackets, log level, and module path removed)
 Variables: {
   "user_id": "3790d414-165b-11f0-8ee4-96dac6adf53a",
   "user_currency": "USD"
+}
+
+Input: "[DEBUG traceback::event] message=\"Processing order\" order_id=123 status=\"pending\""
+Static: "Processing order"
+Variables: {
+  "order_id": 123,
+  "status": "pending"
 }`
             }],
             model: this.model,