feat: implement usage tracking for AI requests

Adds local token usage tracking for all AI providers. Usage data is captured during stream processing and stored in a local database. Includes a new Usage tab in the Providers dashboard to visualize historical token consumption. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
2026-01-06 16:44:14 +02:00
parent 5645d0a0fb
commit 9c3d6f4ad8
12 changed files with 7460 additions and 81 deletions
--- a/open-sse/handlers/chatCore.js
+++ b/open-sse/handlers/chatCore.js
@@ -8,6 +8,46 @@ import { createRequestLogger } from "../utils/requestLogger.js";
 import { getModelTargetFormat, PROVIDER_ID_TO_ALIAS } from "../config/providerModels.js";
 import { createErrorResult, parseUpstreamError, formatProviderError } from "../utils/error.js";
 import { handleBypassRequest } from "../utils/bypassHandler.js";
+import { saveRequestUsage } from "@/lib/usageDb.js";
+
+/**
+ * Extract usage from non-streaming response body
+ * Handles different provider response formats
+ */
+function extractUsageFromResponse(responseBody, provider) {
+  if (!responseBody) return null;
+
+  // OpenAI format
+  if (responseBody.usage) {
+    return {
+      prompt_tokens: responseBody.usage.prompt_tokens || 0,
+      completion_tokens: responseBody.usage.completion_tokens || 0,
+      cached_tokens: responseBody.usage.prompt_tokens_details?.cached_tokens,
+      reasoning_tokens: responseBody.usage.completion_tokens_details?.reasoning_tokens
+    };
+  }
+
+  // Claude format
+  if (responseBody.usage?.input_tokens !== undefined || responseBody.usage?.output_tokens !== undefined) {
+    return {
+      prompt_tokens: responseBody.usage.input_tokens || 0,
+      completion_tokens: responseBody.usage.output_tokens || 0,
+      cache_read_input_tokens: responseBody.usage.cache_read_input_tokens,
+      cache_creation_input_tokens: responseBody.usage.cache_creation_input_tokens
+    };
+  }
+
+  // Gemini format
+  if (responseBody.usageMetadata) {
+    return {
+      prompt_tokens: responseBody.usageMetadata.promptTokenCount || 0,
+      completion_tokens: responseBody.usageMetadata.candidatesTokenCount || 0,
+      reasoning_tokens: responseBody.usageMetadata.thoughtsTokenCount
+    };
+  }
+
+  return null;
+}

 /**
 * Core chat handler - shared between SSE and Worker
@@ -20,8 +60,9 @@ import { handleBypassRequest } from "../utils/bypassHandler.js";
 * @param {function} options.onCredentialsRefreshed - Callback when credentials are refreshed
 * @param {function} options.onRequestSuccess - Callback when request succeeds (to clear error status)
 * @param {function} options.onDisconnect - Callback when client disconnects
+ * @param {string} options.connectionId - Connection ID for usage tracking
 */
