@probelabs/probe 0.6.0-rc224 → 0.6.0-rc225

package/cjs/index.cjs

@@ -39321,11 +39321,12 @@ function createTools(configOptions) {
   return tools2;
 }
 function parseXmlToolCallWithThinking(xmlString, validTools) {
-  const { cleanedXmlString, recoveryResult } = processXmlWithThinkingAndRecovery(xmlString, validTools);
+  const { cleanedXmlString, recoveryResult, thinkingContent } = processXmlWithThinkingAndRecovery(xmlString, validTools);
   if (recoveryResult) {
-    return recoveryResult;
+    return { ...recoveryResult, thinkingContent };
   }
-  return parseXmlToolCall(cleanedXmlString, validTools);
+  const toolCall = parseXmlToolCall(cleanedXmlString, validTools);
+  return toolCall ? { ...toolCall, thinkingContent } : null;
 }
 var import_crypto2, implementToolDefinition, listFilesToolDefinition, searchFilesToolDefinition, listSkillsToolDefinition, useSkillToolDefinition, readImageToolDefinition;
 var init_tools = __esm({
@@ -85596,25 +85597,27 @@ function parseXmlMcpToolCall(xmlString, mcpToolNames = []) {
 function parseHybridXmlToolCall(xmlString, nativeTools = [], mcpBridge = null) {
   const nativeResult = parseNativeXmlToolWithThinking(xmlString, nativeTools);
   if (nativeResult) {
-    return { ...nativeResult, type: "native" };
+    const { thinkingContent, ...rest } = nativeResult;
+    return { ...rest, type: "native", thinkingContent };
   }
   if (mcpBridge) {
     const mcpResult = parseXmlMcpToolCall(xmlString, mcpBridge.getToolNames());
     if (mcpResult) {
-      return { ...mcpResult, type: "mcp" };
+      const { thinkingContent } = processXmlWithThinkingAndRecovery(xmlString, []);
+      return { ...mcpResult, type: "mcp", thinkingContent };
     }
   }
   return null;
 }
 function parseNativeXmlToolWithThinking(xmlString, validTools) {
-  const { cleanedXmlString, recoveryResult } = processXmlWithThinkingAndRecovery(xmlString, validTools);
+  const { cleanedXmlString, recoveryResult, thinkingContent } = processXmlWithThinkingAndRecovery(xmlString, validTools);
   if (recoveryResult) {
-    return recoveryResult;
+    return { ...recoveryResult, thinkingContent };
   }
   for (const toolName of validTools) {
     const result = parseNativeXmlTool(cleanedXmlString, toolName);
     if (result) {
-      return result;
+      return { ...result, thinkingContent };
     }
   }
   return null;
@@ -96075,6 +96078,181 @@ var init_ProbeAgent = __esm({
       _filterMcpTools(mcpToolNames) {
         return mcpToolNames.filter((toolName) => this._isMcpToolAllowed(toolName));
       }
+      /**
+       * Check if tracer is AppTracer (expects sessionId as first param) vs SimpleAppTracer
+       * @returns {boolean} - True if tracer is AppTracer style (requires sessionId)
+       * @private
+       */
+      _isAppTracerStyle() {
+        return this.tracer && typeof this.tracer.sessionSpans !== "undefined";
+      }
+      /**
+       * Record an error classification event for telemetry
+       * Provides unified error recording across all error types
+       * @param {string} errorType - Error type (wrapped_tool, unrecognized_tool, no_tool_call, circuit_breaker)
+       * @param {string} message - Error message
+       * @param {Object} context - Additional context data
+       * @param {number} iteration - Current iteration number
+       * @private
+       */
+      _recordErrorTelemetry(errorType, message, context, iteration) {
+        if (!this.tracer) return;
+        if (this._isAppTracerStyle() && typeof this.tracer.recordErrorClassification === "function") {
+          this.tracer.recordErrorClassification(this.sessionId, iteration, errorType, {
+            message,
+            context
+          });
+        } else if (typeof this.tracer.recordErrorEvent === "function") {
+          this.tracer.recordErrorEvent(errorType, {
+            message,
+            context: { ...context, iteration }
+          });
+        } else {
+          this.tracer.addEvent(`error.${errorType}`, {
+            "error.type": errorType,
+            "error.message": message,
+            "error.recoverable": errorType !== "circuit_breaker",
+            "error.context": JSON.stringify(context).substring(0, 1e3),
+            "iteration": iteration
+          });
+        }
+      }
+      /**
+       * Record AI thinking content for telemetry
+       * @param {string} thinkingContent - The thinking content
+       * @param {number} iteration - Current iteration number
+       * @private
+       */
+      _recordThinkingTelemetry(thinkingContent, iteration) {
+        if (!this.tracer || !thinkingContent) return;
+        if (this._isAppTracerStyle() && typeof this.tracer.recordThinkingContent === "function") {
+          this.tracer.recordThinkingContent(this.sessionId, iteration, thinkingContent);
+        } else if (typeof this.tracer.recordThinkingContent === "function") {
+          this.tracer.recordThinkingContent(thinkingContent, { iteration });
+        } else {
+          this.tracer.addEvent("ai.thinking", {
+            "ai.thinking.content": thinkingContent.substring(0, 5e4),
+            "ai.thinking.length": thinkingContent.length,
+            "iteration": iteration
+          });
+        }
+      }
+      /**
+       * Record AI tool decision for telemetry
+       * @param {string} toolName - The tool name
+       * @param {Object} params - Tool parameters
+       * @param {number} responseLength - Length of AI response
+       * @param {number} iteration - Current iteration number
+       * @private
+       */
+      _recordToolDecisionTelemetry(toolName, params, responseLength, iteration) {
+        if (!this.tracer) return;
+        if (this._isAppTracerStyle() && typeof this.tracer.recordAIToolDecision === "function") {
+          this.tracer.recordAIToolDecision(this.sessionId, iteration, toolName, params);
+        } else if (typeof this.tracer.recordToolDecision === "function") {
+          this.tracer.recordToolDecision(toolName, params, {
+            iteration,
+            "ai.tool_decision.raw_response_length": responseLength
+          });
+        } else {
+          this.tracer.addEvent("ai.tool_decision", {
+            "ai.tool_decision.name": toolName,
+            "ai.tool_decision.params": JSON.stringify(params || {}).substring(0, 2e3),
+            "ai.tool_decision.raw_response_length": responseLength,
+            "iteration": iteration
+          });
+        }
+      }
+      /**
+       * Record tool result for telemetry
+       * @param {string} toolName - The tool name
+       * @param {string|Object} result - Tool result
+       * @param {boolean} success - Whether tool succeeded
+       * @param {number} durationMs - Execution duration in milliseconds
+       * @param {number} iteration - Current iteration number
+       * @private
+       */
+      _recordToolResultTelemetry(toolName, result, success, durationMs, iteration) {
+        if (!this.tracer) return;
+        if (this._isAppTracerStyle() && typeof this.tracer.recordToolResult === "function") {
+          this.tracer.recordToolResult(this.sessionId, iteration, toolName, result, success, durationMs);
+        } else if (typeof this.tracer.recordToolResult === "function") {
+          this.tracer.recordToolResult(toolName, result, success, durationMs, { iteration });
+        } else {
+          const resultStr = typeof result === "string" ? result : JSON.stringify(result || "");
+          this.tracer.addEvent("tool.result", {
+            "tool.name": toolName,
+            "tool.result": resultStr.substring(0, 1e4),
+            "tool.result.length": resultStr.length,
+            "tool.duration_ms": durationMs,
+            "tool.success": success,
+            "iteration": iteration
+          });
+        }
+      }
+      /**
+       * Record MCP tool lifecycle event for telemetry
+       * @param {string} phase - 'start' or 'end'
+       * @param {string} toolName - MCP tool name
+       * @param {Object} params - Tool parameters (for start) or null (for end)
+       * @param {number} iteration - Current iteration number
+       * @param {Object} [endData] - Additional data for end phase (result, success, durationMs, error)
+       * @private
+       */
+      _recordMcpToolTelemetry(phase, toolName, params, iteration, endData = null) {
+        if (!this.tracer) return;
+        if (phase === "start") {
+          if (this._isAppTracerStyle() && typeof this.tracer.recordMcpToolStart === "function") {
+            this.tracer.recordMcpToolStart(this.sessionId, iteration, toolName, "mcp", params);
+          } else if (typeof this.tracer.recordMcpToolStart === "function") {
+            this.tracer.recordMcpToolStart(toolName, "mcp", params, { iteration });
+          } else {
+            this.tracer.addEvent("mcp.tool.start", {
+              "mcp.tool.name": toolName,
+              "mcp.tool.server": "mcp",
+              "mcp.tool.params": JSON.stringify(params || {}).substring(0, 2e3),
+              "iteration": iteration
+            });
+          }
+        } else if (phase === "end" && endData) {
+          const { result, success, durationMs, error: error2 } = endData;
+          if (this._isAppTracerStyle() && typeof this.tracer.recordMcpToolEnd === "function") {
+            this.tracer.recordMcpToolEnd(this.sessionId, iteration, toolName, "mcp", result, success, durationMs, error2);
+          } else if (typeof this.tracer.recordMcpToolEnd === "function") {
+            this.tracer.recordMcpToolEnd(toolName, "mcp", result, success, durationMs, error2, { iteration });
+          } else {
+            const resultStr = typeof result === "string" ? result : JSON.stringify(result || "");
+            this.tracer.addEvent("mcp.tool.end", {
+              "mcp.tool.name": toolName,
+              "mcp.tool.server": "mcp",
+              "mcp.tool.result": resultStr.substring(0, 1e4),
+              "mcp.tool.result.length": resultStr.length,
+              "mcp.tool.duration_ms": durationMs,
+              "mcp.tool.success": success,
+              "mcp.tool.error": error2,
+              "iteration": iteration
+            });
+          }
+        }
+      }
+      /**
+       * Record iteration lifecycle event for telemetry
+       * @param {string} phase - 'end' (start is already handled elsewhere)
+       * @param {number} iteration - Current iteration number
+       * @param {Object} data - Additional iteration data
+       * @private
+       */
+      _recordIterationTelemetry(phase, iteration, data2 = {}) {
+        if (!this.tracer) return;
+        if (typeof this.tracer.recordIterationEvent === "function") {
+          this.tracer.recordIterationEvent(phase, iteration, data2);
+        } else {
+          this.tracer.addEvent(`iteration.${phase}`, {
+            "iteration": iteration,
+            ...data2
+          });
+        }
+      }
       /**
        * Initialize the agent asynchronously (must be called after constructor)
        * This method initializes MCP and merges MCP tools into the tool list, and loads history from storage
@@ -98029,8 +98207,12 @@ You are working with a repository located at: ${searchDirectory}
             }
             const nativeTools = validTools;
             const parsedTool = this.mcpBridge && !options._disableTools ? parseHybridXmlToolCall(assistantResponseContent, nativeTools, this.mcpBridge) : parseXmlToolCallWithThinking(assistantResponseContent, validTools);
+            if (parsedTool?.thinkingContent) {
+              this._recordThinkingTelemetry(parsedTool.thinkingContent, currentIteration);
+            }
             if (parsedTool) {
               const { toolName, params } = parsedTool;
+              this._recordToolDecisionTelemetry(toolName, params, assistantResponseContent.length, currentIteration);
               if (this.debug) console.log(`[DEBUG] Parsed tool call: ${toolName} with params:`, params);
               if (toolName === "attempt_completion") {
                 completionAttempted = true;
@@ -98107,6 +98289,8 @@ You are working with a repository located at: ${searchDirectory}
               } else {
                 const { type } = parsedTool;
                 if (type === "mcp" && this.mcpBridge && this.mcpBridge.isMcpTool(toolName)) {
+                  const mcpStartTime = Date.now();
+                  this._recordMcpToolTelemetry("start", toolName, params, currentIteration);
                   try {
                     if (this.debug) {
                       console.error(`
@@ -98136,6 +98320,13 @@ You are working with a repository located at: ${searchDirectory}
                     } catch (truncateError) {
                       console.error(`[WARN] Tool output truncation failed: ${truncateError.message}`);
                     }
+                    const mcpDurationMs = Date.now() - mcpStartTime;
+                    this._recordMcpToolTelemetry("end", toolName, null, currentIteration, {
+                      result: toolResultContent,
+                      success: true,
+                      durationMs: mcpDurationMs,
+                      error: null
+                    });
                     if (this.debug) {
                       const preview = toolResultContent.length > 500 ? toolResultContent.substring(0, 500) + "..." : toolResultContent;
                       console.error(`[DEBUG] ========================================`);
@@ -98149,6 +98340,13 @@ You are working with a repository located at: ${searchDirectory}
 ${toolResultContent}
 </tool_result>` });
                   } catch (error2) {
+                    const mcpDurationMs = Date.now() - mcpStartTime;
+                    this._recordMcpToolTelemetry("end", toolName, null, currentIteration, {
+                      result: null,
+                      success: false,
+                      durationMs: mcpDurationMs,
+                      error: error2.message
+                    });
                     console.error(`Error executing MCP tool ${toolName}:`, error2);
                     if (this.debug) {
                       console.error(`[DEBUG] ========================================`);
@@ -98240,6 +98438,7 @@ ${errorXml}
                       return await this.toolImplementations[toolName].execute(toolParams);
                     };
                     let toolResult;
+                    const toolStartTime = Date.now();
                     try {
                       if (this.tracer) {
                         toolResult = await this.tracer.withSpan("tool.call", executeToolCall, {
@@ -98250,6 +98449,8 @@ ${errorXml}
                       } else {
                         toolResult = await executeToolCall();
                       }
+                      const toolDurationMs = Date.now() - toolStartTime;
+                      this._recordToolResultTelemetry(toolName, toolResult, true, toolDurationMs, currentIteration);
                       if (this.debug) {
                         const resultPreview = typeof toolResult === "string" ? toolResult.length > 500 ? toolResult.substring(0, 500) + "..." : toolResult : toolResult ? JSON.stringify(toolResult, null, 2).substring(0, 500) + "..." : "No Result";
                         console.error(`[DEBUG] ========================================`);
@@ -98306,6 +98507,20 @@ ${toolResultContent}
                       role: "user",
                       content: toolResultMessage
                     });
+                    if (this.tracer) {
+                      if (typeof this.tracer.recordConversationTurn === "function") {
+                        this.tracer.recordConversationTurn("assistant", assistantResponseContent, {
+                          iteration: currentIteration,
+                          has_tool_call: true,
+                          tool_name: toolName
+                        });
+                        this.tracer.recordConversationTurn("tool_result", toolResultContent, {
+                          iteration: currentIteration,
+                          tool_name: toolName,
+                          tool_success: true
+                        });
+                      }
+                    }
                     if (this.debug) {
                       console.log(`[DEBUG] Tool ${toolName} executed successfully. Result length: ${typeof toolResult === "string" ? toolResult.length : JSON.stringify(toolResult).length}`);
                     }
@@ -98376,6 +98591,7 @@ ${errorXml}
                 if (this.debug) {
                   console.log(`[DEBUG] Detected wrapped tool '${wrappedToolName}' in assistant response - wrong XML format.`);
                 }
+                this._recordErrorTelemetry("wrapped_tool", "Tool call wrapped in markdown", { toolName: wrappedToolName }, currentIteration);
                 const toolError = new ParameterError(
                   `Tool '${wrappedToolName}' found but in WRONG FORMAT - do not wrap tools in other XML tags.`,
                   {
@@ -98401,6 +98617,7 @@ ${formatErrorForAI(toolError)}
                 if (this.debug) {
                   console.log(`[DEBUG] Detected unrecognized tool '${unrecognizedTool}' in assistant response.`);
                 }
+                this._recordErrorTelemetry("unrecognized_tool", `Unknown tool: ${unrecognizedTool}`, { toolName: unrecognizedTool, validTools }, currentIteration);
                 const toolError = new ParameterError(`Tool '${unrecognizedTool}' is not available in this context.`, {
                   suggestion: `Available tools: ${validTools.join(", ")}. Please use one of these tools instead.`
                 });
@@ -98408,6 +98625,7 @@ ${formatErrorForAI(toolError)}
 ${formatErrorForAI(toolError)}
 </tool_result>`;
               } else {
+                this._recordErrorTelemetry("no_tool_call", "AI response did not contain tool call", { responsePreview: assistantResponseContent.substring(0, 500) }, currentIteration);
                 if (currentIteration >= maxIterations) {
                   let cleanedResponse = assistantResponseContent;
                   cleanedResponse = cleanedResponse.replace(/<thinking>[\s\S]*?<\/thinking>/gi, "").trim();
@@ -98483,6 +98701,7 @@ Note: <attempt_complete></attempt_complete> reuses your PREVIOUS assistant messa
                   sameFormatErrorCount++;
                   if (sameFormatErrorCount >= MAX_REPEATED_FORMAT_ERRORS) {
                     const errorDesc = isWrapped ? "wrapped tool format" : unrecognizedTool;
+                    this._recordErrorTelemetry("circuit_breaker", "Format error limit exceeded", { formatErrorCount: sameFormatErrorCount, errorCategory }, currentIteration);
                     console.error(`[ERROR] Format error category '${errorCategory}' repeated ${sameFormatErrorCount} times. Breaking loop early to prevent infinite iteration.`);
                     finalResult = `Error: Unable to complete request. The AI model repeatedly used incorrect tool call format (${errorDesc}). Please try rephrasing your question or using a different model.`;
                     break;
@@ -98496,6 +98715,10 @@ Note: <attempt_complete></attempt_complete> reuses your PREVIOUS assistant messa
                 sameFormatErrorCount = 0;
               }
             }
+            this._recordIterationTelemetry("end", currentIteration, {
+              "iteration.completed": completionAttempted,
+              "iteration.message_count": currentMessages.length
+            });
             if (currentMessages.length > MAX_HISTORY_MESSAGES) {
               const messagesBefore = currentMessages.length;
               const systemMsg = currentMessages[0];
@@ -102836,6 +103059,183 @@ var init_simpleTelemetry = __esm({
           console.log("[Attributes]", attributes);
         }
       }
+      /**
+       * Hash content for deduplication/comparison purposes
+       * @param {string} content - The content to hash
+       * @returns {string} - Hex string hash
+       */
+      hashContent(content) {
+        let hash = 0;
+        const len = Math.min(content.length, 1e3);
+        for (let i5 = 0; i5 < len; i5++) {
+          hash = (hash << 5) - hash + content.charCodeAt(i5);
+          hash |= 0;
+        }
+        return hash.toString(16);
+      }
+      /**
+       * Record a conversation turn (assistant response or tool result)
+       * @param {string} role - The role (assistant, tool_result)
+       * @param {string} content - The turn content
+       * @param {Object} metadata - Additional metadata
+       */
+      recordConversationTurn(role, content, metadata = {}) {
+        if (!this.isEnabled()) return;
+        this.addEvent(`conversation.turn.${role}`, {
+          "session.id": this.sessionId,
+          "conversation.role": role,
+          "conversation.content": content.substring(0, 1e4),
+          "conversation.content.length": content.length,
+          "conversation.content.hash": this.hashContent(content),
+          ...metadata
+        });
+      }
+      /**
+       * Record error events with classification
+       * @param {string} errorType - The type of error (wrapped_tool, unrecognized_tool, no_tool_call, circuit_breaker, etc.)
+       * @param {Object} errorDetails - Error details including message, stack, context
+       */
+      recordErrorEvent(errorType, errorDetails = {}) {
+        if (!this.isEnabled()) return;
+        this.addEvent(`error.${errorType}`, {
+          "session.id": this.sessionId,
+          "error.type": errorType,
+          "error.message": errorDetails.message?.substring(0, 1e3) || null,
+          "error.stack": errorDetails.stack?.substring(0, 2e3) || null,
+          "error.recoverable": errorDetails.recoverable ?? true,
+          "error.context": JSON.stringify(errorDetails.context || {}).substring(0, 1e3),
+          ...Object.fromEntries(
+            Object.entries(errorDetails).filter(([k5]) => !["message", "stack", "context", "recoverable"].includes(k5)).map(([k5, v5]) => [`error.${k5}`, v5])
+          )
+        });
+      }
+      /**
+       * Record AI thinking/reasoning content
+       * @param {string} thinkingContent - The thinking content from AI response
+       * @param {Object} metadata - Additional metadata
+       */
+      recordThinkingContent(thinkingContent, metadata = {}) {
+        if (!this.isEnabled() || !thinkingContent) return;
+        this.addEvent("ai.thinking", {
+          "session.id": this.sessionId,
+          "ai.thinking.content": thinkingContent.substring(0, 5e4),
+          "ai.thinking.length": thinkingContent.length,
+          "ai.thinking.hash": this.hashContent(thinkingContent),
+          ...metadata
+        });
+      }
+      /**
+       * Record AI tool call decision
+       * @param {string} toolName - The tool name AI decided to call
+       * @param {Object} params - The parameters AI provided
+       * @param {Object} metadata - Additional metadata
+       */
+      recordToolDecision(toolName, params, metadata = {}) {
+        if (!this.isEnabled()) return;
+        this.addEvent("ai.tool_decision", {
+          "session.id": this.sessionId,
+          "ai.tool_decision.name": toolName,
+          "ai.tool_decision.params": JSON.stringify(params || {}).substring(0, 2e3),
+          ...metadata
+        });
+      }
+      /**
+       * Record tool result after execution
+       * @param {string} toolName - The tool that was executed
+       * @param {string|Object} result - The tool result
+       * @param {boolean} success - Whether the tool succeeded
+       * @param {number} durationMs - Execution duration in milliseconds
+       * @param {Object} metadata - Additional metadata
+       */
+      recordToolResult(toolName, result, success, durationMs, metadata = {}) {
+        if (!this.isEnabled()) return;
+        const resultStr = typeof result === "string" ? result : JSON.stringify(result);
+        this.addEvent("tool.result", {
+          "session.id": this.sessionId,
+          "tool.name": toolName,
+          "tool.result": resultStr.substring(0, 1e4),
+          "tool.result.length": resultStr.length,
+          "tool.result.hash": this.hashContent(resultStr),
+          "tool.duration_ms": durationMs,
+          "tool.success": success,
+          ...metadata
+        });
+      }
+      /**
+       * Record MCP tool execution start
+       * @param {string} toolName - MCP tool name
+       * @param {string} serverName - MCP server name
+       * @param {Object} params - Tool parameters
+       * @param {Object} metadata - Additional metadata
+       */
+      recordMcpToolStart(toolName, serverName, params, metadata = {}) {
+        if (!this.isEnabled()) return;
+        this.addEvent("mcp.tool.start", {
+          "session.id": this.sessionId,
+          "mcp.tool.name": toolName,
+          "mcp.tool.server": serverName || "unknown",
+          "mcp.tool.params": JSON.stringify(params || {}).substring(0, 2e3),
+          ...metadata
+        });
+      }
+      /**
+       * Record MCP tool execution end
+       * @param {string} toolName - MCP tool name
+       * @param {string} serverName - MCP server name
+       * @param {string|Object} result - Tool result
+       * @param {boolean} success - Whether succeeded
+       * @param {number} durationMs - Execution duration
+       * @param {string} errorMessage - Error message if failed
+       * @param {Object} metadata - Additional metadata
+       */
+      recordMcpToolEnd(toolName, serverName, result, success, durationMs, errorMessage = null, metadata = {}) {
+        if (!this.isEnabled()) return;
+        const resultStr = typeof result === "string" ? result : JSON.stringify(result || "");
+        this.addEvent("mcp.tool.end", {
+          "session.id": this.sessionId,
+          "mcp.tool.name": toolName,
+          "mcp.tool.server": serverName || "unknown",
+          "mcp.tool.result": resultStr.substring(0, 1e4),
+          "mcp.tool.result.length": resultStr.length,
+          "mcp.tool.duration_ms": durationMs,
+          "mcp.tool.success": success,
+          "mcp.tool.error": errorMessage,
+          ...metadata
+        });
+      }
+      /**
+       * Record iteration lifecycle event
+       * @param {string} eventType - start or end
+       * @param {number} iteration - Iteration number
+       * @param {Object} data - Additional data
+       */
+      recordIterationEvent(eventType, iteration, data2 = {}) {
+        if (!this.isEnabled()) return;
+        this.addEvent(`iteration.${eventType}`, {
+          "session.id": this.sessionId,
+          "iteration": iteration,
+          ...data2
+        });
+      }
+      /**
+       * Record per-turn token breakdown
+       * @param {number} iteration - Iteration number
+       * @param {Object} tokenData - Token metrics
+       */
+      recordTokenTurn(iteration, tokenData = {}) {
+        if (!this.isEnabled()) return;
+        this.addEvent("tokens.turn", {
+          "session.id": this.sessionId,
+          "iteration": iteration,
+          "tokens.input": tokenData.inputTokens || 0,
+          "tokens.output": tokenData.outputTokens || 0,
+          "tokens.total": (tokenData.inputTokens || 0) + (tokenData.outputTokens || 0),
+          "tokens.cache_read": tokenData.cacheReadTokens || 0,
+          "tokens.cache_write": tokenData.cacheWriteTokens || 0,
+          "tokens.context_used": tokenData.contextTokens || 0,
+          "tokens.context_remaining": tokenData.maxContextTokens ? tokenData.maxContextTokens - (tokenData.contextTokens || 0) : null
+        });
+      }
       async withSpan(spanName, fn, attributes = {}) {
         if (!this.isEnabled()) {
           return fn();

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@probelabs/probe",
-  "version": "0.6.0-rc224",
+  "version": "0.6.0-rc225",
   "description": "Node.js wrapper for the probe code search tool",
   "main": "src/index.js",
   "module": "src/index.js",