Cyrene_v2/internal/runtime/model/usage_test.go

package model

import (
	"bytes"
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"strings"
	"testing"

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

// 本文件钉住 usage/缓存可观测（design/04 附录 A.3 第 4 条）：
//   - 两种线格式的缓存字段都解析进 Response，**缺字段是容错路径**（不 panic、不影响正文）；
//   - 每次模型调用后记一条结构化用量日志，字段齐全且**不含任何凭据**。

// TestOpenAIUsage_ParsesCacheAndReasoningTokens 钉住 OpenAI 兼容形状：DeepSeek 式的
// prompt_cache_hit_tokens / prompt_cache_miss_tokens 与 completion_tokens_details
// .reasoning_tokens 都进 Response。
func TestOpenAIUsage_ParsesCacheAndReasoningTokens(t *testing.T) {
	t.Parallel()
	body := `{
	  "id": "cmpl-1",
	  "model": "deepseek-chat",
	  "choices": [{"index":0,"message":{"role":"assistant","content":"你好"},"finish_reason":"stop"}],
	  "usage": {
	    "prompt_tokens": 1000, "completion_tokens": 200, "total_tokens": 1200,
	    "prompt_cache_hit_tokens": 900, "prompt_cache_miss_tokens": 100,
	    "completion_tokens_details": {"reasoning_tokens": 150}
	  }
	}`
	srv := newCompletionServer(t, http.StatusOK, "application/json", body, nil)
	ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "deepseek-chat"})

	resp, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{})
	if err != nil {
		t.Fatalf("Chat: %v", err)
	}
	if resp.PromptTokens != 1000 || resp.CompletionTokens != 200 || resp.TotalTokens != 1200 {
		t.Fatalf("用量 = (%d,%d,%d), want (1000,200,1200)", resp.PromptTokens, resp.CompletionTokens, resp.TotalTokens)
	}
	if resp.CacheHitTokens != 900 || resp.CacheMissTokens != 100 {
		t.Fatalf("缓存两档 = (%d,%d), want (900,100)", resp.CacheHitTokens, resp.CacheMissTokens)
	}
	if resp.ReasoningTokens != 150 {
		t.Fatalf("ReasoningTokens = %d, want 150", resp.ReasoningTokens)
	}
	if got := resp.CacheHitRate(); got != 0.9 {
		t.Fatalf("命中率 = %v, want 0.9", got)
	}
}

// TestOpenAIUsage_FallsBackToCachedTokensDetails 钉住另一条 OpenAI 兼容形态：
// 只有 prompt_tokens_details.cached_tokens（OpenAI / DashScope）时也能得到命中数，
// 且未命中数按提供商自己的不变式（prompt = hit + miss）补齐，而不是把命中率算成 100%。
func TestOpenAIUsage_FallsBackToCachedTokensDetails(t *testing.T) {
	t.Parallel()
	body := `{
	  "id": "cmpl-2", "model": "gpt-test",
	  "choices": [{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],
	  "usage": {"prompt_tokens": 400, "completion_tokens": 10, "total_tokens": 410,
	            "prompt_tokens_details": {"cached_tokens": 300}}
	}`
	srv := newCompletionServer(t, http.StatusOK, "application/json", body, nil)
	ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"})

	resp, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{})
	if err != nil {
		t.Fatalf("Chat: %v", err)
	}
	if resp.CacheHitTokens != 300 || resp.CacheMissTokens != 100 {
		t.Fatalf("缓存两档 = (%d,%d), want (300,100)", resp.CacheHitTokens, resp.CacheMissTokens)
	}
}

// TestUsage_MissingFieldsAreTolerated 钉住缺字段容错：用量块只有最基本的三个数时
// 不 panic、不报错，缓存命中数为零，未命中数按提供商自己的不变式（prompt = hit +
// miss）补齐 —— 也就是「这次输入全都没命中缓存」，命中率因此是 0 而不是编造值。
func TestUsage_MissingFieldsAreTolerated(t *testing.T) {
	t.Parallel()
	srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, nil)
	ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"})

	resp, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{})
	if err != nil {
		t.Fatalf("Chat: %v", err)
	}
	if resp.CacheHitTokens != 0 || resp.CacheMissTokens != 5 || resp.ReasoningTokens != 0 {
		t.Fatalf("缺失字段 = (%d,%d,%d), want (0,5,0)", resp.CacheHitTokens, resp.CacheMissTokens, resp.ReasoningTokens)
	}
	if resp.CacheHitRate() != 0 {
		t.Fatalf("无缓存命中时的命中率 = %v, want 0", resp.CacheHitRate())
	}
	if resp.Content != "Hello!" {
		t.Fatalf("正文受用量字段影响: %q", resp.Content)
	}
}

