Modify

[KleinerSource]

Summary by CodeRabbit

Release Notes

• New Features
  • Added token balance endpoint to retrieve token status, remaining quota, and usage amount.
  • Introduced streaming passthrough API endpoint for chat requests with usage tracking.
  • Enhanced token status information with human-readable status descriptions (enabled, disabled, expired, exhausted).

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Caution

Review failed

The pull request is closed.

Walkthrough

This PR introduces a new token balance retrieval endpoint and a passthrough relay handler for streaming requests. The balance endpoint authenticates via Bearer token and returns token quota/status information. The passthrough handler manages chat-stream requests with retry logic, quota pre-consumption, and usage extraction from upstream responses.

Changes

Cohort / File(s)	Summary
Token Balance API controller/balance.go, router/api-router.go	Adds GetTokenBalance handler that authenticates via Bearer token, retrieves token data from DB (strips sk- prefix), computes quota/amount based on unlimited flag, maps token status to human-readable text, and returns JSON response. New route group /usage/api with /balance endpoint applies CriticalRateLimit and TokenAuth middleware.
Passthrough Relay Handler controller/relay_passthrough.go, relay/passthrough_handler.go, router/relay-router.go	Introduces RelayPassthrough controller with ChatStreamRequest type for encrypted payload handling. Manages request validation, channel/adaptor selection, model mapping, quota pre-consumption with rollback, retry loop (bounded by limit), and deferred error serialization. PassthroughHelperWithUsage extracts usage from streaming/non-streaming responses, handles Server-Sent Events parsing, and falls back to estimated tokens. New POST /chat-stream route with token auth and model rate limiting middleware.
Supporting Utilities relay/common/relay_info.go, service/token_counter.go	Adds InitFirstResponseTime method to RelayInfo for manual timing initialization. Exposes GetImageTokenForPassthrough wrapper function delegating to internal image token counting with passthrough flag.

Sequence Diagram(s)

sequenceDiagram
    participant Client
    participant RelayPassthrough
    participant Service
    participant ChannelAdaptor
    participant UpstreamAPI
    participant PassthroughHandler

    Client->>RelayPassthrough: POST /chat-stream (encrypted payload)
    RelayPassthrough->>RelayPassthrough: Validate model & data
    RelayPassthrough->>RelayPassthrough: Perform sensitive content checks
    
    loop Retry Loop (bounded)
        RelayPassthrough->>Service: getPassthroughChannel()
        Service-->>RelayPassthrough: Channel (with adaptor)
        RelayPassthrough->>RelayPassthrough: Apply model mapping
        RelayPassthrough->>Service: calculateInputTokens()
        Service-->>RelayPassthrough: Token count
        RelayPassthrough->>Service: PreConsumeQuota()
        alt Quota Available
            Service-->>RelayPassthrough: Success (set defer rollback)
        else Insufficient Quota
            Service-->>RelayPassthrough: Error
            RelayPassthrough->>RelayPassthrough: Return error
        end
        
        RelayPassthrough->>PassthroughHandler: PassthroughHelperWithUsage()
        PassthroughHandler->>ChannelAdaptor: DoRequest(request body)
        ChannelAdaptor->>UpstreamAPI: Forward request
        UpstreamAPI-->>ChannelAdaptor: Response (streaming/non-streaming)
        ChannelAdaptor-->>PassthroughHandler: Response
        
        alt Streaming Response
            PassthroughHandler->>PassthroughHandler: Stream to client while extracting usage
            PassthroughHandler->>PassthroughHandler: Parse Server-Sent Events (data: lines)
            PassthroughHandler->>PassthroughHandler: Extract usage from event
        else Non-Streaming Response
            PassthroughHandler->>PassthroughHandler: Extract usage from response JSON
            PassthroughHandler->>Client: Send response
        end
        
        PassthroughHandler-->>RelayPassthrough: Usage (or estimated tokens on fallback)
        
        alt Success
            RelayPassthrough->>Service: postPassthroughConsumeQuotaWithUsage()
            Service-->>RelayPassthrough: Confirm consumption
            RelayPassthrough-->>Client: Success response
            RelayPassthrough->>RelayPassthrough: Exit retry loop
        else Retry Eligible Error
            RelayPassthrough->>RelayPassthrough: Log retry attempt
            RelayPassthrough->>Service: ReturnPreConsumedQuota() (defer)
        end
    end

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45-60 minutes

Possibly related PRs

• feat: query usage of token #1161 — Adds token retrieval via Bearer token pattern used in GetTokenBalance for the same token-loading logic and router area.
• feat: gemini-2.5-flash-image-preview 文本和图片输出计费 #1677 — Modifies image token counting for passthrough models; the main PR exposes GetImageTokenForPassthrough which this PR directly affects.
• Alpha #1775 — Changes PreConsumeQuota and ReturnPreConsumedQuota signatures that RelayPassthrough and passthrough_handler.go depend on.

Suggested reviewers

• Calcium-Ion
• creamlike1024

Poem

🐰 A token hops through balance lanes,
With streams of data, usage chains,
Retry loops that gently bound,
And quotas counted, safe and sound,
Through passthroughs wide, the data flows! ✨

✨ Finishing touches

• 📝 Generate docstrings

---

📜 Recent review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between b808b96 and 5f23da4.

📒 Files selected for processing (7)

• controller/balance.go
• controller/relay_passthrough.go
• relay/common/relay_info.go
• relay/passthrough_handler.go
• router/api-router.go
• router/relay-router.go
• service/token_counter.go

---

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 and usage tips.}