-export async function handleChatCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect, clientRawRequest }) {
+export async function handleChatCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect, clientRawRequest, connectionId }) {
  const { provider, model } = modelInfo;

  const sourceFormat = detectFormat(body);
@@ -220,12 +261,29 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred
  // Non-streaming response
  if (!stream) {
    const responseBody = await providerResponse.json();
-    
+
    // Notify success - caller can clear error status if needed
    if (onRequestSuccess) {
      await onRequestSuccess();
    }
-    
+
+    // Log usage for non-streaming responses
+    const usage = extractUsageFromResponse(responseBody, provider);
+    if (usage) {
+      const msg = `[${new Date().toLocaleTimeString("en-US", { hour12: false, hour: "2-digit", minute: "2-digit" })}] 📊 [USAGE] ${provider.toUpperCase()} | in=${usage.prompt_tokens || 0} | out=${usage.completion_tokens || 0}${connectionId ? ` | account=${connectionId.slice(0, 8)}...` : ""}`;
+      console.log(`${COLORS.green}${msg}${COLORS.reset}`);
+
+      saveRequestUsage({
+        provider: provider || "unknown",
+        model: model || "unknown",
+        tokens: usage,
+        timestamp: new Date().toISOString(),
+        connectionId: connectionId || undefined
+      }).catch(err => {
+        console.error("Failed to save usage stats:", err.message);
+      });
+    }
+
    return {
      success: true,
      response: new Response(JSON.stringify(responseBody), {
@@ -254,9 +312,9 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred
  // Create transform stream with logger for streaming response
  let transformStream;
  if (needsTranslation(targetFormat, sourceFormat)) {
-    transformStream = createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider, reqLogger, toolNameMap);
+    transformStream = createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider, reqLogger, toolNameMap, model, connectionId);
  } else {
-    transformStream = createPassthroughStreamWithLogger(provider, reqLogger);
+    transformStream = createPassthroughStreamWithLogger(provider, reqLogger, model, connectionId);
  }

  // Pipe response through transform with disconnect detection
--- a/open-sse/handlers/responsesHandler.js
+++ b/open-sse/handlers/responsesHandler.js
@@ -17,9 +17,10 @@ import { createResponsesApiTransformStream } from "../transformer/responsesTrans
 * @param {function} options.onCredentialsRefreshed - Callback when credentials are refreshed
 * @param {function} options.onRequestSuccess - Callback when request succeeds
 * @param {function} options.onDisconnect - Callback when client disconnects
+ * @param {string} options.connectionId - Connection ID for usage tracking
 * @returns {Promise<{success: boolean, response?: Response, status?: number, error?: string}>}
 */
-export async function handleResponsesCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect }) {
+export async function handleResponsesCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect, connectionId }) {
  // Convert Responses API format to Chat Completions format
  const convertedBody = convertResponsesApiFormat(body);
  
@@ -34,7 +35,8 @@ export async function handleResponsesCore({ body, modelInfo, credentials, log, o
    log,
    onCredentialsRefreshed,
    onRequestSuccess,
-    onDisconnect
+    onDisconnect,
+    connectionId
  });

  if (!result.success || !result.response) {
--- a/open-sse/utils/stream.js
+++ b/open-sse/utils/stream.js
@@ -1,5 +1,6 @@
 import { translateResponse, initState } from "../translator/index.js";
 import { FORMATS } from "../translator/formats.js";
+import { saveRequestUsage } from "@/lib/usageDb.js";

 // Get HH:MM:SS timestamp
 function getTimeString() {
@@ -48,27 +49,39 @@ export const COLORS = {
 };

 // Log usage with cache info (green color)
-function logUsage(provider, usage) {
+function logUsage(provider, usage, model = null, connectionId = null) {
  if (!usage) return;
-  
+
  const p = provider?.toUpperCase() || "UNKNOWN";
  const inTokens = usage.prompt_tokens || 0;
  const outTokens = usage.completion_tokens || 0;
-  
+
  let msg = `[${getTimeString()}] 📊 [USAGE] ${p} | in=${inTokens} | out=${outTokens}`;
-  
+  if (connectionId) msg += ` | account=${connectionId.slice(0, 8)}...`;
+
  if (usage.cache_creation_input_tokens) msg += ` | cache_write=${usage.cache_creation_input_tokens}`;
  if (usage.cache_read_input_tokens) msg += ` | cache_read=${usage.cache_read_input_tokens}`;
  if (usage.cached_tokens) msg += ` | cached=${usage.cached_tokens}`;
  if (usage.reasoning_tokens) msg += ` | reasoning=${usage.reasoning_tokens}`;
-  
+
  console.log(`${COLORS.green}${msg}${COLORS.reset}`);
+
+  // Save to DB
+  saveRequestUsage({
+    provider: provider || "unknown",
+    model: model || "unknown",
+    tokens: usage,
+    timestamp: new Date().toISOString(),
+    connectionId: connectionId || undefined
+  }).catch(err => {
+    console.error("Failed to save usage stats:", err.message);
+  });
 }

 // Parse SSE data line
 function parseSSELine(line) {
  if (!line || !line.startsWith("data:")) return null;
-  
+
  const data = line.slice(5).trim();
  if (data === "[DONE]") return { done: true };

@@ -91,17 +104,17 @@ function parseSSELine(line) {
 */
 export function formatSSE(data, sourceFormat) {
  if (data.done) return "data: [DONE]\n\n";
-  
+
  // OpenAI Responses API format: has event field
  if (data.event && data.data) {
    return `event: ${data.event}\ndata: ${JSON.stringify(data.data)}\n\n`;
  }
-  
+
  // Claude format: include event prefix
  if (sourceFormat === FORMATS.CLAUDE && data.type) {
    return `event: ${data.type}\ndata: ${JSON.stringify(data)}\n\n`;
  }
-  
+
  return `data: ${JSON.stringify(data)}\n\n`;
 }

@@ -121,22 +134,26 @@ const STREAM_MODE = {
 * @param {string} options.sourceFormat - Client format (for translate mode)
 * @param {string} options.provider - Provider name
 * @param {object} options.reqLogger - Request logger instance
+ * @param {string} options.model - Model name
+ * @param {string} options.connectionId - Connection ID for usage tracking
 */
 export function createSSEStream(options = {}) {
-  const { 
-    mode = STREAM_MODE.TRANSLATE, 
-    targetFormat, 
-    sourceFormat, 
-    provider = null, 
+  const {
+    mode = STREAM_MODE.TRANSLATE,
+    targetFormat,
+    sourceFormat,
+    provider = null,
    reqLogger = null,
-    toolNameMap = null
+    toolNameMap = null,
+    model = null,
+    connectionId = null
  } = options;

  const decoder = new TextDecoder();
  const encoder = new TextEncoder();
  let buffer = "";
  let usage = null;
-  
+
  // State for translate mode
  const state = mode === STREAM_MODE.TRANSLATE ? { ...initState(sourceFormat), provider, toolNameMap } : null;

@@ -151,7 +168,7 @@ export function createSSEStream(options = {}) {

      for (const line of lines) {
        const trimmed = line.trim();
-        
+
        // Passthrough mode: normalize and forward
        if (mode === STREAM_MODE.PASSTHROUGH) {
          if (trimmed.startsWith("data:") && trimmed.slice(5).trim() !== "[DONE]") {
@@ -216,7 +233,7 @@ export function createSSEStream(options = {}) {
            reqLogger?.appendConvertedChunk?.(output);
            controller.enqueue(encoder.encode(output));
          }
-          if (usage) logUsage(provider, usage);
+          if (usage) logUsage(provider, usage, model, connectionId);
          return;
        }

@@ -250,7 +267,7 @@ export function createSSEStream(options = {}) {
        reqLogger?.appendConvertedChunk?.(doneOutput);
        controller.enqueue(encoder.encode(doneOutput));

-        if (state?.usage) logUsage(state.provider || targetFormat, state.usage);
+        if (state?.usage) logUsage(state.provider || targetFormat, state.usage, model, connectionId);
      } catch (error) {
        console.log("Error in flush:", error);
      }
@@ -259,22 +276,25 @@ export function createSSEStream(options = {}) {
 }

 // Convenience functions for backward compatibility
-export function createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider = null, reqLogger = null, toolNameMap = null) {
-  return createSSEStream({ 
-    mode: STREAM_MODE.TRANSLATE, 
-    targetFormat, 
-    sourceFormat, 
-    provider, 
+export function createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider = null, reqLogger = null, toolNameMap = null, model = null, connectionId = null) {
+  return createSSEStream({
+    mode: STREAM_MODE.TRANSLATE,
+    targetFormat,
+    sourceFormat,
+    provider,
    reqLogger,
-    toolNameMap
+    toolNameMap,
+    model,
+    connectionId
  });
 }

-export function createPassthroughStreamWithLogger(provider = null, reqLogger = null) {
-  return createSSEStream({ 
-    mode: STREAM_MODE.PASSTHROUGH, 
-    provider, 
-    reqLogger 
+export function createPassthroughStreamWithLogger(provider = null, reqLogger = null, model = null, connectionId = null) {
+  return createSSEStream({
+    mode: STREAM_MODE.PASSTHROUGH,
+    provider,
+    reqLogger,
+    model,
+    connectionId
  });
 }
-