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 是一个 OpenAI 兼容的 /chat/completions 端点，记录
// 每个请求的模型名，并或失败或成功。它支撑
// 本地级测试（裁决 03.4.3）：该级的意义在于请求以本地模型名
// 到达本地端点，而该级关闭时绝不如此。
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 验证第三级降级
// （primary → fallback → local）只在启用且完全
// 配置好时才加入链。
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 验证本地级在主
// 步骤与降级步骤都失败后真的接管请求，且线上用
// 本地模型名与本地端点自己的 base URL/API key（而非组的）。
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 钉住默认关闭的红线：该级
// 禁用时绝不能触碰本地端点，即便其它所有
// 步骤都失败。
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 钉住链不会丢弃
// 路由器的决策：主步骤发送请求中被路由的模型，而
// 降级步骤仍钉住自己配置的模型（裁决 03.4.1）。
func TestChain_Chat_PrimaryStepKeepsRoutedModel(t *testing.T) {
	t.Parallel()
	// 一个端点同时服务两个步骤（工厂让每一步都针对
	// 组的 base URL 构建）：除降级模型外它都失败，所以
	// 记录下来的模型名显示了每一步实际发送了哪个模型。
	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)
	}
}