486 lines · 21.9 KB · Go
RawDownload
package executor import (	"context"	"errors"	"io"	"net/http"	"strings"	"testing" 	"github.com/gin-gonic/gin"	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	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") // midSystemLegacyPayload is a caller body pairing a legacy model with a// mid-conversation role=system turn. The turn ends the array, so the shape is// rejected by the model rather than by Anthropic's ordering rule.func midSystemLegacyPayload(model string) []byte {	return []byte(`{"model":"` + model + `","max_tokens":32,` +		`"system":[{"type":"text","text":"Top rule"}],` +		`"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]},` +		`{"role":"system","content":[{"type":"text","text":"Mid rule"}]}],` +		`"metadata":{"user_id":"{\"device_id\":\"0000000000000000000000000000000000000000000000000000000000000000\",\"account_uuid\":\"\",\"session_id\":\"11111111-2222-4333-8444-555555555555\"}"}}`)} // midSystemUpstream intercepts the transport instead of standing up a test// server, so the executor keeps the default https://api.anthropic.com base URL.// The guard only fires on Anthropic's first-party origin, which a httptest// server address would not satisfy.type midSystemUpstream struct {	body    []byte	called  bool	headers http.Header} func (u *midSystemUpstream) context(t *testing.T, headers http.Header) context.Context {	t.Helper()	gin.SetMode(gin.TestMode)	ginCtx, _ := gin.CreateTestContext(nil)	ginCtx.Request = httptest_NewRequest()	ginCtx.Request.Header = headers.Clone()	if ginCtx.Request.Header == nil {		ginCtx.Request.Header = make(http.Header)	}	transport := roundTripperFunc(func(req *http.Request) (*http.Response, error) {		payload, errRead := io.ReadAll(req.Body)		if errRead != nil {			t.Fatal(errRead)		}		u.body = payload		u.called = true		u.headers = req.Header.Clone()		contentType := "application/json"		responseBody := `{"id":"msg_1","type":"message","role":"assistant","model":"m","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`		if strings.Contains(req.URL.Path, "count_tokens") {			responseBody = `{"input_tokens":18}`		} else if gjson.GetBytes(payload, "stream").Bool() {			contentType = "text/event-stream"			// A translated caller aggregates the stream back into one message, so			// the stub has to complete the block and report a stop reason.			responseBody = strings.Join([]string{				`event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"m","content":[],"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":0}}}`,				`event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,				`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}`,				`event: content_block_stop` + "\n" + `data: {"type":"content_block_stop","index":0}`,				`event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}`,				`event: message_stop` + "\n" + `data: {"type":"message_stop"}`,			}, "\n\n") + "\n\n"		}		return &http.Response{			StatusCode: http.StatusOK,			Header:     http.Header{"Content-Type": []string{contentType}},			Body:       io.NopCloser(strings.NewReader(responseBody)),			Request:    req,		}, nil	})	ctx := context.WithValue(context.Background(), "gin", ginCtx)	return context.WithValue(ctx, "cliproxy.roundtripper", http.RoundTripper(transport))} func httptest_NewRequest() *http.Request {	req, _ := http.NewRequest(http.MethodPost, "http://example.invalid/", nil)	return req} func midSystemAuth() *cliproxyauth.Auth {	// No base_url, so the executor keeps Anthropic's first-party origin.	return &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "key-123", "cloak_mode": "always"}}} func midSystemConfig() *config.Config {	return &config.Config{ClaudeKey: []config.ClaudeKey{{APIKey: "key-123"}}}} func assertMidSystemRejected(t *testing.T, err error, upstream *midSystemUpstream) {	t.Helper()	if err == nil {		t.Fatal("error = nil, want the legacy pairing rejected")	}	if upstream.called {		t.Fatalf("upstream must not be called for a guaranteed rejection; got %s", upstream.body)	}	var statusCoder interface{ StatusCode() int }	if !errors.As(err, &statusCoder) || statusCoder.StatusCode() != http.StatusBadRequest {		t.Fatalf("error = %v, want a 400 status error", err)	}	var scoped interface{ IsRequestScoped() bool }	if !errors.As(err, &scoped) || !scoped.IsRequestScoped() {		t.Fatalf("error = %v, want a request-scoped error so no credential is retried", err)	}	if !strings.Contains(err.Error(), "role 'system' is not supported on this model") {		t.Fatalf("error = %v, want Anthropic's wording preserved", err)	}} // Every executor path that can send the pairing to Anthropic must answer it// locally instead of spending an upstream call on a guaranteed 400.func TestClaudeExecutor_LegacyMidSystemMessageRejectedOnEveryUpstreamPath(t *testing.T) {	for _, test := range []struct {		name  string		model string		send  func(t *testing.T, ex *ClaudeExecutor, ctx context.Context, model string) error	}{		{name: "execute", model: "claude-haiku-4-5-20251001", send: sendMidSystemExecute},		{name: "execute stream", model: "claude-haiku-4-5-20251001", send: sendMidSystemStream},		{name: "count tokens", model: "claude-haiku-4-5-20251001", send: sendMidSystemCountTokens},		{name: "execute legacy sonnet", model: "claude-sonnet-4-6", send: sendMidSystemExecute},	} {		t.Run(test.name, func(t *testing.T) {			upstream := &midSystemUpstream{}			ex := NewClaudeExecutor(midSystemConfig())			err := test.send(t, ex, upstream.context(t, nil), test.model)			assertMidSystemRejected(t, err, upstream)		})	}} func sendMidSystemExecute(t *testing.T, ex *ClaudeExecutor, ctx context.Context, model string) error {	t.Helper()	_, err := ex.Execute(ctx, midSystemAuth(), cliproxyexecutor.Request{		Model: model, Payload: midSystemLegacyPayload(model),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude})	return err} func sendMidSystemStream(t *testing.T, ex *ClaudeExecutor, ctx context.Context, model string) error {	t.Helper()	result, err := ex.ExecuteStream(ctx, midSystemAuth(), cliproxyexecutor.Request{		Model: model, Payload: midSystemLegacyPayload(model),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude})	if err != nil {		return err	}	for chunk := range result.Chunks {		if chunk.Err != nil {			return chunk.Err		}	}	return nil} func sendMidSystemCountTokens(t *testing.T, ex *ClaudeExecutor, ctx context.Context, model string) error {	t.Helper()	_, err := ex.CountTokens(ctx, midSystemAuth(), cliproxyexecutor.Request{		Model: model, Payload: midSystemLegacyPayload(model),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude})	return err} // Payload rules run long after translation and can rewrite model and messages,// so the guard has to read the finished body rather than an intermediate one.func TestClaudeExecutor_PayloadOverrideCannotSmuggleLegacyMidSystemMessage(t *testing.T) {	upstream := &midSystemUpstream{}	cfg := midSystemConfig()	cfg.Payload.Override = []config.PayloadRule{{		Models: []config.PayloadModelRule{{Name: "*"}},		Params: map[string]any{"model": "claude-haiku-4-5-20251001"},	}}	ex := NewClaudeExecutor(cfg) 	// The caller addresses a model that accepts the turn; only the payload rule	// turns it into the rejected pairing.	_, err := ex.Execute(upstream.context(t, nil), midSystemAuth(), cliproxyexecutor.Request{		Model: "claude-sonnet-5", Payload: midSystemLegacyPayload("claude-sonnet-5"),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude})	assertMidSystemRejected(t, err, upstream)} // A caller may already have the exact role=system turn that cloaking would// otherwise insert. The message-count proof must keep that turn caller-owned,// so a later legacy model rewrite is rejected instead of silently consuming it.func TestClaudeExecutor_PayloadOverrideDoesNotClaimMatchingCallerTurn(t *testing.T) {	upstream := &midSystemUpstream{}	cfg := midSystemConfig()	cfg.Payload.Override = []config.PayloadRule{{		Models: []config.PayloadModelRule{{Name: "*"}},		Params: map[string]any{"model": "claude-haiku-4-5-20251001"},	}}	ex := NewClaudeExecutor(cfg)	payload := []byte(`{"model":"claude-sonnet-5","max_tokens":32,` +		`"system":[{"type":"text","text":"Same rule"}],` +		`"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]},` +		`{"role":"system","content":[{"type":"text","text":"Same rule"}]}]}`) 	_, err := ex.Execute(upstream.context(t, nil), midSystemAuth(), cliproxyexecutor.Request{		Model: "claude-sonnet-5", Payload: payload,	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude})	assertMidSystemRejected(t, err, upstream)} // Cloaking relocates a caller's system prompt into a role=system turn for models// that accept one. A payload rule can then rewrite the model to one that does// not. Because the caller never wrote that turn, CPA must reconcile its own// placement through the legacy reminder path instead of returning 400.func TestClaudeExecutor_PayloadOverrideReconcilesRelocatedSystemPrompt(t *testing.T) {	for _, test := range []struct {		name string		send func(t *testing.T, ex *ClaudeExecutor, ctx context.Context, payload []byte) error	}{		{name: "execute", send: func(t *testing.T, ex *ClaudeExecutor, ctx context.Context, payload []byte) error {			_, err := ex.Execute(ctx, midSystemAuth(), cliproxyexecutor.Request{				Model: "claude-sonnet-5", Payload: payload,			}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude})			return err		}},		{name: "execute stream", send: func(t *testing.T, ex *ClaudeExecutor, ctx context.Context, payload []byte) error {			result, err := ex.ExecuteStream(ctx, midSystemAuth(), cliproxyexecutor.Request{				Model: "claude-sonnet-5", Payload: payload,			}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude})			if err != nil {				return err			}			for chunk := range result.Chunks {				if chunk.Err != nil {					return chunk.Err				}			}			return nil		}},	} {		t.Run(test.name, func(t *testing.T) {			upstream := &midSystemUpstream{}			cfg := midSystemConfig()			cfg.Payload.Override = []config.PayloadRule{{				Models: []config.PayloadModelRule{{Name: "*"}},				Params: map[string]any{"model": "claude-haiku-4-5-20251001"},			}}			ex := NewClaudeExecutor(cfg) 			// Only a top-level system prompt: the caller never writes a			// role=system turn.			payload := []byte(`{"model":"claude-sonnet-5","max_tokens":32,` +				`"system":[{"type":"text","text":"Caller top"}],` +				`"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) 			if err := test.send(t, ex, upstream.context(t, nil), payload); err != nil {				t.Fatalf("request error = %v, want CPA's inserted turn reconciled", err)			}			if !upstream.called {				t.Fatal("expected reconciled request to reach upstream")			}			if got := gjson.GetBytes(upstream.body, "model").String(); got != "claude-haiku-4-5-20251001" {				t.Fatalf("upstream model = %q, want payload override preserved", got)			}			if gjson.GetBytes(upstream.body, `messages.#(role=="system")`).Exists() {				t.Fatalf("reconciled body still carries role=system; body=%s", upstream.body)			}			if !strings.Contains(gjson.GetBytes(upstream.body, "messages.0.content").Raw, "<system-reminder>") ||				!strings.Contains(gjson.GetBytes(upstream.body, "messages.0.content").Raw, "Caller top") {				t.Fatalf("caller system prompt was not replayed as a legacy reminder; body=%s", upstream.body)			}		})	}} // A confirmed native caller owns its wire. It gates the turn on the model// itself, so CPA forwards the body untouched and lets the upstream answer.func TestClaudeExecutor_ConfirmedNativeLegacyMidSystemMessageForwarded(t *testing.T) {	upstream := &midSystemUpstream{}	ex := NewClaudeExecutor(midSystemConfig())	headers := claudeNativeHelperHeaders("claude-code-20250219,"+claudeNativeHelperCoreBetas, "gzip, deflate, br, zstd") 	if _, err := ex.Execute(upstream.context(t, headers), midSystemAuth(), cliproxyexecutor.Request{		Model:   "claude-haiku-4-5-20251001",		Payload: midSystemLegacyPayload("claude-haiku-4-5-20251001"),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude, Headers: headers}); err != nil {		t.Fatalf("Execute() error = %v, want the native body forwarded", err)	}	if !upstream.called {		t.Fatal("expected the native request to reach the upstream")	}	if !gjson.GetBytes(upstream.body, `messages.#(role=="system")`).Exists() {		t.Fatalf("confirmed native caller lost its role=system turn; body=%s", upstream.body)	}} // The rejection was measured against api.anthropic.com. A third-party gateway// may map the same model ID onto something that accepts the turn, and answering// locally would also stop failover to another credential or base URL.func TestClaudeExecutor_LegacyMidSystemMessageForwardedToThirdPartyGateway(t *testing.T) {	upstream := &midSystemUpstream{}	ex := NewClaudeExecutor(&config.Config{ClaudeKey: []config.ClaudeKey{{		APIKey: "key-123", BaseURL: "https://gateway.example",	}}})	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"api_key": "key-123", "base_url": "https://gateway.example",	}} 	if _, err := ex.Execute(upstream.context(t, nil), auth, cliproxyexecutor.Request{		Model:   "claude-haiku-4-5-20251001",		Payload: midSystemLegacyPayload("claude-haiku-4-5-20251001"),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}); err != nil {		t.Fatalf("Execute() error = %v, want a third-party gateway to decide for itself", err)	}	if !upstream.called {		t.Fatal("expected the request to reach the third-party gateway")	}} // A model outside claudeLegacySystemReminderModels stays optimistic, matching// how checkSystemInstructions treats unknown and future IDs.func TestClaudeExecutor_SupportedModelMidSystemMessageForwarded(t *testing.T) {	for _, model := range []string{"claude-sonnet-5", "claude-sonnet-9"} {		t.Run(model, func(t *testing.T) {			upstream := &midSystemUpstream{}			ex := NewClaudeExecutor(midSystemConfig())			if _, err := ex.Execute(upstream.context(t, nil), midSystemAuth(), cliproxyexecutor.Request{				Model: model, Payload: midSystemLegacyPayload(model),			}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}); err != nil {				t.Fatalf("Execute() error = %v, want the request forwarded", err)			}			if !upstream.called {				t.Fatal("expected the request to reach the upstream")			}		})	}} // The opt-in rescues the pairing by folding the turn into the system slot, so// the guard must run after it rather than rejecting the request outright.func TestClaudeExecutor_LegacyMidSystemMessageOptInStillRebuilds(t *testing.T) {	upstream := &midSystemUpstream{}	ex := NewClaudeExecutor(midSystemConfig())	auth := midSystemAuth()	auth.Attributes["rebuild_mid_system_message"] = "true" 	if _, err := ex.Execute(upstream.context(t, nil), auth, cliproxyexecutor.Request{		Model:   "claude-haiku-4-5-20251001",		Payload: midSystemLegacyPayload("claude-haiku-4-5-20251001"),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}); err != nil {		t.Fatalf("Execute() error = %v, want the opt-in rebuild to rescue the request", err)	}	if !upstream.called {		t.Fatal("expected the rebuilt request to reach the upstream")	}	if gjson.GetBytes(upstream.body, `messages.#(role=="system")`).Exists() {		t.Fatalf("opt-in rebuild left a role=system turn; body=%s", upstream.body)	}} // The pairing must never originate inside CPA. A non-Claude caller reaches the// Claude executor through a translator, and every translator hoists system// content into the top-level system field, so no translated body can carry a// role=system turn to a legacy model. This pins that guarantee: the guard is for// callers that speak Claude natively, never for a translated request.func TestTranslatedRequestNeverPairsLegacyModelWithMidSystemMessage(t *testing.T) {	const legacyModel = "claude-haiku-4-5-20251001"	for _, test := range []struct {		name    string		format  sdktranslator.Format		payload string	}{		{name: "openai chat with a mid conversation system message", format: sdktranslator.FormatOpenAI,			payload: `{"model":"` + legacyModel + `","messages":[{"role":"system","content":"Top rule"},{"role":"user","content":"hi"},{"role":"system","content":"Mid rule"},{"role":"assistant","content":"ok"},{"role":"user","content":"go"}]}`},		{name: "openai chat ending on a system message", format: sdktranslator.FormatOpenAI,			payload: `{"model":"` + legacyModel + `","messages":[{"role":"user","content":"hi"},{"role":"system","content":"Mid rule"}]}`},		{name: "gemini with a system instruction", format: sdktranslator.FormatGemini,			payload: `{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"systemInstruction":{"parts":[{"text":"Top rule"}]}}`},		{name: "openai responses with instructions", format: sdktranslator.FormatOpenAIResponse,			payload: `{"model":"` + legacyModel + `","instructions":"Top rule","input":[{"role":"user","content":[{"type":"input_text","text":"hi"}]}]}`},		{name: "interactions with a system instruction", format: sdktranslator.FormatInteractions,			payload: `{"model":"` + legacyModel + `","system_instruction":"Top rule","input":[{"type":"user_input","content":[{"type":"text","text":"hi"}]}]}`},	} {		t.Run(test.name, func(t *testing.T) {			upstream := &midSystemUpstream{}			ex := NewClaudeExecutor(midSystemConfig()) 			if _, err := ex.Execute(upstream.context(t, nil), midSystemAuth(), cliproxyexecutor.Request{				Model: legacyModel, Payload: []byte(test.payload),			}, cliproxyexecutor.Options{SourceFormat: test.format}); err != nil {				t.Fatalf("Execute() error = %v, want the translated request forwarded", err)			}			if !upstream.called {				t.Fatal("expected the translated request to reach the upstream")			}			// Without these the subject under test could drift away: a body that			// no longer addresses the legacy model, or that lost the caller's			// turns, would satisfy the role assertions for the wrong reason.			if got := gjson.GetBytes(upstream.body, "model").String(); got != legacyModel {				t.Fatalf("upstream model = %q, want the legacy model %q under test", got, legacyModel)			}			if got := len(gjson.GetBytes(upstream.body, "messages").Array()); got == 0 {				t.Fatalf("upstream messages are empty, so the role assertions prove nothing; body=%s", upstream.body)			}			for _, role := range gjson.GetBytes(upstream.body, "messages.#.role").Array() {				if strings.EqualFold(role.String(), "system") {					t.Fatalf("translated body carries a system role; body=%s", upstream.body)				}			}		})	}} func TestClaudePayloadHasMidSystemMessage(t *testing.T) {	for _, test := range []struct {		name    string		payload string		want    bool	}{		{name: "mid conversation turn", want: true,			payload: `{"messages":[{"role":"user","content":"a"},{"role":"system","content":"s"}]}`},		{name: "role casing is ignored", want: true,			payload: `{"messages":[{"role":"SySTeM","content":"s"}]}`},		{name: "surrounding whitespace is ignored", want: true,			payload: `{"messages":[{"role":" system ","content":"s"}]}`},		{name: "only user and assistant turns",			payload: `{"messages":[{"role":"user","content":"a"},{"role":"assistant","content":"b"}]}`},		{name: "top level system field is not a turn",			payload: `{"system":[{"type":"text","text":"s"}],"messages":[{"role":"user","content":"a"}]}`},		{name: "messages missing", payload: `{"model":"claude-haiku-4-5"}`},		{name: "messages is not an array", payload: `{"messages":"system"}`},		{name: "messages holds a bare string", payload: `{"messages":["system"]}`},		{name: "system appears only in content", payload: `{"messages":[{"role":"user","content":"role: system"}]}`},	} {		t.Run(test.name, func(t *testing.T) {			if got := claudePayloadHasMidSystemMessage([]byte(test.payload)); got != test.want {				t.Fatalf("claudePayloadHasMidSystemMessage = %v, want %v", got, test.want)			}		})	}} func TestValidateClaudeMidSystemMessageModel(t *testing.T) {	const turn = `,"messages":[{"role":"user","content":"a"},{"role":"system","content":"s"}]}`	for _, test := range []struct {		name       string		payload    string		confirmed  bool		thirdParty bool		wantError  bool	}{		{name: "legacy model is rejected", wantError: true,			payload: `{"model":"claude-haiku-4-5-20251001"` + turn},		{name: "vendor prefixed legacy model is rejected", wantError: true,			payload: `{"model":"anthropic/claude-sonnet-4-6"` + turn},		{name: "model casing is ignored", wantError: true,			payload: `{"model":"Claude-Haiku-4-5-20251001"` + turn},		{name: "confirmed native keeps the passthrough", confirmed: true,			payload: `{"model":"claude-haiku-4-5-20251001"` + turn},		{name: "third party gateway decides for itself", thirdParty: true,			payload: `{"model":"claude-haiku-4-5-20251001"` + turn},		{name: "supported model is forwarded",			payload: `{"model":"claude-sonnet-5"` + turn},		{name: "unknown model stays optimistic",			payload: `{"model":"claude-sonnet-9"` + turn},		{name: "legacy model without the turn is forwarded",			payload: `{"model":"claude-haiku-4-5-20251001","messages":[{"role":"user","content":"a"}]}`},	} {		t.Run(test.name, func(t *testing.T) {			err := validateClaudeMidSystemMessageModel([]byte(test.payload), test.confirmed, !test.thirdParty)			if test.wantError != (err != nil) {				t.Fatalf("validateClaudeMidSystemMessageModel error = %v, want error %v", err, test.wantError)			}			if err == nil {				return			}			if !strings.Contains(err.Error(), gjson.Get(test.payload, "model").String()) {				t.Fatalf("error = %v, want the offending model named", err)			}		})	}}