Cyrene_v2/internal/runtime/model/budget_test.go

package model

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/config/runtimeconfig"
)

// newFallbackDispatchServer serves two chain steps from one endpoint: the
// "main" model fails so the chain moves on, every other model succeeds.
func newFallbackDispatchServer(t *testing.T, ch chan<- capturedRequest) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]any
		_ = json.NewDecoder(r.Body).Decode(&body)
		ch <- capturedRequest{method: r.Method, path: r.URL.Path, body: body}
		if body["model"] == "main" {
			w.WriteHeader(http.StatusInternalServerError)
			_, _ = io.WriteString(w, `{"error":"down"}`)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, okCompletion)
	}))
	t.Cleanup(srv.Close)
	return srv
}

func chatOnce(t *testing.T, chain *Chain) {
	t.Helper()
	if _, err := chain.Chat(context.Background(),
		Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); err != nil {
		t.Fatalf("Chain.Chat error: %v", err)
	}
}

// TestChain_Chat_GroupBudgetAppliesToPrimary asserts the active group's own
// maxTokens is applied to a request that carries no explicit budget.
func TestChain_Chat_GroupBudgetAppliesToPrimary(t *testing.T) {
	t.Parallel()
	ch := make(chan capturedRequest, 1)
	srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
	chain := NewChain(Step{Model: "m", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "m"})}).
		WithTokenBudget(4096, 8192)

	chatOnce(t, chain)
	if got := (<-ch).body["max_tokens"]; got != float64(4096) {
		t.Errorf("body.max_tokens = %v, want the group budget 4096", got)
	}
}

// TestChain_Chat_ExplicitRequestWinsOverGroupBudget asserts an explicit
// Request.MaxTokens is never overridden (the small review/judge budgets).
func TestChain_Chat_ExplicitRequestWinsOverGroupBudget(t *testing.T) {
	t.Parallel()
	ch := make(chan capturedRequest, 1)
	srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
	chain := NewChain(Step{Model: "m", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "m"})}).
		WithTokenBudget(4096, 8192)

	if _, err := chain.Chat(context.Background(),
		Request{Messages: []Message{TextMessage(RoleUser, "hi")}, MaxTokens: 64}, Options{}); err != nil {
		t.Fatalf("Chain.Chat error: %v", err)
	}
	if got := (<-ch).body["max_tokens"]; got != float64(64) {
		t.Errorf("body.max_tokens = %v, want the explicit 64", got)
	}
}

// TestChain_Chat_FallbackEndAppliesEnvBudget asserts the last-resort budget only
// reaches the final fallback step, never the primary.
func TestChain_Chat_FallbackEndAppliesEnvBudget(t *testing.T) {
	t.Parallel()
	primaryCh := make(chan capturedRequest, 1)
	primary := newCompletionServer(t, http.StatusInternalServerError, "application/json", `{"error":"down"}`, primaryCh)
	backupCh := make(chan capturedRequest, 1)
	backup := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, backupCh)
	chain := NewChain(
		Step{Model: "main", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: primary.URL, Model: "main"})},
		Step{Model: "backup", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: backup.URL, Model: "backup"})},
	).WithTokenBudget(0, 3000)

	chatOnce(t, chain)
	if got := (<-primaryCh).body["max_tokens"]; got != nil {
		t.Errorf("primary body.max_tokens = %v, want omitted (no budget on the primary)", got)
	}
	if got := (<-backupCh).body["max_tokens"]; got != float64(3000) {
		t.Errorf("fallback body.max_tokens = %v, want the env fallback budget 3000", got)
	}
}

