Repository navigation
Conversation
…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
force-pushed
the
fix/quick-answer-stream-usage
branch
from
September 30, 2026 04:31
90a6159 to
7d8ef6d
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 answerDonemarker on the finish_reason chunk and skipped everything after it (answerCompleted).completeQuickAnswerTurnthen filledAgentCompleteDatawithMessageIDonly, sohandleCompleteomittedusagefrom the SSEcompleteevent and from the persisted message — while the agent (smart-reasoning) path, whose engine accumulatesTurnUsageand completes only after the whole stream, kept working. The reportedtotal_duration_ms: 0is the same gap (nothing filled it).This PR closes the gap on both streaming paths:
chat_completion_stream.go: the answerDonemarker 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; theStreamEndedEarlyErrorand truncation handling are unchanged.consumeFallbackStream(the model-fallback path): same treatment — it previouslybreaked 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.AgentFinalAnswerDatagains aUsagefield (interface{}, same convention asAgentCompleteData.Usageso the event package stays free of atypesimport). 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.completeQuickAnswerTurnnow fillsUsageand a realTotalDurationMs(anchored at stream setup), so:completeevent carriesusageat both the top level anddata.usage, exactly like the agent path;Message.Usage(previously always empty for quick-answer turns) also feeds the history context-token calibration;AgentCompleteDataevent with no further changes.Behaviour note: the
completeevent (and the SSE answerdonemarker) now fires when the model stream fully closes rather than atfinish_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
Related Issue
Fixes #3865
Testing
go build ./...,go veton the four changed packages,gofmt -lon 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).TestStreamCarriesUsageOnDoneMarker— usage on the EOF-sentinel chunk reaches the deferredDonemarker;TestStreamUsageOnFinishChunkAloneStillCarried— usage on the finish_reason chunk itself, stream closes right after;TestConsumeFallbackStreamCarriesUsageOnDone— the fallback path defersDoneand carries usage too.Donemarker is now a separate trailing event (content events are unchanged):TestStreamReportsEmptyLengthTruncation,TestStreamMarksPartialLengthTruncation,TestConsumeFallbackStreamTruncation(assertions rewritten to be order-based rather than count-based).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...HEADpassesgolangci-lint run --new-from-rev=origin/main ./...) — local v2.10.1 panics on Go export data, noted abovewebsite-docs/, Swagger annotations, etc.)Screenshots / Recordings
N/A — SSE payload addition; no UI changes (the frontend already reads
usageoff thecompleteevent uniformly for both modes).