Merge pull request #245 from BaaRai/feat/noa-budget-wiring

feat(agent): size NOA to the configured window and refuse irreducible requests

agent/noa.go

@@ -39,10 +39,17 @@ func enableNoa(opts *agentcore.Options, enabled func() bool, archiveRoot, sessio
 	if opts.OnWarn == nil {
 		opts.OnWarn = onWarn
 	}
-	if err := noaadapter.Enable(opts, noaadapter.Options{
-		ArchiveBaseDir: filepath.Join(archiveRoot, "noa"),
-		SessionID:      sessionID,
-		OnWarn:         onWarn,
+	// noa sizes its thresholds to the window this agent is configured for, read
+	// before Compaction is cleared below; without it noa assumes its own default.
+	window := defaultWindowK * 1000
+	if opts.Compaction != nil && opts.Compaction.ContextWindow > 0 {
+		window = opts.Compaction.ContextWindow
+	}
+	if err := noaEnable(opts, noaadapter.Options{
+		ArchiveBaseDir:    filepath.Join(archiveRoot, "noa"),
+		SessionID:         sessionID,
+		OnWarn:            onWarn,
+		ModelContextLimit: window,
 	}); err != nil {
 		if onWarn != nil {
 			onWarn("failed to enable noa compression, falling back to built-in compaction: " + err.Error())
@@ -52,4 +59,10 @@ func enableNoa(opts *agentcore.Options, enabled func() bool, archiveRoot, sessio
 	// Compactor overrides Compaction, but agentcore warns on every run when both are
 	// set; clear it explicitly.
 	opts.Compaction = nil
+	// noa can only trim what it may drop; a request whose protected part alone
+	// exceeds the window is refused here instead of by the provider.
+	opts.Provider = &noaBudgetProvider{Provider: opts.Provider, window: window}
 }
+
+// noaEnable is noaadapter.Enable; tests replace it to observe the options.
+var noaEnable = noaadapter.Enable

agent/noa_budget.go

@@ -0,0 +1,44 @@
+package agent
+
+import (
+	"context"
+	"fmt"
+	"iter"
+
+	"github.com/Autumn-27/artex/sidequestion"
+	"github.com/Autumn-27/norma/llm"
+)
+
+// noaRequestOverhead covers request framing the input estimate does not count.
+const noaRequestOverhead = 512
+
+// noaBudgetProvider refuses, before any transport, a request that cannot fit the
+// window noa was configured with. noa trims what it is allowed to drop; when
+// the protected part (system prompt, recent turns, the output reservation)
+// alone is over the window, sending it only buys a provider rejection.
+type noaBudgetProvider struct {
+	llm.Provider
+	window int
+}
+
+func (p *noaBudgetProvider) check(req llm.CompletionRequest) error {
+	estimate := sidequestion.EstimateInputTokens(req) + max(req.MaxTokens, 0) + noaRequestOverhead
+	if estimate > p.window {
+		return fmt.Errorf("Norma NOA request estimate %d exceeds configured budget %d; protected context cannot fit, increase the configured window or start a new session", estimate, p.window)
+	}
+	return nil
+}
+
+func (p *noaBudgetProvider) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
+	if err := p.check(req); err != nil {
+		return llm.Message{}, "", llm.Usage{}, err
+	}
+	return p.Provider.Complete(ctx, req)
+}
+
+func (p *noaBudgetProvider) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
+	if err := p.check(req); err != nil {
+		return func(yield func(llm.StreamEvent, error) bool) { yield(llm.StreamEvent{}, err) }
+	}
+	return p.Provider.Stream(ctx, req)
+}

agent/noa_budget_test.go

@@ -0,0 +1,134 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"iter"
+	"strings"
+	"testing"
+
+	"github.com/Autumn-27/norma/agentcore"
+	"github.com/Autumn-27/norma/llm"
+	"github.com/Autumn-27/norma/noaadapter"
+)
+
+// countingProvider answers with one text turn and counts the calls it receives.
+type countingProvider struct{ calls int }
+
+func (p *countingProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
+	p.calls++
+	return func(yield func(llm.StreamEvent, error) bool) {
+		if yield(llm.StreamEvent{Type: llm.SETextDelta, Text: "ok"}, nil) {
+			yield(llm.StreamEvent{Type: llm.SEMessageStop}, nil)
+		}
+	}
+}
+
+func (p *countingProvider) Complete(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
+	p.calls++
+	return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("ok")}}, "end_turn", llm.Usage{}, nil
+}
+
+// withNoaEnable swaps the noaadapter.Enable seam for the duration of a test.
+func withNoaEnable(t *testing.T, fn func(*agentcore.Options, noaadapter.Options) error) {
+	t.Helper()
+	prev := noaEnable
+	noaEnable = fn
+	t.Cleanup(func() { noaEnable = prev })
+}
+
+func TestEnableNoaUsesConfiguredWindow(t *testing.T) {
+	for _, tc := range []struct {
+		name       string
+		compaction func() *agentcore.Options
+		want       int
+	}{
+		{"configured", func() *agentcore.Options { return &agentcore.Options{Compaction: compactionConfig(32000)} }, 32000},
+		{"default", func() *agentcore.Options { return &agentcore.Options{} }, defaultWindowK * 1000},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var got noaadapter.Options
+			withNoaEnable(t, func(opts *agentcore.Options, o noaadapter.Options) error {
+				got = o
+				return noaadapter.Enable(opts, o)
+			})
+			inner := &countingProvider{}
+			opts := tc.compaction()
+			opts.Provider = inner
+			enableNoa(opts, func() bool { return true }, t.TempDir(), "noa-window", nil)
+
+			if got.ModelContextLimit != tc.want {
+				t.Fatalf("noaadapter.Options.ModelContextLimit=%d, want %d", got.ModelContextLimit, tc.want)
+			}
+			bp, ok := opts.Provider.(*noaBudgetProvider)
+			if !ok || bp.window != tc.want || bp.Provider != inner {
+				t.Fatalf("provider %T %+v, want *noaBudgetProvider{window %d} around the original", opts.Provider, opts.Provider, tc.want)
+			}
+			if opts.Compaction != nil || opts.Compactor == nil {
+				t.Fatalf("Compaction=%v Compactor=%v; noa must replace the built-in compaction", opts.Compaction, opts.Compactor)
+			}
+		})
+	}
+}
+
+func irreducibleRequest() llm.CompletionRequest {
+	return llm.CompletionRequest{System: []string{strings.Repeat("protected system prompt ", 8000)}, MaxTokens: 8000}
+}
+
+func TestNoaBudgetRefusesIrreducibleRequestBeforeTransport(t *testing.T) {
+	t.Run("atomic", func(t *testing.T) {
+		inner := &countingProvider{}
+		p := &noaBudgetProvider{Provider: inner, window: 32000}
+		_, _, _, err := p.Complete(t.Context(), irreducibleRequest())
+		if err == nil || !strings.Contains(err.Error(), "exceeds configured budget 32000") || inner.calls != 0 {
+			t.Fatalf("err=%v provider calls=%d", err, inner.calls)
+		}
+		if _, _, _, err := p.Complete(t.Context(), llm.CompletionRequest{System: []string{"short"}, MaxTokens: 1000}); err != nil || inner.calls != 1 {
+			t.Fatalf("a request that fits was refused or not sent: err=%v calls=%d", err, inner.calls)
+		}
+	})
+	t.Run("stream", func(t *testing.T) {
+		inner := &countingProvider{}
+		p := &noaBudgetProvider{Provider: inner, window: 32000}
+		var events int
+		var err error
+		for ev, e := range p.Stream(t.Context(), irreducibleRequest()) {
+			if e != nil {
+				err = e
+				break
+			}
+			_ = ev
+			events++
+		}
+		if err == nil || !strings.Contains(err.Error(), "exceeds configured budget 32000") || events != 0 || inner.calls != 0 {
+			t.Fatalf("err=%v events=%d provider calls=%d", err, events, inner.calls)
+		}
+		for _, e := range p.Stream(t.Context(), llm.CompletionRequest{System: []string{"short"}, MaxTokens: 1000}) {
+			if e != nil {
+				t.Fatal(e)
+			}
+		}
+		if inner.calls != 1 {
+			t.Fatalf("a request that fits was not streamed: calls=%d", inner.calls)
+		}
+	})
+}
+
+func TestEnableNoaFailureKeepsBuiltInCompaction(t *testing.T) {
+	withNoaEnable(t, func(*agentcore.Options, noaadapter.Options) error { return errors.New("archive dir unusable") })
+	inner := &countingProvider{}
+	compaction := compactionConfig(32000)
+	opts := &agentcore.Options{Compaction: compaction, Provider: inner}
+	var warns []string
+	enableNoa(opts, func() bool { return true }, t.TempDir(), "noa-fail", func(m string) { warns = append(warns, m) })
+
+	if len(warns) != 1 || !strings.Contains(warns[0], "falling back to built-in compaction") || !strings.Contains(warns[0], "archive dir unusable") {
+		t.Fatalf("warnings %q", warns)
+	}
+	if opts.Compaction != compaction || opts.Compactor != nil {
+		t.Fatalf("Compaction=%v Compactor=%v; the built-in compaction must stay", opts.Compaction, opts.Compactor)
+	}
+	if opts.Provider != llm.Provider(inner) {
+		t.Fatalf("provider %T was wrapped although noa is off", opts.Provider)
+	}
+}

