ituoke/internal/agent/model_test.go

package agent

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/zalando/go-keyring"
)

type memorySecrets map[string]string

func (m memorySecrets) Get(id string) (string, error) {
	v, ok := m[id]
	if !ok {
		return "", keyring.ErrNotFound
	}
	return v, nil
}
func (m memorySecrets) Set(id, value string) error { m[id] = value; return nil }
func (m memorySecrets) Delete(id string) error     { delete(m, id); return nil }

func testService(t *testing.T) (*Service, memorySecrets) {
	t.Helper()
	secrets := memorySecrets{}
	return newService(filepath.Join(t.TempDir(), "model.json"), secrets), secrets
}
func testConfig(base string) ModelConfig {
	c := defaultModelConfig()
	c.Enabled = true
	c.BaseURL = base
	c.Model = "test-model"
	c.ContextTurns = 1
	c.MaxTokens = 32
	return c
}
func reply(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "application/json")
	fmt.Fprint(w, `{"choices":[{"message":{"content":"测试回答"}}]}`)
}

func TestSaveConfigKeepsCredentialsOutOfFilesAndResponses(t *testing.T) {
	s, secrets := testService(t)
	c := testConfig("https://example.com/v1/")
	view, err := s.SaveModelConfig(ModelConfigInput{Config: c, APIKey: "test-secret"})
	if err != nil {
		t.Fatal(err)
	}
	if !view.HasAPIKey || view.Config.BaseURL != "https://example.com/v1" {
		t.Fatalf("bad view: %+v", view)
	}
	raw, err := os.ReadFile(s.configPath)
	if err != nil {
		t.Fatal(err)
	}
	encoded, _ := json.Marshal(view)
	if strings.Contains(string(raw), "test-secret") || strings.Contains(string(encoded), "test-secret") {
		t.Fatal("credential escaped secret store")
	}
	reloaded := newService(s.configPath, secrets)
	got, err := reloaded.GetModelConfig()
	if err != nil || !got.HasAPIKey || got.Config.Model != c.Model {
		t.Fatalf("reload: %+v %v", got, err)
	}
	changed := got.Config
	changed.BaseURL = "https://different.example/v1"
	if _, err := reloaded.SaveModelConfig(ModelConfigInput{Config: changed}); err == nil {
		t.Fatal("credential was reused at another endpoint")
	}
	got.Config.Enabled = false
	cleared, err := reloaded.SaveModelConfig(ModelConfigInput{Config: got.Config, ClearAPIKey: true})
	if err != nil || cleared.HasAPIKey {
		t.Fatalf("clear: %+v %v", cleared, err)
	}
	if _, err := secrets.Get(credentialID(got.Config)); !errors.Is(err, keyring.ErrNotFound) {
		t.Fatal("credential was not deleted")
	}
}

func TestSaveFailureRollsBackCredential(t *testing.T) {
	s, secrets := testService(t)
	c := testConfig("https://example.com/v1")
	secrets[credentialID(c)] = "old-key"
	blocked := filepath.Join(t.TempDir(), "not-a-directory")
	os.WriteFile(blocked, []byte("x"), 0600)
	s.configPath = filepath.Join(blocked, "model.json")
	if _, err := s.SaveModelConfig(ModelConfigInput{Config: c, APIKey: "new-key"}); err == nil {
		t.Fatal("expected failed config write")
	}
	if secrets[credentialID(c)] != "old-key" {
		t.Fatal("credential changed after failed save")
	}
}

func TestDraftModelTestAndListDoNotPersist(t *testing.T) {
	requests := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requests++
		if r.Header.Get("Authorization") != "Bearer draft-key" {
			t.Error("missing request credential")
		}
		switch r.URL.Path {
		case "/v1/models":
			fmt.Fprint(w, `{"data":[{"id":"z"},{"id":"a"},{"id":"z"}]}`)
		case "/v1/chat/completions":
			reply(w)
		default:
			t.Errorf("unexpected path %s", r.URL.Path)
			w.WriteHeader(404)
		}
	}))
	defer server.Close()
	s, secrets := testService(t)
	draft := ModelConfigInput{Config: testConfig(server.URL + "/v1"), APIKey: "draft-key"}
	if result := s.TestModel(draft); !result.OK {
		t.Fatal(result.Message)
	}
	draft.Config.Model = ""
	models, err := s.ListModels(draft)
	if err != nil || strings.Join(models, ",") != "a,z" {
		t.Fatalf("models: %v %v", models, err)
	}
	if requests != 2 || len(secrets) != 0 || len(s.history) != 0 {
		t.Fatal("draft test changed persistent state")
	}
	if _, err := os.Stat(s.configPath); !errors.Is(err, os.ErrNotExist) {
		t.Fatal("draft was written to disk")
	}
}

