diff --git a/internal/runtime/executor/usage_helpers.go b/internal/runtime/executor/usage_helpers.go index 5669d9bc..78581b82 100644 --- a/internal/runtime/executor/usage_helpers.go +++ b/internal/runtime/executor/usage_helpers.go @@ -275,6 +275,20 @@ func parseClaudeStreamUsage(line []byte) (usage.Detail, bool) { return detail, true } +func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail { + detail := usage.Detail{ + InputTokens: node.Get("promptTokenCount").Int(), + OutputTokens: node.Get("candidatesTokenCount").Int(), + ReasoningTokens: node.Get("thoughtsTokenCount").Int(), + TotalTokens: node.Get("totalTokenCount").Int(), + CachedTokens: node.Get("cachedContentTokenCount").Int(), + } + if detail.TotalTokens == 0 { + detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens + } + return detail +} + func parseGeminiCLIUsage(data []byte) usage.Detail { usageNode := gjson.ParseBytes(data) node := usageNode.Get("response.usageMetadata") @@ -284,16 +298,7 @@ func parseGeminiCLIUsage(data []byte) usage.Detail { if !node.Exists() { return usage.Detail{} } - detail := usage.Detail{ - InputTokens: node.Get("promptTokenCount").Int(), - OutputTokens: node.Get("candidatesTokenCount").Int(), - ReasoningTokens: node.Get("thoughtsTokenCount").Int(), - TotalTokens: node.Get("totalTokenCount").Int(), - } - if detail.TotalTokens == 0 { - detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens - } - return detail + return parseGeminiFamilyUsageDetail(node) } func parseGeminiUsage(data []byte) usage.Detail { @@ -305,16 +310,7 @@ func parseGeminiUsage(data []byte) usage.Detail { if !node.Exists() { return usage.Detail{} } - detail := usage.Detail{ - InputTokens: node.Get("promptTokenCount").Int(), - OutputTokens: node.Get("candidatesTokenCount").Int(), - ReasoningTokens: node.Get("thoughtsTokenCount").Int(), - TotalTokens: node.Get("totalTokenCount").Int(), - } - if detail.TotalTokens == 0 { - detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens - } - return detail + return parseGeminiFamilyUsageDetail(node) } func parseGeminiStreamUsage(line []byte) (usage.Detail, bool) { @@ -329,16 +325,7 @@ func parseGeminiStreamUsage(line []byte) (usage.Detail, bool) { if !node.Exists() { return usage.Detail{}, false } - detail := usage.Detail{ - InputTokens: node.Get("promptTokenCount").Int(), - OutputTokens: node.Get("candidatesTokenCount").Int(), - ReasoningTokens: node.Get("thoughtsTokenCount").Int(), - TotalTokens: node.Get("totalTokenCount").Int(), - } - if detail.TotalTokens == 0 { - detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens - } - return detail, true + return parseGeminiFamilyUsageDetail(node), true } func parseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) { @@ -353,16 +340,7 @@ func parseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) { if !node.Exists() { return usage.Detail{}, false } - detail := usage.Detail{ - InputTokens: node.Get("promptTokenCount").Int(), - OutputTokens: node.Get("candidatesTokenCount").Int(), - ReasoningTokens: node.Get("thoughtsTokenCount").Int(), - TotalTokens: node.Get("totalTokenCount").Int(), - } - if detail.TotalTokens == 0 { - detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens - } - return detail, true + return parseGeminiFamilyUsageDetail(node), true } func parseAntigravityUsage(data []byte) usage.Detail { @@ -377,16 +355,7 @@ func parseAntigravityUsage(data []byte) usage.Detail { if !node.Exists() { return usage.Detail{} } - detail := usage.Detail{ - InputTokens: node.Get("promptTokenCount").Int(), - OutputTokens: node.Get("candidatesTokenCount").Int(), - ReasoningTokens: node.Get("thoughtsTokenCount").Int(), - TotalTokens: node.Get("totalTokenCount").Int(), - } - if detail.TotalTokens == 0 { - detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens - } - return detail + return parseGeminiFamilyUsageDetail(node) } func parseAntigravityStreamUsage(line []byte) (usage.Detail, bool) { @@ -404,16 +373,7 @@ func parseAntigravityStreamUsage(line []byte) (usage.Detail, bool) { if !node.Exists() { return usage.Detail{}, false } - detail := usage.Detail{ - InputTokens: node.Get("promptTokenCount").Int(), - OutputTokens: node.Get("candidatesTokenCount").Int(), - ReasoningTokens: node.Get("thoughtsTokenCount").Int(), - TotalTokens: node.Get("totalTokenCount").Int(), - } - if detail.TotalTokens == 0 { - detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens - } - return detail, true + return parseGeminiFamilyUsageDetail(node), true } var stopChunkWithoutUsage sync.Map