ModelRouter/internal/lua/usage_cost_test.go

package lua

import (
	"encoding/json"
	"testing"
)

// OpenCode Zen/Go 在 usage 里报告输出有多少是思考 token，并在**顶层**用一个
// **字符串**给出本次费用。前者决定「completion_tokens 里多少是可见回答」，
// 后者是唯一的费用信号。两者都曾被适配器丢掉。
func TestOpenCodeCostAndReasoningPassthrough(t *testing.T) {
	for _, name := range []string{"opencodego", "opencodezen"} {
		vm := NewVM(freshAdapterDir(t))
		if err := vm.Start(); err != nil {
			t.Fatal(err)
		}

		// ---- 非流式 ----
		resp := `{"id":"x","object":"chat.completion","model":"deepseek-v4.1-flash",
			"choices":[{"index":0,"message":{"role":"assistant","content":"hi","reasoning_content":"think"},
			"finish_reason":"stop"}],
			"usage":{"prompt_tokens":37,"completion_tokens":40,"total_tokens":77,
			"prompt_cache_hit_tokens":0,"prompt_cache_miss_tokens":37,
			"prompt_tokens_details":{"cached_tokens":0},
			"completion_tokens_details":{"reasoning_tokens":40}},
			"cost":"0"}`
		out, err := vm.Transform(name, "transform_response", resp)
		if err != nil {
			t.Fatalf("%s transform_response: %v", name, err)
		}
		var r struct {
			TokenUsage struct {
				Completion int `json:"completion"`
				Details    *struct {
					ReasoningTokens int `json:"reasoning_tokens"`
				} `json:"completion_tokens_details"`
			} `json:"token_usage"`
			Cost string `json:"cost"`
		}
		if err := json.Unmarshal([]byte(out), &r); err != nil {
			t.Fatalf("%s unmarshal: %v (%s)", name, err, out)
		}
		if r.Cost != "0" {
			t.Errorf("%s: cost 未透传（上游给的是字符串 \"0\"）: %q", name, r.Cost)
		}
		if r.TokenUsage.Details == nil || r.TokenUsage.Details.ReasoningTokens != 40 {
			t.Errorf("%s: completion_tokens_details.reasoning_tokens 未透传: %s", name, out)
		}
		// 前提校验：reasoning_tokens 确实 ≤ completion_tokens（全为思考）。
		if r.TokenUsage.Details != nil && r.TokenUsage.Details.ReasoningTokens > r.TokenUsage.Completion {
			t.Errorf("%s: reasoning 不应超过 completion", name)
		}

		// ---- 流式：非零费用也要透传 ----
		chunk := `{"id":"x","object":"chat.completion.chunk","choices":[],
			"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,
			"completion_tokens_details":{"reasoning_tokens":4}},"cost":"0.0012"}`
		co, err := vm.Transform(name, "transform_stream_chunk", chunk)
		if err != nil {
			t.Fatalf("%s transform_stream_chunk: %v", name, err)
		}
		var c struct {
			Cost  string `json:"cost"`
			Usage *struct {
				Details *struct {
					ReasoningTokens int `json:"reasoning_tokens"`
				} `json:"completion_tokens_details"`
			} `json:"usage"`
		}
		if err := json.Unmarshal([]byte(co), &c); err != nil {
			t.Fatalf("%s unmarshal chunk: %v (%s)", name, err, co)
		}
		if c.Usage == nil || c.Usage.Details == nil || c.Usage.Details.ReasoningTokens != 4 {
			t.Errorf("%s: 流式 reasoning_tokens 未透传: %s", name, co)
		}
		if c.Cost != "0.0012" {
			t.Errorf("%s: 流式 cost 未透传: %s", name, co)
		}

		// ---- 独立的 cost 帧（上游就这么发：choices 为空、只有 cost）----
		only := `{"choices":[],"cost":"0"}`
		oo, err := vm.Transform(name, "transform_stream_chunk", only)
		if err != nil {
			t.Fatalf("%s cost-only chunk: %v", name, err)
		}
		if oo == "" {
			t.Fatalf("%s: 纯 cost 帧被整个丢弃（客户端将看不到费用）", name)
		}
		var oc struct {
			Cost string `json:"cost"`
		}
		if err := json.Unmarshal([]byte(oo), &oc); err != nil {
			t.Fatalf("%s unmarshal cost chunk: %v (%s)", name, err, oo)
		}
		if oc.Cost != "0" {
			t.Errorf("%s: 纯 cost 帧内容丢失: %s", name, oo)
		}

		// ---- 回归：没有 cost / 没有 details 时不得凭空造字段 ----
		bare := `{"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":"stop"}],
			"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`
		bo, _ := vm.Transform(name, "transform_response", bare)
		if got := bo; len(got) > 0 && containsAll(got, `"cost"`) {
			t.Errorf("%s: 上游没给 cost 却透传了该字段: %s", name, got)
		}

		vm.Stop()
	}
}

func containsAll(s string, subs ...string) bool {
	for _, sub := range subs {
		found := false
		for i := 0; i+len(sub) <= len(s); i++ {
			if s[i:i+len(sub)] == sub {
				found = true
				break
			}
		}
		if !found {
			return false
		}
	}
	return true
}

// 流式与非流式对同一上游切片的产出必须一致：命中缓存时上游同时给
// prompt_tokens_details.cached_tokens 和独立的 hit/miss 字段，两边都要透传。
func TestOpenCodeStreamCacheFieldsMatchNonStream(t *testing.T) {
	for _, name := range []string{"opencodego", "opencodezen"} {
		vm := NewVM(freshAdapterDir(t))
		if err := vm.Start(); err != nil {
			t.Fatal(err)
		}
		type cache struct {
			Details *struct {
				CachedTokens int `json:"cached_tokens"`
			} `json:"prompt_tokens_details"`
			Hit  int `json:"prompt_cache_hit_tokens"`
			Miss int `json:"prompt_cache_miss_tokens"`
		}
		usage := `"usage":{"prompt_tokens":5732,"completion_tokens":20,"total_tokens":5752,
			"prompt_cache_hit_tokens":5504,"prompt_cache_miss_tokens":228,
			"prompt_tokens_details":{"cached_tokens":5504}}`

		// 非流式
		nout, _ := vm.Transform(name, "transform_response",
			`{"choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],`+usage+`}`)
		var nr struct {
			U cache `json:"token_usage"`
		}
		if err := json.Unmarshal([]byte(nout), &nr); err != nil {
			t.Fatalf("%s: %v", name, err)
		}
		// 流式
		sout, _ := vm.Transform(name, "transform_stream_chunk",
			`{"choices":[],`+usage+`}`)
		var sr struct {
			U cache `json:"usage"`
		}
		if err := json.Unmarshal([]byte(sout), &sr); err != nil {
			t.Fatalf("%s chunk: %v (%s)", name, err, sout)
		}

		if nr.U.Hit != 5504 || nr.U.Miss != 228 {
			t.Errorf("%s: 非流式 hit/miss 异常: %+v", name, nr.U)
		}
		if sr.U.Hit != nr.U.Hit || sr.U.Miss != nr.U.Miss {
			t.Errorf("%s: 流式 hit/miss 与非流式不一致 (stream %d/%d vs non-stream %d/%d): %s",
				name, sr.U.Hit, sr.U.Miss, nr.U.Hit, nr.U.Miss, sout)
		}
		if sr.U.Details == nil || sr.U.Details.CachedTokens != 5504 {
			t.Errorf("%s: 流式 cached_tokens 丢失: %s", name, sout)
		}
		vm.Stop()
	}
}