internal/app/codex_session_cache_test.go

package app

import (
	"context"
	"net/http"
	"regexp"
	"strings"
	"testing"
	"time"

	"ccLoad/internal/model"
	"ccLoad/internal/protocol"
)

var uuidPattern = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`)

func assertFieldOrder(t *testing.T, body string, fields ...string) {
	t.Helper()
	prev := -1
	for _, field := range fields {
		idx := strings.Index(body, field)
		if idx < 0 {
			t.Fatalf("field %s missing in %s", field, body)
		}
		if idx <= prev {
			t.Fatalf("field order broken at %s in %s", field, body)
		}
		prev = idx
	}
}

func resetCodexSessionCache() {
	codexSessionMu.Lock()
	codexSessionMap = make(map[string]codexSessionEntry)
	codexSessionMu.Unlock()
}

func TestResolveCodexSessionHint_AnthropicWithUserID(t *testing.T) {
	resetCodexSessionCache()

	rc := &requestContext{
		clientProtocol:   protocol.Anthropic,
		upstreamProtocol: protocol.Codex,
		originalModel:    "gpt-5-codex",
		originalBody:     []byte(`{"metadata":{"user_id":"abc123"}}`),
	}
	id1 := resolveCodexSessionHint(rc, nil, "", nil)
	if !uuidPattern.MatchString(id1) {
		t.Fatalf("expected UUID, got %q", id1)
	}
	// 相同 user_id 再次调用应返回同一 UUID（命中缓存）
	id2 := resolveCodexSessionHint(rc, nil, "", nil)
	if id1 != id2 {
		t.Fatalf("expected cached session id, got %q vs %q", id1, id2)
	}
}

func TestResolveCodexSessionHint_AnthropicDifferentModelsOrUsers(t *testing.T) {
	resetCodexSessionCache()

	mkCtx := func(model, userID string) *requestContext {
		return &requestContext{
			clientProtocol:   protocol.Anthropic,
			upstreamProtocol: protocol.Codex,
			originalModel:    model,
			originalBody:     []byte(`{"metadata":{"user_id":"` + userID + `"}}`),
		}
	}
	idA := resolveCodexSessionHint(mkCtx("m1", "u1"), nil, "", nil)
	idB := resolveCodexSessionHint(mkCtx("m1", "u2"), nil, "", nil)
	idC := resolveCodexSessionHint(mkCtx("m2", "u1"), nil, "", nil)
	if idA == idB || idA == idC || idB == idC {
		t.Fatalf("expected distinct UUIDs for distinct buckets; got %s %s %s", idA, idB, idC)
	}
}

func TestResolveCodexSessionHint_AnthropicMissingUserID(t *testing.T) {
	resetCodexSessionCache()

	rc := &requestContext{
		clientProtocol:   protocol.Anthropic,
		upstreamProtocol: protocol.Codex,
		originalBody:     []byte(`{}`),
	}
	if got := resolveCodexSessionHint(rc, nil, "", nil); got != "" {
		t.Fatalf("expected empty when metadata.user_id, session header and apiKey all missing, got %q", got)
	}

	// 头 X-Claude-Code-Session-Id 存在时优先于 apiKey
	h := http.Header{}
	h.Set("X-Claude-Code-Session-Id", "sid-1")
	s1 := resolveCodexSessionHint(rc, nil, "sk-abc", h)
	s2 := resolveCodexSessionHint(rc, nil, "sk-xyz", h)
	s3 := resolveCodexSessionHint(rc, nil, "sk-abc", nil)
	if !uuidPattern.MatchString(s1) {
		t.Fatalf("expected UUID from session header, got %q", s1)
	}
	if s1 != s2 {
		t.Fatalf("expected session-id to dominate apiKey, got %q vs %q", s1, s2)
	}
	if s1 == s3 {
		t.Fatalf("session-id UUID should differ from apiKey UUID")
	}

	// Claude Code 客户端无 user_id 且无 session 头时 fallback 到 apiKey 稳定 UUID
	a1 := resolveCodexSessionHint(rc, nil, "sk-abc", nil)
	a2 := resolveCodexSessionHint(rc, nil, "sk-abc", nil)
	b := resolveCodexSessionHint(rc, nil, "sk-xyz", nil)
	if !uuidPattern.MatchString(a1) {
		t.Fatalf("expected UUID fallback, got %q", a1)
	}
	if a1 != a2 {
		t.Fatalf("expected deterministic fallback UUID, got %q vs %q", a1, a2)
	}
	if a1 == b {
		t.Fatalf("expected different UUIDs for different apiKeys")
	}
}

func TestResolveCodexSessionHint_CodexPassthrough(t *testing.T) {
	rc := &requestContext{
		clientProtocol:   protocol.Codex,
		upstreamProtocol: protocol.Codex,
	}
	body := []byte(`{"prompt_cache_key":"existing-uuid","model":"gpt-5-codex"}`)
	if got := resolveCodexSessionHint(rc, body, "", nil); got != "existing-uuid" {
		t.Fatalf("expected passthrough prompt_cache_key, got %q", got)
	}

	if got := resolveCodexSessionHint(rc, []byte(`{"model":"x"}`), "", nil); got != "" {
		t.Fatalf("expected empty when codex body has no prompt_cache_key, got %q", got)
	}
}

func TestResolveCodexSessionHint_OpenAIDeterministic(t *testing.T) {
	rc := &requestContext{
		clientProtocol:   protocol.OpenAI,
		upstreamProtocol: protocol.Codex,
	}
	a1 := resolveCodexSessionHint(rc, nil, "sk-abc", nil)
	a2 := resolveCodexSessionHint(rc, nil, "sk-abc", nil)
	b := resolveCodexSessionHint(rc, nil, "sk-xyz", nil)
	if a1 == "" || !uuidPattern.MatchString(a1) {
		t.Fatalf("expected UUID for openai key, got %q", a1)
	}
	if a1 != a2 {
		t.Fatalf("expected deterministic UUID for same key, got %q vs %q", a1, a2)
	}
	if a1 == b {
		t.Fatalf("expected different UUIDs for different keys")
	}
	if got := resolveCodexSessionHint(rc, nil, "", nil); got != "" {
		t.Fatalf("expected empty for empty apiKey, got %q", got)
	}
}

func TestResolveCodexSessionHint_NonCodexUpstream(t *testing.T) {
	rc := &requestContext{
		clientProtocol:   protocol.Anthropic,
		upstreamProtocol: protocol.Anthropic,
		originalBody:     []byte(`{"metadata":{"user_id":"abc"}}`),
	}
	if got := resolveCodexSessionHint(rc, nil, "", nil); got != "" {
		t.Fatalf("expected empty when upstream is not Codex, got %q", got)
	}
}

func TestInjectCodexPromptCacheKey(t *testing.T) {
	body := []byte(`{"model":"gpt-5-codex","input":[]}`)
	out := injectCodexPromptCacheKey(body, "deadbeef")
	if readCodexPromptCacheKey(out) != "deadbeef" {
		t.Fatalf("expected prompt_cache_key injected, got %s", out)
	}

	// 已存在非空值时不覆盖
	preset := []byte(`{"prompt_cache_key":"kept","model":"x"}`)
	out2 := injectCodexPromptCacheKey(preset, "new")
	if readCodexPromptCacheKey(out2) != "kept" {
		t.Fatalf("expected existing prompt_cache_key preserved, got %s", out2)
	}

	// 空 body / 空 id / 非 JSON 原样返回
	if got := injectCodexPromptCacheKey(nil, "x"); got != nil {
		t.Fatalf("expected nil body preserved, got %s", got)
	}
	if got := injectCodexPromptCacheKey(body, ""); string(got) != string(body) {
		t.Fatalf("expected body unchanged for empty id")
	}
	raw := []byte(`not json`)
	if got := injectCodexPromptCacheKey(raw, "x"); string(got) != string(raw) {
		t.Fatalf("expected non-json body unchanged")
	}
}

func TestInjectCodexPromptCacheKey_PreservesExistingFieldOrder(t *testing.T) {
	body := []byte(`{"model":"gpt-5-codex","instructions":"keep this prefix","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]}],"stream":true}`)

	out := injectCodexPromptCacheKey(body, "stable-session")
	got := string(out)

	for _, want := range []string{`"model"`, `"instructions"`, `"input"`, `"stream"`, `"prompt_cache_key"`} {
		if !strings.Contains(got, want) {
			t.Fatalf("expected %s in injected body: %s", want, got)
		}
	}
	assertFieldOrder(t, got, `"model"`, `"instructions"`, `"input"`, `"stream"`, `"prompt_cache_key"`)
	if !strings.HasPrefix(got, `{"model":"gpt-5-codex","instructions":"keep this prefix","input":`) {
		t.Fatalf("prompt_cache_key injection reordered cache-sensitive prefix: %s", got)
	}
}

