Sakurasan
99829ef745
- Auth 加载的完整用户放进上下文，resolveUser 直接复用，省掉每请求 一次重复的 users 查询 - last_used_at 是纯展示字段，改为距上次更新满一分钟才写，高 QPS 下 每请求一次 DB 写降到每密钥每分钟至多一次
104 lines
4.4 KiB
Go
104 lines
4.4 KiB
Go
package proxy
|
||
|
||
import (
|
||
"encoding/json"
|
||
"testing"
|
||
)
|
||
|
||
// 复现线上火山方舟 qwen 流：data:{...} 无空格（省略 data: 后的空格）。
|
||
func TestSSEDataPayloadNoSpace(t *testing.T) {
|
||
lines := []string{
|
||
`data:{"message":{"model":"qwen3.8-flash","id":"msg_1","role":"assistant","type":"message","content":[],"usage":{"input_tokens":31626,"output_tokens":0}},"type":"message_start"}`,
|
||
`data:{"delta":{"type":"text_delta","text":"你好"},"type":"content_block_delta","index":0}`,
|
||
`data:{"delta":{"type":"text_delta","text":"!"},"type":"content_block_delta","index":0}`,
|
||
`data:{"delta":{"stop_reason":"end_turn"},"type":"message_delta","usage":{"cache_creation":{"ephemeral_5m_input_tokens":33065},"output_tokens":8,"cache_creation_input_tokens":33065,"input_tokens":8,"cache_read_input_tokens":0}}`,
|
||
}
|
||
var out string
|
||
for _, l := range lines {
|
||
out += sseContentText([]byte(l))
|
||
}
|
||
if out != "你好!" {
|
||
t.Fatalf("outputText=%q, want %q", out, "你好!")
|
||
}
|
||
// message_delta 的 delta.usage 应能提取（output_tokens=8）
|
||
u := scanUsage([]byte(lines[3]))
|
||
if u == nil {
|
||
t.Fatal("scanUsage returned nil for message_delta with usage")
|
||
}
|
||
var sh usageShape
|
||
if err := json.Unmarshal(u, &sh); err != nil {
|
||
t.Fatalf("unmarshal usage: %v", err)
|
||
}
|
||
if sh.OutputTokens != 8 {
|
||
t.Fatalf("output_tokens=%d, want 8", sh.OutputTokens)
|
||
}
|
||
}
|
||
|
||
// 兼容带空格的单 data: 行（标准 SSE）与 event:+data: 多行块。
|
||
func TestSSEDataPayloadSpacedAndMultiLine(t *testing.T) {
|
||
// 标准：data: {...}
|
||
if got := sseContentText([]byte(`data: {"delta":{"type":"text_delta","text":"hi"},"type":"content_block_delta","index":0}`)); got != "hi" {
|
||
t.Fatalf("spaced single line: got %q, want hi", got)
|
||
}
|
||
// 多行块：event: message_delta\ndata: {...}
|
||
block := []byte("event: message_delta\ndata: {\"delta\":{\"type\":\"text_delta\",\"text\":\"yo\"},\"type\":\"content_block_delta\",\"index\":0}\n")
|
||
if got := sseContentText(block); got != "yo" {
|
||
t.Fatalf("multiline block: got %q, want yo", got)
|
||
}
|
||
}
|
||
|
||
// 复现线上 qwen（dashscope）messages 流式缓存场景：
|
||
// message_start.usage.input_tokens 是总输入，message_delta.usage.input_tokens 是非缓存输入
|
||
// 且带 cache_read/cache_creation，是最终计费口径。合并后：
|
||
//
|
||
// in(落库)=input+cache_read+cache_creation，计价 in 只算非缓存部分。
|
||
//
|
||
// 此前 delta 的 input 覆盖 start 的 input 导致总输入丢失（31790 → 8）。
|
||
func TestUsageSinkMessageDeltaAuthoritative(t *testing.T) {
|
||
sink := &usageSink{}
|
||
// message_start：总输入 31790
|
||
start := json.RawMessage(`{"input_tokens":31790,"output_tokens":0}`)
|
||
sink.push(start)
|
||
if got := sink.us.InputTokens; got != 31790 {
|
||
t.Fatalf("after start: input=%d, want 31790", got)
|
||
}
|
||
// message_delta：非缓存输入 8 + 缓存写 33229（最终口径，整体替换）
|
||
delta := json.RawMessage(`{"output_tokens":8,"cache_creation_input_tokens":33229,"input_tokens":8,"cache_read_input_tokens":0}`)
|
||
sink.push(delta)
|
||
s := sink.Shape()
|
||
if s.InputTokens != 8 || s.CacheCreationInputTokens != 33229 || s.OutputTokens != 8 {
|
||
t.Fatalf("after delta: %+v, want input=8 cache_create=33229 output=8", s)
|
||
}
|
||
// finishUsage 口径：落库 input = 8 + 0 + 33229 = 33237（总量），计价 in=8、cacheCreate=33229
|
||
in := s.InputTokens + s.CacheReadInputTokens + s.CacheCreationInputTokens
|
||
if in != 33237 {
|
||
t.Fatalf("total input=%d, want 33237", in)
|
||
}
|
||
}
|
||
|
||
// 缓存命中场景（id=55）：delta input=76 非缓存 + cache_read=33229 + cache_creation=17。
|
||
func TestUsageSinkCacheHitMerge(t *testing.T) {
|
||
sink := &usageSink{}
|
||
sink.push(json.RawMessage(`{"input_tokens":31862,"output_tokens":0}`))
|
||
sink.push(json.RawMessage(`{"output_tokens":32,"cache_creation_input_tokens":17,"input_tokens":76,"cache_read_input_tokens":33229}`))
|
||
s := sink.Shape()
|
||
total := s.InputTokens + s.CacheReadInputTokens + s.CacheCreationInputTokens
|
||
if total != 33322 {
|
||
t.Fatalf("total input=%d, want 33322 (76+33229+17)", total)
|
||
}
|
||
if s.OutputTokens != 32 {
|
||
t.Fatalf("output=%d, want 32", s.OutputTokens)
|
||
}
|
||
}
|
||
|
||
// OpenAI chat 末块（无 cache 字段）仍走零值不覆盖合并，不受整体替换影响。
|
||
func TestUsageSinkChatLastChunkStillMerges(t *testing.T) {
|
||
sink := &usageSink{}
|
||
sink.push(json.RawMessage(`{"prompt_tokens":65,"completion_tokens":0}`))
|
||
sink.push(json.RawMessage(`{"prompt_tokens":65,"completion_tokens":82}`))
|
||
s := sink.Shape()
|
||
if s.PromptTokens != 65 || s.CompletionTokens != 82 {
|
||
t.Fatalf("chat merge broken: %+v", s)
|
||
}
|
||
}
|