noema/internal/llm/router/router_test.go

package router

import (
	"context"
	"errors"
	"log/slog"
	"testing"
	"time"

	"github.com/andydixon/noema/internal/llm"
	"github.com/andydixon/noema/internal/minds"
	"github.com/andydixon/noema/internal/security/secrets"
	"github.com/andydixon/noema/internal/storage"
	"github.com/andydixon/noema/migrations"
	"github.com/andydixon/noema/pkg/ids"
)

func TestRouterOverrideBudgetRecordingAndReplay(t *testing.T) {
	pool := storage.TestPool(t, migrations.FS)
	ctx := context.Background()
	st := NewStore(pool)
	sec, _ := secrets.New("")
	r := NewRouter(st, sec, slog.New(slog.DiscardHandler), nil, 2)
	m, _ := minds.NewStore(pool).Create(ctx, minds.CreateParams{Slug: "r"})

	if _, err := r.Complete(ctx, llm.Request{Role: llm.RoleJudge, Messages: []llm.Message{{Role: "user", Content: "x"}}}); !errors.Is(err, llm.ErrNoModel) {
		t.Fatalf("expected no model, got %v", err)
	}
	fake := &llm.Fake{Script: []string{"first", "second"}}
	r.Override = map[llm.Role]llm.LanguageModel{llm.RoleJudge: fake}
	resp, err := r.Complete(ctx, llm.Request{Role: llm.RoleJudge, MindID: m.ID, Messages: []llm.Message{{Role: "user", Content: "x"}}})
	if err != nil || resp.Text != "first" || resp.CallID.IsZero() {
		t.Fatalf("override call: %v %+v", err, resp)
	}
	calls, _ := st.ListCalls(ctx, CallFilter{MindID: m.ID})
	if len(calls) != 1 || calls[0].Role != llm.RoleJudge || calls[0].PromptTokens == 0 {
		t.Fatalf("call not recorded: %+v", calls)
	}
	// Replay serves the recording for an identical request without calling the model.
	r.Replay = true
	resp2, err := r.Complete(ctx, llm.Request{Role: llm.RoleJudge, MindID: m.ID, Messages: []llm.Message{{Role: "user", Content: "x"}}})
	if err != nil || !resp2.Replayed || resp2.Text != "first" || len(fake.Calls) != 1 {
		t.Fatalf("replay wrong: %v %+v calls=%d", err, resp2, len(fake.Calls))
	}
	r.Replay = false
	// Budget: exhaust and expect refusal.
	r.DailyTokenBudget = 1
	// Override path bypasses configured routing; budget applies to configured models, so configure one.
	r.Override = nil
	p, err := st.CreateProvider(ctx, ProviderInput{Name: "local", Kind: llm.KindOllama, Endpoint: "http://127.0.0.1:1", AllowPrivate: true, Enabled: true})
	if err != nil {
		t.Fatal(err)
	}
	mdl, _ := st.CreateModel(ctx, ModelInput{ProviderID: p.ID, Name: "m", Kind: "chat", Enabled: true})
	st.SetRole(ctx, RoleAssignment{Role: llm.RoleJudge, ModelID: &mdl.ID}) //nolint:errcheck
	r.Invalidate()
	if _, err := r.Complete(ctx, llm.Request{Role: llm.RoleJudge, Messages: []llm.Message{{Role: "user", Content: "x"}}}); !errors.Is(err, llm.ErrBudgetExceeded) {
		t.Fatalf("expected budget error, got %v", err)
	}
	r.DailyTokenBudget = 0
	// Provider unreachable → provider error recorded with error text.
	_, err = r.Complete(ctx, llm.Request{Role: llm.RoleJudge, Messages: []llm.Message{{Role: "user", Content: "x"}}, Timeout: time.Second})
	if err == nil {
		t.Fatal("unreachable provider should fail")
	}
	errs, _ := st.ListCalls(ctx, CallFilter{ErrorsOnly: true})
	if len(errs) != 1 {
		t.Fatalf("error not recorded: %d", len(errs))
	}
	// Provider input validation: credentials in headers are rejected; key refs kept on update.
	if _, err := st.CreateProvider(ctx, ProviderInput{Name: "bad", Kind: llm.KindOpenAI, Headers: map[string]string{"Authorization": "x"}}); err == nil {
		t.Fatal("credential header accepted")
	}
	p2, _ := st.UpdateProvider(ctx, p.ID, ProviderInput{Name: "local", Kind: llm.KindOllama, Endpoint: p.Endpoint, APIKeyRef: "env:X", Enabled: true, AllowPrivate: true})
	p3, _ := st.UpdateProvider(ctx, p.ID, ProviderInput{Name: "local", Kind: llm.KindOllama, Endpoint: p.Endpoint, APIKeyRef: "", Enabled: true, AllowPrivate: true})
	if p2.APIKeyRef != "env:X" || p3.APIKeyRef != "env:X" {
		t.Fatalf("key reference handling wrong: %q %q", p2.APIKeyRef, p3.APIKeyRef)
	}
	if h := Hash("m", llm.Request{Role: "a"}); h != Hash("m", llm.Request{Role: "a"}) || h == Hash("m2", llm.Request{Role: "a"}) {
		t.Fatal("hash not stable/discriminating")
	}
	_ = ids.Nil
}