db/llm_usage.go

@@ -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
 		}

db/llm_usage_cost_test.go

@@ -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)
+	}
+}

go.mod

@@ -3,7 +3,7 @@ module github.com/Autumn-27/artex
 go 1.26.10
 
 require (
-	github.com/Autumn-27/norma v0.4.3
+	github.com/Autumn-27/norma v0.4.9
 	github.com/golang-jwt/jwt/v5 v5.3.1
 	github.com/google/uuid v1.6.0
 	github.com/jackc/pgx/v5 v5.10.0
@@ -61,4 +61,4 @@ require (
 	modernc.org/memory v1.11.0 // indirect
 )
 
-replace github.com/Autumn-27/norma => github.com/BaaRai/Norma v0.4.8
+replace github.com/Autumn-27/norma => github.com/BaaRai/Norma v0.4.9

go.sum

@@ -1,5 +1,5 @@
-github.com/BaaRai/Norma v0.4.8 h1:k3at2aYSJ6aIDjEo7tC4VDSLBHXw2HobYjpKChVe6hI=
-github.com/BaaRai/Norma v0.4.8/go.mod h1:BudWrf7iHl96FgHqFr/HwN3mKYmGjJDt+xL2ZHXfFZw=
+github.com/BaaRai/Norma v0.4.9 h1:veWwo9e+OgmszutYEuglRfnvMDKThX9ilaFaXmOEGxY=
+github.com/BaaRai/Norma v0.4.9/go.mod h1:BudWrf7iHl96FgHqFr/HwN3mKYmGjJDt+xL2ZHXfFZw=
 github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro=
 github.com/andybalholm/brotli v1.2.1/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
 github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3dyBCFEj5IhUbnKptjxatkF07cF2ak3yi77so=

llmrec/llmrec.go

@@ -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

llmrec/usage_echo_test.go

@@ -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)
+			}
+		})
+	}
+}

