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=