feat(agent): report and attribute per-turn LLM token usage

Accumulate token usage across every LLM round of an agent turn and
expose the aggregate everywhere a turn is observable:

- AgentState.TurnUsage accumulates each round's reported usage,
  including the final-answer synthesis call; TokenUsage.Accumulate
  keeps the cache counters as prompt subsets and recomputes the
  combined cache status (unsupported survives only while every call
  was itself unsupported).
- The agent completion event carries the aggregate, and the complete
  stream event exposes it both as data.usage and as a typed
  StreamEvent.Usage that buildStreamResponse promotes into the
  previously-unused StreamResponse.Usage wire field.
- The assistant message persists it as a nullable usage column on
  both migration targets (postgres 000085 jsonb, sqlite 000012 text),
  for web and IM channels alike, so history reads still carry the
  turn's cost after the live stream is gone.
- "[LLM Usage]" log lines gain ", session_id=..., principal=type:id"
  when the context carries them; calls outside sessions stay
  byte-identical.

Usage from failed retry attempts is not accumulated: the per-call
usage log remains the ground truth, while the turn aggregate serves
as the operational signal.
This commit is contained in:
ochan.kwon
2026-08-21 22:50:15 +09:00
committed by lyingbug
parent 412dcc41c6
commit 894cbec48a
20 changed files with 393 additions and 10 deletions
+1
View File
@@ -559,6 +559,7 @@ func (e *AgentEngine) runReActIteration(
response = resp
if response.Usage.TotalTokens > 0 {
e.lastUsage = response.Usage
state.TurnUsage.Accumulate(response.Usage)
logger.Debugf(ctx, "[Agent][Round-%d] Usage: prompt=%d, completion=%d, total=%d",
round, response.Usage.PromptTokens,
response.Usage.CompletionTokens, response.Usage.TotalTokens)
+18
View File
@@ -144,6 +144,12 @@ Now generate the final answer:`, query, imageRequirement)
})
}
// The synthesis call is often the largest of the turn — fold its usage
// into the turn aggregate like every ReAct round.
if llmResult.Usage != nil {
state.TurnUsage.Accumulate(*llmResult.Usage)
}
// Safety net: strip any residual <think> blocks that may have leaked through
fullAnswer := agenttools.StripThinkBlocks(llmResult.Content)
logger.Infof(ctx, "[Agent][FinalAnswer] Final answer generated: %d characters", len(fullAnswer))
@@ -195,6 +201,7 @@ func (e *AgentEngine) emitCompletionEvent(
FinalAnswer: state.FinalAnswer,
KnowledgeRefs: knowledgeRefsInterface,
AgentSteps: state.RoundSteps, // Include detailed execution steps for message storage
Usage: turnUsage(state),
TotalSteps: len(state.RoundSteps),
TotalDurationMs: time.Since(startTime).Milliseconds(),
MessageID: messageID, // Include message ID for proper message update
@@ -203,3 +210,14 @@ func (e *AgentEngine) emitCompletionEvent(
logger.Infof(ctx, "Agent execution completed in %d rounds", state.CurrentRound)
}
// turnUsage returns the turn's aggregated LLM usage, or nil when no round
// reported usage so the field stays absent from the completion event and the
// persisted message alike.
func turnUsage(state *types.AgentState) *types.TokenUsage {
if state == nil || state.TurnUsage.TotalTokens == 0 {
return nil
}
usage := state.TurnUsage
return &usage
}
+33
View File
@@ -0,0 +1,33 @@
package agent
import (
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestTurnUsageNilWhenNothingReported(t *testing.T) {
if turnUsage(nil) != nil {
t.Fatal("nil state must yield no usage")
}
if turnUsage(&types.AgentState{}) != nil {
t.Fatal("a turn whose rounds reported no usage must omit the field entirely")
}
}
func TestTurnUsageCopiesTheAggregate(t *testing.T) {
state := &types.AgentState{}
state.TurnUsage.Accumulate(types.TokenUsage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120})
usage := turnUsage(state)
if usage == nil || usage.TotalTokens != 120 {
t.Fatalf("aggregate not propagated: %+v", usage)
}
// The returned pointer must be a copy: later state mutation must not
// reach an event that has already been emitted.
state.TurnUsage.Accumulate(types.TokenUsage{PromptTokens: 1, TotalTokens: 1})
if usage.TotalTokens != 120 {
t.Fatalf("emitted usage must be detached from state: %+v", usage)
}
}
@@ -27,13 +27,13 @@ var versionedSQLiteColumns = map[string][]string{
"tenants": {"api_principal_config"}, // 000064
"users": {"is_system_admin"}, // 000053
"knowledges": {"pending_subtasks_count"}, // 000056
"messages": {"attachments"}, // 000034
"messages": {"attachments", "usage"}, // 000034, 000085
"tenant_invitations": {"token", "accepted_count"}, // 000054
"embed_channels": {"allow_memory"}, // 000060
"mcp_oauth_tokens": {"principal_type", "principal_id"}, // 000064
}
const expectedSQLiteMigrationVersion = 11
const expectedSQLiteMigrationVersion = 12
func TestSQLiteMigrationsCreateVersionedSchema(t *testing.T) {
repoRoot := sqliteRepoRoot(t)
+1
View File
@@ -139,6 +139,7 @@ type AgentCompleteData struct {
FinalAnswer string `json:"final_answer"`
KnowledgeRefs []interface{} `json:"knowledge_refs,omitempty"` // []*types.SearchResult
AgentSteps interface{} `json:"agent_steps,omitempty"` // []types.AgentStep - detailed execution steps
Usage interface{} `json:"usage,omitempty"` // *types.TokenUsage - LLM token usage aggregated over the turn
TotalDurationMs int64 `json:"total_duration_ms"`
MessageID string `json:"message_id,omitempty"` // Assistant message ID
RequestID string `json:"request_id,omitempty"`
@@ -646,6 +646,12 @@ func (h *AgentStreamHandler) handleComplete(ctx context.Context, evt event.Event
}
}
// Persist the turn's aggregated LLM usage with the message so history
// reads still carry it after the live stream is gone.
if usage, ok := data.Usage.(*types.TokenUsage); ok && usage != nil {
h.assistantMessage.Usage = usage
}
// Drain skill-generated files from the sandbox into persistent
// storage. Best-effort: any failure is logged and the turn is
// persisted without artifacts. Collect is a no-op when either the
@@ -728,6 +734,13 @@ func (h *AgentStreamHandler) handleComplete(ctx context.Context, evt event.Event
if len(h.assistantMessage.Artifacts) > 0 {
completeData["artifacts"] = publicArtifactViews(h.assistantMessage.Artifacts)
}
// Carry the turn's aggregated LLM usage both inside data (map consumers)
// and on the typed event field, which buildStreamResponse promotes to the
// response's top-level usage.
turnUsage, _ := data.Usage.(*types.TokenUsage)
if turnUsage != nil {
completeData["usage"] = turnUsage
}
if err := h.streamManager.AppendEvent(h.ctx, h.sessionID, h.assistantMessageID, interfaces.StreamEvent{
ID: evt.ID,
Type: types.ResponseTypeComplete,
@@ -735,6 +748,7 @@ func (h *AgentStreamHandler) handleComplete(ctx context.Context, evt event.Event
Done: true,
Timestamp: time.Now(),
Data: completeData,
Usage: turnUsage,
}); err != nil {
logger.GetLogger(h.ctx).Errorf("Append complete event to stream failed: %v", err)
}
+1
View File
@@ -194,6 +194,7 @@ func buildStreamResponse(evt interfaces.StreamEvent, requestID string) *types.St
Content: evt.Content,
Done: evt.Done,
Data: evt.Data,
Usage: evt.Usage,
}
// Extract session_id and assistant_message_id for agent_query events
@@ -0,0 +1,50 @@
package session
import (
"encoding/json"
"testing"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestBuildStreamResponsePromotesUsageOnCompleteEvents(t *testing.T) {
usage := &types.TokenUsage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120}
response := buildStreamResponse(interfaces.StreamEvent{
Type: types.ResponseTypeComplete,
Done: true,
Data: map[string]interface{}{"total_steps": 3, "usage": usage},
Usage: usage,
}, "req-1")
require.NotNil(t, response.Usage)
assert.Equal(t, 120, response.Usage.TotalTokens)
assert.Equal(t, usage, response.Data["usage"])
}
func TestBuildStreamResponseLeavesUsageNilWhenAbsent(t *testing.T) {
response := buildStreamResponse(interfaces.StreamEvent{
Type: types.ResponseTypeAnswer,
Data: map[string]interface{}{"event_id": "e-1"},
}, "req-1")
assert.Nil(t, response.Usage)
}
func TestStreamEventUsageSurvivesJSONRoundTrip(t *testing.T) {
// The stream manager persists events as JSON (Redis) before the SSE loop
// reads them back — the typed usage must survive that round trip.
usage := &types.TokenUsage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120}
usage.SetPromptCacheUsage(80, 0, 20, true)
raw, err := json.Marshal(interfaces.StreamEvent{Type: types.ResponseTypeComplete, Usage: usage})
require.NoError(t, err)
var restored interfaces.StreamEvent
require.NoError(t, json.Unmarshal(raw, &restored))
require.NotNil(t, restored.Usage)
assert.Equal(t, *usage, *restored.Usage)
}
+3
View File
@@ -701,6 +701,9 @@ func applyIMCompleteDataToMessage(msg *types.Message, data event.AgentCompleteDa
}
msg.IsCompleted = true
msg.AgentDurationMs = data.TotalDurationMs
if usage, ok := data.Usage.(*types.TokenUsage); ok && usage != nil {
msg.Usage = usage
}
if len(data.KnowledgeRefs) > 0 {
refs := make([]*types.SearchResult, 0, len(data.KnowledgeRefs))
collectIMKnowledgeReferences(&refs, data.KnowledgeRefs)
+25 -2
View File
@@ -2,6 +2,7 @@ package chat
import (
"context"
"strings"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
@@ -16,8 +17,30 @@ func logUsage(ctx context.Context, model string, u *types.TokenUsage) {
}
purpose, prefixFingerprint := types.LLMCallMetadataFromContext(ctx)
logger.Infof(ctx,
"[LLM Usage] model=%s, purpose=%s, prompt_prefix=%s, prompt_tokens=%d, completion_tokens=%d, total_tokens=%d, cached_tokens=%d, cache_read_tokens=%d, cache_write_tokens=%d, cache_miss_tokens=%d, cache_reported=%t, cache_status=%s",
"[LLM Usage] model=%s, purpose=%s, prompt_prefix=%s, prompt_tokens=%d, completion_tokens=%d, "+
"total_tokens=%d, cached_tokens=%d, cache_read_tokens=%d, cache_write_tokens=%d, "+
"cache_miss_tokens=%d, cache_reported=%t, cache_status=%s%s",
model, purpose, prefixFingerprint, u.PromptTokens, u.CompletionTokens, u.TotalTokens,
u.CachedTokens, u.CacheReadTokens, u.CacheWriteTokens, u.CacheMissTokens,
u.CacheReported, u.CacheStatus)
u.CacheReported, u.CacheStatus, usageAttribution(ctx))
}
// usageAttribution renders the ", session_id=…, principal=…" suffix that
// attributes a usage line to the session and terminal principal that
// triggered the call. Calls that run outside a session or without a resolved
// principal (document parsing, title generation, background jobs) render an
// empty suffix, keeping their lines byte-identical to before.
func usageAttribution(ctx context.Context) string {
var b strings.Builder
if sessionID, ok := types.SessionIDFromContext(ctx); ok && sessionID != "" {
b.WriteString(", session_id=")
b.WriteString(sessionID)
}
if principal, ok := types.PrincipalFromContext(ctx); ok {
b.WriteString(", principal=")
b.WriteString(principal.Type)
b.WriteString(":")
b.WriteString(principal.ID)
}
return b.String()
}
@@ -0,0 +1,34 @@
package chat
import (
"context"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestUsageAttributionEmptyOutsideSessions(t *testing.T) {
// Calls without session or principal context (document parsing, title
// generation, background jobs) must keep the usage line byte-identical.
if got := usageAttribution(context.Background()); got != "" {
t.Fatalf("expected empty suffix, got %q", got)
}
}
func TestUsageAttributionCarriesSessionAndPrincipal(t *testing.T) {
ctx := types.WithSessionID(context.Background(), "sess-123")
ctx = types.WithPrincipal(ctx, types.Principal{Type: types.PrincipalAPIExternalUser, ID: "10000:42"})
got := usageAttribution(ctx)
want := ", session_id=sess-123, principal=api_external_user:10000:42"
if got != want {
t.Fatalf("got %q want %q", got, want)
}
}
func TestUsageAttributionSessionOnly(t *testing.T) {
ctx := types.WithSessionID(context.Background(), "sess-9")
if got := usageAttribution(ctx); got != ", session_id=sess-9" {
t.Fatalf("got %q", got)
}
}
+1
View File
@@ -256,6 +256,7 @@ type AgentState struct {
IsComplete bool `json:"is_complete"` // Whether agent has finished
FinalAnswer string `json:"final_answer"` // The final answer to the query
KnowledgeRefs []*SearchResult `json:"knowledge_refs"` // Collected knowledge references
TurnUsage TokenUsage `json:"turn_usage"` // LLM token usage accumulated across every round of this turn
}
// FunctionDefinition represents a function definition for LLM function calling
+73
View File
@@ -73,6 +73,79 @@ func (u *TokenUsage) MarkPromptCacheUnsupported() {
u.CacheStatus = PromptCacheStatusUnsupported
}
// Accumulate adds another call's usage into u, preserving the subset
// semantics: prompt/completion/total and every cache counter sum
// independently (cache counters stay subsets of the prompt count and are
// never folded into it). CacheReported ORs, and the cache status is
// recomputed from the combined counters so a single cache hit anywhere in
// the accumulated calls reads as a hit.
func (u *TokenUsage) Accumulate(other TokenUsage) {
if u == nil {
return
}
u.PromptTokens += other.PromptTokens
u.CompletionTokens += other.CompletionTokens
u.TotalTokens += other.TotalTokens
u.CachedTokens += other.CachedTokens
u.CacheReadTokens += other.CacheReadTokens
u.CacheWriteTokens += other.CacheWriteTokens
u.CacheMissTokens += other.CacheMissTokens
u.CacheReported = u.CacheReported || other.CacheReported
switch {
case !u.CacheReported:
u.CacheStatus = mergeUnreportedCacheStatus(u.CacheStatus, other.CacheStatus)
case u.CacheReadTokens > 0:
u.CacheStatus = PromptCacheStatusHit
default:
u.CacheStatus = PromptCacheStatusMiss
}
}
// mergeUnreportedCacheStatus folds the statuses of never-reported usage.
// "unsupported" survives only while every accumulated call was itself
// classified unsupported — the first accumulation adopts the incoming
// classification, and any later call that is not known-unsupported (an
// unreported or unclassified one) degrades the aggregate to unreported.
func mergeUnreportedCacheStatus(accumulated, incoming PromptCacheStatus) PromptCacheStatus {
if accumulated == "" {
if incoming == "" {
return PromptCacheStatusUnreported
}
return incoming
}
if accumulated == PromptCacheStatusUnsupported && incoming == PromptCacheStatusUnsupported {
return PromptCacheStatusUnsupported
}
return PromptCacheStatusUnreported
}
// Value persists the usage as a jsonb column (assistant messages carry the
// turn's aggregate); nil writes SQL NULL. Mirrors the nullable-pointer
// pattern APIPrincipalConfig uses.
func (u *TokenUsage) Value() (driver.Value, error) {
if u == nil {
return nil, nil
}
return json.Marshal(u)
}
// Scan restores a jsonb usage column; NULL leaves the receiver zero-valued.
func (u *TokenUsage) Scan(value interface{}) error {
if value == nil {
return nil
}
var b []byte
switch v := value.(type) {
case []byte:
b = v
case string:
b = []byte(v)
default:
return nil
}
return json.Unmarshal(b, u)
}
// LLMToolCall represents a function/tool call from the LLM
type LLMToolCall struct {
ID string `json:"id"`
+7 -6
View File
@@ -9,12 +9,13 @@ import (
// StreamEvent represents a single event in the stream
type StreamEvent struct {
ID string `json:"id"` // Unique event ID
Type types.ResponseType `json:"type"` // Event type (thinking, tool_call, tool_result, references, complete, etc.)
Content string `json:"content"` // Event content (chunk for streaming events)
Done bool `json:"done"` // Whether this event is done
Timestamp time.Time `json:"timestamp"` // When this event occurred
Data map[string]interface{} `json:"data,omitempty"` // Additional event data (references, metadata, etc.)
ID string `json:"id"` // Unique event ID
Type types.ResponseType `json:"type"` // Event type (thinking, tool_call, complete, etc.)
Content string `json:"content"` // Event content (chunk for streaming events)
Done bool `json:"done"` // Whether this event is done
Timestamp time.Time `json:"timestamp"` // When this event occurred
Data map[string]interface{} `json:"data,omitempty"` // Additional event data (references, metadata, etc.)
Usage *types.TokenUsage `json:"usage,omitempty"` // LLM token usage aggregated over the turn (complete events)
}
// StreamManager stream manager interface - minimal append-only design
+4
View File
@@ -285,6 +285,10 @@ type Message struct {
IsFallback bool `json:"is_fallback,omitempty"`
// Agent total execution duration in milliseconds (from query start to answer start)
AgentDurationMs int64 `json:"agent_duration_ms,omitempty" gorm:"column:agent_duration_ms;default:0"`
// LLM token usage aggregated across every round of the turn that produced this
// assistant message. Persisted so history reads can attribute cost after the
// live stream is gone; NULL (nil) for user messages and pre-feature rows.
Usage *TokenUsage `json:"usage,omitempty" gorm:"type:jsonb;column:usage"`
// RenderedContent stores the full RAG-augmented user message (with retrieved context)
// sent to the LLM. Used to preserve retrieval context across conversation turns.
// Empty for non-retrieval intents or assistant messages.
+119
View File
@@ -0,0 +1,119 @@
package types
import (
"encoding/json"
"testing"
)
func TestTokenUsageAccumulateSumsEveryCounter(t *testing.T) {
var turn TokenUsage
first := TokenUsage{PromptTokens: 1000, CompletionTokens: 200, TotalTokens: 1200}
first.SetPromptCacheUsage(800, 0, 200, true)
second := TokenUsage{PromptTokens: 1500, CompletionTokens: 300, TotalTokens: 1800}
second.SetPromptCacheUsage(0, 1500, 1500, true)
turn.Accumulate(first)
turn.Accumulate(second)
if turn.PromptTokens != 2500 || turn.CompletionTokens != 500 || turn.TotalTokens != 3000 {
t.Fatalf("token sums wrong: %+v", turn)
}
if turn.CacheReadTokens != 800 || turn.CacheWriteTokens != 1500 || turn.CacheMissTokens != 1700 {
t.Fatalf("cache sums wrong: %+v", turn)
}
if turn.CachedTokens != turn.CacheReadTokens {
t.Fatalf("legacy alias diverged from cache reads: %+v", turn)
}
if !turn.CacheReported || turn.CacheStatus != PromptCacheStatusHit {
t.Fatalf("a hit anywhere in the turn must read as a hit: %+v", turn)
}
}
func TestTokenUsageAccumulateStatusFollowsCombinedCounters(t *testing.T) {
var unreported TokenUsage
unreported.Accumulate(TokenUsage{PromptTokens: 10, TotalTokens: 10})
if unreported.CacheReported || unreported.CacheStatus != PromptCacheStatusUnreported {
t.Fatalf("all-unreported accumulation must stay unreported: %+v", unreported)
}
var missOnly TokenUsage
reportedMiss := TokenUsage{PromptTokens: 10, TotalTokens: 10}
reportedMiss.SetPromptCacheUsage(0, 0, 10, true)
missOnly.Accumulate(reportedMiss)
if missOnly.CacheStatus != PromptCacheStatusMiss {
t.Fatalf("reported without reads must read as miss: %+v", missOnly)
}
}
func TestTokenUsageAccumulatePreservesUnsupported(t *testing.T) {
var unsupported TokenUsage
call := TokenUsage{PromptTokens: 10, TotalTokens: 10}
call.MarkPromptCacheUnsupported()
unsupported.Accumulate(call)
unsupported.Accumulate(call)
if unsupported.CacheStatus != PromptCacheStatusUnsupported {
t.Fatalf("all-unsupported accumulation must stay unsupported: %+v", unsupported)
}
// Mixing in a merely-unreported call degrades the aggregate: the turn no
// longer proves every provider path was incapable of reporting.
unsupported.Accumulate(TokenUsage{PromptTokens: 5, TotalTokens: 5})
if unsupported.CacheStatus != PromptCacheStatusUnreported {
t.Fatalf("unsupported+unreported mix must read unreported: %+v", unsupported)
}
}
func TestTokenUsageAccumulateOnNilReceiverIsNoOp(t *testing.T) {
var u *TokenUsage
u.Accumulate(TokenUsage{PromptTokens: 1, TotalTokens: 1}) // must not panic
}
func TestTokenUsageValueScanRoundTrip(t *testing.T) {
original := &TokenUsage{PromptTokens: 42, CompletionTokens: 7, TotalTokens: 49}
original.SetPromptCacheUsage(30, 0, 12, true)
value, err := original.Value()
if err != nil {
t.Fatalf("Value failed: %v", err)
}
raw, ok := value.([]byte)
if !ok {
t.Fatalf("Value must produce bytes, got %T", value)
}
var restored TokenUsage
if err := restored.Scan(raw); err != nil {
t.Fatalf("Scan failed: %v", err)
}
if restored != *original {
t.Fatalf("round trip diverged: got %+v want %+v", restored, *original)
}
// json.Marshal must agree with Value so API responses and the persisted
// column carry the same shape.
direct, err := json.Marshal(original)
if err != nil {
t.Fatalf("json.Marshal failed: %v", err)
}
if string(direct) != string(raw) {
t.Fatalf("Value diverged from json.Marshal: %s vs %s", raw, direct)
}
}
func TestTokenUsageValueNilAndScanNull(t *testing.T) {
var u *TokenUsage
value, err := u.Value()
if err != nil || value != nil {
t.Fatalf("nil usage must persist as SQL NULL, got (%v, %v)", value, err)
}
restored := TokenUsage{PromptTokens: 5}
if err := restored.Scan(nil); err != nil {
t.Fatalf("Scan(NULL) failed: %v", err)
}
if restored.PromptTokens != 5 {
t.Fatalf("Scan(NULL) must leave the receiver untouched: %+v", restored)
}
}
@@ -0,0 +1 @@
ALTER TABLE messages DROP COLUMN usage;
@@ -0,0 +1,4 @@
-- Mirrors versioned migration 000085_message_usage:
-- per-turn LLM token usage persisted with the assistant message.
ALTER TABLE messages ADD COLUMN usage TEXT;
@@ -0,0 +1 @@
ALTER TABLE messages DROP COLUMN IF EXISTS usage;
@@ -0,0 +1 @@
ALTER TABLE messages ADD COLUMN IF NOT EXISTS usage JSONB;