// TestOpenAIStreamUsage_ParsesCacheFields 钉住流式末块的用量块：缓存的 hit/miss 与
// 思维链 token 一样原样折进聚合结果（流式是最常用的路径，漏了它等于没做）。
func TestOpenAIStreamUsage_ParsesCacheFields(t *testing.T) {
	t.Parallel()
	lines := []string{
		`data: {"model":"deepseek-chat","choices":[{"index":0,"delta":{"content":"你好"}}]}`,
		`data: {"model":"deepseek-chat","choices":[],"usage":{"prompt_tokens":500,"completion_tokens":50,"total_tokens":550,"prompt_cache_hit_tokens":480,"prompt_cache_miss_tokens":20,"completion_tokens_details":{"reasoning_tokens":30}}}`,
		`data: [DONE]`,
	}
	srv := newCompletionServer(t, http.StatusOK, "text/event-stream", strings.Join(lines, "\n")+"\n", nil)
	ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "deepseek-chat"})

	resp, err := ad.ChatStream(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, nil)
	if err != nil {
		t.Fatalf("ChatStream: %v", err)
	}
	if resp.PromptTokens != 500 || resp.CacheHitTokens != 480 || resp.CacheMissTokens != 20 || resp.ReasoningTokens != 30 {
		t.Fatalf("流式用量 = %+v, want prompt 500 / hit 480 / miss 20 / reasoning 30", resp)
	}
	if resp.Content != "你好" {
		t.Fatalf("流式正文 = %q, want 你好", resp.Content)
	}
}

// TestAnthropicUsage_ParsesPromptCacheFields 钉住 Anthropic 形状：命中 =
// cache_read_input_tokens，未命中 = input_tokens + cache_creation_input_tokens
// （Anthropic 的 input_tokens 不含缓存读到的部分，只取它会把命中率算成 0）。
func TestAnthropicUsage_ParsesPromptCacheFields(t *testing.T) {
	t.Parallel()
	body := `{
	  "id": "msg_1", "model": "claude-test", "stop_reason": "end_turn",
	  "content": [{"type":"text","text":"你好"}],
	  "usage": {"input_tokens": 100, "output_tokens": 20,
	            "cache_creation_input_tokens": 300, "cache_read_input_tokens": 700}
	}`
	srv := newAnthropicServer(t, http.StatusOK, "application/json", body, nil)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	resp, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{})
	if err != nil {
		t.Fatalf("Chat: %v", err)
	}
	if resp.CacheHitTokens != 700 || resp.CacheMissTokens != 400 {
		t.Fatalf("缓存两档 = (%d,%d), want (700,400)", resp.CacheHitTokens, resp.CacheMissTokens)
	}
	// 既有口径不变：input_tokens 仍然只算 PromptTokens。
	if resp.PromptTokens != 100 || resp.TotalTokens != 120 {
		t.Fatalf("用量 = (%d,...,%d), want prompt 100 / total 120", resp.PromptTokens, resp.TotalTokens)
	}
	if resp.ReasoningTokens != 0 {
		t.Fatalf("Anthropic 不单独上报思维链 token，ReasoningTokens = %d, want 0", resp.ReasoningTokens)
	}
}

