Cyrene_v2/internal/runtime/model/local_rung_test.go

package model

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

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

// recordingServer is an OpenAI-compatible /chat/completions endpoint that records
// the model name of every request and either fails or succeeds. It backs the
// local-rung tests (裁决 03.4.3): the point of the rung is that a request reaches
// the local endpoint with the local model name, and never when the rung is off.
type recordingServer struct {
	mu     sync.Mutex
	models []string
}

func (s *recordingServer) modelsSeen() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]string(nil), s.models...)
}

func newRecordingServer(t *testing.T, fail bool) (*httptest.Server, *recordingServer) {
	t.Helper()
	rec := &recordingServer{}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body struct {
			Model string `json:"model"`
		}
		_ = json.NewDecoder(r.Body).Decode(&body)
		rec.mu.Lock()
		rec.models = append(rec.models, body.Model)
		rec.mu.Unlock()
		if fail {
			w.WriteHeader(http.StatusInternalServerError)
			_, _ = io.WriteString(w, `{"error":"down"}`)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = io.WriteString(w, okCompletion)
	}))
	t.Cleanup(srv.Close)
	return srv, rec
}

// TestNewChatChainWithLog_LocalRungAppended verifies the third degradation level
// (primary → fallback → local) only joins the chain when it is enabled AND fully
// configured.
func TestNewChatChainWithLog_LocalRungAppended(t *testing.T) {
	t.Parallel()
	cases := []struct {
		name  string
		local config.LocalLLMConfig
		want  int
	}{
		{"off", config.LocalLLMConfig{BaseURL: "http://127.0.0.1:1/v1", Model: "local"}, 2},
		{"enabled but no model", config.LocalLLMConfig{Enabled: true, BaseURL: "http://127.0.0.1:1/v1"}, 2},
		{"enabled but no base url", config.LocalLLMConfig{Enabled: true, Model: "local"}, 2},
		{"enabled and configured", config.LocalLLMConfig{Enabled: true, BaseURL: "http://127.0.0.1:1/v1", Model: "local"}, 3},
		{"enabled with unknown provider", config.LocalLLMConfig{Enabled: true, Provider: "gemini", BaseURL: "http://127.0.0.1:1/v1", Model: "local"}, 2},
	}
	for _, c := range cases {
		c := c
		t.Run(c.name, func(t *testing.T) {
			t.Parallel()
			chain, err := NewChatChainWithLog(config.LLMConfig{
				Provider: ProviderOpenAI, BaseURL: "http://127.0.0.1:1/v1", Model: "main", Fallback: "fb",
				Local: c.local,
			}, nil)
			if err != nil {
				t.Fatalf("NewChatChainWithLog: %v", err)
			}
			if len(chain.steps) != c.want {
				t.Fatalf("steps = %d, want %d", len(chain.steps), c.want)
			}
			if c.want == 3 && chain.steps[2].Model != "local" {
				t.Errorf("steps[2].Model = %q, want local", chain.steps[2].Model)
			}
		})
	}
}

// TestChain_Chat_LocalRungServesAfterFallback verifies the local rung really takes
// over the request after the primary and fallback steps fail, with the local model
// name on the wire and the local endpoint's own base URL/API key (not the group's).
func TestChain_Chat_LocalRungServesAfterFallback(t *testing.T) {
	t.Parallel()
	mainSrv, _ := newRecordingServer(t, true)
	localSrv, localRec := newRecordingServer(t, false)

	chain, err := NewChatChainWithLog(config.LLMConfig{
		Provider: ProviderOpenAI, BaseURL: mainSrv.URL, APIKey: "sk-main", Model: "main", Fallback: "fb",
		Local: config.LocalLLMConfig{Enabled: true, BaseURL: localSrv.URL, APIKey: "sk-local", Model: "local-model"},
	}, nil)
	if err != nil {
		t.Fatalf("NewChatChainWithLog: %v", err)
	}
	resp, err := chain.Chat(context.Background(),
		Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Purpose: "chat"}, Options{})
	if err != nil {
		t.Fatalf("Chain.Chat: %v", err)
	}
	if resp.Content != "Hello!" {
		t.Errorf("resp.Content = %q, want Hello!", resp.Content)
	}
	seen := localRec.modelsSeen()
	if len(seen) != 1 || seen[0] != "local-model" {
		t.Fatalf("local endpoint models = %v, want [local-model]", seen)
	}
}

// TestChain_Chat_LocalRungOffNeverCalled pins the default-off red line: with the
// rung disabled the local endpoint must never be touched, even when every other
// step fails.
func TestChain_Chat_LocalRungOffNeverCalled(t *testing.T) {
	t.Parallel()
	mainSrv, _ := newRecordingServer(t, true)
	localSrv, localRec := newRecordingServer(t, false)

	chain, err := NewChatChainWithLog(config.LLMConfig{
		Provider: ProviderOpenAI, BaseURL: mainSrv.URL, Model: "main", Fallback: "fb",
		Local: config.LocalLLMConfig{BaseURL: localSrv.URL, Model: "local-model"},
	}, nil)
	if err != nil {
		t.Fatalf("NewChatChainWithLog: %v", err)
	}
	if _, err := chain.Chat(context.Background(),
		Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrUpstream) {
		t.Fatalf("err = %v, want ErrUpstream", err)
	}
	if seen := localRec.modelsSeen(); len(seen) != 0 {
		t.Fatalf("local endpoint called %v, want never called while the rung is off", seen)
	}
}

// TestChain_Chat_PrimaryStepKeepsRoutedModel pins that the chain does not discard
// the router's decision: the primary step sends the request's routed model, and
// the fallback step still pins its own configured model (裁决 03.4.1).
func TestChain_Chat_PrimaryStepKeepsRoutedModel(t *testing.T) {
	t.Parallel()
	// One endpoint serves both steps (the factory builds every step against the
	// group's base URL): it fails for anything but the fallback model, so the
	// recorded model names show which model each step actually sent.
	rec := &recordingServer{}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body struct {
			Model string `json:"model"`
		}
		_ = json.NewDecoder(r.Body).Decode(&body)
		rec.mu.Lock()
		rec.models = append(rec.models, body.Model)
		rec.mu.Unlock()
		if body.Model != "group-fb" {
			w.WriteHeader(http.StatusInternalServerError)
			_, _ = io.WriteString(w, `{"error":"down"}`)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = io.WriteString(w, okCompletion)
	}))
	t.Cleanup(srv.Close)

	chain, err := NewChatChainWithLog(config.LLMConfig{
		Provider: ProviderOpenAI, BaseURL: srv.URL, Model: "group-main", Fallback: "group-fb",
	}, nil)
	if err != nil {
		t.Fatalf("NewChatChainWithLog: %v", err)
	}
	if _, err := chain.Chat(context.Background(), Request{
		Model:    "persona-main",
		Messages: []Message{TextMessage(RoleUser, "hi")},
	}, Options{}); err != nil {
		t.Fatalf("Chain.Chat: %v", err)
	}
	seen := rec.modelsSeen()
	if len(seen) != 2 || seen[0] != "persona-main" || seen[1] != "group-fb" {
		t.Fatalf("models seen = %v, want [persona-main group-fb] (routed model first, step model on fallback)", seen)
	}
}