Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 38 additions & 4 deletions src/app/v1/_lib/proxy/response-handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1560,6 +1560,10 @@ export function parseUsageFromResponseText(
let messageStartUsage: UsageMetrics | null = null;
let messageDeltaUsage: UsageMetrics | null = null;

// Gemini SSE: usageMetadata 需要 last-wins(完整 token 计数仅在最后事件中)
let lastGeminiUsage: UsageMetrics | null = null;
let lastGeminiUsageRecord: Record<string, unknown> | null = null;

const mergeUsageMetrics = (base: UsageMetrics | null, patch: UsageMetrics): UsageMetrics => {
if (!base) {
return { ...patch };
Expand Down Expand Up @@ -1633,18 +1637,37 @@ export function parseUsageFromResponseText(
}

// 非 Claude 格式的 SSE 处理(Gemini 等)
// 注意:Gemini SSE 流中,usageMetadata 在每个事件中都可能存在,
// 但只有最后一个事件包含完整的 token 计数(candidatesTokenCount、thoughtsTokenCount 等)
// 因此需要持续更新,使用最后一个有效值
if (!messageStartUsage && !messageDeltaUsage) {
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

condition !messageStartUsage && !messageDeltaUsage may prevent Gemini last-wins logic from running if a provider returns both Claude-style events AND Gemini-style usageMetadata in the same response

if a malformed response contains both message_start events and usageMetadata fields, the Gemini SSE handling (lines 1647-1671) will be skipped, causing Gemini usage data to be ignored

Prompt To Fix With AI
This is a comment left during a code review.
Path: src/app/v1/_lib/proxy/response-handler.ts
Line: 1643:1643

Comment:
condition `!messageStartUsage && !messageDeltaUsage` may prevent Gemini last-wins logic from running if a provider returns both Claude-style events AND Gemini-style `usageMetadata` in the same response

if a malformed response contains both `message_start` events and `usageMetadata` fields, the Gemini SSE handling (lines 1647-1671) will be skipped, causing Gemini usage data to be ignored

How can I resolve this? If you propose a fix, please make it concise.

// Standard usage fields (data.usage)
// Standard usage fields (data.usage) - 仍使用 first-wins 策略
applyUsageValue(data.usage, `sse.${event.event}.usage`);

// Gemini usageMetadata
applyUsageValue(data.usageMetadata, `sse.${event.event}.usageMetadata`);
// Gemini usageMetadata - 改为 last-wins 策略
// 跳过 applyUsageValue(它是 first-wins),直接更新
if (data.usageMetadata && typeof data.usageMetadata === "object") {
const extracted = extractUsageMetrics(data.usageMetadata);
if (extracted) {
// 持续更新,最后一个有效值会覆盖之前的
lastGeminiUsage = extracted;
lastGeminiUsageRecord = data.usageMetadata as Record<string, unknown>;
}
}

// Handle response wrapping in SSE
if (!usageMetrics && data.response && typeof data.response === "object") {
const responseObj = data.response as Record<string, unknown>;
applyUsageValue(responseObj.usage, `sse.${event.event}.response.usage`);
applyUsageValue(responseObj.usageMetadata, `sse.${event.event}.response.usageMetadata`);

// response.usageMetadata 也使用 last-wins 策略
if (responseObj.usageMetadata && typeof responseObj.usageMetadata === "object") {
const extracted = extractUsageMetrics(responseObj.usageMetadata);
if (extracted) {
lastGeminiUsage = extracted;
lastGeminiUsageRecord = responseObj.usageMetadata as Record<string, unknown>;
}
}
}
Comment on lines +1647 to 1671
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic for extracting and updating lastGeminiUsage from usageMetadata is duplicated for data.usageMetadata and responseObj.usageMetadata. This can be refactored to improve code clarity and maintainability.

Specifically, the following pattern is repeated:

if (metadataSource && typeof metadataSource === 'object') {
  const extracted = extractUsageMetrics(metadataSource);
  if (extracted) {
    lastGeminiUsage = extracted;
    lastGeminiUsageRecord = metadataSource as Record<string, unknown>;
  }
}

Consider extracting this into a helper function within parseUsageFromResponseText to avoid repetition. For example:

const applyLastGeminiUsage = (value: unknown) => {
  if (value && typeof value === 'object') {
    const extracted = extractUsageMetrics(value);
    if (extracted) {
      lastGeminiUsage = extracted;
      lastGeminiUsageRecord = value as Record<string, unknown>;
    }
  }
};

This would make the code in the loop more concise and easier to maintain.

}
}
Expand All @@ -1665,6 +1688,17 @@ export function parseUsageFromResponseText(
usage: usageMetrics,
});
}

// Gemini SSE 处理:使用最后一个有效的 usageMetadata
// 仅当 Claude SSE 没有提供 usage 且 applyUsageValue 也没有找到时才使用
if (!usageMetrics && lastGeminiUsage) {
usageMetrics = adjustUsageForProviderType(lastGeminiUsage, providerType);
usageRecord = lastGeminiUsageRecord;
logger.debug("[ResponseHandler] Final usage from Gemini SSE (last event)", {
providerType,
usage: usageMetrics,
});
}
}

return { usageRecord, usageMetrics };
Expand Down
Loading