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

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=