openteam/server/internal/proxy/usage_refine_test.go

package proxy

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/openteam/server/internal/store"
	"github.com/shopspring/decimal"
)

// runPath 走完整链路（Auth → Handle），路径与请求体由调用方给。
func runPath(t *testing.T, h *harness, path, body string) *httptest.ResponseRecorder {
	t.Helper()
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(body))
	req.Header.Set("Authorization", "Bearer "+h.key)
	c.Request = req
	h.gw.Auth(c)
	if !c.IsAborted() {
		h.gw.Handle(c)
	}
	return w
}

// usage 细化：reasoning token 与上游实际应答模型落库；计费口径不变
// （reasoning 已含在 completion_tokens 内，不重复计费）。
func TestUsageReasoningAndServedModelRecorded(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"model":"gpt-x-actual","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,"completion_tokens_details":{"reasoning_tokens":3}}}`))
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 0, true, nil)
	w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK {
		t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
	}
	h.rec.Close()
	var logs []store.UsageLog
	h.db.Find(&logs)
	if len(logs) != 1 {
		t.Fatalf("应有一条记账，实际 %d", len(logs))
	}
	l := logs[0]
	if l.ServedModel != "gpt-x-actual" {
		t.Fatalf("served_model 应记录上游应答的模型名，实际 %q", l.ServedModel)
	}
	if l.ReasoningTokens != 3 {
		t.Fatalf("reasoning_tokens 应为 3，实际 %d", l.ReasoningTokens)
	}
	// 计费不含 reasoning 加成：10×0.15/1e6 + 5×0.6/1e6 = 4.5e-6
	if !l.Cost.Equal(decimal.NewFromFloat(4.5e-6)) {
		t.Fatalf("cost 应为 4.5e-6（reasoning 不重复计费），实际 %s", l.Cost)
	}
}

// Anthropic 缓存写细分：1h 与 5m 分开计价（1h = 2× 输入价，5m = 1.25×）。
func TestAnthropicCache1hBillingSplit(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"id":"msg_1","model":"claude-x","role":"assistant","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":100,"output_tokens":10,"cache_read_input_tokens":200,"cache_creation_input_tokens":300,"cache_creation":{"ephemeral_5m_input_tokens":100,"ephemeral_1h_input_tokens":200}}}`))
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 0, true, nil)
	// 渠道改为 Anthropic 形态，模型定价：输入 1、输出 2、缓存读 0.1、缓存写未配置（1.25× 输入）
	var ch1 store.Channel
	if err := h.db.Where("name = ?", "c1").First(&ch1).Error; err != nil {
		t.Fatal(err)
	}
	ch1.Formats = []string{store.FormatMessages}
	if err := h.db.Save(&ch1).Error; err != nil {
		t.Fatal(err)
	}
	if err := h.db.Model(&store.Model{}).Where("name = ?", "m1").
		Updates(map[string]any{"input_price": 1.0, "output_price": 2.0, "cache_read_price": 0.1}).Error; err != nil {
		t.Fatal(err)
	}

	w := runPath(t, h, "/v1/messages", `{"model":"m1","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK {
		t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
	}
	h.rec.Close()
	var logs []store.UsageLog
	h.db.Find(&logs)
	if len(logs) != 1 {
		t.Fatalf("应有一条记账，实际 %d", len(logs))
	}
	l := logs[0]
	if l.CacheWrite1hTokens != 200 {
		t.Fatalf("cache_write_1h_tokens 应为 200，实际 %d", l.CacheWrite1hTokens)
	}
	if l.ServedModel != "claude-x" {
		t.Fatalf("served_model 应为 claude-x，实际 %q", l.ServedModel)
	}
	// 100/1e6*1 + 10/1e6*2 + 200/1e6*0.1 + 100/1e6*1.25 + 200/1e6*(1.25*1.6) = 6.65e-4
	million := decimal.NewFromInt(1e6)
	want := decimal.NewFromInt(100).Div(million).Add(decimal.NewFromInt(10).Mul(decimal.NewFromFloat(2)).Div(million)).
		Add(decimal.NewFromInt(200).Mul(decimal.NewFromFloat(0.1)).Div(million)).
		Add(decimal.NewFromInt(100).Mul(decimal.NewFromFloat(1.25)).Div(million)).
		Add(decimal.NewFromInt(200).Mul(decimal.NewFromFloat(2.0)).Div(million))
	if !l.Cost.Equal(want) {
		t.Fatalf("cost 应为 %s（1h/5m 分开计价），实际 %s", want, l.Cost)
	}
	// 落库 input 为含缓存的总量口径
	if l.InputTokens != 600 || l.OutputTokens != 10 {
		t.Fatalf("in=%d out=%d，期望 600/10", l.InputTokens, l.OutputTokens)
	}
	if !strings.Contains(w.Body.String(), `"text":"hi"`) {
		t.Fatalf("响应应原样透传：%s", w.Body.String())
	}
}