// TestAnthropicUsage_MissingCacheFieldsAreTolerated 钉住 Anthropic 的缺字段容错：
// 未启用提示缓存的端点只给 input/output 两个数，解析照常、缓存为零。
func TestAnthropicUsage_MissingCacheFieldsAreTolerated(t *testing.T) {
	t.Parallel()
	srv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, nil)
	ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})

	resp, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{})
	if err != nil {
		t.Fatalf("Chat: %v", err)
	}
	if resp.CacheHitTokens != 0 || resp.CacheMissTokens != 5 {
		// anthropicTextResponse 的 usage 是 input_tokens=5 / output_tokens=2：
		// 未命中数 = input + cache_creation = 5（缓存字段缺失时即为输入本身）。
		t.Fatalf("缓存两档 = (%d,%d), want (0,5)", resp.CacheHitTokens, resp.CacheMissTokens)
	}
}

// TestUsageLog_HasAllFieldsAndNoCredentials 钉住每次模型调用后那条结构化日志：
// prompt / completion / total / reasoning / cacheHit / cacheMiss / hitRate 一个不少，
// 且**绝不出现凭据**（apiKey / token）。
func TestUsageLog_HasAllFieldsAndNoCredentials(t *testing.T) {
	t.Parallel()
	const secret = "sk-super-secret-credential"
	body := `{
	  "id": "cmpl-1", "model": "deepseek-chat",
	  "choices": [{"index":0,"message":{"role":"assistant","content":"你好"},"finish_reason":"stop"}],
	  "usage": {"prompt_tokens": 1000, "completion_tokens": 200, "total_tokens": 1200,
	            "prompt_cache_hit_tokens": 900, "prompt_cache_miss_tokens": 100,
	            "completion_tokens_details": {"reasoning_tokens": 150}}
	}`
	srv := newCompletionServer(t, http.StatusOK, "application/json", body, nil)
	log, buf := captureJSONLog()
	ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, APIKey: secret, Model: "deepseek-chat"},
		WithOpenAILogger(log))

	if _, err := ad.Chat(context.Background(), Request{
		Purpose:  "chat",
		Messages: []Message{TextMessage(RoleUser, "hi")},
	}, Options{}); err != nil {
		t.Fatalf("Chat: %v", err)
	}

	rec := usageLogRecord(t, buf.String())
	if rec["msg"] != usageLogEvent {
		t.Fatalf("事件名 = %v, want %q", rec["msg"], usageLogEvent)
	}
	for field, want := range map[string]float64{
		"prompt": 1000, "completion": 200, "total": 1200, "reasoning": 150,
		"cacheHit": 900, "cacheMiss": 100, "hitRate": 0.9,
	} {
		if got, ok := rec[field].(float64); !ok || got != want {
			t.Fatalf("日志字段 %s = %v, want %v（完整记录：%v）", field, rec[field], want, rec)
		}
	}
	if rec["model"] != "deepseek-chat" || rec["purpose"] != "chat" {
		t.Fatalf("日志缺少模型/用途：%v", rec)
	}
	raw := buf.String()
	for _, credential := range []string{secret, "Authorization", "apiKey", "Bearer"} {
		if strings.Contains(raw, credential) {
			t.Fatalf("用量日志里出现了凭据相关内容 %q：%s", credential, raw)
		}
	}
}

// captureJSONLog 返回输出 JSON 记录的 logger 与缓冲：用量日志要按**字段**断言，
// 因此不能像 reasoning 用例那样只看文本子串。
func captureJSONLog() (*slog.Logger, *bytes.Buffer) {
	buf := &bytes.Buffer{}
	return slog.New(slog.NewJSONHandler(buf, &slog.HandlerOptions{Level: slog.LevelDebug})), buf
}

// usageLogRecord 从捕获的日志文本里取出唯一一条用量记录（找不到或不止一条就报错）。
func usageLogRecord(t *testing.T, raw string) map[string]any {
	t.Helper()
	var found []map[string]any
	for _, line := range strings.Split(strings.TrimSpace(raw), "\n") {
		line = strings.TrimSpace(line)
		if line == "" || !strings.Contains(line, `"msg":"`+usageLogEvent+`"`) {
			continue
		}
		rec := map[string]any{}
		if err := json.Unmarshal([]byte(line), &rec); err != nil {
			t.Fatalf("用量日志不是合法 JSON 行: %v（%q）", err, line)
		}
		found = append(found, rec)
	}
	if len(found) != 1 {
		t.Fatalf("用量日志条数 = %d, want 1（每次调用恰好一条）：%s", len(found), raw)
	}
	return found[0]
}