Compare commits

..
Author SHA1 Message Date
BaaRai ce096cdba1 fix(llmrec,sidequestion): fold echoed stream usage with MergeEcho
Since Norma v0.4.7 message_delta echoes the input side reported by message_start. The recorder and side answers added both events, counting input and cache twice in llm_usage and side-answer usage. message_start is still added; message_delta is now folded with llm.Usage.MergeEcho, as llm.Accumulator does.
2026-11-15 17:51:53 +01:00
BaaRai 50f6d07640 Merge pull request #241 from BaaRai/fix/task-llm-usage
fix(server): sum task LLM usage across retries and failover
2026-11-15 17:48:11 +01:00
4 changed files with 180 additions and 7 deletions
+6 -6
View File
@@ -195,11 +195,11 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S
return
}
// Always track usage (cheap); accumulate text/thinking only for bodies.
// Anthropic (and the other providers) split token usage across events:
// message_start carries ONLY the input side (input + cache), message_delta
// ONLY the output. They must be FOLDED with Add — overwriting on delta
// would zero out the input already counted at start (mirrors the SDK's own
// llm.Accumulator; see norma/llm/accumulate.go).
// message_start carries the input side (input + cache) and is added.
// Since Norma v0.4.7 message_delta also echoes that input side next to the
// output, so it is folded with llm.Usage.MergeEcho (max on echoed fields,
// sum on output), as llm.Accumulator does; Add would count the input and
// cache twice. An output-only delta folds the same either way.
switch ev.Type {
case llm.SETextDelta:
if recordBodies {
@@ -215,7 +215,7 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S
if ev.StopReason != "" {
stopReason = ev.StopReason
}
usage.Add(ev.Usage)
usage.MergeEcho(ev.Usage)
}
if !yield(ev, err) {
return
+118
View File
@@ -0,0 +1,118 @@
package llmrec
import (
"context"
"iter"
"testing"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/norma/llm"
"github.com/google/uuid"
)
// echoStreamProvider streams one message_start and one message_delta with the
// given (already normalized) usage, as Norma's providers emit them.
type echoStreamProvider struct {
completeProvider
start, delta llm.Usage
}
func (p echoStreamProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
return func(yield func(llm.StreamEvent, error) bool) {
for _, ev := range []llm.StreamEvent{
{Type: llm.SEMessageStart, Usage: p.start},
{Type: llm.SETextDelta, Text: "answer"},
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: p.delta},
{Type: llm.SEMessageStop},
} {
if !yield(ev, nil) {
return
}
}
}
}
// Norma normalizes Anthropic usage so InputTokens includes the cache read and
// write subsets, and since v0.4.7 message_delta echoes that input side.
var usageEchoCases = []struct {
name string
start, delta llm.Usage
want llm.Usage
}{
{
name: "full_echo",
start: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
delta: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3, OutputTokens: 5},
want: llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3},
},
{
// Raw start: input 500, cache read 200, cache creation 40 → normalized 740.
// Raw delta: input 500, output 77, no cache fields → normalized 500.
name: "partial_echo",
start: llm.Usage{InputTokens: 740, CacheReadTokens: 200, CacheWriteTokens: 40},
delta: llm.Usage{InputTokens: 500, OutputTokens: 77},
want: llm.Usage{InputTokens: 740, OutputTokens: 77, CacheReadTokens: 200, CacheWriteTokens: 40},
},
{
name: "no_echo",
start: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
delta: llm.Usage{OutputTokens: 5},
want: llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3},
},
}
func drainRecorderStream(t *testing.T, ctx context.Context, r *Recorder) {
t.Helper()
for _, err := range r.Stream(ctx, llm.CompletionRequest{}) {
if err != nil {
t.Fatal(err)
}
}
}
// The metered usage of a streamed call counts an echoed input side once.
func TestStreamUsageEchoCountedOnce(t *testing.T) {
for _, tc := range usageEchoCases {
t.Run(tc.name, func(t *testing.T) {
var got []llm.Usage
ctx := WithUsageSink(t.Context(), func(u llm.Usage) { got = append(got, u) })
drainRecorderStream(t, ctx, Wrap(echoStreamProvider{start: tc.start, delta: tc.delta}, nil, "fixture", "profile", "", "", func() bool { return false }))
if len(got) != 1 || got[0] != tc.want {
t.Fatalf("metered %+v, want one call with %+v", got, tc.want)
}
})
}
}
// The llm_usage ledger gets one row per streamed call, with the echo counted once.
func TestStreamUsageEchoLedgerRow(t *testing.T) {
dsn, _, err := db.DSN()
if err != nil {
t.Skipf("no database config (%v)", err)
}
pg, err := db.Open(dsn)
if err != nil {
t.Skipf("postgres unavailable (%v)", err)
}
defer pg.Close()
if err := pg.EnsureLLMUsageTable(); err != nil {
t.Fatal(err)
}
for _, tc := range usageEchoCases {
t.Run(tc.name, func(t *testing.T) {
profile := "usage-echo-" + uuid.NewString()
t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM llm_usage WHERE profile_name=$1`, profile) })
drainRecorderStream(t, t.Context(), Wrap(echoStreamProvider{start: tc.start, delta: tc.delta}, pg, "fixture", profile, "", "", func() bool { return false }))
var rows int
var got llm.Usage
err := pg.QueryRow(`SELECT count(*), COALESCE(sum(input_tokens),0), COALESCE(sum(output_tokens),0), COALESCE(sum(cache_read),0), COALESCE(sum(cache_write),0)
FROM llm_usage WHERE profile_name=$1`, profile).Scan(&rows, &got.InputTokens, &got.OutputTokens, &got.CacheReadTokens, &got.CacheWriteTokens)
if err != nil {
t.Fatal(err)
}
if rows != 1 || got != tc.want {
t.Fatalf("ledger rows=%d usage=%+v, want 1 row with %+v", rows, got, tc.want)
}
})
}
}
+4 -1
View File
@@ -32,8 +32,11 @@ func (s SideQuestionService) Answer(ctx context.Context, req llm.CompletionReque
out.Text += ev.Text
case llm.SEToolUseStart:
out.ToolUse = true
case llm.SEMessageStart, llm.SEMessageDelta:
case llm.SEMessageStart:
out.Usage.Add(ev.Usage)
case llm.SEMessageDelta:
// message_delta may echo message_start's input side: MergeEcho, not Add.
out.Usage.MergeEcho(ev.Usage)
case llm.SEMessageStop:
complete = true
}
+52
View File
@@ -0,0 +1,52 @@
package sidequestion
import (
"context"
"iter"
"testing"
"github.com/Autumn-27/norma/llm"
)
// A streamed side answer counts message_delta's echo of the input side once.
func TestAnswerStreamUsageEchoCountedOnce(t *testing.T) {
for _, tc := range []struct {
name string
start, delta llm.Usage
want llm.Usage
}{
{"full_echo",
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3, OutputTokens: 5},
llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3}},
{"partial_echo",
llm.Usage{InputTokens: 740, CacheReadTokens: 200, CacheWriteTokens: 40},
llm.Usage{InputTokens: 500, OutputTokens: 77},
llm.Usage{InputTokens: 740, OutputTokens: 77, CacheReadTokens: 200, CacheWriteTokens: 40}},
{"no_echo",
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
llm.Usage{OutputTokens: 5},
llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3}},
} {
t.Run(tc.name, func(t *testing.T) {
p := fakeProvider{stream: func(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
return func(yield func(llm.StreamEvent, error) bool) {
for _, ev := range []llm.StreamEvent{
{Type: llm.SEMessageStart, Usage: tc.start},
{Type: llm.SETextDelta, Text: "answer"},
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: tc.delta},
{Type: llm.SEMessageStop},
} {
if !yield(ev, nil) {
return
}
}
}
}}
out, err := (SideQuestionService{Provider: p}).Answer(t.Context(), llm.CompletionRequest{}, true, nil)
if err != nil || out.Text != "answer" || out.Usage != tc.want {
t.Fatalf("answer=%+v err=%v, want usage %+v", out, err, tc.want)
}
})
}
}