fix(task): record token usage in async billing logs

[ElevenZhou]

Summary

• add prompt/completion token fields to async task billing log params
• persist async task billing token usage into logs.prompt_tokens and logs.completion_tokens
• include token usage in quota export for async task billing logs

Why

Async task settlement can recalculate quota from upstream total_tokens, but the follow-up consume/refund log only stored the token count in the human-readable content string. As a result, the admin log token column and token-based exports showed zero tokens for these settlement logs even though billing used the token count correctly.

Test

• Not run: local environment does not have the Go toolchain installed.

Summary by CodeRabbit

• New Features

  • Billing and quota records now include separate prompt and completion token counts for improved usage tracking.
  • Token-based quota recalculation now uses the recorded token breakdown for more accurate accounting.

• Bug Fixes

  • Fixed quota settlement so token usage is captured consistently when billing logs are created.
  • Improved how usage data is stored for downstream reporting and export flows.

[coderabbitai]

Walkthrough

Adds PromptTokens and CompletionTokens fields to RecordTaskBillingLogParams, persists them on Log records, and computes TokenUsed from their sum. Introduces RecalculateTaskQuotaWithUsage to carry token counts through settlement; RecalculateTaskQuota and RecalculateTaskQuotaByTokens now delegate to it.

Changes

Task Billing Token Usage Tracking

Layer / File(s)	Summary
Log data model and persistence model/log.go	RecordTaskBillingLogParams gains PromptTokens/CompletionTokens; RecordTaskBillingLog persists these into Log and computes LogQuotaData.TokenUsed as their sum.
Quota recalculation propagation service/task_billing.go	New RecalculateTaskQuotaWithUsage helper accepts token counts and records them via RecordTaskBillingLog; RecalculateTaskQuota delegates to it, and RecalculateTaskQuotaByTokens calls it with 0 prompt tokens and totalTokens completion tokens.

Estimated code review effort: 2 (Simple) | ~10 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant RecalculateTaskQuota
  participant RecalculateTaskQuotaWithUsage
  participant RecordTaskBillingLog
  participant LogStore

  Caller->>RecalculateTaskQuota: request settlement
  RecalculateTaskQuota->>RecalculateTaskQuotaWithUsage: delegate (promptTokens, completionTokens)
  RecalculateTaskQuotaWithUsage->>RecordTaskBillingLog: record(params incl. token counts)
  RecordTaskBillingLog->>LogStore: persist Log with PromptTokens, CompletionTokens
  RecordTaskBillingLog->>LogStore: compute TokenUsed = PromptTokens + CompletionTokens

Loading

Possibly related PRs

• QuantumNous/new-api#5684: Both PRs modify the same task billing logging pipeline (RecordTaskBillingLogParams and the settlement path in model/log.go and service/task_billing.go), one adding token fields and the other adding node-attribution fields.

Poem

A rabbit counts each token twice,
Prompt and completion, oh so precise! 🐇
Through quota's path they now do flow,
Logged and summed for all to know.
Hop, hop, hooray—the ledger's tight! 📊

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: recording token usage in async task billing logs.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

service/task_billing.go (1)

192-253: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add test coverage for the new token-usage propagation.

The referenced test file doesn't appear to assert PromptTokens/CompletionTokens values on the resulting Log record for RecalculateTaskQuotaWithUsage/RecalculateTaskQuotaByTokens. Since the PR's core motivation is fixing zero-token display for these settlement logs, a test verifying log.PromptTokens/log.CompletionTokens are populated as expected would guard against regressions.

Also applies to: 258-308

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@service/task_billing.go` around lines 192 - 253, Add test coverage around
RecalculateTaskQuotaWithUsage and RecalculateTaskQuotaByTokens to assert token
usage is propagated into the billing log. In the relevant test(s), inspect the
Log record created by RecordTaskBillingLog and verify PromptTokens and
CompletionTokens are populated with the expected values rather than left at
zero, using the existing RecalculateTaskQuotaWithUsage helper and any shared
billing-log assertions to locate the record. Ensure the test covers the
settlement path where quota is recalculated so this regression is caught in
future changes.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@service/task_billing.go`:
- Around line 192-253: Add test coverage around RecalculateTaskQuotaWithUsage
and RecalculateTaskQuotaByTokens to assert token usage is propagated into the
billing log. In the relevant test(s), inspect the Log record created by
RecordTaskBillingLog and verify PromptTokens and CompletionTokens are populated
with the expected values rather than left at zero, using the existing
RecalculateTaskQuotaWithUsage helper and any shared billing-log assertions to
locate the record. Ensure the test covers the settlement path where quota is
recalculated so this regression is caught in future changes.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: cf31de21-93bb-4ad4-9407-8389588a1c81

📥 Commits

Reviewing files that changed from the base of the PR and between 52858ad and e134b5b.

📒 Files selected for processing (2)

• model/log.go
• service/task_billing.go