harrison/ein-ai-proxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: track usage for all providers, not just when tokens extracted

Browse Source 
- Rewrote extractUsageFromSSE() to parse data: lines directly instead of
  buggy indexOf-based event/data pairing that silently failed for Z.ai
  and DeepSeek SSE formats
- Always call recordUsage() on successful responses even when token
  extraction yields 0 — ensures request counts are tracked
- Added console.warn when token extraction fails for debugging
This commit is contained in:
William Stuckey
2026-03-27 11:38:59 -05:00
parent c01e1a6e87
commit f31a89f10c
1 changed files with 56 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
proxy.mjs
+56 -40
View File
@@ -455,33 +455,45 @@ function extractUsageFromSSE(text, provider) {
let inputTokens = 0; let inputTokens = 0;
let outputTokens = 0; let outputTokens = 0;
// Split by "data:" lines (SSE format) // Parse every "data:" line for usage info (provider-agnostic approach)
const lines = text.split("\n"); const lines = text.split("\n");
for (const line of lines) { for (let i = 0; i < lines.length; i++) {
const trimmed = line.trim(); const trimmed = lines[i].trim();
if (trimmed.startsWith("event: message_start")) { if (!trimmed.startsWith("data:")) continue;
// Find the next data line
const dataIndex = lines.indexOf(line) + 1; const dataStr = trimmed.slice(5).trim();
if (dataIndex < lines.length && lines[dataIndex].trim().startsWith("data:")) { if (!dataStr || dataStr === "[DONE]") continue;
try {
const dataStr = lines[dataIndex].trim().slice(5).trim(); try {
const data = JSON.parse(dataStr); const data = JSON.parse(dataStr);
if (data.message?.usage?.input_tokens) {
inputTokens = data.message.usage.input_tokens; // Anthropic message_start: { type: "message_start", message: { usage: { input_tokens } } }
} if (data.type === "message_start" && data.message?.usage?.input_tokens) {
} catch {} inputTokens = data.message.usage.input_tokens;
} }
} else if (trimmed.startsWith("event: message_delta")) {
const dataIndex = lines.indexOf(line) + 1; // Anthropic message_delta: { type: "message_delta", usage: { output_tokens } }
if (dataIndex < lines.length && lines[dataIndex].trim().startsWith("data:")) { if (data.type === "message_delta" && data.usage?.output_tokens) {
try { outputTokens = data.usage.output_tokens;
const dataStr = lines[dataIndex].trim().slice(5).trim();
const data = JSON.parse(dataStr);
if (data.usage?.output_tokens) {
outputTokens = data.usage.output_tokens;
}
} catch {}
} }
// Some providers put usage at the top level of non-typed events
if (data.usage) {
if (data.usage.input_tokens && data.type !== "message_start") {
inputTokens = inputTokens || data.usage.input_tokens;
}
if (data.usage.output_tokens && data.type !== "message_delta") {
outputTokens = outputTokens || data.usage.output_tokens;
}
}
// Some providers put usage in a final message object
if (data.message?.usage && data.type !== "message_start") {
inputTokens = inputTokens || data.message.usage.input_tokens || 0;
outputTokens = outputTokens || data.message.usage.output_tokens || 0;
}
} catch {
// Not valid JSON — skip (e.g. SSE comment lines)
} }
} }
@@ -561,8 +573,10 @@ async function handleMessages(req, res) {
// Extract usage from the accumulated SSE text // Extract usage from the accumulated SSE text
const { inputTokens, outputTokens } = extractUsageFromSSE(fullText, route.provider); const { inputTokens, outputTokens } = extractUsageFromSSE(fullText, route.provider);
if (inputTokens > 0 || outputTokens > 0) { // Always record usage — even if token extraction fails, we count the request
recordUsage(requestedModel, route.provider, actualModel, inputTokens, outputTokens); recordUsage(requestedModel, route.provider, actualModel, inputTokens, outputTokens);
if (inputTokens === 0 && outputTokens === 0) {
console.warn(`[usage] No tokens extracted from ${route.provider} SSE stream for ${requestedModel} (stream length: ${fullText.length} chars)`);
} }
providerStatus[route.provider] = { status: "ok", lastSeen: new Date().toISOString(), lastError: null }; providerStatus[route.provider] = { status: "ok", lastSeen: new Date().toISOString(), lastError: null };
} catch (streamErr) { } catch (streamErr) {
@@ -576,20 +590,22 @@ async function handleMessages(req, res) {
const responseBody = await upstreamRes.text(); const responseBody = await upstreamRes.text();
const responseHeaders = { "content-type": upstreamRes.headers.get("content-type") || "application/json" }; const responseHeaders = { "content-type": upstreamRes.headers.get("content-type") || "application/json" };
// Try to extract usage from non-streaming response // Record usage from non-streaming response
if (statusCode === 200 && responseHeaders["content-type"]?.includes("application/json")) { if (statusCode === 200) {
try { let extractedInput = 0;
const json = JSON.parse(responseBody); let extractedOutput = 0;
if (json.usage?.input_tokens || json.usage?.output_tokens) { if (responseHeaders["content-type"]?.includes("application/json")) {
recordUsage( try {
requestedModel, const json = JSON.parse(responseBody);
route.provider, extractedInput = json.usage?.input_tokens || 0;
actualModel, extractedOutput = json.usage?.output_tokens || 0;
json.usage.input_tokens, } catch {}
json.usage.output_tokens }
); // Always record — even if token extraction fails, we count the request
} recordUsage(requestedModel, route.provider, actualModel, extractedInput, extractedOutput);
} catch {} if (extractedInput === 0 && extractedOutput === 0) {
console.warn(`[usage] No tokens extracted from ${route.provider} non-streaming response for ${requestedModel}`);
}
} }
// Update provider status based on HTTP status // Update provider status based on HTTP status