Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
18 kB · 605 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606package executor
import ( "bytes" "context" "encoding/base64" "io" "net/http" "net/http/httptest" "testing"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" internalsignature "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" _ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "github.com/tidwall/gjson" "google.golang.org/protobuf/encoding/protowire")
const testClaudeCAISSample = "CAISqwIKiAEIEBgCKkBHRlRBsNiptQUWfPoOhuQKwi5LnncZVO9bB5jqOs76D7uBtgktML0zqJtNmLHXHHcgD6lk4MQu4QBXzFd1lbC3Mg5jbGF1ZGUtZmFibGUtNTgBQgh0aGlua2luZ1okZDk3NDM5NzUtNGJiMC00OTM2LTllMjgtZDViMGQyMWJkYzQ4EgxCGh+XVFFFeySAjtAaDL/A1LltGu6MMJ+eXSIwsN0oBpDrqLv22UBfkMnTotnIbkvkOyb9xZHgigG6OZVHaI3gThm+maLKmgO5PrFLKlDFYp+YZksy/wKwszJlnLTPzAK+NUlfzagOE1ymtZTXhAYK260XyFYmg/te/C231+Fr/hoX+EJoUBnrn0gD7hqMISOT+TaFEuOXYsN517GfaxgB"
func testNativeGemini3ThoughtSignature() string { inner := protowire.AppendTag(nil, 1, protowire.BytesType) inner = protowire.AppendBytes(inner, []byte{0x01, 0x0c, 0x39, 0xd6, 0xc7, 0x34}) encoded := protowire.AppendTag(nil, 2, protowire.BytesType) encoded = protowire.AppendBytes(encoded, inner) return base64.StdEncoding.EncodeToString(encoded)}
func claudeRequestWithThinkingSignature(sig string) (cliproxyexecutor.Request, cliproxyexecutor.Options) { req := cliproxyexecutor.Request{ Model: "gemini-2.5-flash", Payload: []byte(`{ "model": "claude-3-7-sonnet-20250219", "messages": [ { "role": "assistant", "content": [ {"type": "thinking", "thinking": "Let me think...", "signature": "` + sig + `"}, {"type": "text", "text": "Here is the response."} ] }, { "role": "user", "content": [ {"type": "text", "text": "Follow up question."} ] } ] }`), Metadata: map[string]any{ "cliproxy.resolved_api_key_model_info": ®istry.ModelInfo{IsCompat: true}, }, } opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatClaude, } return req, opts}
func TestGeminiExecutorExecute_SanitizesClaudeCAISSignature(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}`)) })) defer server.Close()
executor := NewGeminiExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Attributes: map[string]string{ "api_key": "test-api-key", "base_url": server.URL, }, }
req, opts := claudeRequestWithThinkingSignature(testClaudeCAISSample)
_, err := executor.Execute(context.Background(), auth, req, opts) if err != nil { t.Fatalf("Execute() error = %v", err) }
if bytes.Contains(upstreamBody, []byte(testClaudeCAISSample)) { t.Fatalf("upstream request leaked raw Claude CAIS signature: %s", upstreamBody) }
contents := gjson.GetBytes(upstreamBody, "contents").Array() for _, content := range contents { if content.Get("role").String() == "model" { for _, part := range content.Get("parts").Array() { if sig := part.Get("thoughtSignature").String(); sig == testClaudeCAISSample { t.Fatalf("model part thoughtSignature contains raw Claude CAIS signature: %s", upstreamBody) } } } }}
func TestGeminiExecutorExecuteStream_SanitizesClaudeCAISSignature(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"chunk\"}]}}]}\n\n")) })) defer server.Close()
executor := NewGeminiExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Attributes: map[string]string{ "api_key": "test-api-key", "base_url": server.URL, }, }
req, opts := claudeRequestWithThinkingSignature(testClaudeCAISSample)
res, err := executor.ExecuteStream(context.Background(), auth, req, opts) if err != nil { t.Fatalf("ExecuteStream() error = %v", err) } for range res.Chunks { }
if bytes.Contains(upstreamBody, []byte(testClaudeCAISSample)) { t.Fatalf("upstream stream request leaked raw Claude CAIS signature: %s", upstreamBody) }}
func TestGeminiExecutorCountTokens_SanitizesClaudeCAISSignature(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"totalTokens": 42}`)) })) defer server.Close()
executor := NewGeminiExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Attributes: map[string]string{ "api_key": "test-api-key", "base_url": server.URL, }, }
req, opts := claudeRequestWithThinkingSignature(testClaudeCAISSample)
_, err := executor.CountTokens(context.Background(), auth, req, opts) if err != nil { t.Fatalf("CountTokens() error = %v", err) }
if bytes.Contains(upstreamBody, []byte(testClaudeCAISSample)) { t.Fatalf("upstream countTokens request leaked raw Claude CAIS signature: %s", upstreamBody) }}
func TestGeminiExecutorExecute_FunctionCall_ReplacesClaudeSignatureWithBypass(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`)) })) defer server.Close()
executor := NewGeminiExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Attributes: map[string]string{ "api_key": "test-api-key", "base_url": server.URL, }, }
reqPayload := []byte(`{ "contents": [ { "role": "model", "parts": [ { "functionCall": {"name": "search", "args": {"q": "go"}}, "thoughtSignature": "` + testClaudeCAISSample + `" } ] }, { "role": "user", "parts": [ { "functionResponse": {"name": "search", "response": {"result": "found"}} } ] } ] }`)
req := cliproxyexecutor.Request{ Model: "gemini-2.5-flash", Payload: reqPayload, }
opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatGemini, }
_, err := executor.Execute(context.Background(), auth, req, opts) if err != nil { t.Fatalf("Execute() error = %v", err) }
gotSig := gjson.GetBytes(upstreamBody, "contents.1.parts.0.thoughtSignature").String() if gotSig != internalsignature.GeminiSkipThoughtSignatureValidator { t.Fatalf("first functionCall thoughtSignature = %q, want bypass sentinel %q; upstreamBody=%s", gotSig, internalsignature.GeminiSkipThoughtSignatureValidator, upstreamBody) }}
func TestGeminiExecutorExecute_PreservesNativeGeminiSignature(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`)) })) defer server.Close()
nativeSig := testNativeGemini3ThoughtSignature() executor := NewGeminiExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Attributes: map[string]string{ "api_key": "test-api-key", "base_url": server.URL, }, }
reqPayload := []byte(`{ "contents": [ { "role": "model", "parts": [ { "functionCall": {"name": "search", "args": {"q": "go"}}, "thoughtSignature": "` + nativeSig + `" } ] }, { "role": "user", "parts": [ { "functionResponse": {"name": "search", "response": {"result": "found"}} } ] } ] }`)
req := cliproxyexecutor.Request{ Model: "gemini-2.5-flash", Payload: reqPayload, }
opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatGemini, }
_, err := executor.Execute(context.Background(), auth, req, opts) if err != nil { t.Fatalf("Execute() error = %v", err) }
gotSig := gjson.GetBytes(upstreamBody, "contents.1.parts.0.thoughtSignature").String() if gotSig != nativeSig { t.Fatalf("thoughtSignature = %q, want preserved native signature %q; upstreamBody=%s", gotSig, nativeSig, upstreamBody) }}
func TestGeminiExecutorExecute_UnsignedRequestNotCorrupted(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`)) })) defer server.Close()
executor := NewGeminiExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Attributes: map[string]string{ "api_key": "test-api-key", "base_url": server.URL, }, }
reqPayload := []byte(`{ "contents": [ { "role": "user", "parts": [{"text": "Hello world"}] } ] }`)
req := cliproxyexecutor.Request{ Model: "gemini-2.5-flash", Payload: reqPayload, }
opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatGemini, }
_, err := executor.Execute(context.Background(), auth, req, opts) if err != nil { t.Fatalf("Execute() error = %v", err) }
text := gjson.GetBytes(upstreamBody, "contents.0.parts.0.text").String() if text != "Hello world" { t.Fatalf("text = %q, want 'Hello world'; upstreamBody=%s", text, upstreamBody) }}
func geminiRequestWithThinkingSignature(sig string) (cliproxyexecutor.Request, cliproxyexecutor.Options) { req := cliproxyexecutor.Request{ Model: "gemini-2.5-flash", Payload: []byte(`{ "contents": [ { "role": "model", "parts": [ {"text": "Let me think...", "thought": true, "thoughtSignature": "` + sig + `"}, {"text": "Here is the response."} ] }, { "role": "user", "parts": [ {"text": "Follow up question."} ] } ] }`), } opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatGemini, } return req, opts}
func TestGeminiVertexExecutorExecute_GeminiPayload_SanitizesClaudeCAISSignature(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`)) })) defer server.Close()
executor := NewGeminiVertexExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "vertex", Attributes: map[string]string{ "api_key": "test-vertex-key", "base_url": server.URL, }, }
req, opts := geminiRequestWithThinkingSignature(testClaudeCAISSample)
_, err := executor.Execute(context.Background(), auth, req, opts) if err != nil { t.Fatalf("Execute() error = %v", err) }
if bytes.Contains(upstreamBody, []byte(testClaudeCAISSample)) { t.Fatalf("vertex upstream request leaked raw Claude CAIS signature: %s", upstreamBody) }}
func TestGeminiVertexExecutorExecuteStream_GeminiPayload_SanitizesClaudeCAISSignature(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"chunk\"}]}}]}\n\n")) })) defer server.Close()
executor := NewGeminiVertexExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "vertex", Attributes: map[string]string{ "api_key": "test-vertex-key", "base_url": server.URL, }, }
req, opts := geminiRequestWithThinkingSignature(testClaudeCAISSample)
res, err := executor.ExecuteStream(context.Background(), auth, req, opts) if err != nil { t.Fatalf("ExecuteStream() error = %v", err) } for range res.Chunks { }
if bytes.Contains(upstreamBody, []byte(testClaudeCAISSample)) { t.Fatalf("vertex stream upstream request leaked raw Claude CAIS signature: %s", upstreamBody) }}
func TestGeminiVertexExecutorCountTokens_GeminiPayload_SanitizesClaudeCAISSignature(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"totalTokens": 42}`)) })) defer server.Close()
executor := NewGeminiVertexExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "vertex", Attributes: map[string]string{ "api_key": "test-vertex-key", "base_url": server.URL, }, }
req, opts := geminiRequestWithThinkingSignature(testClaudeCAISSample)
ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) _, err := executor.CountTokens(ctx, auth, req, opts) if err != nil { t.Fatalf("CountTokens() error = %v", err) } if !cliproxyexecutor.UpstreamAttempted(ctx) { t.Fatal("CountTokens() did not mark the HTTP request as an upstream attempt") }
if bytes.Contains(upstreamBody, []byte(testClaudeCAISSample)) { t.Fatalf("vertex countTokens upstream request leaked raw Claude CAIS signature: %s", upstreamBody) }}
func TestGeminiVertexExecutorExecute_PreservesNativeGeminiSignature(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`)) })) defer server.Close()
nativeSig := testNativeGemini3ThoughtSignature() executor := NewGeminiVertexExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "vertex", Attributes: map[string]string{ "api_key": "test-vertex-key", "base_url": server.URL, }, }
reqPayload := []byte(`{ "contents": [ { "role": "model", "parts": [ { "functionCall": {"name": "search", "args": {"q": "go"}}, "thoughtSignature": "` + nativeSig + `" } ] }, { "role": "user", "parts": [ { "functionResponse": {"name": "search", "response": {"result": "found"}} } ] } ] }`)
req := cliproxyexecutor.Request{ Model: "gemini-2.5-flash", Payload: reqPayload, }
opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatGemini, }
_, err := executor.Execute(context.Background(), auth, req, opts) if err != nil { t.Fatalf("Execute() error = %v", err) }
gotSig := gjson.GetBytes(upstreamBody, "contents.1.parts.0.thoughtSignature").String() if gotSig != nativeSig { t.Fatalf("thoughtSignature = %q, want preserved native signature %q; upstreamBody=%s", gotSig, nativeSig, upstreamBody) }}
func TestGeminiVertexExecutorExecute_UnsignedRequestNotCorrupted(t *testing.T) { var upstreamBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, errRead := io.ReadAll(r.Body) if errRead != nil { t.Fatalf("read request body: %v", errRead) } upstreamBody = append([]byte(nil), body...) w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}]}`)) })) defer server.Close()
executor := NewGeminiVertexExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "vertex", Attributes: map[string]string{ "api_key": "test-vertex-key", "base_url": server.URL, }, }
reqPayload := []byte(`{ "contents": [ { "role": "user", "parts": [{"text": "Hello world"}] } ] }`)
req := cliproxyexecutor.Request{ Model: "gemini-2.5-flash", Payload: reqPayload, }
opts := cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatGemini, }
_, err := executor.Execute(context.Background(), auth, req, opts) if err != nil { t.Fatalf("Execute() error = %v", err) }
text := gjson.GetBytes(upstreamBody, "contents.0.parts.0.text").String() if text != "Hello world" { t.Fatalf("text = %q, want 'Hello world'; upstreamBody=%s", text, upstreamBody) }}