fix(proxy): estimate ensemble panel + judge sub-call usage on missing-usage backends

[jarvis9443]

Follow-up to #794 (AISIX-Cloud#1074). An ensemble request fans out to one usage event per sub-call — each panel member plus the judge — but none of those events went through the #1074 token-estimation fallback. Against a backend that reports no usage (a relay/transit upstream), each sub-call recorded silent zeros. #794 wired only the streaming judge (via the SSE Drop guard); the non-streaming panel + judge, and the panel members on the streaming path, were still unflagged.

What was missing

sub-call emit site	before
non-streaming panel members (emit_panel_member)	❌ zeros
non-streaming judge (emit_subcalls)	❌ zeros
streaming-path panel members (on_complete)	❌ zeros
streaming judge	✅ #794

The panel members' answer texts and the judge's synthesis request lived only inside ensemble.rs, so the emit sites had nothing to estimate from — the refactor #794 explicitly deferred here.

Change

• PanelOutcome carries est_output_text (the member's answer text) and EnsembleOutcome carries judge_req (the synthesis request), so every emit site has the request+output it needs.
• One shared helper estimate_subcall_tokens(req, model, usage, output_text) applies the feat(usage): estimate tokens locally when the upstream reports no usage #794 per-field or-semantics — a non-zero upstream value always wins; estimation fills only zeros (prompt from the sub-call's own request, completion from its answer text) — and returns the usage_estimated flag. Routing all four emit sites through it keeps the family from drifting again.

This closes the whole bug class, not just the non-streaming path named in the issue: the streaming-path panel members shared the identical gap and are fixed in the same PR.

Estimation feeds telemetry only (usage events → exporters / TPM / metrics). The client-facing response and SSE stream are never rewritten with synthesised usage. No CP or schema change — the sub-call events already carry usage_estimated on the existing /dp/telemetry wire.

LiteLLM has no ensemble (panel/judge) construct emitting per-sub-call usage, so there is no baseline to compare against here; the underlying token-counting semantics are unchanged from #794 (already LiteLLM-aligned).

Tests

• E2E (ensemble-usage-estimation-796-e2e.test.ts): a real ensemble (2 panel members + judge) over usage-less mock upstreams, non-streaming and streaming. Asserts, per-record, that every sub-call (both panel members + judge) ships its own usage_estimated-flagged event. Verified failing before the fix (non-streaming: nothing flagged; streaming: only the judge flagged) and passing after.
• Rust unit: estimate_subcall_tokens or-semantics — fully-missing usage fills+flags, reported usage is preserved unflagged, and a reported prompt + missing completion fills only the completion.

Fixes #796

Summary by CodeRabbit

• Bug Fixes

  • Improved token usage estimation for ensemble panel and judge calls when upstream usage data is missing or incomplete.
  • Ensured estimated usage is consistently recorded for both streaming and non-streaming ensemble requests.
  • Preserved client-facing usage behavior while improving telemetry accuracy.

• Tests

  • Added coverage for missing, complete, and partially missing usage data.
  • Added end-to-end validation for usage estimation across ensemble streaming modes.

[coderabbitai]

Warning

Review limit reached

You’ve reached a temporary PR review limit under our Fair Usage Limits Policy.

Your recent review volume is higher than typical usage, so adaptive limits are currently applied.

Next review available in: 30 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 47ae2f22-ee0f-4458-9615-9fa97982d467

📥 Commits

Reviewing files that changed from the base of the PR and between b2d7a85 and c90b7e5.

📒 Files selected for processing (1)

• crates/aisix-proxy/src/chat.rs

📝 Walkthrough

Walkthrough

Changes

The ensemble proxy now centralizes fallback token estimation for missing or partial upstream usage data. Non-streaming and streaming panel/judge telemetry records include estimated counters and usage_estimated, with unit and end-to-end coverage for both execution modes.

Ensemble usage estimation

Layer / File(s)	Summary
Token estimation contract crates/aisix-proxy/src/chat.rs	Adds shared estimation logic for missing prompt or completion counters and applies it to non-streaming panel and judge events with unit tests.
Ensemble telemetry data flow crates/aisix-proxy/src/ensemble.rs, crates/aisix-proxy/src/chat.rs	Carries panel answer text and the judge request into streaming completion callbacks for consistent usage estimation.
End-to-end telemetry validation tests/e2e/src/cases/ensemble-usage-estimation-796-e2e.test.ts	Tests non-streaming and streaming ensembles against usage-less upstream responses and verifies usage_estimated SLS records.

Estimated code review effort: 3 (Moderate) | ~30 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant EnsembleProxy
  participant PanelBackends
  participant JudgeBackend
  participant SLSTelemetry
  Client->>EnsembleProxy: ensemble chat completion
  EnsembleProxy->>PanelBackends: panel sub-calls
  PanelBackends-->>EnsembleProxy: answers with missing usage
  EnsembleProxy->>JudgeBackend: judge sub-call
  JudgeBackend-->>EnsembleProxy: synthesized answer
  EnsembleProxy->>SLSTelemetry: estimated sub-call usage with usage_estimated
  EnsembleProxy-->>Client: ensemble response or stream

Loading

Possibly related PRs

• api7/aisix#794: Introduced the local ensemble token-estimation machinery extended by this change.

Suggested reviewers: moonming

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
E2e Test Quality Review	⚠️ Warning	FAIL: the ensemble estimators still pass display names (member.model/outcome.judge_model) into estimate_subcall_tokens, so alias-backed models can get the wrong tokenizer.	Resolve and pass each sub-call’s upstream model into all four estimation sites, then add an E2E case where display_name differs from model_name to catch alias/tokenizer skew.

✅ Passed checks (5 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 fix: ensemble sub-call usage estimation for missing-usage backends.
Linked Issues check	✅ Passed	The PR implements the #796 non-streaming panel and judge estimation flow, sets usage_estimated, and adds the required usage-less ensemble test.
Out of Scope Changes check	✅ Passed	The added streaming-path panel estimation and unit tests are consistent with the PR objectives and do not appear unrelated.
Security Check	✅ Passed	No security issues found: the PR only threads internal telemetry fields into emit_usage_event; no new logging, auth, DB writes, or secret/PII exposure paths were added.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/ensemble-usage-estimation-796

---

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

[jarvis9443]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 1

🤖 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.

Inline comments:
In `@crates/aisix-proxy/src/chat.rs`:
- Around line 2608-2612: Update all estimate_subcall_tokens call sites to pass
each sub-call’s resolved upstream model rather than ensemble display names such
as member.model, member.attempt_model, or outcome.judge_model. Thread the
resolved model through the streaming panel telemetry path as well, preserving
the existing tokenizer and usage-estimation behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 67f543aa-ba46-468a-994b-a96356dc1619

📥 Commits

Reviewing files that changed from the base of the PR and between 47e6252 and b2d7a85.

📒 Files selected for processing (3)

• crates/aisix-proxy/src/chat.rs
• crates/aisix-proxy/src/ensemble.rs
• tests/e2e/src/cases/ensemble-usage-estimation-796-e2e.test.ts

[jarvis9443]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.