func TestChatUsesSavedModelAndBoundedHistory(t *testing.T) {
	var bodies []map[string]json.RawMessage
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]json.RawMessage
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Error(err)
		}
		if r.Header.Get("Authorization") != "Bearer saved-key" {
			t.Error("saved key not used")
		}
		bodies = append(bodies, body)
		reply(w)
	}))
	defer server.Close()
	s, _ := testService(t)
	c := testConfig(server.URL + "/v1")
	if _, err := s.Chat("phone", "hello"); err == nil {
		t.Fatal("unconfigured chat was accepted")
	}
	if _, err := s.SaveModelConfig(ModelConfigInput{Config: c, APIKey: "saved-key"}); err != nil {
		t.Fatal(err)
	}
	for _, msg := range []string{"first", "second", "third"} {
		if _, err := s.Chat("phone", msg); err != nil {
			t.Fatal(err)
		}
	}
	for i, body := range bodies {
		if _, ok := body["temperature"]; ok {
			t.Fatal("temperature was sent without opting in")
		}
		if string(body["max_completion_tokens"]) != "32" {
			t.Fatal("OpenAI token limit field missing")
		}
		var messages []modelMessage
		json.Unmarshal(body["messages"], &messages)
		want := 4
		if i == 0 {
			want = 2
		}
		if len(messages) != want {
			t.Fatalf("turn %d: got %d messages want %d", i, len(messages), want)
		}
		if i == 2 && messages[1].Content != "second" {
			t.Fatalf("old history was not trimmed: %+v", messages)
		}
	}
	c.ContextTurns = 0
	c.Provider = "custom"
	c.UseTemperature = true
	c.Temperature = 0.5
	if _, err := s.SaveModelConfig(ModelConfigInput{Config: c}); err != nil {
		t.Fatal(err)
	}
	if _, err := s.Chat("phone", "fourth"); err != nil {
		t.Fatal(err)
	}
	last := bodies[len(bodies)-1]
	if string(last["temperature"]) != "0.5" || string(last["max_tokens"]) != "32" {
		t.Fatal("custom parameters not applied")
	}
	if len(s.history["phone"]) != 0 {
		t.Fatal("zero context retained history")
	}
}

func TestProviderErrorsAreUsefulAndRedacted(t *testing.T) {
	for _, status := range []int{401, 404, 429, 500} {
		t.Run(fmt.Sprint(status), func(t *testing.T) {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(status)
				fmt.Fprint(w, `{"error":{"message":"rejected secret-value"}}`)
			}))
			defer server.Close()
			s, _ := testService(t)
			result := s.TestModel(ModelConfigInput{Config: testConfig(server.URL), APIKey: "secret-value"})
			if result.OK || !strings.Contains(result.Message, fmt.Sprint(status)) || strings.Contains(result.Message, "secret-value") {
				t.Fatalf("unsafe/errorless result: %+v", result)
			}
		})
	}
}

func TestModelRequestsDoNotFollowRedirects(t *testing.T) {
	hit := false
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hit = true; reply(w) }))
	defer target.Close()
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, target.URL, http.StatusTemporaryRedirect)
	}))
	defer server.Close()
	s, _ := testService(t)
	result := s.TestModel(ModelConfigInput{Config: testConfig(server.URL), APIKey: "secret"})
	if result.OK || hit {
		t.Fatal("request/credentials followed redirect")
	}
}

func TestConfigValidation(t *testing.T) {
	for _, base := range []string{"ftp://example.com", "https://key@example.com", "https://example.com?key=secret", "http://public.example.com"} {
		if _, err := normalizeConfig(testConfig(base), true); err == nil {
			t.Errorf("accepted unsafe/invalid base %s", base)
		}
	}
	c := testConfig("http://127.0.0.1:1234/v1/chat/completions/")
	got, err := normalizeConfig(c, true)
	if err != nil || got.BaseURL != "http://127.0.0.1:1234/v1" {
		t.Fatalf("endpoint normalization: %+v %v", got, err)
	}
	c.ContextTurns = 31
	if _, err := normalizeConfig(c, true); err == nil {
		t.Fatal("unbounded context accepted")
	}
}