mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 00:48:00 +00:00
Anthropic and OpenAI disagree on what the prompt count means:
Anthropic: input_tokens EXCLUDES cached blocks; cache_read_input_tokens and
cache_creation_input_tokens are separate, additive, billed input.
OpenAI: prompt_tokens INCLUDES its cached_tokens subset.
anthropic.lua mapped input_tokens straight onto prompt, so a cache-heavy turn
was doubly wrong: the billed prompt was undercounted by the entire cache
portion, and cached_tokens could exceed prompt_tokens — a cache hit rate above
100% for any client that divides one by the other. cache_creation_input_tokens
was never read at all, so a cache-write turn silently lost those billed tokens.
Worse, the streaming path dropped the cache split entirely: message_delta
carries the FINAL usage and only mapped input/output, so every streamed
response reported no cache information even when the upstream sent it.
All three counts are now summed into prompt, with the read half exposed as
prompt_tokens_details.cached_tokens plus the DeepSeek-legacy hit/miss pair, via
one shared map_usage() used by transform_response, message_start and
message_delta. A reported zero stays distinguishable from "never reported": the
split is emitted whenever either cache field is present, and omitted entirely
when the upstream mentions neither (justwoker reports only input/output plus its
own cost fields, so its output is byte-identical to before). map_usage returns
nil for a countless object, preserving "no usage in this chunk means say
nothing" rather than reporting zeros.
message_start's placeholder count is still emitted: justwoker reports 160 there
and the real 6931 in message_delta, and the gateway's mergeUsage lets the later
non-zero value win.
2005 lines
67 KiB
Go
2005 lines
67 KiB
Go
package lua
|
||
|
||
import (
|
||
"encoding/json"
|
||
"os"
|
||
"path/filepath"
|
||
"regexp"
|
||
"strings"
|
||
"sync"
|
||
"testing"
|
||
"time"
|
||
)
|
||
|
||
// freshAdapterDir returns a path under a temp dir that does not exist yet, so
|
||
// VM.Start() treats it as first-run and seeds the bundled adapters.
|
||
func freshAdapterDir(t *testing.T) string {
|
||
t.Helper()
|
||
return filepath.Join(t.TempDir(), "adapters")
|
||
}
|
||
|
||
func TestLoadBundledAdapters(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
adapters := vm.ListAdapters()
|
||
if len(adapters) == 0 {
|
||
t.Fatal("no adapters loaded")
|
||
}
|
||
names := map[string]bool{}
|
||
for _, a := range adapters {
|
||
names[a.Name] = true
|
||
}
|
||
for _, want := range []string{"openai", "deepseek", "anthropic", "gemini", "ollama", "kimicode", "opencode"} {
|
||
if !names[want] {
|
||
t.Errorf("missing adapter %s (got %v)", want, names)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestFirstRunSeedsEveryBundledAdapter(t *testing.T) {
|
||
dir := filepath.Join(t.TempDir(), "adapters")
|
||
vm := NewVM(dir)
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
// every embedded .lua must have been copied out (and no extra, non-.lua files)
|
||
embedded, err := bundledAdapters.ReadDir("adapters")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
seeded := 0
|
||
for _, e := range embedded {
|
||
if e.IsDir() || !strings.HasSuffix(e.Name(), ".lua") {
|
||
continue
|
||
}
|
||
seeded++
|
||
if _, err := os.Stat(filepath.Join(dir, e.Name())); err != nil {
|
||
t.Errorf("embedded adapter %s not seeded", e.Name())
|
||
}
|
||
}
|
||
if seeded == 0 {
|
||
t.Fatal("no embedded adapters found")
|
||
}
|
||
written := 0
|
||
entries, _ := os.ReadDir(dir)
|
||
for _, e := range entries {
|
||
if !e.IsDir() && strings.HasSuffix(e.Name(), ".lua") {
|
||
written++
|
||
}
|
||
}
|
||
if written != seeded {
|
||
t.Fatalf("seeded %d files but %d embedded adapters exist", written, seeded)
|
||
}
|
||
// re-start with the existing dir: authoritative, never rewritten
|
||
vm2 := NewVM(dir)
|
||
if err := vm2.Start(); err != nil {
|
||
t.Fatalf("second start: %v", err)
|
||
}
|
||
defer vm2.Stop()
|
||
}
|
||
|
||
func TestTransformRequest(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
out, err := vm.Transform("openai", "transform_request", `{"model":"x","disable_thinking":true,"messages":[]}`)
|
||
if err != nil {
|
||
t.Fatalf("transform: %v", err)
|
||
}
|
||
if strings.Contains(out, "disable_thinking") {
|
||
t.Fatalf("disable_thinking not stripped: %s", out)
|
||
}
|
||
}
|
||
|
||
func TestBuildHeadersFallbackStatic(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
hdrs, err := vm.BuildHeaders("anthropic", nil)
|
||
if err != nil {
|
||
t.Fatalf("build headers: %v", err)
|
||
}
|
||
if hdrs["anthropic-version"] != "2023-06-01" {
|
||
t.Fatalf("static header missing: %v", hdrs)
|
||
}
|
||
}
|
||
|
||
func TestOpenCodeAdapterFingerprint(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
hdrs, err := vm.BuildHeaders("opencode", map[string]interface{}{
|
||
"api_key": "public",
|
||
"timestamp": 1700000000,
|
||
"body": `{"model":"x"}`,
|
||
"url": "https://opencode.ai/zen/v1/chat/completions",
|
||
"method": "POST",
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("build headers: %v", err)
|
||
}
|
||
// zen 按客户端指纹路由请求池:UA 必须是 opencode 真实格式,
|
||
// 且必须携带 x-opencode-* 身份头,否则带 tools 的请求会被上游
|
||
// 判为匿名客户端并返回 network_error 空流(表现为空回复)。
|
||
if ua := hdrs["User-Agent"]; !strings.HasPrefix(ua, "opencode/") ||
|
||
!strings.Contains(ua, "ai-sdk/provider-utils/") {
|
||
t.Fatalf("opencode adapter must send the real opencode client UA, got %q", ua)
|
||
}
|
||
for _, h := range []string{"x-opencode-client", "x-opencode-project", "x-opencode-session", "x-opencode-request"} {
|
||
if hdrs[h] == "" {
|
||
t.Fatalf("opencode adapter must send %q (zen client fingerprint), got %v", h, hdrs)
|
||
}
|
||
}
|
||
hdrs2, _ := vm.BuildHeaders("opencode", map[string]interface{}{"api_key": "public", "timestamp": 1700000001, "body": `{"model":"y"}`, "method": "POST"})
|
||
if hdrs["x-opencode-request"] == hdrs2["x-opencode-request"] {
|
||
t.Fatalf("x-opencode-request must vary per request, got %q twice", hdrs["x-opencode-request"])
|
||
}
|
||
if hdrs["Authorization"] != "" {
|
||
t.Fatalf("opencode adapter must not hardcode Authorization (config api_key supplies it), got %q", hdrs["Authorization"])
|
||
}
|
||
// passthrough behaves like the openai adapter
|
||
out, err := vm.Transform("opencode", "transform_request", `{"model":"x","disable_thinking":true,"extra_body":{},"messages":[{"role":"user","content":"hi"}]}`)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if strings.Contains(out, "disable_thinking") || strings.Contains(out, "extra_body") {
|
||
t.Fatalf("opencode transform_request must strip provider fields: %s", out)
|
||
}
|
||
}
|
||
|
||
func TestOpenCodeNormalizesRoles(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
// zen 上游只接受 system/user/assistant/tool/latest_reminder;
|
||
// OpenAI 系客户端发 `developer`(Claude Code 还会用 `function`)必须归一化成 system。
|
||
body := `{"model":"x","messages":[
|
||
{"role":"developer","content":"you are helpful"},
|
||
{"role":"system","content":"sys"},
|
||
{"role":"user","content":"hi"},
|
||
{"role":"assistant","content":"ok"},
|
||
{"role":"tool","tool_call_id":"t1","content":"{\"ok\":true}"}
|
||
]}`
|
||
out, err := vm.Transform("opencode", "transform_request", body)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if strings.Contains(out, `"developer"`) {
|
||
t.Fatalf("developer role must be normalized: %s", out)
|
||
}
|
||
if !strings.Contains(out, `"system"`) || !strings.Contains(out, `"user"`) {
|
||
t.Fatalf("allowed roles must survive: %s", out)
|
||
}
|
||
if !strings.Contains(out, `"assistant"`) {
|
||
t.Fatalf("assistant role must survive: %s", out)
|
||
}
|
||
if !strings.Contains(out, `"tool"`) {
|
||
t.Fatalf("tool role must survive: %s", out)
|
||
}
|
||
}
|
||
|
||
func TestOpenCodeStripsMultiModalParts(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
body := `{"model":"x","messages":[
|
||
{"role":"user","content":[
|
||
{"type":"text","text":"describe"},
|
||
{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}},
|
||
{"type":"input_audio","input_audio":{"data":"QQ","format":"wav"}}
|
||
]},
|
||
{"role":"user","content":[{"type":"image_url","image_url":{"url":"https://ex.com/a.png"}}]},
|
||
{"role":"assistant","content":"plain"}
|
||
]}`
|
||
out, err := vm.Transform("opencode", "transform_request", body)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if strings.Contains(out, "image_url") || strings.Contains(out, "input_audio") {
|
||
t.Fatalf("opencode must strip multimodal parts: %s", out)
|
||
}
|
||
if !strings.Contains(out, "describe") {
|
||
t.Fatalf("text parts must survive: %s", out)
|
||
}
|
||
// image-only message dropped: only the text message and plain assistant message remain
|
||
if got := strings.Count(out, `"role"`); got != 2 {
|
||
t.Fatalf("image-only message must be dropped, got %d messages: %s", got, out)
|
||
}
|
||
}
|
||
|
||
func TestBuildHeadersCustomHook(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
hdrs, err := vm.BuildHeaders("kimicode", map[string]interface{}{
|
||
"timestamp": int64(12345),
|
||
"api_key": "k",
|
||
"body": "{}",
|
||
"method": "POST",
|
||
"url": "http://x/chat",
|
||
"source": map[string]interface{}{"meta": map[string]interface{}{
|
||
"app_id": "app-9", "app_secret": "s", "api_key": "k",
|
||
}},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("build headers: %v", err)
|
||
}
|
||
if hdrs["X-App-Id"] != "app-9" {
|
||
t.Fatalf("x-app-id = %q", hdrs["X-App-Id"])
|
||
}
|
||
if hdrs["X-App-Sign"] == "" {
|
||
t.Fatal("expected signature header")
|
||
}
|
||
if hdrs["X-Timestamp"] != "12345" {
|
||
t.Fatalf("timestamp = %q", hdrs["X-Timestamp"])
|
||
}
|
||
}
|
||
|
||
func TestDisableThinkingPassthrough(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
body := `{"model":"x","disable_thinking":true,"messages":[{"role":"user","content":"hi"}]}`
|
||
out, err := vm.Transform("deepseek", "transform_request", body)
|
||
if err != nil {
|
||
t.Fatalf("deepseek transform: %v", err)
|
||
}
|
||
var req struct {
|
||
ExtraBody map[string]interface{} `json:"extra_body"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &req); err != nil {
|
||
t.Fatalf("unmarshal: %v (%s)", err, out)
|
||
}
|
||
thinking, ok := req.ExtraBody["thinking"].(map[string]interface{})
|
||
if !ok {
|
||
t.Fatalf("deepseek should emit extra_body.thinking on disable_thinking: %s", out)
|
||
}
|
||
if thinking["type"] != "disabled" {
|
||
t.Fatalf("thinking.type = %v", thinking["type"])
|
||
}
|
||
|
||
// anthropic: thinking is opt-in (Issue 3) — never emitted by default, and
|
||
// disable_thinking is not a trigger either.
|
||
out2, err := vm.Transform("anthropic", "transform_request", body)
|
||
if err != nil {
|
||
t.Fatalf("anthropic transform: %v", err)
|
||
}
|
||
if strings.Contains(out2, "thinking") {
|
||
t.Fatalf("anthropic should drop thinking when disable_thinking: %s", out2)
|
||
}
|
||
out3, err := vm.Transform("anthropic", "transform_request", `{"model":"x","messages":[{"role":"user","content":"hi"}]}`)
|
||
if err != nil {
|
||
t.Fatalf("anthropic transform: %v", err)
|
||
}
|
||
if strings.Contains(out3, "thinking") {
|
||
t.Fatalf("anthropic must not enable thinking by default: %s", out3)
|
||
}
|
||
// opt-in path: extra_body.thinking is forwarded verbatim
|
||
out4, err := vm.Transform("anthropic", "transform_request",
|
||
`{"model":"x","messages":[{"role":"user","content":"hi"}],"extra_body":{"thinking":{"type":"enabled","budget_tokens":2048}}}`)
|
||
if err != nil {
|
||
t.Fatalf("anthropic transform: %v", err)
|
||
}
|
||
if !strings.Contains(out4, `"type":"enabled"`) || !strings.Contains(out4, "2048") {
|
||
t.Fatalf("anthropic should forward extra_body.thinking: %s", out4)
|
||
}
|
||
}
|
||
|
||
func TestMultimodalTransform(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
body := `{"model":"x","messages":[{"role":"user","content":[
|
||
{"type":"text","text":"what is this?"},
|
||
{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}},
|
||
{"type":"image_url","image_url":{"url":"https://ex.com/a.png"}}
|
||
]}]}`
|
||
|
||
// anthropic: image_url -> image block (base64/url), text preserved
|
||
out, err := vm.Transform("anthropic", "transform_request", body)
|
||
if err != nil {
|
||
t.Fatalf("anthropic: %v", err)
|
||
}
|
||
for _, want := range []string{`"media_type":"image/png"`, `"data":"QUJD"`, `"type":"url","url":"https://ex.com/a.png"`, `"what is this?"`} {
|
||
if !strings.Contains(out, want) {
|
||
t.Fatalf("anthropic multimodal missing %s: %s", want, out)
|
||
}
|
||
}
|
||
|
||
// gemini: image_url -> inline_data
|
||
gout, err := vm.Transform("gemini", "transform_request", body)
|
||
if err != nil {
|
||
t.Fatalf("gemini: %v", err)
|
||
}
|
||
var g struct {
|
||
Contents []struct {
|
||
Parts []map[string]interface{} `json:"parts"`
|
||
} `json:"contents"`
|
||
}
|
||
if err := json.Unmarshal([]byte(gout), &g); err != nil {
|
||
t.Fatalf("gemini unmarshal: %v", err)
|
||
}
|
||
if len(g.Contents) == 0 {
|
||
t.Fatalf("gemini no contents")
|
||
}
|
||
var found bool
|
||
for _, p := range g.Contents[0].Parts {
|
||
if v, ok := p["inline_data"].(map[string]interface{}); ok && v["data"] == "QUJD" && v["mime_type"] == "image/png" {
|
||
found = true
|
||
}
|
||
}
|
||
if !found {
|
||
t.Fatalf("gemini missing inline_data image: %s", gout)
|
||
}
|
||
|
||
// ollama: image_url -> images base64 array
|
||
out, err = vm.Transform("ollama", "transform_request", body)
|
||
if err != nil {
|
||
t.Fatalf("ollama: %v", err)
|
||
}
|
||
if !strings.Contains(out, `"images":["QUJD"]`) {
|
||
t.Fatalf("ollama multimodal missing images: %s", out)
|
||
}
|
||
|
||
// openai passthrough keeps the content array intact
|
||
po, _ := vm.Transform("openai", "transform_request", body)
|
||
if !strings.Contains(po, `"image_url"`) || !strings.Contains(po, `,QUJD"`) {
|
||
t.Fatalf("openai passthrough lost multimodal content: %s", po)
|
||
}
|
||
}
|
||
|
||
func TestOpenCodeAdapterNormalizesDeveloperRole(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
raw := `{"model":"zen","messages":[` +
|
||
`{"role":"developer","content":"be concise"},` +
|
||
`{"role":"user","content":"hi"},` +
|
||
`{"role":"assistant","content":"hello"},` +
|
||
`{"role":"function","content":"{\"a\":1}"},` +
|
||
`{"role":"latest_reminder","content":"remind"}]}`
|
||
out, err := vm.Transform("opencode", "transform_request", raw)
|
||
if err != nil {
|
||
t.Fatalf("transform_request: %v", err)
|
||
}
|
||
var req struct {
|
||
Messages []struct {
|
||
Role string `json:"role"`
|
||
} `json:"messages"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &req); err != nil {
|
||
t.Fatalf("output not JSON: %v\n%s", err, out)
|
||
}
|
||
want := []string{"system", "user", "assistant", "system", "latest_reminder"}
|
||
if len(req.Messages) != len(want) {
|
||
t.Fatalf("got %d messages, want %d: %s", len(req.Messages), len(want), out)
|
||
}
|
||
for i, w := range want {
|
||
if req.Messages[i].Role != w {
|
||
t.Errorf("message %d role = %q, want %q", i, req.Messages[i].Role, w)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestOpenCodeAdapterKeepsWhitelistedRolesAndDropsMultimodal(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
raw := `{"messages":[` +
|
||
`{"role":"system","content":"sys"},` +
|
||
`{"role":"user","content":[{"type":"text","text":"keep"},{"type":"image_url","url":"x"}]},` +
|
||
`{"role":"user","content":[{"type":"image_url","url":"x"}]}]}`
|
||
out, err := vm.Transform("opencode", "transform_request", raw)
|
||
if err != nil {
|
||
t.Fatalf("transform_request: %v", err)
|
||
}
|
||
var req struct {
|
||
Messages []struct {
|
||
Role string `json:"role"`
|
||
Content any `json:"content"`
|
||
} `json:"messages"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &req); err != nil {
|
||
t.Fatalf("output not JSON: %v\n%s", err, out)
|
||
}
|
||
if len(req.Messages) != 2 {
|
||
t.Fatalf("got %d messages, want 2 (multimodal-only dropped): %s", len(req.Messages), out)
|
||
}
|
||
if req.Messages[0].Role != "system" || req.Messages[1].Role != "user" {
|
||
t.Fatalf("unexpected roles: %s", out)
|
||
}
|
||
}
|
||
|
||
func TestAdaptersPassFinishReason(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
// OpenAI-style chunk with a real finish reason must surface finish_reason
|
||
// and done=true; an empty-string finish_reason (sensenova sends "" on
|
||
// every chunk) must NOT terminate the stream.
|
||
chunk := `{"choices":[{"index":0,"finish_reason":"tool_calls","delta":{"content":""}}]}`
|
||
empty := `{"choices":[{"index":0,"finish_reason":"","delta":{"content":"x"}}]}`
|
||
for _, name := range []string{"openai", "deepseek", "github", "groq", "kimicode", "mistral", "opencode"} {
|
||
out, err := vm.Transform(name, "transform_stream_chunk", chunk)
|
||
if err != nil {
|
||
t.Fatalf("%s: %v", name, err)
|
||
}
|
||
if !strings.Contains(out, `"finish_reason":"tool_calls"`) {
|
||
t.Fatalf("%s: finish_reason lost: %s", name, out)
|
||
}
|
||
if !strings.Contains(out, `"done":true`) {
|
||
t.Fatalf("%s: done not set on real finish: %s", name, out)
|
||
}
|
||
out, err = vm.Transform(name, "transform_stream_chunk", empty)
|
||
if err != nil {
|
||
t.Fatalf("%s: %v", name, err)
|
||
}
|
||
if strings.Contains(out, `"done":true`) {
|
||
t.Fatalf("%s: empty finish_reason must not end stream: %s", name, out)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestAnthropicToolRoundTrip verifies the OpenAI->Anthropic request mapping for
|
||
// a full agent tool-call round (Issue 1): assistant tool_calls become
|
||
// tool_use content blocks and role:"tool" results become user tool_result
|
||
// blocks merged into ONE user message. This is the regression that made agent
|
||
// clients (dsh/Claude Code/Cursor) repeatedly re-invoke the same tool.
|
||
func TestAnthropicToolRoundTrip(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
round1 := `{"model":"x","tools":[{"type":"function","function":{"name":"calc","description":"multiply","parameters":{"type":"object","properties":{"a":{"type":"integer"},"b":{"type":"integer"}},"required":["a","b"]}}}],"messages":[{"role":"user","content":"what is 17*23?"}]}`
|
||
out1, err := vm.Transform("anthropic", "transform_request", round1)
|
||
if err != nil {
|
||
t.Fatalf("round1 transform: %v", err)
|
||
}
|
||
var r1 struct {
|
||
Tools []struct {
|
||
Name string `json:"name"`
|
||
InputSchema map[string]interface{} `json:"input_schema"`
|
||
} `json:"tools"`
|
||
Messages []map[string]interface{} `json:"messages"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out1), &r1); err != nil {
|
||
t.Fatalf("unmarshal r1: %v (%s)", err, out1)
|
||
}
|
||
if len(r1.Tools) != 1 || r1.Tools[0].Name != "calc" {
|
||
t.Fatalf("tools not mapped: %s", out1)
|
||
}
|
||
|
||
// Round 2: assistant tool_calls + tool result
|
||
round2 := `{"model":"x","messages":[
|
||
{"role":"user","content":"what is 17*23?"},
|
||
{"role":"assistant","content":"","tool_calls":[
|
||
{"id":"call_1","type":"function","function":{"name":"calc","arguments":"{\"a\":17,\"b\":23}"}}
|
||
]},
|
||
{"role":"tool","tool_call_id":"call_1","content":"391"}
|
||
]}`
|
||
out2, err := vm.Transform("anthropic", "transform_request", round2)
|
||
if err != nil {
|
||
t.Fatalf("round2 transform: %v", err)
|
||
}
|
||
var r2 struct {
|
||
Messages []struct {
|
||
Role string `json:"role"`
|
||
Content []struct {
|
||
Type string `json:"type"`
|
||
ID string `json:"id"`
|
||
Name string `json:"name"`
|
||
Input map[string]interface{} `json:"input"`
|
||
ToolUseID string `json:"tool_use_id"`
|
||
ContentText string `json:"content"`
|
||
} `json:"content"`
|
||
} `json:"messages"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out2), &r2); err != nil {
|
||
t.Fatalf("unmarshal r2: %v (%s)", err, out2)
|
||
}
|
||
// Message 2 (index 1) must be assistant with one tool_use block
|
||
am := r2.Messages[1]
|
||
if am.Role != "assistant" {
|
||
t.Fatalf("msg[1].role = %q, want assistant", am.Role)
|
||
}
|
||
var toolUse struct {
|
||
Type string `json:"type"`
|
||
ID string `json:"id"`
|
||
Name string `json:"name"`
|
||
Input map[string]interface{} `json:"input"`
|
||
}
|
||
for _, b := range am.Content {
|
||
if b.Type == "tool_use" {
|
||
toolUse = struct {
|
||
Type string `json:"type"`
|
||
ID string `json:"id"`
|
||
Name string `json:"name"`
|
||
Input map[string]interface{} `json:"input"`
|
||
}{b.Type, b.ID, b.Name, b.Input}
|
||
}
|
||
}
|
||
if toolUse.ID != "call_1" || toolUse.Name != "calc" {
|
||
t.Fatalf("tool_use not mapped: %+v", am.Content)
|
||
}
|
||
if toolUse.Input["a"] != float64(17) || toolUse.Input["b"] != float64(23) {
|
||
t.Fatalf("tool_use input args not decoded from JSON string: %+v", toolUse.Input)
|
||
}
|
||
// Message 3 (index 2) must be user with one tool_result block
|
||
um := r2.Messages[2]
|
||
if um.Role != "user" {
|
||
t.Fatalf("msg[2].role = %q, want user (tool_result)", um.Role)
|
||
}
|
||
if len(um.Content) != 1 || um.Content[0].Type != "tool_result" || um.Content[0].ToolUseID != "call_1" {
|
||
t.Fatalf("tool_result not mapped: %+v", um.Content)
|
||
}
|
||
|
||
// Round 3: consecutive tool results must merge into ONE user message
|
||
round3 := `{"model":"x","messages":[
|
||
{"role":"user","content":"do both"},
|
||
{"role":"assistant","content":"","tool_calls":[
|
||
{"id":"c1","type":"function","function":{"name":"calc","arguments":"{\"a\":1,\"b\":2}"}},
|
||
{"id":"c2","type":"function","function":{"name":"calc","arguments":"{\"a\":3,\"b\":4}"}}
|
||
]},
|
||
{"role":"tool","tool_call_id":"c1","content":"2"},
|
||
{"role":"tool","tool_call_id":"c2","content":"12"}
|
||
]}`
|
||
out3, err := vm.Transform("anthropic", "transform_request", round3)
|
||
if err != nil {
|
||
t.Fatalf("round3 transform: %v", err)
|
||
}
|
||
var r3 struct {
|
||
Messages []map[string]interface{} `json:"messages"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out3), &r3); err != nil {
|
||
t.Fatalf("unmarshal r3: %v (%s)", err, out3)
|
||
}
|
||
// messages: user, assistant(tool_use x2), user(tool_result x2 merged)
|
||
if len(r3.Messages) != 3 {
|
||
t.Fatalf("round3 len(messages) = %d, want 3 (merged tool results): %s", len(r3.Messages), out3)
|
||
}
|
||
last := r3.Messages[3-1]
|
||
blocks := last["content"].([]interface{})
|
||
if len(blocks) != 2 {
|
||
t.Fatalf("last user message content blocks = %d, want 2 merged tool_results: %s", len(blocks), out3)
|
||
}
|
||
}
|
||
|
||
// TestOpenAITransformErrorEnvelopes covers Issue 6: the openai adapter must
|
||
// condense every shape of upstream error body that real OpenAI-compatible
|
||
// gateways emit, not just the {error:{message}} envelope. Unhandled shapes
|
||
// used to fall through to Go's generic "unknown error" and dump raw HTML /
|
||
// JSON into the server log.
|
||
func TestOpenAITransformErrorEnvelopes(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
cases := []struct {
|
||
name string
|
||
status int
|
||
body string
|
||
want string
|
||
}{
|
||
{"standard openai envelope", 429,
|
||
`{"error":{"message":"Rate limit reached","type":"rate_limit"}}`,
|
||
"Rate limit reached"},
|
||
{"error as bare string", 400,
|
||
`{"error":"bad request"}`,
|
||
"bad request"},
|
||
{"flat numeric code (qijiar/siliconflow style)", 400,
|
||
`{"code":20012,"message":"Model does not exist. Please check it carefully.","data":null}`,
|
||
"20012: Model does not exist. Please check it carefully."},
|
||
{"flat string code (remotezen style)", 401,
|
||
`{"code":"INVALID_API_KEY","message":"Invalid API key"}`,
|
||
"INVALID_API_KEY: Invalid API key"},
|
||
{"fastapi detail", 422,
|
||
`{"detail":"validation failed"}`,
|
||
"validation failed"},
|
||
{"nginx html error page", 413,
|
||
`<html> <head><title>413 Request Entity Too Large</title></head> <body> <center><h1>413 Request Entity Too Large</h1></center> <hr><center>nginx/1.18.0 (Ubuntu)</center> </body> </html>`,
|
||
"413 Request Entity Too Large"},
|
||
{"plain text body", 502,
|
||
"upstream connect error",
|
||
"upstream connect error"},
|
||
}
|
||
for _, tc := range cases {
|
||
t.Run(tc.name, func(t *testing.T) {
|
||
got, ok, err := vm.TransformError("openai", tc.status, tc.body)
|
||
if err != nil {
|
||
t.Fatalf("TransformError: %v", err)
|
||
}
|
||
if !ok {
|
||
t.Fatalf("hook returned no reason for %s body: %s", tc.name, tc.body)
|
||
}
|
||
if got != tc.want {
|
||
t.Fatalf("reason = %q, want %q", got, tc.want)
|
||
}
|
||
})
|
||
}
|
||
|
||
// A body carrying no usable message must fall through (ok=false) so the
|
||
// Go-side generic condenser stays in charge instead of inventing text.
|
||
t.Run("no usable message falls through", func(t *testing.T) {
|
||
if _, ok, _ := vm.TransformError("openai", 500, `{"foo":"bar"}`); ok {
|
||
t.Fatal("expected fallthrough for a body with no message field")
|
||
}
|
||
})
|
||
}
|
||
|
||
// TestAdaptersReportZeroCacheHit covers the "distinguish missed from not
|
||
// reported" contract on BOTH adapter paths: whenever an upstream reports a
|
||
// cache field, the adapter must emit prompt_tokens_details even when the hit
|
||
// count is 0, so the gateway can record cache_reported=true and the UI shows
|
||
// 0% instead of "—". Adapters that dropped the 0 case made a reported miss
|
||
// indistinguishable from an upstream that never reported cache info.
|
||
func TestAdaptersReportZeroCacheHit(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
// OpenAI-shaped upstreams: usage.prompt_tokens_details.cached_tokens = 0
|
||
openaiLike := []string{"openai", "deepseek", "sensenova", "opencode",
|
||
"agentrouter", "github", "groq", "kimicode", "mistral"}
|
||
respBody := `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],
|
||
"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12,
|
||
"prompt_tokens_details":{"cached_tokens":0}}}`
|
||
streamBody := `{"choices":[],"usage":{"prompt_tokens":10,"completion_tokens":2,
|
||
"total_tokens":12,"prompt_tokens_details":{"cached_tokens":0}}}`
|
||
for _, name := range openaiLike {
|
||
out, err := vm.Transform(name, "transform_response", respBody)
|
||
if err != nil {
|
||
t.Fatalf("%s transform_response: %v", name, err)
|
||
}
|
||
if !strings.Contains(out, "prompt_tokens_details") {
|
||
t.Fatalf("%s: zero cached_tokens dropped in transform_response: %s", name, out)
|
||
}
|
||
out, err = vm.Transform(name, "transform_stream_chunk", streamBody)
|
||
if err != nil {
|
||
t.Fatalf("%s transform_stream_chunk: %v", name, err)
|
||
}
|
||
if !strings.Contains(out, "prompt_tokens_details") {
|
||
t.Fatalf("%s: zero cached_tokens dropped in transform_stream_chunk: %s", name, out)
|
||
}
|
||
}
|
||
|
||
// gemini: usageMetadata.cachedContentTokenCount = 0
|
||
gResp := `{"candidates":[{"content":{"parts":[{"text":"hi"}]},"finishReason":"STOP"}],
|
||
"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,
|
||
"totalTokenCount":12,"cachedContentTokenCount":0}}`
|
||
out, err := vm.Transform("gemini", "transform_response", gResp)
|
||
if err != nil {
|
||
t.Fatalf("gemini transform_response: %v", err)
|
||
}
|
||
if !strings.Contains(out, "prompt_tokens_details") {
|
||
t.Fatalf("gemini: cachedContentTokenCount not normalized in transform_response: %s", out)
|
||
}
|
||
gStream := `{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,
|
||
"totalTokenCount":12,"cachedContentTokenCount":0}}`
|
||
out, err = vm.Transform("gemini", "transform_stream_chunk", gStream)
|
||
if err != nil {
|
||
t.Fatalf("gemini transform_stream_chunk: %v", err)
|
||
}
|
||
if !strings.Contains(out, "prompt_tokens_details") {
|
||
t.Fatalf("gemini: zero cachedContentTokenCount dropped in stream: %s", out)
|
||
}
|
||
|
||
// anthropic: usage.cache_read_input_tokens = 0
|
||
aResp := `{"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn",
|
||
"usage":{"input_tokens":10,"output_tokens":2,"cache_read_input_tokens":0}}`
|
||
out, err = vm.Transform("anthropic", "transform_response", aResp)
|
||
if err != nil {
|
||
t.Fatalf("anthropic transform_response: %v", err)
|
||
}
|
||
if !strings.Contains(out, "prompt_tokens_details") {
|
||
t.Fatalf("anthropic: zero cache_read_input_tokens dropped in response: %s", out)
|
||
}
|
||
aStream := `{"type":"message_start","message":{"usage":{"input_tokens":10,
|
||
"output_tokens":2,"cache_read_input_tokens":0}}}`
|
||
out, err = vm.Transform("anthropic", "transform_stream_chunk", aStream)
|
||
if err != nil {
|
||
t.Fatalf("anthropic transform_stream_chunk: %v", err)
|
||
}
|
||
if !strings.Contains(out, "prompt_tokens_details") {
|
||
t.Fatalf("anthropic: zero cache_read_input_tokens dropped in stream: %s", out)
|
||
}
|
||
|
||
// Negative control: an upstream that reports NO cache field at all must
|
||
// not fabricate details (that would flip "not reported" into a fake 0%).
|
||
noCache := `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],
|
||
"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`
|
||
out, err = vm.Transform("openai", "transform_response", noCache)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if strings.Contains(out, "prompt_tokens_details") {
|
||
t.Fatalf("openai fabricated cache details when upstream reported none: %s", out)
|
||
}
|
||
}
|
||
|
||
// ---- elastic pool sizing (plan 阶段 2) ----
|
||
|
||
// poolOf reaches into the VM for one adapter's pool. Tests assert on internal
|
||
// counters because the whole point of the change is that they no longer grow
|
||
// monotonically.
|
||
func poolOf(t *testing.T, v *VM, name string) *adapterPool {
|
||
t.Helper()
|
||
p := v.pool(name)
|
||
if p == nil {
|
||
t.Fatalf("adapter %q not loaded", name)
|
||
}
|
||
return p
|
||
}
|
||
|
||
func poolCounts(p *adapterPool) (created, idle, inUse int) {
|
||
p.mu.Lock()
|
||
defer p.mu.Unlock()
|
||
return p.created, len(p.idle), p.inUse
|
||
}
|
||
|
||
// TestPoolStartsEmpty is the startup-memory claim: loading adapters must not
|
||
// boot a single Lua state, no matter how large the configured ceiling is.
|
||
func TestPoolStartsEmpty(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 64, "deepseek": 32})
|
||
for _, st := range vm.PoolStats() {
|
||
if st.Created != 0 || st.Idle != 0 || st.InUse != 0 {
|
||
t.Fatalf("adapter %q booted eagerly: created=%d idle=%d in_use=%d",
|
||
st.Name, st.Created, st.Idle, st.InUse)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestGrowStepFollowsMaxConcurrent pins requirement 3: the growth step is
|
||
// derived from the adapter's maximum concurrency.
|
||
func TestGrowStepFollowsMaxConcurrent(t *testing.T) {
|
||
p := newAdapterPool("t", "return {name='t'}", staticInfo{name: "t"})
|
||
for _, tc := range []struct{ max, want int }{
|
||
{1, 1}, {4, 1}, {8, 1}, {9, 2}, {16, 2}, {32, 4}, {64, 8}, {200, 8},
|
||
} {
|
||
p.setMax(tc.max)
|
||
p.mu.Lock()
|
||
got := p.growStepLocked()
|
||
p.mu.Unlock()
|
||
if got != tc.want {
|
||
t.Fatalf("max=%d: grow step = %d, want %d", tc.max, got, tc.want)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestShrinkStepFollowsConnections pins requirement 4: the shrink step is
|
||
// driven by the live connection count — no connections means collapse, busy
|
||
// means give up one state at a time.
|
||
func TestShrinkStepFollowsConnections(t *testing.T) {
|
||
p := newAdapterPool("t", "return {name='t'}", staticInfo{name: "t"})
|
||
p.setMax(32)
|
||
p.used = true
|
||
|
||
cases := []struct {
|
||
created, idle, inUse, want int
|
||
why string
|
||
}{
|
||
{created: 9, idle: 9, inUse: 0, want: 8, why: "no connections: collapse to the resident floor in one round"},
|
||
{created: 9, idle: 8, inUse: 1, want: 4, why: "1 connection: halve the slack"},
|
||
{created: 9, idle: 6, inUse: 3, want: 2, why: "3 connections: quarter the slack"},
|
||
{created: 9, idle: 2, inUse: 7, want: 1, why: "busy: release one state at a time"},
|
||
{created: 4, idle: 0, inUse: 4, want: 0, why: "everything checked out: nothing to reclaim"},
|
||
{created: 1, idle: 1, inUse: 0, want: 0, why: "at the resident floor: keep the warm state"},
|
||
}
|
||
for _, tc := range cases {
|
||
p.mu.Lock()
|
||
p.created, p.inUse = tc.created, tc.inUse
|
||
p.idle = make([]*worker, tc.idle)
|
||
got := p.shrinkStepLocked()
|
||
p.mu.Unlock()
|
||
if got != tc.want {
|
||
t.Fatalf("%s: created=%d idle=%d inUse=%d -> step %d, want %d",
|
||
tc.why, tc.created, tc.idle, tc.inUse, got, tc.want)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestPoolGrowsOnDemandAndReclaims is the end-to-end elasticity loop: boot on
|
||
// demand, stay under the ceiling, then collapse back to the resident floor.
|
||
func TestPoolGrowsOnDemandAndReclaims(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 16})
|
||
p := poolOf(t, vm, "openai")
|
||
|
||
if created, _, _ := poolCounts(p); created != 0 {
|
||
t.Fatalf("pool must start empty, created=%d", created)
|
||
}
|
||
|
||
// concurrent load: hold several workers at once so the pool must grow
|
||
const hold = 6
|
||
var wg sync.WaitGroup
|
||
release := make(chan struct{})
|
||
for i := 0; i < hold; i++ {
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
w, err := p.acquire()
|
||
if err != nil {
|
||
t.Errorf("acquire: %v", err)
|
||
return
|
||
}
|
||
<-release
|
||
p.release(w)
|
||
}()
|
||
}
|
||
// wait until all six are checked out
|
||
deadline := time.Now().Add(5 * time.Second)
|
||
for {
|
||
_, _, inUse := poolCounts(p)
|
||
if inUse >= hold {
|
||
break
|
||
}
|
||
if time.Now().After(deadline) {
|
||
t.Fatalf("workers never checked out, in_use=%d", inUse)
|
||
}
|
||
time.Sleep(5 * time.Millisecond)
|
||
}
|
||
created, _, inUse := poolCounts(p)
|
||
if inUse != hold {
|
||
t.Fatalf("in_use = %d, want %d", inUse, hold)
|
||
}
|
||
if created < hold {
|
||
t.Fatalf("created = %d, must cover %d concurrent checkouts", created, hold)
|
||
}
|
||
if created > 16 {
|
||
t.Fatalf("created = %d exceeds the ceiling 16", created)
|
||
}
|
||
close(release)
|
||
wg.Wait()
|
||
|
||
// all connections closed: the janitor collapses the pool to the floor.
|
||
// shrinkGraceRounds rounds are required, and one extra round proves the
|
||
// floor is stable rather than shrinking to zero.
|
||
for i := 0; i < shrinkGraceRounds+2; i++ {
|
||
vm.ReclaimIdleNow()
|
||
}
|
||
created, idle, inUse := poolCounts(p)
|
||
if inUse != 0 {
|
||
t.Fatalf("in_use = %d after release, want 0", inUse)
|
||
}
|
||
if created != residentWorkers || idle != residentWorkers {
|
||
t.Fatalf("after reclaim: created=%d idle=%d, want %d/%d (resident floor)",
|
||
created, idle, residentWorkers, residentWorkers)
|
||
}
|
||
}
|
||
|
||
// TestReclaimNeedsGraceRounds: a single quiet round must not tear the pool
|
||
// down — a gap between requests is not the end of a load period.
|
||
func TestReclaimNeedsGraceRounds(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 8})
|
||
p := poolOf(t, vm, "openai")
|
||
|
||
// build slack: 4 states created, all returned
|
||
var ws []*worker
|
||
for i := 0; i < 4; i++ {
|
||
w, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire: %v", err)
|
||
}
|
||
ws = append(ws, w)
|
||
}
|
||
for _, w := range ws {
|
||
p.release(w)
|
||
}
|
||
created, _, _ := poolCounts(p)
|
||
if created != 4 {
|
||
t.Fatalf("created = %d, want 4", created)
|
||
}
|
||
|
||
if n := vm.ReclaimIdleNow(); n != 0 {
|
||
t.Fatalf("first quiet round must not reclaim (grace), closed %d", n)
|
||
}
|
||
if c, _, _ := poolCounts(p); c != 4 {
|
||
t.Fatalf("created = %d after the grace round, want 4", c)
|
||
}
|
||
if n := vm.ReclaimIdleNow(); n == 0 {
|
||
t.Fatal("second quiet round must reclaim")
|
||
}
|
||
}
|
||
|
||
// TestGrowthResetsGrace: growth (real contention) must cancel a pending shrink
|
||
// decision, but ordinary sequential traffic must NOT — otherwise a busy gateway
|
||
// that always has a request in flight would pin every state a past burst
|
||
// created, which is exactly the leak this change removes.
|
||
func TestGrowthResetsGrace(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 8})
|
||
p := poolOf(t, vm, "openai")
|
||
|
||
var ws []*worker
|
||
for i := 0; i < 4; i++ {
|
||
w, _ := p.acquire()
|
||
ws = append(ws, w)
|
||
}
|
||
for _, w := range ws {
|
||
p.release(w)
|
||
}
|
||
if n := vm.ReclaimIdleNow(); n != 0 {
|
||
t.Fatalf("grace round reclaimed %d", n)
|
||
}
|
||
|
||
// a lone sequential request reuses a warm state: no growth, so the pending
|
||
// shrink decision stands and the second round reclaims
|
||
w, _ := p.acquire()
|
||
p.release(w)
|
||
if n := vm.ReclaimIdleNow(); n == 0 {
|
||
t.Fatal("sequential traffic must not keep resetting the grace counter")
|
||
}
|
||
|
||
// now rebuild slack and prove that GROWTH does reset it
|
||
ws = ws[:0]
|
||
for i := 0; i < 4; i++ {
|
||
w, _ := p.acquire() // concurrent holds force created to grow
|
||
ws = append(ws, w)
|
||
}
|
||
for _, w := range ws {
|
||
p.release(w)
|
||
}
|
||
if n := vm.ReclaimIdleNow(); n != 0 {
|
||
t.Fatalf("growth must reset the grace counter, reclaimed %d", n)
|
||
}
|
||
}
|
||
|
||
// TestPoolCeilingIsRespectedUnderContention: acquires beyond the ceiling block
|
||
// instead of booting unbounded states.
|
||
func TestPoolCeilingIsRespectedUnderContention(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 2})
|
||
p := poolOf(t, vm, "openai")
|
||
|
||
w1, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire 1: %v", err)
|
||
}
|
||
w2, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire 2: %v", err)
|
||
}
|
||
got := make(chan struct{})
|
||
go func() {
|
||
w3, err := p.acquire()
|
||
if err == nil {
|
||
p.release(w3)
|
||
}
|
||
close(got)
|
||
}()
|
||
select {
|
||
case <-got:
|
||
t.Fatal("a third acquire must block at the ceiling of 2")
|
||
case <-time.After(100 * time.Millisecond):
|
||
}
|
||
if created, _, _ := poolCounts(p); created != 2 {
|
||
t.Fatalf("created = %d, must not exceed the ceiling 2", created)
|
||
}
|
||
p.release(w1)
|
||
select {
|
||
case <-got:
|
||
case <-time.After(2 * time.Second):
|
||
t.Fatal("releasing a state must unblock the waiter")
|
||
}
|
||
p.release(w2)
|
||
}
|
||
|
||
// TestLoweredCeilingReleasesStates: shrinking max_concurrent in the config must
|
||
// actually give memory back rather than leaving orphaned states parked.
|
||
func TestLoweredCeilingReleasesStates(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 8})
|
||
p := poolOf(t, vm, "openai")
|
||
|
||
var ws []*worker
|
||
for i := 0; i < 6; i++ {
|
||
w, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire: %v", err)
|
||
}
|
||
ws = append(ws, w)
|
||
}
|
||
if created, _, _ := poolCounts(p); created != 6 {
|
||
t.Fatalf("created = %d, want 6", created)
|
||
}
|
||
// operator lowers the ceiling while the states are still checked out
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 2})
|
||
for _, w := range ws {
|
||
p.release(w)
|
||
}
|
||
created, idle, _ := poolCounts(p)
|
||
if created > 2 || idle > 2 {
|
||
t.Fatalf("lowering the ceiling must drop the excess: created=%d idle=%d", created, idle)
|
||
}
|
||
}
|
||
|
||
// TestPoolStatsSurfacesAlgorithm checks the observability payload the status API
|
||
// exposes.
|
||
func TestPoolStatsSurfacesAlgorithm(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 32})
|
||
p := poolOf(t, vm, "openai")
|
||
w, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire: %v", err)
|
||
}
|
||
|
||
var st PoolStats
|
||
for _, s := range vm.PoolStats() {
|
||
if s.Name == "openai" {
|
||
st = s
|
||
}
|
||
}
|
||
if st.Max != 32 {
|
||
t.Fatalf("Max = %d, want 32", st.Max)
|
||
}
|
||
if st.InUse != 1 || st.Created != 1 {
|
||
t.Fatalf("InUse=%d Created=%d, want 1/1", st.InUse, st.Created)
|
||
}
|
||
if st.GrowStep != 4 {
|
||
t.Fatalf("GrowStep = %d, want ceil(32/8)=4", st.GrowStep)
|
||
}
|
||
if st.PeakInUse != 1 {
|
||
t.Fatalf("PeakInUse = %d, want 1", st.PeakInUse)
|
||
}
|
||
p.release(w)
|
||
|
||
// stats must be sorted by name for a stable UI
|
||
names := []string{}
|
||
for _, s := range vm.PoolStats() {
|
||
names = append(names, s.Name)
|
||
}
|
||
for i := 1; i < len(names); i++ {
|
||
if names[i-1] > names[i] {
|
||
t.Fatalf("PoolStats not sorted: %v", names)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestHotReloadKeepsCeiling: replacing an adapter's code must not reset its
|
||
// elastic ceiling back to 1.
|
||
func TestHotReloadKeepsCeiling(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 24})
|
||
if err := vm.LoadAdapterSource("openai", "return {name='openai', version='9'}"); err != nil {
|
||
t.Fatalf("reload: %v", err)
|
||
}
|
||
p := poolOf(t, vm, "openai")
|
||
p.mu.Lock()
|
||
maxW := p.maxW
|
||
p.mu.Unlock()
|
||
if maxW != 24 {
|
||
t.Fatalf("ceiling after hot reload = %d, want 24", maxW)
|
||
}
|
||
}
|
||
|
||
// TestConcurrentTransformStress exercises the grow/shrink paths together under
|
||
// real hook calls, and asserts the accounting stays consistent (no leaked
|
||
// states, no negative counters).
|
||
func TestConcurrentTransformStress(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 16})
|
||
|
||
var wg sync.WaitGroup
|
||
for i := 0; i < 40; i++ {
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
raw := `{"model":"m","messages":[{"role":"user","content":"hi"}]}`
|
||
if _, err := vm.Transform("openai", "transform_request", raw); err != nil {
|
||
t.Errorf("transform: %v", err)
|
||
}
|
||
}()
|
||
}
|
||
wg.Wait()
|
||
|
||
p := poolOf(t, vm, "openai")
|
||
created, idle, inUse := poolCounts(p)
|
||
if inUse != 0 {
|
||
t.Fatalf("in_use = %d after the stress run, want 0 (leaked checkout)", inUse)
|
||
}
|
||
if created != idle {
|
||
t.Fatalf("created=%d idle=%d must match when nothing is checked out", created, idle)
|
||
}
|
||
if created > 16 {
|
||
t.Fatalf("created = %d exceeds the ceiling", created)
|
||
}
|
||
for i := 0; i < shrinkGraceRounds+2; i++ {
|
||
vm.ReclaimIdleNow()
|
||
}
|
||
created, idle, _ = poolCounts(p)
|
||
if created != residentWorkers || idle != residentWorkers {
|
||
t.Fatalf("post-stress reclaim: created=%d idle=%d, want %d/%d",
|
||
created, idle, residentWorkers, residentWorkers)
|
||
}
|
||
}
|
||
|
||
// TestStopIsIdempotent guards the janitor's stop channel against a double close.
|
||
func TestStopIsIdempotent(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
vm.Stop()
|
||
vm.Stop() // must not panic
|
||
}
|
||
|
||
// TestSequentialTrafficDoesNotBatch: a single caller looping must keep reusing
|
||
// one warm state — batch prewarm is for genuine concurrency only.
|
||
func TestSequentialTrafficDoesNotBatch(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 64}) // grow step would be 8
|
||
p := poolOf(t, vm, "openai")
|
||
for i := 0; i < 20; i++ {
|
||
w, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire %d: %v", i, err)
|
||
}
|
||
p.release(w)
|
||
}
|
||
created, _, _ := poolCounts(p)
|
||
if created != 1 {
|
||
t.Fatalf("sequential traffic must reuse one state, created=%d", created)
|
||
}
|
||
}
|
||
|
||
// TestContentionBatchPrewarms: when callers are actually QUEUED behind a full
|
||
// pool, the pool warms extra states for them instead of booting one per request.
|
||
// The batch is bounded by the queue depth, not by the ceiling, so it can never
|
||
// warm states nobody is waiting for.
|
||
func TestContentionBatchPrewarms(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
// ceiling 32 -> step cap 4; queue depth will be the real bound
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 32})
|
||
p := poolOf(t, vm, "openai")
|
||
|
||
// Fill the pool to its ceiling-independent state: hold 1, then queue 3 more
|
||
// so waiting == 3 when the next miss happens.
|
||
p.mu.Lock()
|
||
p.maxW = 1 // force the next acquires to queue
|
||
p.mu.Unlock()
|
||
w1, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire 1: %v", err)
|
||
}
|
||
done := make(chan *worker, 3)
|
||
for i := 0; i < 3; i++ {
|
||
go func() {
|
||
w, err := p.acquire()
|
||
if err != nil {
|
||
done <- nil
|
||
return
|
||
}
|
||
done <- w
|
||
}()
|
||
}
|
||
// wait until all three are queued
|
||
deadline := time.Now().Add(5 * time.Second)
|
||
for {
|
||
p.mu.Lock()
|
||
w := p.waiting
|
||
p.mu.Unlock()
|
||
if w >= 3 {
|
||
break
|
||
}
|
||
if time.Now().After(deadline) {
|
||
t.Fatalf("waiters never queued, waiting=%d", w)
|
||
}
|
||
time.Sleep(5 * time.Millisecond)
|
||
}
|
||
|
||
// raise the ceiling: the next miss sees waiting=3 and prewarms a batch
|
||
p.setMax(32)
|
||
got := make([]*worker, 0, 3)
|
||
for i := 0; i < 3; i++ {
|
||
select {
|
||
case w := <-done:
|
||
if w == nil {
|
||
t.Fatal("a queued acquire failed")
|
||
}
|
||
got = append(got, w)
|
||
case <-time.After(5 * time.Second):
|
||
t.Fatal("queued acquires never completed after the ceiling was raised")
|
||
}
|
||
}
|
||
created, _, _ := poolCounts(p)
|
||
if created < 4 {
|
||
t.Fatalf("queued demand must be satisfied, created=%d want >=4", created)
|
||
}
|
||
if created > 32 {
|
||
t.Fatalf("prewarm exceeded the ceiling, created=%d", created)
|
||
}
|
||
p.release(w1)
|
||
for _, w := range got {
|
||
p.release(w)
|
||
}
|
||
}
|
||
|
||
// TestNoPrewarmWithoutWaiters is the production regression this bound fixes: two
|
||
// concurrent requests against a wide adapter (ceiling 76, step 8) must NOT warm
|
||
// eight states. Warming for nobody meant the janitor discarded the surplus a
|
||
// minute later, so RSS oscillated between ~46 MB and ~56 MB forever.
|
||
func TestNoPrewarmWithoutWaiters(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 76})
|
||
p := poolOf(t, vm, "openai")
|
||
|
||
// two overlapping requests: the second misses while the first holds its
|
||
// state, but nobody is blocked, so no batch may be warmed
|
||
w1, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire 1: %v", err)
|
||
}
|
||
w2, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire 2: %v", err)
|
||
}
|
||
time.Sleep(100 * time.Millisecond) // give any background prewarm a chance
|
||
created, _, _ := poolCounts(p)
|
||
if created != 2 {
|
||
t.Fatalf("two overlapping requests must create exactly 2 states, got %d", created)
|
||
}
|
||
p.release(w1)
|
||
p.release(w2)
|
||
}
|
||
|
||
// TestBurstLeftoverIsReclaimedUnderSteadyTraffic reproduces what production
|
||
// showed right after the first deploy: a burst grew the openai pool to 9 states,
|
||
// then steady low-concurrency traffic kept arriving and the pool never shrank,
|
||
// because any checkout reset the shrink grace counter. Slack must be reclaimed
|
||
// even while the gateway keeps serving requests one at a time.
|
||
func TestBurstLeftoverIsReclaimedUnderSteadyTraffic(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatalf("start: %v", err)
|
||
}
|
||
defer vm.Stop()
|
||
vm.ConfigureConcurrency(map[string]int{"openai": 32})
|
||
p := poolOf(t, vm, "openai")
|
||
|
||
// burst: 9 concurrent holds
|
||
var ws []*worker
|
||
for i := 0; i < 9; i++ {
|
||
w, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("acquire: %v", err)
|
||
}
|
||
ws = append(ws, w)
|
||
}
|
||
for _, w := range ws {
|
||
p.release(w)
|
||
}
|
||
if created, _, _ := poolCounts(p); created < 9 {
|
||
t.Fatalf("burst should have grown the pool, created=%d", created)
|
||
}
|
||
|
||
// steady state: one request between every janitor round, forever
|
||
for round := 0; round < 8; round++ {
|
||
w, err := p.acquire()
|
||
if err != nil {
|
||
t.Fatalf("round %d acquire: %v", round, err)
|
||
}
|
||
p.release(w)
|
||
vm.ReclaimIdleNow()
|
||
}
|
||
created, idle, inUse := poolCounts(p)
|
||
if inUse != 0 {
|
||
t.Fatalf("in_use = %d, want 0", inUse)
|
||
}
|
||
if created > residentWorkers+idleHeadroom {
|
||
t.Fatalf("burst leftovers were never reclaimed under steady traffic: created=%d idle=%d",
|
||
created, idle)
|
||
}
|
||
}
|
||
|
||
// TestNonStreamToolCallsPreserved is the regression for the bug that killed
|
||
// tool-using conversations on their SECOND request: several adapters handled
|
||
// tool_calls in transform_stream_chunk but dropped them in
|
||
// transform_response. A client then received finish_reason:"tool_calls" with no
|
||
// tool_calls payload, replayed an assistant message whose function
|
||
// name/arguments were empty, and the upstream rejected the next turn with
|
||
//
|
||
// 400 invalid tool_call function, function/name/arguments cannot be empty
|
||
//
|
||
// Every OpenAI-shaped adapter must forward non-streaming tool calls.
|
||
func TestNonStreamToolCallsPreserved(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
// upstream shape: arguments is a JSON *string*, as OpenAI sends it
|
||
body := `{"choices":[{"index":0,"message":{"role":"assistant","content":"",` +
|
||
`"tool_calls":[{"index":0,"id":"call_abc","type":"function",` +
|
||
`"function":{"name":"get_weather","arguments":"{\"city\":\"Beijing\"}"}}]},` +
|
||
`"finish_reason":"tool_calls"}],` +
|
||
`"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}`
|
||
|
||
for _, name := range []string{
|
||
"openai", "sensenova", "trae", "deepseek", "github",
|
||
"groq", "kimicode", "mistral", "opencode", "agentrouter",
|
||
} {
|
||
out, err := vm.Transform(name, "transform_response", body)
|
||
if err != nil {
|
||
t.Fatalf("%s transform_response: %v", name, err)
|
||
}
|
||
var got struct {
|
||
FinishReason string `json:"finish_reason"`
|
||
ToolCalls []struct {
|
||
ID string `json:"id"`
|
||
Type string `json:"type"`
|
||
Name string `json:"name"`
|
||
Arguments map[string]interface{} `json:"arguments"`
|
||
} `json:"tool_calls"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||
t.Fatalf("%s: unmarshal %v (%s)", name, err, out)
|
||
}
|
||
if len(got.ToolCalls) != 1 {
|
||
t.Fatalf("%s: dropped non-streaming tool_calls (client would replay an empty function): %s", name, out)
|
||
}
|
||
tc := got.ToolCalls[0]
|
||
if tc.Name != "get_weather" {
|
||
t.Errorf("%s: tool name lost: %s", name, out)
|
||
}
|
||
if tc.ID != "call_abc" {
|
||
t.Errorf("%s: tool id lost: %s", name, out)
|
||
}
|
||
// arguments must be decoded into an object, not left as a string
|
||
if tc.Arguments == nil || tc.Arguments["city"] != "Beijing" {
|
||
t.Errorf("%s: arguments not decoded: %s", name, out)
|
||
}
|
||
if got.FinishReason != "tool_calls" {
|
||
t.Errorf("%s: finish_reason lost: %s", name, out)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestGeminiNonStreamToolCalls covers Gemini's native shape (functionCall parts
|
||
// under candidates[].content.parts). The streaming path handled these; the
|
||
// non-streaming path silently returned text only.
|
||
func TestGeminiNonStreamToolCalls(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
body := `{"candidates":[{"content":{"parts":[` +
|
||
`{"text":"checking"},` +
|
||
`{"functionCall":{"name":"get_weather","args":{"city":"Beijing"}}}` +
|
||
`]},"finishReason":"STOP"}],` +
|
||
`"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":4,"totalTokenCount":12}}`
|
||
|
||
out, err := vm.Transform("gemini", "transform_response", body)
|
||
if err != nil {
|
||
t.Fatalf("gemini transform_response: %v", err)
|
||
}
|
||
var got struct {
|
||
Content string `json:"content"`
|
||
FinishReason string `json:"finish_reason"`
|
||
ToolCalls []struct {
|
||
Name string `json:"name"`
|
||
Arguments map[string]interface{} `json:"arguments"`
|
||
} `json:"tool_calls"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||
t.Fatalf("unmarshal: %v (%s)", err, out)
|
||
}
|
||
if len(got.ToolCalls) != 1 || got.ToolCalls[0].Name != "get_weather" {
|
||
t.Fatalf("gemini dropped non-streaming functionCall: %s", out)
|
||
}
|
||
if got.ToolCalls[0].Arguments["city"] != "Beijing" {
|
||
t.Errorf("gemini args lost: %s", out)
|
||
}
|
||
// a tool call must not be reported as a plain stop
|
||
if got.FinishReason != "tool_calls" {
|
||
t.Errorf("gemini finish_reason should become tool_calls: %s", out)
|
||
}
|
||
}
|
||
|
||
// TestOllamaNonStreamToolCalls covers Ollama's non-streaming shape
|
||
// (message.tool_calls with an already-decoded arguments object). The adapter
|
||
// initialized tool_calls to an empty table and never filled it.
|
||
func TestOllamaNonStreamToolCalls(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
body := `{"message":{"role":"assistant","content":"",` +
|
||
`"tool_calls":[{"function":{"name":"get_weather","arguments":{"city":"Beijing"}}}]},` +
|
||
`"done_reason":"stop","prompt_eval_count":7,"eval_count":3}`
|
||
|
||
out, err := vm.Transform("ollama", "transform_response", body)
|
||
if err != nil {
|
||
t.Fatalf("ollama transform_response: %v", err)
|
||
}
|
||
var got struct {
|
||
FinishReason string `json:"finish_reason"`
|
||
ToolCalls []struct {
|
||
Name string `json:"name"`
|
||
Arguments map[string]interface{} `json:"arguments"`
|
||
} `json:"tool_calls"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||
t.Fatalf("unmarshal: %v (%s)", err, out)
|
||
}
|
||
if len(got.ToolCalls) != 1 || got.ToolCalls[0].Name != "get_weather" {
|
||
t.Fatalf("ollama dropped non-streaming tool_calls: %s", out)
|
||
}
|
||
if got.ToolCalls[0].Arguments["city"] != "Beijing" {
|
||
t.Errorf("ollama args lost: %s", out)
|
||
}
|
||
if got.FinishReason != "tool_calls" {
|
||
t.Errorf("ollama finish_reason should become tool_calls: %s", out)
|
||
}
|
||
}
|
||
|
||
// TestTraeTextToolCallRecovery covers the trae-specific failure: its upstream
|
||
// (trae-local-api's OpenAI endpoint) never reads the request's `tools` array, so
|
||
// the relayed model prints a <tool_call> block as TEXT, leaves
|
||
// message.tool_calls null and reports finish_reason:"stop". An OpenAI client
|
||
// then sees an ordinary completion and its agent loop ends mid-conversation.
|
||
// The adapter must recover the structured call, strip the block from content,
|
||
// and correct finish_reason.
|
||
func TestTraeTextToolCallRecovery(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
cases := []struct {
|
||
name string
|
||
body string
|
||
}{
|
||
{
|
||
"tool_call tag with arguments",
|
||
`{"choices":[{"message":{"role":"assistant","content":"<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Beijing\"}}\n</tool_call>"},"finish_reason":"stop"}]}`,
|
||
},
|
||
{
|
||
// the same codebase's Anthropic endpoint prompts for <toolcall> and "params"
|
||
"toolcall tag with params",
|
||
`{"choices":[{"message":{"role":"assistant","content":"<toolcall>{\"name\": \"get_weather\", \"params\": {\"city\": \"Beijing\"}}</toolcall>"},"finish_reason":"stop"}]}`,
|
||
},
|
||
{
|
||
"prose around the block is preserved",
|
||
`{"choices":[{"message":{"role":"assistant","content":"Let me check.\n<tool_call>{\"name\":\"get_weather\",\"arguments\":{\"city\":\"Beijing\"}}</tool_call>"},"finish_reason":"stop"}]}`,
|
||
},
|
||
}
|
||
|
||
for _, c := range cases {
|
||
out, err := vm.Transform("trae", "transform_response", c.body)
|
||
if err != nil {
|
||
t.Fatalf("%s: transform: %v", c.name, err)
|
||
}
|
||
var got struct {
|
||
Content string `json:"content"`
|
||
FinishReason string `json:"finish_reason"`
|
||
ToolCalls []struct {
|
||
Name string `json:"name"`
|
||
Arguments map[string]interface{} `json:"arguments"`
|
||
} `json:"tool_calls"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &got); err != nil {
|
||
t.Fatalf("%s: unmarshal %v (%s)", c.name, err, out)
|
||
}
|
||
if len(got.ToolCalls) != 1 {
|
||
t.Fatalf("%s: text tool call not recovered: %s", c.name, out)
|
||
}
|
||
if got.ToolCalls[0].Name != "get_weather" {
|
||
t.Errorf("%s: name wrong: %s", c.name, out)
|
||
}
|
||
if got.ToolCalls[0].Arguments["city"] != "Beijing" {
|
||
t.Errorf("%s: args wrong: %s", c.name, out)
|
||
}
|
||
if got.FinishReason != "tool_calls" {
|
||
t.Errorf("%s: finish_reason must be corrected to tool_calls: %s", c.name, out)
|
||
}
|
||
if strings.Contains(got.Content, "<tool_call") || strings.Contains(got.Content, "<toolcall") {
|
||
t.Errorf("%s: raw block left in user-visible content: %s", c.name, out)
|
||
}
|
||
}
|
||
|
||
// a plain text answer must be untouched
|
||
plain := `{"choices":[{"message":{"role":"assistant","content":"just text"},"finish_reason":"stop"}]}`
|
||
out, err := vm.Transform("trae", "transform_response", plain)
|
||
if err != nil {
|
||
t.Fatalf("plain: %v", err)
|
||
}
|
||
if strings.Contains(out, "tool_calls") {
|
||
t.Errorf("plain answer must not gain tool_calls: %s", out)
|
||
}
|
||
if !strings.Contains(out, `"finish_reason":"stop"`) {
|
||
t.Errorf("plain answer finish_reason changed: %s", out)
|
||
}
|
||
}
|
||
|
||
// TestToolIDSanitize locks the tool-call id sanitizer in both adapters that
|
||
// talk to Claude upstreams. Anthropic requires tool_use.id /
|
||
// tool_result.tool_use_id to match ^[a-zA-Z0-9_-]{1,64}$ and rejects the WHOLE
|
||
// request otherwise (REQUEST_BODY_INVALID / "Invalid tool use format"), while
|
||
// OpenAI has no such rule — so an OpenAI-compatible model can mint "bash:0"
|
||
// (observed from moonshotai/kimi-k3) and poison a client's history, killing
|
||
// every Claude slot on replay. tabitoken/扇贝 are Claude-behind-OpenAI, so
|
||
// openai.lua needs the same treatment as anthropic.lua.
|
||
func TestToolIDSanitize(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
legal := regexp.MustCompile(`^[a-zA-Z0-9_-]{1,64}$`)
|
||
long := strings.Repeat("a", 200)
|
||
|
||
body := func(id string) string {
|
||
b, err := json.Marshal(map[string]interface{}{
|
||
"model": "x",
|
||
"messages": []interface{}{
|
||
map[string]interface{}{"role": "user", "content": "run whoami"},
|
||
map[string]interface{}{"role": "assistant", "content": "", "tool_calls": []interface{}{
|
||
map[string]interface{}{"id": id, "type": "function",
|
||
"function": map[string]interface{}{"name": "bash", "arguments": `{"command":"whoami"}`}},
|
||
}},
|
||
map[string]interface{}{"role": "tool", "tool_call_id": id, "content": "root"},
|
||
},
|
||
})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return string(b)
|
||
}
|
||
|
||
// ---- anthropic: tool_use.id and tool_result.tool_use_id must agree ----
|
||
for _, id := range []string{"bash:0", "call_ok_123", long, "toolu~sig1:AB+/=="} {
|
||
out, err := vm.Transform("anthropic", "transform_request", body(id))
|
||
if err != nil {
|
||
t.Fatalf("anthropic transform %q: %v", id, err)
|
||
}
|
||
var r struct {
|
||
Messages []struct {
|
||
Content []struct {
|
||
Type string `json:"type"`
|
||
ID string `json:"id"`
|
||
ToolUseID string `json:"tool_use_id"`
|
||
} `json:"content"`
|
||
} `json:"messages"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &r); err != nil {
|
||
t.Fatalf("anthropic unmarshal %q: %v (%s)", id, err, out)
|
||
}
|
||
var useID, resID string
|
||
for _, m := range r.Messages {
|
||
for _, c := range m.Content {
|
||
switch c.Type {
|
||
case "tool_use":
|
||
useID = c.ID
|
||
case "tool_result":
|
||
resID = c.ToolUseID
|
||
}
|
||
}
|
||
}
|
||
if !legal.MatchString(useID) {
|
||
t.Errorf("anthropic %q: tool_use.id %q violates Anthropic's id pattern", id, useID)
|
||
}
|
||
if useID != resID {
|
||
t.Errorf("anthropic %q: tool_use.id %q != tool_result.tool_use_id %q (unpaired call)", id, useID, resID)
|
||
}
|
||
if legal.MatchString(id) && useID != id {
|
||
t.Errorf("anthropic %q: already-legal id must pass through untouched, got %q", id, useID)
|
||
}
|
||
}
|
||
|
||
// ---- openai: same rewrite, and both sites must stay in sync ----
|
||
for _, id := range []string{"bash:0", "call_ok_123", long} {
|
||
out, err := vm.Transform("openai", "transform_request", body(id))
|
||
if err != nil {
|
||
t.Fatalf("openai transform %q: %v", id, err)
|
||
}
|
||
var r struct {
|
||
Messages []struct {
|
||
Role string `json:"role"`
|
||
ToolCalls []struct {
|
||
ID string `json:"id"`
|
||
} `json:"tool_calls"`
|
||
ToolCallID string `json:"tool_call_id"`
|
||
} `json:"messages"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &r); err != nil {
|
||
t.Fatalf("openai unmarshal %q: %v (%s)", id, err, out)
|
||
}
|
||
var callID, resID string
|
||
for _, m := range r.Messages {
|
||
if len(m.ToolCalls) > 0 {
|
||
callID = m.ToolCalls[0].ID
|
||
}
|
||
if m.Role == "tool" {
|
||
resID = m.ToolCallID
|
||
}
|
||
}
|
||
if !legal.MatchString(callID) {
|
||
t.Errorf("openai %q: tool_calls[0].id %q still illegal for Claude-behind-OpenAI upstreams", id, callID)
|
||
}
|
||
if callID != resID {
|
||
t.Errorf("openai %q: tool_calls id %q != tool_call_id %q (unpaired call)", id, callID, resID)
|
||
}
|
||
if legal.MatchString(id) && callID != id {
|
||
t.Errorf("openai %q: already-legal id must pass through untouched, got %q", id, callID)
|
||
}
|
||
}
|
||
|
||
// ---- distinct dirty ids must NOT collapse into one (digest suffix) ----
|
||
seen := map[string]string{}
|
||
for _, id := range []string{"a:b", "a_b", "a-b", "a b", "a/b"} {
|
||
out, err := vm.Transform("anthropic", "transform_request", body(id))
|
||
if err != nil {
|
||
t.Fatalf("collision probe %q: %v", id, err)
|
||
}
|
||
var r struct {
|
||
Messages []struct {
|
||
Content []struct {
|
||
Type string `json:"type"`
|
||
ID string `json:"id"`
|
||
} `json:"content"`
|
||
} `json:"messages"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &r); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
got := ""
|
||
for _, m := range r.Messages {
|
||
for _, c := range m.Content {
|
||
if c.Type == "tool_use" {
|
||
got = c.ID
|
||
}
|
||
}
|
||
}
|
||
if prev, dup := seen[got]; dup {
|
||
t.Errorf("ids %q and %q both map to %q — a collision makes a duplicate tool_use", prev, id, got)
|
||
}
|
||
seen[got] = id
|
||
}
|
||
|
||
// ---- determinism: same input, same output ----
|
||
a, _ := vm.Transform("anthropic", "transform_request", body("bash:0"))
|
||
b, _ := vm.Transform("anthropic", "transform_request", body("bash:0"))
|
||
if a != b {
|
||
t.Errorf("sanitizer is not deterministic:\n%s\n%s", a, b)
|
||
}
|
||
}
|
||
|
||
// TestToolIDSanitizeResponse locks the OUTBOUND half of the sanitizer. The
|
||
// request-side rewrite alone is not enough: an id minted by a permissive
|
||
// OpenAI-compatible upstream (moonshotai/kimi-k3 returns "bash:0") is handed to
|
||
// the client, stored in its history, and replayed forever after. Sanitizing on
|
||
// the way out means a poisoned id never enters a client session in the first
|
||
// place — both for non-streaming responses and for the first fragment of a
|
||
// streamed tool call (later argument fragments carry no id).
|
||
func TestToolIDSanitizeResponse(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
legal := regexp.MustCompile(`^[a-zA-Z0-9_-]{1,64}$`)
|
||
|
||
// ---- openai non-streaming ----
|
||
resp := `{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[
|
||
{"id":"bash:0","type":"function","function":{"name":"bash","arguments":"{\"command\":\"whoami\"}"}}]},
|
||
"finish_reason":"tool_calls"}]}`
|
||
out, err := vm.Transform("openai", "transform_response", resp)
|
||
if err != nil {
|
||
t.Fatalf("openai transform_response: %v", err)
|
||
}
|
||
var r struct {
|
||
ToolCalls []struct {
|
||
ID string `json:"id"`
|
||
} `json:"tool_calls"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &r); err != nil {
|
||
t.Fatalf("unmarshal: %v (%s)", err, out)
|
||
}
|
||
if len(r.ToolCalls) != 1 {
|
||
t.Fatalf("tool_calls lost: %s", out)
|
||
}
|
||
if !legal.MatchString(r.ToolCalls[0].ID) {
|
||
t.Errorf("non-stream response leaks illegal id %q to the client", r.ToolCalls[0].ID)
|
||
}
|
||
nonStreamID := r.ToolCalls[0].ID
|
||
|
||
// ---- openai streaming: first fragment carries the id ----
|
||
chunk := `{"choices":[{"delta":{"tool_calls":[
|
||
{"index":0,"id":"bash:0","type":"function","function":{"name":"bash","arguments":""}}]},
|
||
"finish_reason":null}]}`
|
||
out, err = vm.Transform("openai", "transform_stream_chunk", chunk)
|
||
if err != nil {
|
||
t.Fatalf("openai transform_stream_chunk: %v", err)
|
||
}
|
||
var c struct {
|
||
ToolCalls []struct {
|
||
ID string `json:"id"`
|
||
} `json:"tool_calls"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &c); err != nil {
|
||
t.Fatalf("unmarshal chunk: %v (%s)", err, out)
|
||
}
|
||
if len(c.ToolCalls) != 1 {
|
||
t.Fatalf("stream tool_calls lost: %s", out)
|
||
}
|
||
if !legal.MatchString(c.ToolCalls[0].ID) {
|
||
t.Errorf("stream response leaks illegal id %q to the client", c.ToolCalls[0].ID)
|
||
}
|
||
// Streaming and non-streaming must agree, otherwise a client that mixes
|
||
// modes within one session produces unpaired tool calls.
|
||
if c.ToolCalls[0].ID != nonStreamID {
|
||
t.Errorf("stream id %q != non-stream id %q for the same input", c.ToolCalls[0].ID, nonStreamID)
|
||
}
|
||
|
||
// ---- a later argument fragment has no id and must stay id-less ----
|
||
argChunk := `{"choices":[{"delta":{"tool_calls":[
|
||
{"index":0,"function":{"arguments":"{\"command\":\"whoami\"}"}}]},"finish_reason":null}]}`
|
||
out, err = vm.Transform("openai", "transform_stream_chunk", argChunk)
|
||
if err != nil {
|
||
t.Fatalf("arg fragment: %v", err)
|
||
}
|
||
if strings.Contains(out, `"id"`) {
|
||
t.Errorf("argument fragment gained an id (breaks index-based accumulation): %s", out)
|
||
}
|
||
|
||
// ---- anthropic response side: tool_use id and stream start ----
|
||
aResp := `{"content":[{"type":"tool_use","id":"toolu~sig1:AB+/==","name":"bash","input":{"command":"whoami"}}],"stop_reason":"tool_use"}`
|
||
out, err = vm.Transform("anthropic", "transform_response", aResp)
|
||
if err != nil {
|
||
t.Fatalf("anthropic transform_response: %v", err)
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &r); err != nil {
|
||
t.Fatalf("unmarshal anthropic: %v (%s)", err, out)
|
||
}
|
||
if len(r.ToolCalls) != 1 || !legal.MatchString(r.ToolCalls[0].ID) {
|
||
t.Errorf("anthropic response leaks illegal id: %s", out)
|
||
}
|
||
|
||
aChunk := `{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu~sig1:AB+/==","name":"bash"}}`
|
||
out, err = vm.Transform("anthropic", "transform_stream_chunk", aChunk)
|
||
if err != nil {
|
||
t.Fatalf("anthropic stream chunk: %v", err)
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &c); err != nil {
|
||
t.Fatalf("unmarshal anthropic chunk: %v (%s)", err, out)
|
||
}
|
||
if len(c.ToolCalls) != 1 || !legal.MatchString(c.ToolCalls[0].ID) {
|
||
t.Errorf("anthropic stream leaks illegal id: %s", out)
|
||
}
|
||
}
|
||
|
||
// TestAnthropicUsageCacheMapping pins the Anthropic→OpenAI usage conversion.
|
||
// The two APIs disagree on what the prompt count means: Anthropic's
|
||
// input_tokens EXCLUDES cached blocks (cache_read_input_tokens and
|
||
// cache_creation_input_tokens are separate, additive, billed input), while
|
||
// OpenAI's prompt_tokens INCLUDES its cached_tokens subset. Mapping
|
||
// input_tokens straight onto prompt undercounted the billed prompt by the whole
|
||
// cache portion and could report cached_tokens > prompt_tokens (a hit rate over
|
||
// 100%); cache_creation_input_tokens was dropped entirely.
|
||
func TestAnthropicUsageCacheMapping(t *testing.T) {
|
||
vm := NewVM(freshAdapterDir(t))
|
||
if err := vm.Start(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer vm.Stop()
|
||
|
||
type usage struct {
|
||
Prompt int `json:"prompt"`
|
||
Completion int `json:"completion"`
|
||
Total int `json:"total"`
|
||
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"`
|
||
}
|
||
|
||
// ---- non-streaming, cache read + cache write reported ----
|
||
resp := `{"content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn",
|
||
"usage":{"input_tokens":1200,"cache_read_input_tokens":40000,
|
||
"cache_creation_input_tokens":500,"output_tokens":80}}`
|
||
out, err := vm.Transform("anthropic", "transform_response", resp)
|
||
if err != nil {
|
||
t.Fatalf("transform_response: %v", err)
|
||
}
|
||
// Each case decodes into a FRESH value: json.Unmarshal leaves fields absent
|
||
// from the payload untouched, so a reused struct would carry the previous
|
||
// case's prompt_tokens_details into a response that has none.
|
||
decode := func(out string) usage {
|
||
t.Helper()
|
||
var r struct {
|
||
TokenUsage usage `json:"token_usage"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &r); err != nil {
|
||
t.Fatalf("unmarshal: %v (%s)", err, out)
|
||
}
|
||
return r.TokenUsage
|
||
}
|
||
decodeChunk := func(out string) (usage, string) {
|
||
t.Helper()
|
||
var c struct {
|
||
Usage usage `json:"usage"`
|
||
FinishReason string `json:"finish_reason"`
|
||
}
|
||
if err := json.Unmarshal([]byte(out), &c); err != nil {
|
||
t.Fatalf("unmarshal chunk: %v (%s)", err, out)
|
||
}
|
||
return c.Usage, c.FinishReason
|
||
}
|
||
|
||
u := decode(out)
|
||
if u.Prompt != 1200+40000+500 {
|
||
t.Errorf("prompt must sum fresh+cache_read+cache_creation (41700), got %d", u.Prompt)
|
||
}
|
||
if u.Total != u.Prompt+u.Completion {
|
||
t.Errorf("total %d != prompt %d + completion %d", u.Total, u.Prompt, u.Completion)
|
||
}
|
||
if u.Details == nil || u.Details.CachedTokens != 40000 {
|
||
t.Errorf("cached_tokens must carry cache_read (40000): %s", out)
|
||
}
|
||
if u.Details != nil && u.Details.CachedTokens > u.Prompt {
|
||
t.Errorf("cached_tokens %d > prompt %d implies a hit rate above 100%%",
|
||
u.Details.CachedTokens, u.Prompt)
|
||
}
|
||
if u.Hit != 40000 || u.Miss != 1200+500 {
|
||
t.Errorf("hit/miss split wrong: hit=%d miss=%d (want 40000/1700)", u.Hit, u.Miss)
|
||
}
|
||
|
||
// ---- a reported zero hit must stay distinguishable from "not reported" ----
|
||
zero := `{"content":[],"stop_reason":"end_turn","usage":{"input_tokens":100,"cache_read_input_tokens":0,"output_tokens":5}}`
|
||
out, err = vm.Transform("anthropic", "transform_response", zero)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if decode(out).Details == nil {
|
||
t.Errorf("a reported cache_read of 0 must still emit prompt_tokens_details: %s", out)
|
||
}
|
||
|
||
// An upstream that never mentions caching must not gain a fabricated split
|
||
// (justwoker reports only input_tokens/output_tokens plus its own cost fields).
|
||
none := `{"content":[],"stop_reason":"end_turn","usage":{"input_tokens":6931,"output_tokens":1}}`
|
||
out, err = vm.Transform("anthropic", "transform_response", none)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
plain := decode(out)
|
||
if plain.Details != nil {
|
||
t.Errorf("upstream reported no cache fields; details must be absent: %s", out)
|
||
}
|
||
if plain.Prompt != 6931 || plain.Total != 6932 {
|
||
t.Errorf("plain usage mismapped: %s", out)
|
||
}
|
||
|
||
// ---- streaming: message_delta carries the FINAL usage and must map cache ----
|
||
delta := `{"type":"message_delta","delta":{"stop_reason":"end_turn"},
|
||
"usage":{"input_tokens":1200,"cache_read_input_tokens":40000,"output_tokens":80}}`
|
||
out, err = vm.Transform("anthropic", "transform_stream_chunk", delta)
|
||
if err != nil {
|
||
t.Fatalf("stream delta: %v", err)
|
||
}
|
||
du, fin := decodeChunk(out)
|
||
if du.Prompt != 41200 {
|
||
t.Errorf("stream final prompt must include cache_read: got %d", du.Prompt)
|
||
}
|
||
if du.Details == nil || du.Details.CachedTokens != 40000 {
|
||
t.Errorf("stream final usage dropped the cache split: %s", out)
|
||
}
|
||
if fin != "stop" {
|
||
t.Errorf("finish_reason regressed: %s", out)
|
||
}
|
||
|
||
// message_start's placeholder count must not suppress the later real one:
|
||
// justwoker reports input_tokens=160 at message_start and 6931 at
|
||
// message_delta, and the gateway's mergeUsage lets the later value win.
|
||
start := `{"type":"message_start","message":{"usage":{"input_tokens":160,"output_tokens":1}}}`
|
||
out, err = vm.Transform("anthropic", "transform_stream_chunk", start)
|
||
if err != nil {
|
||
t.Fatalf("message_start: %v", err)
|
||
}
|
||
if su, _ := decodeChunk(out); su.Prompt != 160 {
|
||
t.Errorf("message_start usage lost: %s", out)
|
||
}
|
||
|
||
// A usage-less chunk must stay silent rather than reporting zeros.
|
||
quiet := `{"type":"message_start","message":{}}`
|
||
out, err = vm.Transform("anthropic", "transform_stream_chunk", quiet)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if strings.Contains(out, "usage") {
|
||
t.Errorf("usage-less chunk must not emit a usage object: %s", out)
|
||
}
|
||
}
|