diff --git a/proxy.mjs b/proxy.mjs index df0642c..20626aa 100644 --- a/proxy.mjs +++ b/proxy.mjs @@ -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