// TestChain_Chat_GroupBudgetWinsOverFallbackEnv asserts a group maxTokens beats
// the last-resort env budget even on the final fallback step.
func TestChain_Chat_GroupBudgetWinsOverFallbackEnv(t *testing.T) {
	t.Parallel()
	primaryCh := make(chan capturedRequest, 1)
	primary := newCompletionServer(t, http.StatusInternalServerError, "application/json", `{"error":"down"}`, primaryCh)
	backupCh := make(chan capturedRequest, 1)
	backup := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, backupCh)
	chain := NewChain(
		Step{Model: "main", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: primary.URL, Model: "main"})},
		Step{Model: "backup", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: backup.URL, Model: "backup"})},
	).WithTokenBudget(2000, 999)

	chatOnce(t, chain)
	if got := (<-backupCh).body["max_tokens"]; got != float64(2000) {
		t.Errorf("fallback body.max_tokens = %v, want the group budget 2000", got)
	}
}

// TestChain_Chat_SingleStepIgnoresFallbackEnv asserts a chain with no fallback
// model (one step = the primary) never gets the last-resort budget, so the
// adapter default is preserved.
func TestChain_Chat_SingleStepIgnoresFallbackEnv(t *testing.T) {
	t.Parallel()
	ch := make(chan capturedRequest, 1)
	srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
	chain := NewChain(Step{Model: "m", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "m"})}).
		WithTokenBudget(0, 3000)

	chatOnce(t, chain)
	if got := (<-ch).body["max_tokens"]; got != nil {
		t.Errorf("body.max_tokens = %v, want omitted (single-step chain is not a fallback)", got)
	}
}

// TestRuntimeLLM_GroupMaxTokensEndToEnd proves the models.json maxTokens field
// reaches the provider request body through the real runtime stack.
func TestRuntimeLLM_GroupMaxTokensEndToEnd(t *testing.T) {
	ch := make(chan capturedRequest, 1)
	srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
	path := filepath.Join(t.TempDir(), "models.json")
	if err := runtimeconfig.SaveModels(path, []runtimeconfig.Group{
		{ID: "openai", Provider: ProviderOpenAI, BaseURL: srv.URL, APIKey: "sk-1", Model: "m", MaxTokens: 1234},
	}); err != nil {
		t.Fatalf("SaveModels: %v", err)
	}
	m := NewRuntimeLLM(path, rtEnv(), nil)
	if got := m.Groups()[0].MaxTokens; got != 1234 {
		t.Fatalf("Groups()[0].MaxTokens = %d, want 1234", got)
	}
	if _, err := m.Chat(context.Background(),
		Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); err != nil {
		t.Fatalf("RuntimeLLM.Chat error: %v", err)
	}
	if got := (<-ch).body["max_tokens"]; got != float64(1234) {
		t.Errorf("body.max_tokens = %v, want the group budget 1234", got)
	}
}

// TestRuntimeLLM_FallbackEndAppliesEnvBudget proves LLM_FALLBACK_MAX_TOKENS is
// wired from config through buildSnapshot into the final fallback step.
func TestRuntimeLLM_FallbackEndAppliesEnvBudget(t *testing.T) {
	ch := make(chan capturedRequest, 2)
	srv := newFallbackDispatchServer(t, ch)
	path := filepath.Join(t.TempDir(), "models.json")
	if err := runtimeconfig.SaveModels(path, []runtimeconfig.Group{
		{ID: "openai", Provider: ProviderOpenAI, BaseURL: srv.URL, APIKey: "sk-1", Model: "main", Fallback: "backup"},
	}); err != nil {
		t.Fatalf("SaveModels: %v", err)
	}
	env := rtEnv()
	env.FallbackMaxTokens = 2500
	m := NewRuntimeLLM(path, env, nil)
	if _, err := m.Chat(context.Background(),
		Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); err != nil {
		t.Fatalf("RuntimeLLM.Chat error: %v", err)
	}
	primary := <-ch
	backup := <-ch
	if got := primary.body["max_tokens"]; got != nil {
		t.Errorf("primary body.max_tokens = %v, want omitted", got)
	}
	if got := backup.body["max_tokens"]; got != float64(2500) {
		t.Errorf("fallback body.max_tokens = %v, want LLM_FALLBACK_MAX_TOKENS 2500", got)
	}
}