func TestExtractAnthropicUserID(t *testing.T) {
	cases := []struct {
		name string
		body []byte
		want string
	}{
		{"happy", []byte(`{"metadata":{"user_id":"  abc123  "}}`), "abc123"},
		{"missing metadata", []byte(`{"model":"x"}`), ""},
		{"missing user_id", []byte(`{"metadata":{}}`), ""},
		{"empty body", nil, ""},
		{"invalid json", []byte(`not json`), ""},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if got := extractAnthropicUserID(tc.body); got != tc.want {
				t.Fatalf("want %q, got %q", tc.want, got)
			}
		})
	}
}

func TestBuildProxyRequest_CodexSessionInjection_Anthropic(t *testing.T) {
	resetCodexSessionCache()
	srv := newInMemoryServer(t)

	cfg := &model.Config{
		ID:          1,
		Name:        "codex-ch",
		URL:         "https://api.example.com",
		ChannelType: "openai",
	}

	originalBody := []byte(`{"metadata":{"user_id":"claude-code-user-42"}}`)
	translatedBody := []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]}]}`)

	reqCtx := &requestContext{
		ctx:              context.Background(),
		startTime:        time.Now(),
		clientProtocol:   protocol.Anthropic,
		upstreamProtocol: protocol.Codex,
		originalModel:    "gpt-5-codex",
		originalBody:     originalBody,
		translatedBody:   translatedBody,
	}

	req, err := srv.buildProxyRequest(
		reqCtx,
		cfg,
		"sk-test",
		http.MethodPost,
		translatedBody,
		http.Header{"Content-Type": []string{"application/json"}},
		"",
		"/v1/responses",
		cfg.URL,
	)
	if err != nil {
		t.Fatalf("buildProxyRequest failed: %v", err)
	}

	sid := req.Header.Get("Session_id")
	if !uuidPattern.MatchString(sid) {
		t.Fatalf("Session_id header missing or invalid: %q", sid)
	}

	bodyReader, _ := req.GetBody()
	defer func() { _ = bodyReader.Close() }()
	buf := make([]byte, 4096)
	n, _ := bodyReader.Read(buf)
	if key := readCodexPromptCacheKey(buf[:n]); key != sid {
		t.Fatalf("expected body prompt_cache_key == Session_id header, got body=%q header=%q", key, sid)
	}
}

func TestBuildProxyRequest_CodexSessionInjection_NonCodexUpstreamSkipped(t *testing.T) {
	resetCodexSessionCache()
	srv := newInMemoryServer(t)

	cfg := &model.Config{
		ID:          1,
		Name:        "anthropic-ch",
		URL:         "https://api.example.com",
		ChannelType: "anthropic",
	}

	reqCtx := &requestContext{
		ctx:              context.Background(),
		startTime:        time.Now(),
		clientProtocol:   protocol.Anthropic,
		upstreamProtocol: protocol.Anthropic,
		originalModel:    "claude-3",
		originalBody:     []byte(`{"metadata":{"user_id":"u1"}}`),
	}

	req, err := srv.buildProxyRequest(
		reqCtx,
		cfg,
		"sk-test",
		http.MethodPost,
		[]byte(`{"model":"claude-3"}`),
		http.Header{"Content-Type": []string{"application/json"}},
		"",
		"/v1/messages",
		cfg.URL,
	)
	if err != nil {
		t.Fatalf("buildProxyRequest failed: %v", err)
	}

	if got := req.Header.Get("Session_id"); got != "" {
		t.Fatalf("expected no Session_id for non-Codex upstream, got %q", got)
	}
}

func TestBuildProxyRequest_CodexSessionInjection_ClientHeaderNotOverwritten(t *testing.T) {
	resetCodexSessionCache()
	srv := newInMemoryServer(t)

	cfg := &model.Config{
		ID:          1,
		Name:        "codex-ch",
		URL:         "https://api.example.com",
		ChannelType: "openai",
	}

	reqCtx := &requestContext{
		ctx:              context.Background(),
		startTime:        time.Now(),
		clientProtocol:   protocol.Codex,
		upstreamProtocol: protocol.Codex,
		originalModel:    "gpt-5-codex",
		originalBody:     []byte(`{"prompt_cache_key":"client-supplied","model":"gpt-5-codex"}`),
	}

	req, err := srv.buildProxyRequest(
		reqCtx,
		cfg,
		"sk-test",
		http.MethodPost,
		[]byte(`{"prompt_cache_key":"client-supplied","model":"gpt-5-codex"}`),
		http.Header{
			"Content-Type": []string{"application/json"},
			"Session_id":   []string{"client-session"},
		},
		"",
		"/v1/responses",
		cfg.URL,
	)
	if err != nil {
		t.Fatalf("buildProxyRequest failed: %v", err)
	}

	if got := req.Header.Get("Session_id"); got != "client-session" {
		t.Fatalf("expected client Session_id preserved, got %q", got)
	}
}