Skip to content

fix(qa): report token usage on quick-answer complete events (#3865) - #3899

Open
BaiMoHan wants to merge 1 commit into
Tencent:mainfrom
BaiMoHan:fix/quick-answer-stream-usage
Open

BaiMoHan wants to merge 1 commit into
Tencent:mainfrom
BaiMoHan:fix/quick-answer-stream-usage

Conversation

@BaiMoHan

Copy link
Copy Markdown

Description

Fixes #3865.

Quick-answer (KnowledgeQA) turns reported no token usage: providers deliver usage on the stream's closing chunk — after finish_reason — but the KnowledgeQA stream consumer emitted the answer Done marker on the finish_reason chunk and skipped everything after it (answerCompleted). completeQuickAnswerTurn then filled AgentCompleteData with MessageID only, so handleComplete omitted usage from the SSE complete event and from the persisted message — while the agent (smart-reasoning) path, whose engine accumulates TurnUsage and completes only after the whole stream, kept working. The reported total_duration_ms: 0 is the same gap (nothing filled it).

This PR closes the gap on both streaming paths:

  • chat_completion_stream.go: the answer Done marker is now held back until the model channel closes and rides out with the usage captured from whichever closing chunk carried it (some providers attach usage to the finish_reason chunk, some to a separate EOF sentinel — both are covered, including chunks the duplicate-terminal guard would otherwise skip). Content still streams out immediately, so a mid-stream cancellation cannot lose text; the StreamEndedEarlyError and truncation handling are unchanged.
  • consumeFallbackStream (the model-fallback path): same treatment — it previously breaked out of the loop on the first done chunk, which also dropped any usage. A guard now ignores repeated terminal chunks after completion instead of overwriting the truncation marker.
  • AgentFinalAnswerData gains a Usage field (interface{}, same convention as AgentCompleteData.Usage so the event package stays free of a types import). It is only set when a provider actually reported usage — a typed nil is never boxed, so the field serializes as absent rather than "usage": null.
  • The KnowledgeQA final-answer handler stores that usage on the assistant message; completeQuickAnswerTurn now fills Usage and a real TotalDurationMs (anchored at stream setup), so:
    • the SSE complete event carries usage at both the top level and data.usage, exactly like the agent path;
    • the persisted Message.Usage (previously always empty for quick-answer turns) also feeds the history context-token calibration;
    • IM consumers pick usage up from the same AgentCompleteData event with no further changes.

Behaviour note: the complete event (and the SSE answer done marker) now fires when the model stream fully closes rather than at finish_reason — typically tens of milliseconds later, and strictly more correct, since every stream event is guaranteed to precede the completion.

No breaking changes, no config surface changes.

Type of Change

  • 🐛 Bug fix
  • ✨ New feature
  • 💥 Breaking change
  • 📚 Documentation update
  • 🎨 Refactor
  • ⚡ Performance improvement
  • 🧪 Test
  • 🔧 Configuration / Build / CI

Related Issue

Fixes #3865

Testing

  • go build ./..., go vet on the four changed packages, gofmt -l on the changed files, git diff --check — all clean.
  • golangci-lint run --new-from-rev=origin/main: the locally installed v2.10.1 panics on current Go export data (same as noted in fix(wiki): fan out follow-up triggers so backlogs drain at ingest_max_inflight #3726; CI's pinned v2.12.2 should be authoritative here).
  • New tests:
    • TestStreamCarriesUsageOnDoneMarker — usage on the EOF-sentinel chunk reaches the deferred Done marker;
    • TestStreamUsageOnFinishChunkAloneStillCarried — usage on the finish_reason chunk itself, stream closes right after;
    • TestConsumeFallbackStreamCarriesUsageOnDone — the fallback path defers Done and carries usage too.
  • Updated existing tests whose event counts changed because the Done marker is now a separate trailing event (content events are unchanged): TestStreamReportsEmptyLengthTruncation, TestStreamMarksPartialLengthTruncation, TestConsumeFallbackStreamTruncation (assertions rewritten to be order-based rather than count-based).
  • Regression: go test ./internal/application/service/... ./internal/handler/session/... ./internal/event/... ./internal/agent/... ./internal/im/... — all pass, no pre-existing failures in these packages in this environment.

Checklist

  • git diff --check origin/main...HEAD passes
  • Changed source files are formatted
  • Targeted tests for the changed packages/components pass
  • Diff-scoped lint passes where applicable (for Go: golangci-lint run --new-from-rev=origin/main ./...) — local v2.10.1 panics on Go export data, noted above
  • Full-repository checks were run, or any unrelated/environment-dependent failures are documented above
  • Self-reviewed the code
  • Added/updated tests covering the change
  • Updated related documentation (README, website-docs/, Swagger annotations, etc.)
  • Breaking changes are clearly called out in the description above

Screenshots / Recordings

N/A — SSE payload addition; no UI changes (the frontend already reads usage off the complete event uniformly for both modes).

…3865)

KnowledgeQA turns lost their token usage: providers report it on the
stream's closing chunk (after finish_reason), but the stream consumer
emitted the Done marker on the finish_reason chunk and skipped the
closing chunk, so completeQuickAnswerTurn filled AgentCompleteData with
MessageID only — the SSE complete event and the persisted message both
ended up without usage (and total_duration_ms stayed 0).

Hold the answer Done marker until the channel closes and carry the
captured usage on it; the same applies to the fallback stream, which
also broke out of its loop on the first done chunk. completeQuickAnswerTurn now fills Usage and TotalDurationMs, so
handleComplete publishes usage on the complete event and persists it
with the message — the agent (smart-reasoning) path already did. IM
consumers pick usage up from the same event.

Provider-reported usage is absent only when a provider reports none;
the field stays omitted rather than serializing a typed nil.
@BaiMoHan
BaiMoHan force-pushed the fix/quick-answer-stream-usage branch from 90a6159 to 7d8ef6d Compare September 30, 2026 04:31

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: 快速问答(KnowledgeQA)流式完成事件未返回 usage,智能推理正常返回

1 participant