mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-03 06:35:00 +08:00
fix(executor): sanitize thought signatures in Gemini and Gemini Vertex executors (#5110)
Ensure all Gemini and Gemini Vertex execution and token count pathways invoke SanitizeGeminiRequestThoughtSignatures before dispatching upstream requests. This prevents raw non-Gemini (e.g. Claude CAIS) thought signatures from leaking to upstream Gemini endpoints while preserving valid native protobuf signatures and setting appropriate validator bypass sentinels for function calls. Co-authored-by: W ARELIK <warelik@WARELIK-MB.local>
This commit is contained in:
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
|
||||
internalsignature "github.com/router-for-me/CLIProxyAPI/v7/internal/signature"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
@@ -162,6 +163,7 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
|
||||
body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
|
||||
body = helps.SetStringIfDifferent(body, "model", baseModel)
|
||||
body = capGeminiMaxOutputTokens(body, baseModel)
|
||||
body = internalsignature.SanitizeGeminiRequestThoughtSignatures(body, "contents")
|
||||
|
||||
action := "generateContent"
|
||||
if req.Metadata != nil {
|
||||
@@ -279,6 +281,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
|
||||
body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
|
||||
body = helps.SetStringIfDifferent(body, "model", baseModel)
|
||||
body = capGeminiMaxOutputTokens(body, baseModel)
|
||||
body = internalsignature.SanitizeGeminiRequestThoughtSignatures(body, "contents")
|
||||
body = helps.EnsureGeminiLeadingUserContent(body, "contents")
|
||||
|
||||
baseURL := resolveGeminiBaseURL(auth)
|
||||
@@ -645,6 +648,7 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
||||
translatedReq = helps.SetStringIfDifferent(translatedReq, "model", baseModel)
|
||||
translatedReq = internalsignature.SanitizeGeminiRequestThoughtSignatures(translatedReq, "contents")
|
||||
translatedReq = helps.EnsureGeminiLeadingUserContent(translatedReq, "contents")
|
||||
|
||||
baseURL := resolveGeminiBaseURL(auth)
|
||||
|
||||
601
internal/runtime/executor/gemini_executor_signature_test.go
Normal file
601
internal/runtime/executor/gemini_executor_signature_test.go
Normal file
@@ -0,0 +1,601 @@
|
||||
package 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)
|
||||
|
||||
_, 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("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)
|
||||
}
|
||||
}
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
vertexauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/vertex"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
|
||||
internalsignature "github.com/router-for-me/CLIProxyAPI/v7/internal/signature"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
@@ -342,6 +343,7 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au
|
||||
body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
|
||||
body = helps.SetStringIfDifferent(body, "model", baseModel)
|
||||
body = helps.StripVertexOpenAIResponsesToolCallIDs(body, from.String())
|
||||
body = internalsignature.SanitizeGeminiRequestThoughtSignatures(body, "contents")
|
||||
}
|
||||
|
||||
action := getVertexAction(baseModel, false)
|
||||
@@ -471,6 +473,7 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip
|
||||
body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
|
||||
body = helps.SetStringIfDifferent(body, "model", baseModel)
|
||||
body = helps.StripVertexOpenAIResponsesToolCallIDs(body, from.String())
|
||||
body = internalsignature.SanitizeGeminiRequestThoughtSignatures(body, "contents")
|
||||
|
||||
action := getVertexAction(baseModel, false)
|
||||
if req.Metadata != nil {
|
||||
@@ -590,6 +593,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte
|
||||
body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
|
||||
body = helps.SetStringIfDifferent(body, "model", baseModel)
|
||||
body = helps.StripVertexOpenAIResponsesToolCallIDs(body, from.String())
|
||||
body = internalsignature.SanitizeGeminiRequestThoughtSignatures(body, "contents")
|
||||
|
||||
action := getVertexAction(baseModel, true)
|
||||
body = helps.EnsureGeminiLeadingUserContent(body, "contents")
|
||||
@@ -738,6 +742,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth
|
||||
body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
|
||||
body = helps.SetStringIfDifferent(body, "model", baseModel)
|
||||
body = helps.StripVertexOpenAIResponsesToolCallIDs(body, from.String())
|
||||
body = internalsignature.SanitizeGeminiRequestThoughtSignatures(body, "contents")
|
||||
|
||||
action := getVertexAction(baseModel, true)
|
||||
body = helps.EnsureGeminiLeadingUserContent(body, "contents")
|
||||
@@ -878,6 +883,7 @@ func (e *GeminiVertexExecutor) countTokensWithServiceAccount(ctx context.Context
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
||||
translatedReq = internalsignature.SanitizeGeminiRequestThoughtSignatures(translatedReq, "contents")
|
||||
translatedReq = helps.EnsureGeminiLeadingUserContent(translatedReq, "contents")
|
||||
|
||||
baseURL := vertexBaseURL(location)
|
||||
@@ -970,6 +976,7 @@ func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth *
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
||||
translatedReq = internalsignature.SanitizeGeminiRequestThoughtSignatures(translatedReq, "contents")
|
||||
translatedReq = helps.EnsureGeminiLeadingUserContent(translatedReq, "contents")
|
||||
|
||||
// For API key auth, use simpler URL format without project/location
|
||||
|
||||
Reference in New Issue
Block a user