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 outputTokens = 0;
// Split by "data:" lines (SSE format)
// Parse every "data:" line for usage info (provider-agnostic approach)
const lines = text.split("\n");
for (const line of lines) {
const trimmed = line.trim();
if (trimmed.startsWith("event: message_start")) {
// Find the next data line
const dataIndex = lines.indexOf(line) + 1;
if (dataIndex < lines.length && lines[dataIndex].trim().startsWith("data:")) {
try {
const dataStr = lines[dataIndex].trim().slice(5).trim();
const data = JSON.parse(dataStr);
if (data.message?.usage?.input_tokens) {
inputTokens = data.message.usage.input_tokens;
}
} catch {}
for (let i = 0; i < lines.length; i++) {
const trimmed = lines[i].trim();
if (!trimmed.startsWith("data:")) continue;
const dataStr = trimmed.slice(5).trim();
if (!dataStr || dataStr === "[DONE]") continue;
try {
const data = JSON.parse(dataStr);
// Anthropic message_start: { type: "message_start", message: { usage: { input_tokens } } }
if (data.type === "message_start" && data.message?.usage?.input_tokens) {
inputTokens = data.message.usage.input_tokens;
}
} else if (trimmed.startsWith("event: message_delta")) {
const dataIndex = lines.indexOf(line) + 1;
if (dataIndex < lines.length && lines[dataIndex].trim().startsWith("data:")) {
try {
const dataStr = lines[dataIndex].trim().slice(5).trim();
const data = JSON.parse(dataStr);
if (data.usage?.output_tokens) {
outputTokens = data.usage.output_tokens;
}
} catch {}
// Anthropic message_delta: { type: "message_delta", usage: { output_tokens } }
if (data.type === "message_delta" && data.usage?.output_tokens) {
outputTokens = data.usage.output_tokens;
}
// 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
const { inputTokens, outputTokens } = extractUsageFromSSE(fullText, route.provider);
if (inputTokens > 0 || outputTokens > 0) {
recordUsage(requestedModel, route.provider, actualModel, inputTokens, outputTokens);
// Always record usage — even if token extraction fails, we count the request
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 };
} catch (streamErr) {
@@ -576,20 +590,22 @@ async function handleMessages(req, res) {
const responseBody = await upstreamRes.text();
const responseHeaders = { "content-type": upstreamRes.headers.get("content-type") || "application/json" };
// Try to extract usage from non-streaming response
if (statusCode === 200 && responseHeaders["content-type"]?.includes("application/json")) {
try {
const json = JSON.parse(responseBody);
if (json.usage?.input_tokens || json.usage?.output_tokens) {
recordUsage(
requestedModel,
route.provider,
actualModel,
json.usage.input_tokens,
json.usage.output_tokens
);
}
} catch {}
// Record usage from non-streaming response
if (statusCode === 200) {
let extractedInput = 0;
let extractedOutput = 0;
if (responseHeaders["content-type"]?.includes("application/json")) {
try {
const json = JSON.parse(responseBody);
extractedInput = json.usage?.input_tokens || 0;
extractedOutput = json.usage?.output_tokens || 0;
} catch {}
}
// Always record — even if token extraction fails, we count the request
recordUsage(requestedModel, route.provider, actualModel, extractedInput, extractedOutput);
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