Compare commits

..
Author SHA1 Message Date
BaaRai 60728092e1 fix(db): stop billing cache writes twice in profile cost
Norma normalizes InputTokens to input + cache read + cache write, so both cache buckets are subsets of it. The fresh share billed at the input price is now input minus reads minus writes, still clamped at zero for gateways that report disjoint buckets.
2026-11-15 18:17:33 +01:00
BaaRai ff599c65d3 Merge pull request #243 from BaaRai/fix/stream-usage-echo
fix(llmrec,sidequestion): fold echoed stream usage with MergeEcho
2026-11-15 18:14:20 +01:00
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
6 changed files with 311 additions and 12 deletions
+6 -5
View File
@@ -341,11 +341,12 @@ func applyCostEur(d *DB, out []ProfileUsage) error {
if !ok || (pr.in == 0 && pr.outp == 0 && pr.cr == 0 && pr.cw == 0) {
continue
}
// InputTokens includes the cached subset (Anthropic and OpenAI usage
// semantics), so the fresh share is input minus cache reads; charging
// both buckets at full input price double-counts the cache. Clamped at
// zero for gateways that report the buckets disjointly.
fresh := out[i].InputTokens - out[i].CacheReadTokens
// InputTokens already includes both cached subsets: Norma normalizes it to
// input + cache read + cache write (OpenAI prompt_tokens semantics). The
// fresh share is therefore input minus cache reads minus cache writes;
// charging either subset at full input price as well bills it twice.
// Clamped at zero for gateways that report the buckets disjointly.
fresh := out[i].InputTokens - out[i].CacheReadTokens - out[i].CacheWriteTokens
if fresh < 0 {
fresh = 0
}
+125
View File
@@ -0,0 +1,125 @@
package db
import (
"fmt"
"math"
"testing"
"time"
)
// costProfiles saves a priced profile (EUR per million tokens: input 2.00,
// output 4.00, cache read 0.20, cache write 2.50) and an unpriced one.
func costProfiles(t *testing.T, d *DB) (priced, unpriced string) {
t.Helper()
suffix := time.Now().UnixNano()
priced, unpriced = fmt.Sprintf("t-cost-priced-%d", suffix), fmt.Sprintf("t-cost-unpriced-%d", suffix)
for _, p := range []*LLMProfile{
{Name: priced, Format: "anthropic", Model: "m", APIKey: "k",
PriceInputPerM: 2, PriceOutputPerM: 4, PriceCacheReadPerM: 0.2, PriceCacheWritePerM: 2.5},
{Name: unpriced, Format: "anthropic", Model: "m", APIKey: "k"},
} {
id, err := d.SaveProfile(p)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) })
}
return priced, unpriced
}
func assertCost(t *testing.T, got *float64, want float64) {
t.Helper()
if got == nil {
t.Fatalf("cost=nil, want %.2f", want)
}
if math.Abs(*got-want) > 1e-9 {
t.Fatalf("cost=%.4f, want %.2f", *got, want)
}
}
// InputTokens already includes the cache read and cache write subsets, so only
// the remainder is billed at the input price.
func TestCostSeparatesFreshReadAndWriteBuckets(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Fatal(err)
}
defer d.Close()
priced, unpriced := costProfiles(t, d)
// Every case adds 100k output tokens: 0.1 × 4.00 = 0.40.
for _, tc := range []struct {
name string
in, read, write int
want float64
}{
{"fresh", 1_000_000, 0, 0, 2.40}, // 1M×2.00 + 0.40
{"read_only", 1_000_000, 1_000_000, 0, 0.60}, // 0 fresh + 1M×0.20 + 0.40
{"write_only", 1_000_000, 0, 1_000_000, 2.90}, // 0 fresh + 1M×2.50 + 0.40
{"mixed", 1_000_000, 200_000, 300_000, 2.19}, // 500k×2.00 + 200k×0.20 + 300k×2.50 + 0.40
{"inconsistent_gateway", 100_000, 200_000, 300_000, 1.19}, // fresh clamped to 0 + 0.04 + 0.75 + 0.40
} {
t.Run(tc.name, func(t *testing.T) {
out := []ProfileUsage{
{ProfileName: priced, InputTokens: tc.in, CacheReadTokens: tc.read, CacheWriteTokens: tc.write, OutputTokens: 100_000},
{ProfileName: unpriced, InputTokens: tc.in, CacheReadTokens: tc.read, CacheWriteTokens: tc.write, OutputTokens: 100_000},
{ProfileName: "t-cost-unknown-profile", InputTokens: tc.in, OutputTokens: 100_000},
}
if err := applyCostEur(d, out); err != nil {
t.Fatal(err)
}
assertCost(t, out[0].CostEur, tc.want)
if out[1].CostEur != nil || out[2].CostEur != nil {
t.Fatalf("unpriced usage got a cost: %v %v", out[1].CostEur, out[2].CostEur)
}
})
}
}
// Several ledger rows of one profile are summed before pricing; with no clamp
// involved the profile cost equals the sum of the per-row costs.
func TestUsageByProfileCostAggregatesRows(t *testing.T) {
d, err := Open(testDSN(t))
if err != nil {
t.Fatal(err)
}
defer d.Close()
priced, unpriced := costProfiles(t, d)
t.Cleanup(func() { d.Exec(`DELETE FROM llm_usage WHERE profile_name IN ($1,$2)`, priced, unpriced) })
rows := []LLMUsage{
// fresh 400k×2.00 + 200k×0.20 + 40k×4.00 = 0.80 + 0.04 + 0.16 = 1.00
{InputTokens: 600_000, CacheRead: 200_000, OutputTokens: 40_000},
// fresh 100k×2.00 + 300k×2.50 + 60k×4.00 = 0.20 + 0.75 + 0.24 = 1.19
{InputTokens: 400_000, CacheWrite: 300_000, OutputTokens: 60_000},
}
for _, profile := range []string{priced, unpriced} {
for _, r := range rows {
r.ProfileName, r.Model, r.Status = profile, "m", "ok"
if err := d.InsertLLMUsage(&r); err != nil {
t.Fatal(err)
}
}
}
usage, err := d.UsageByProfile()
if err != nil {
t.Fatal(err)
}
found := 0
for _, p := range usage {
switch p.ProfileName {
case priced:
found++
if p.Calls != 2 || p.InputTokens != 1_000_000 || p.CacheReadTokens != 200_000 || p.CacheWriteTokens != 300_000 || p.OutputTokens != 100_000 {
t.Fatalf("aggregate %+v", p)
}
assertCost(t, p.CostEur, 2.19) // = 1.00 + 1.19, the mixed case
case unpriced:
found++
if p.Calls != 2 || p.CostEur != nil {
t.Fatalf("unpriced aggregate %+v", p)
}
}
}
if found != 2 {
t.Fatalf("profiles missing from UsageByProfile: found %d of 2", found)
}
}
+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)
}
})
}
}