server/task_llm.go

@@ -317,10 +317,14 @@ func (r *taskLLMRuntime) Complete(ctx context.Context, req llm.CompletionRequest
 }
 
 func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) (llm.Message, string, llm.Usage, error) {
+	// total is the usage of the logical call: the sum of every physical attempt
+	// (same-provider retries and profile failover). Each attempt is already billed
+	// once by the recorder below; this sum only feeds the caller's accounting.
+	var total llm.Usage
 	for {
 		selection, err := hooks.current(ctx)
 		if err != nil {
-			return llm.Message{}, "", llm.Usage{}, err
+			return llm.Message{}, "", total, err
 		}
 		var (
 			msg     llm.Message
@@ -333,13 +337,14 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque
 		retries, backoffOf := sameProviderRetryPolicy(selection.retry)
 		for attempt := 0; ; attempt++ {
 			msg, sr, usage, callErr = selection.provider.Complete(ctx, req)
+			total.Add(usage)
 			if callErr != nil && ctx.Err() == nil &&
 				attempt < retries && isRetryableStreamError(callErr) {
 				backoff := backoffOf(attempt)
 				log.Printf("[task-llm] task %s non-streaming call failed, retrying on the same provider after %v (%d/%d): %v",
 					taskID, backoff, attempt+1, retries, callErr)
 				if sleepCtx(ctx, backoff) {
-					break // ctx cancelled during backoff → stop retrying
+					return llm.Message{}, "", total, ctx.Err() // cancelled during backoff → stop
 				}
 				continue
 			}
@@ -347,7 +352,7 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque
 		}
 		if callErr == nil {
 			resetLLMFailStreak(taskID, selection.profileID)
-			return msg, sr, usage, nil
+			return msg, sr, total, nil
 		}
 		// Non-quota failures on a chain profile escalate after a consecutive
 		// streak (#99): same advance path as quota exhaust, truthful reason.
@@ -356,31 +361,31 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque
 				cause := fmt.Errorf("profile #%d failed %d consecutive calls (last: %w)", selection.profileID, llmFailStreakThreshold, callErr)
 				transition, markErr := hooks.exhaust(selection, cause)
 				if markErr != nil {
-					return llm.Message{}, "", llm.Usage{}, fmt.Errorf("mark profile exhausted after %v: %w", callErr, markErr)
+					return llm.Message{}, "", total, fmt.Errorf("mark profile exhausted after %v: %w", callErr, markErr)
 				}
 				if transition.Advanced && !transition.Stale && hooks.transition != nil {
 					hooks.transition(selection, transition, callErr)
 				}
 				if !transition.Stale && transition.NextProfileID == nil {
-					return llm.Message{}, "", llm.Usage{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
+					return llm.Message{}, "", total, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
 				}
 				continue // replay the same logical request on the next profile
 			}
-			return llm.Message{}, "", llm.Usage{}, callErr
+			return llm.Message{}, "", total, callErr
 		}
 		// profileID=0: the explicit chain was cleared; non-quota error: pass through. Neither changes the task failover state.
 		if selection.profileID == 0 || !isQuotaExhaustedError(callErr) {
-			return llm.Message{}, "", llm.Usage{}, callErr
+			return llm.Message{}, "", total, callErr
 		}
 		transition, markErr := hooks.exhaust(selection, callErr)
 		if markErr != nil {
-			return llm.Message{}, "", llm.Usage{}, fmt.Errorf("mark profile quota exhausted after %v: %w", callErr, markErr)
+			return llm.Message{}, "", total, fmt.Errorf("mark profile quota exhausted after %v: %w", callErr, markErr)
 		}
 		if transition.Advanced && !transition.Stale && hooks.transition != nil {
 			hooks.transition(selection, transition, callErr)
 		}
 		if !transition.Stale && transition.NextProfileID == nil {
-			return llm.Message{}, "", llm.Usage{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
+			return llm.Message{}, "", total, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
 		}
 		// No output has been delivered to the caller, so replaying the same logical request on the next profile is safe.
 	}
@@ -397,6 +402,35 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
 			committed := false
 			var pending []llm.StreamEvent
 			var streamErr error
+			// The caller's llm.Accumulator sums usage across the whole logical call
+			// (every attempt and profile): message_start is added, message_delta is
+			// folded with MergeEcho. Folding one attempt's echo against the sum of
+			// earlier attempts would lose input/cache tokens, so each attempt's echoes
+			// are normalized here against that attempt alone and only the increment
+			// is passed on, as a message_start. The recorder below already billed the
+			// physical attempt; this only shapes what the caller adds up.
+			var attemptUsage llm.Usage
+			forward := func(event llm.StreamEvent) bool {
+				switch event.Type {
+				case llm.SEMessageStart:
+					attemptUsage.Add(event.Usage)
+				case llm.SEMessageDelta:
+					next := attemptUsage
+					next.MergeEcho(event.Usage)
+					inc := llm.Usage{
+						InputTokens:      next.InputTokens - attemptUsage.InputTokens,
+						OutputTokens:     next.OutputTokens - attemptUsage.OutputTokens,
+						CacheReadTokens:  next.CacheReadTokens - attemptUsage.CacheReadTokens,
+						CacheWriteTokens: next.CacheWriteTokens - attemptUsage.CacheWriteTokens,
+					}
+					attemptUsage = next
+					if inc != (llm.Usage{}) && !yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: inc}, nil) {
+						return false
+					}
+					event.Usage = llm.Usage{}
+				}
+				return yield(event, nil)
+			}
 			retries, backoffOf := sameProviderRetryPolicy(selection.retry)
 			// Same-provider safe-window retry: a transient failure before committed
 			// (nothing delivered to the caller yet) can be replayed verbatim without
@@ -405,6 +439,7 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
 				committed = false
 				pending = nil
 				streamErr = nil
+				attemptUsage = llm.Usage{}
 				for event, err := range selection.provider.Stream(ctx, req) {
 					if err != nil {
 						streamErr = err
@@ -416,24 +451,34 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
 					}
 					if !committed {
 						for _, buffered := range pending {
-							if !yield(buffered, nil) {
+							if !forward(buffered) {
 								return
 							}
 						}
 						pending = nil
 						committed = true
 					}
-					if !yield(event, nil) {
+					if !forward(event) {
 						return
 					}
 				}
+				if streamErr != nil && !committed {
+					// A failed uncommitted attempt delivers none of its events (they are
+					// never replayed), but the usage it reported survives as one
+					// message_start so the logical call still counts it.
+					if u := pendingUsage(pending); u != (llm.Usage{}) && !yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: u}, nil) {
+						return
+					}
+					pending = nil
+				}
 				if streamErr != nil && !committed && ctx.Err() == nil &&
 					attempt < retries && isRetryableStreamError(streamErr) {
 					backoff := backoffOf(attempt)
 					log.Printf("[task-llm] task %s pre-commit stream failure, retrying on the same provider after %v (%d/%d): %v",
 						taskID, backoff, attempt+1, retries, streamErr)
 					if sleepCtx(ctx, backoff) {
-						break // ctx cancelled during backoff → stop retrying
+						yield(llm.StreamEvent{}, ctx.Err()) // cancelled during backoff → stop
+						return
 					}
 					continue
 				}
@@ -442,7 +487,7 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
 			if streamErr == nil {
 				resetLLMFailStreak(taskID, selection.profileID)
 				for _, buffered := range pending {
-					if !yield(buffered, nil) {
+					if !forward(buffered) {
 						return
 					}
 				}
@@ -474,11 +519,6 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
 			// bundle was still in use. Agent/global fallback errors follow the legacy
 			// behavior and never mutate task failover state.
 			if selection.profileID == 0 || !isQuotaExhaustedError(streamErr) {
-				for _, buffered := range pending {
-					if !yield(buffered, nil) {
-						return
-					}
-				}
 				yield(llm.StreamEvent{}, streamErr)
 				return
 			}
@@ -508,9 +548,25 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
 	}
 }
 
+// pendingUsage is the usage the buffered (uncommitted) events of one attempt
+// reported. Only message_start can be buffered: message_delta commits.
+func pendingUsage(pending []llm.StreamEvent) llm.Usage {
+	var u llm.Usage
+	for _, ev := range pending {
+		if ev.Type == llm.SEMessageStart {
+			u.Add(ev.Usage)
+		}
+	}
+	return u
+}
+
+// streamEventCommitsOutput reports whether delivering event makes the attempt
+// unreplayable. Signed and redacted thinking and tool-use signatures are opaque
+// output the model bound to this exact response, so they commit like text.
 func streamEventCommitsOutput(event llm.StreamEvent) bool {
 	switch event.Type {
-	case llm.SETextDelta, llm.SEThinkingDelta, llm.SEToolInputJSON:
+	case llm.SETextDelta, llm.SEThinkingDelta, llm.SEToolInputJSON,
+		llm.SEThinkingSignature, llm.SERedactedThinking, llm.SEToolUseSignature:
 		return event.Text != ""
 	case llm.SEToolUseStart, llm.SEMessageDelta, llm.SEMessageStop:
 		return true

server/task_llm_usage_test.go

@@ -0,0 +1,286 @@
+package server
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"iter"
+	"testing"
+	"time"
+
+	"github.com/Autumn-27/artex/agent"
+	"github.com/Autumn-27/artex/db"
+	"github.com/Autumn-27/artex/llmrec"
+	"github.com/Autumn-27/norma/llm"
+)
+
+// Usage reported by one failed attempt, and by a successful one. A stream
+// reports the input side in message_start and echoes it, with the output, in
+// message_delta; the atomic path returns the attempt's normalized usage.
+var (
+	failedAttemptUsage = llm.Usage{InputTokens: 11, OutputTokens: 2, CacheReadTokens: 3, CacheWriteTokens: 2}
+	okStartUsage       = llm.Usage{InputTokens: 13, CacheReadTokens: 3, CacheWriteTokens: 2}
+	okDeltaUsage       = llm.Usage{InputTokens: 17, OutputTokens: 5, CacheReadTokens: 4, CacheWriteTokens: 2}
+	okAttemptUsage     = llm.Usage{InputTokens: 17, OutputTokens: 5, CacheReadTokens: 4, CacheWriteTokens: 2}
+)
+
+// usageAttemptProvider fails its first fail attempts (always, when fail < 0)
+// after reporting failedAttemptUsage, then answers "done".
+type usageAttemptProvider struct {
+	fail    int
+	failErr error
+	onFail  func()
+	// extra events a failed stream attempt emits before its error.
+	failEvents []llm.StreamEvent
+	calls      int
+}
+
+func (p *usageAttemptProvider) failing() bool {
+	p.calls++
+	failed := p.fail < 0 || p.calls <= p.fail
+	if failed && p.onFail != nil {
+		p.onFail()
+	}
+	return failed
+}
+
+func (p *usageAttemptProvider) Complete(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
+	if p.failing() {
+		return llm.Message{}, "", failedAttemptUsage, p.failErr
+	}
+	return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("done")}}, "end_turn", okAttemptUsage, nil
+}
+
+func (p *usageAttemptProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
+	return func(yield func(llm.StreamEvent, error) bool) {
+		if p.failing() {
+			events := append([]llm.StreamEvent{{Type: llm.SEMessageStart, Usage: failedAttemptUsage}}, p.failEvents...)
+			for _, ev := range events {
+				if !yield(ev, nil) {
+					return
+				}
+			}
+			yield(llm.StreamEvent{}, p.failErr)
+			return
+		}
+		for _, ev := range []llm.StreamEvent{
+			{Type: llm.SEMessageStart, Usage: okStartUsage},
+			{Type: llm.SETextDelta, Text: "done"},
+			{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: okDeltaUsage},
+			{Type: llm.SEMessageStop},
+		} {
+			if !yield(ev, nil) {
+				return
+			}
+		}
+	}
+}
+
+// logicalResult is what the SDK sees for one logical call.
+type logicalResult struct {
+	usage  llm.Usage
+	text   string
+	events []llm.StreamEvent
+	err    error
+}
+
+func runLogicalCall(ctx context.Context, stream bool, taskID string, hooks taskLLMStreamHooks) logicalResult {
+	if !stream {
+		msg, _, usage, err := completeTaskLLM(ctx, taskID, llm.CompletionRequest{}, hooks)
+		return logicalResult{usage: usage, text: msg.Text(), err: err}
+	}
+	var r logicalResult
+	acc := llm.NewAccumulator()
+	for ev, err := range streamTaskLLM(ctx, taskID, llm.CompletionRequest{}, hooks) {
+		if err != nil {
+			r.err = err
+			break
+		}
+		r.events = append(r.events, ev)
+		acc.Add(ev)
+	}
+	r.usage, r.text = acc.Usage, acc.Message().Text()
+	return r
+}
+
+func singleProfileHooks(p llm.Provider, retry agent.RetryConfig) taskLLMStreamHooks {
+	return taskLLMStreamHooks{
+		current: func(context.Context) (taskLLMSelection, error) {
+			return taskLLMSelection{provider: p, retry: retry}, nil
+		},
+		exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
+			return db.TaskLLMTransition{}, nil
+		},
+	}
+}
+
+func sumUsage(us ...llm.Usage) llm.Usage {
+	var t llm.Usage
+	for _, u := range us {
+		t.Add(u)
+	}
+	return t
+}
+
+var usageModes = []struct {
+	name   string
+	stream bool
+}{{"atomic", false}, {"stream", true}}
+
+func TestTaskLLMRetrySumsAttemptUsage(t *testing.T) {
+	for _, m := range usageModes {
+		t.Run(m.name, func(t *testing.T) {
+			p := &usageAttemptProvider{fail: 1, failErr: errors.New("connection reset by peer")}
+			r := runLogicalCall(t.Context(), m.stream, "p12-retry", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
+			if r.err != nil {
+				t.Fatal(r.err)
+			}
+			want := sumUsage(failedAttemptUsage, okAttemptUsage) // {28 7 7 4}
+			if r.usage != want || p.calls != 2 || r.text != "done" {
+				t.Fatalf("usage=%+v want %+v, calls=%d, text=%q", r.usage, want, p.calls, r.text)
+			}
+		})
+	}
+}
+
+func TestTaskLLMExhaustedRetriesReturnSummedUsage(t *testing.T) {
+	for _, m := range usageModes {
+		t.Run(m.name, func(t *testing.T) {
+			p := &usageAttemptProvider{fail: -1, failErr: errors.New("connection reset by peer")}
+			r := runLogicalCall(t.Context(), m.stream, "p12-exhausted", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
+			if r.err == nil || r.err.Error() != "connection reset by peer" {
+				t.Fatalf("err=%v", r.err)
+			}
+			want := sumUsage(failedAttemptUsage, failedAttemptUsage)
+			if r.usage != want || p.calls != 2 {
+				t.Fatalf("usage=%+v want %+v, calls=%d", r.usage, want, p.calls)
+			}
+		})
+	}
+}
+
+func TestTaskLLMProfileFailoverSumsUsage(t *testing.T) {
+	for _, m := range usageModes {
+		t.Run(m.name, func(t *testing.T) {
+			first := &usageAttemptProvider{fail: -1, failErr: errors.New("openai: status 402: insufficient_quota")}
+			second := &usageAttemptProvider{}
+			active := int64(11)
+			hooks := taskLLMStreamHooks{
+				current: func(context.Context) (taskLLMSelection, error) {
+					if active == 11 {
+						return taskLLMSelection{profileID: 11, provider: first}, nil
+					}
+					return taskLLMSelection{profileID: 22, provider: second}, nil
+				},
+				exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
+					active = 22
+					next := active
+					return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil
+				},
+			}
+			r := runLogicalCall(t.Context(), m.stream, "p12-failover-"+m.name, hooks)
+			if r.err != nil {
+				t.Fatal(r.err)
+			}
+			want := sumUsage(failedAttemptUsage, okAttemptUsage)
+			if r.usage != want || first.calls != 1 || second.calls != 1 || r.text != "done" {
+				t.Fatalf("usage=%+v want %+v, calls first=%d second=%d, text=%q", r.usage, want, first.calls, second.calls, r.text)
+			}
+		})
+	}
+}
+
+func TestTaskLLMCancelDuringBackoffStops(t *testing.T) {
+	for _, m := range usageModes {
+		t.Run(m.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(t.Context())
+			defer cancel()
+			p := &usageAttemptProvider{fail: -1, failErr: errors.New("connection reset by peer")}
+			// Cancel once the retry is waiting out its (long) backoff.
+			p.onFail = func() { time.AfterFunc(20*time.Millisecond, cancel) }
+			r := runLogicalCall(ctx, m.stream, "p12-cancel", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 2, StreamInterval: time.Minute}))
+			if !errors.Is(r.err, context.Canceled) {
+				t.Fatalf("err=%v, want context.Canceled", r.err)
+			}
+			if p.calls != 1 || r.usage != failedAttemptUsage {
+				t.Fatalf("calls=%d usage=%+v", p.calls, r.usage)
+			}
+		})
+	}
+}
+
+func TestTaskLLMUncommittedStreamAttemptKeepsUsageWithoutReplay(t *testing.T) {
+	p := &usageAttemptProvider{
+		fail:    1,
+		failErr: errors.New("anthropic: overloaded_error"),
+		// Buffered, non-committing events of the failed attempt.
+		failEvents: []llm.StreamEvent{{Type: llm.SEThinkingDelta}, {Type: llm.SEToolInputJSON}},
+	}
+	r := runLogicalCall(t.Context(), true, "p12-uncommitted", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
+	if r.err != nil {
+		t.Fatal(r.err)
+	}
+	if want := sumUsage(failedAttemptUsage, okAttemptUsage); r.usage != want {
+		t.Fatalf("usage=%+v want %+v", r.usage, want)
+	}
+	if r.text != "done" {
+		t.Fatalf("text=%q, want the second attempt's only", r.text)
+	}
+	// The failed attempt contributes one usage-only message_start, nothing else.
+	if len(r.events) == 0 || r.events[0].Type != llm.SEMessageStart || r.events[0].Usage != failedAttemptUsage {
+		t.Fatalf("first event %+v", r.events)
+	}
+	for _, ev := range r.events {
+		if ev.Type == llm.SEThinkingDelta || ev.Type == llm.SEToolInputJSON {
+			t.Fatalf("buffered event of the failed attempt was delivered: %+v", r.events)
+		}
+	}
+}
+
+func TestTaskLLMOpaqueThinkingCommitsStream(t *testing.T) {
+	defer withZeroRetryBackoff()()
+	for _, kind := range []llm.StreamEventType{llm.SEThinkingSignature, llm.SERedactedThinking, llm.SEToolUseSignature} {
+		t.Run(string(kind), func(t *testing.T) {
+			if !streamEventCommitsOutput(llm.StreamEvent{Type: kind, Text: "opaque"}) {
+				t.Fatalf("%s does not commit", kind)
+			}
+			// A retryable failure after the opaque block is surfaced, not replayed.
+			p := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: kind, Text: "opaque"}}, err: errors.New("connection reset by peer")}
+			events, err := collectTaskLLMStream(streamTaskLLM(t.Context(), "p12-opaque", llm.CompletionRequest{}, singleProfileHooks(p, agent.RetryConfig{})))
+			if err == nil || p.calls != 1 || len(events) != 1 || events[0].Type != kind {
+				t.Fatalf("calls=%d events=%+v err=%v", p.calls, events, err)
+			}
+		})
+	}
+}
+
+// The recorder below the task router writes one llm_usage row per physical
+// attempt; summing the attempts for the logical call adds no row.
+func TestTaskLLMRetryKeepsOneLedgerRowPerPhysicalCall(t *testing.T) {
+	m, err := NewManager(t.TempDir(), "")
+	if err != nil {
+		t.Skipf("postgres unavailable (%v) — skipping", err)
+	}
+	defer m.Close()
+	pg := m.PG()
+	for _, mode := range usageModes {
+		t.Run(mode.name, func(t *testing.T) {
+			taskID := fmt.Sprintf("p12-ledger-%s-%d", mode.name, time.Now().UnixNano())
+			t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM llm_usage WHERE task_id=$1`, taskID) })
+			p := &usageAttemptProvider{fail: 1, failErr: errors.New("connection reset by peer")}
+			rec := llmrec.Wrap(p, pg, "test-model", "test-profile", "", "", func() bool { return false })
+			ctx := llmrec.WithTaskID(t.Context(), taskID)
+			r := runLogicalCall(ctx, mode.stream, taskID, singleProfileHooks(rec, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
+			if r.err != nil {
+				t.Fatal(r.err)
+			}
+			var rows, failed int
+			if err := pg.QueryRow(`SELECT count(*), count(*) FILTER (WHERE status='error') FROM llm_usage WHERE task_id=$1`, taskID).Scan(&rows, &failed); err != nil {
+				t.Fatal(err)
+			}
+			if rows != 2 || failed != 1 || p.calls != 2 {
+				t.Fatalf("llm_usage rows=%d (error %d), physical calls=%d; want 2 rows, 1 error", rows, failed, p.calls)
+			}
+		})
+	}
+}

sidequestion/service.go

@@ -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
 			}

sidequestion/usage_echo_test.go

@@ -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)
+			}
+		})
+	}
+}