package executor import (	"context"	"fmt"	"net/http"	"net/http/httptest"	"strings"	"sync"	"testing"	"time" 	"github.com/gorilla/websocket"	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	"github.com/router-for-me/CLIProxyAPI/v7/internal/wsrelay"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"	"github.com/tidwall/gjson") func TestAIStudioTranslateRequestPreservesSummaryFromOriginalRequest(t *testing.T) {	executor := NewAIStudioExecutor(&config.Config{}, "aistudio", nil)	req := cliproxyexecutor.Request{		Model:   "gemini-3.6-flash",		Payload: []byte(`{"model":"gemini-3.6-flash","input":"hi"}`),	}	opts := cliproxyexecutor.Options{		SourceFormat:    sdktranslator.FormatOpenAIResponse,		OriginalRequest: []byte(`{"model":"gemini-3.6-flash","reasoning":{"summary":"auto"},"input":"hi"}`),	}	payload, _, err := executor.translateRequest(context.Background(), req, opts, false)	if err != nil {		t.Fatalf("translateRequest() error = %v", err)	}	if !gjson.GetBytes(payload, "generationConfig.thinkingConfig.includeThoughts").Bool() {		t.Fatalf("original request summary intent was lost: %s", payload)	}} func TestAIStudioTranslateRequestNormalizesThinkingLevel(t *testing.T) {	executor := NewAIStudioExecutor(&config.Config{}, "aistudio", nil)	req := cliproxyexecutor.Request{		Model: "gemini-3.7-flash",		Payload: []byte(`{			"contents":[{"role":"user","parts":[{"text":"Human: Hi"}]}],			"model":"gemini-3.7-flash",			"generationConfig":{"thinkingConfig":{"thinkingLevel":"high","includeThoughts":true}}		}`),	}	opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini} 	payload, _, err := executor.translateRequest(context.Background(), req, opts, false)	if err != nil {		t.Fatalf("translateRequest() error = %v", err)	}	if got := gjson.GetBytes(payload, "generationConfig.thinkingConfig.thinkingLevel").String(); got != "HIGH" {		t.Fatalf("thinkingLevel = %q, want HIGH; payload=%s", got, payload)	}	if !gjson.GetBytes(payload, "generationConfig.thinkingConfig.includeThoughts").Bool() {		t.Fatalf("includeThoughts was lost: %s", payload)	}} func TestAIStudioTranslateRequestNormalizesThinkingLevelAfterPayloadOverride(t *testing.T) {	executor := NewAIStudioExecutor(&config.Config{Payload: config.PayloadConfig{Override: []config.PayloadRule{{		Models: []config.PayloadModelRule{{Name: "gemini-3.7-flash", Protocol: "gemini"}},		Params: map[string]any{"generationConfig.thinkingConfig.thinkingLevel": "medium"},	}}}}, "aistudio", nil)	req := cliproxyexecutor.Request{		Model:   "gemini-3.7-flash",		Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}],"model":"gemini-3.7-flash"}`),	}	opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini} 	payload, _, err := executor.translateRequest(context.Background(), req, opts, false)	if err != nil {		t.Fatalf("translateRequest() error = %v", err)	}	if got := gjson.GetBytes(payload, "generationConfig.thinkingConfig.thinkingLevel").String(); got != "MEDIUM" {		t.Fatalf("thinkingLevel = %q, want MEDIUM; payload=%s", got, payload)	}} func TestNormalizeAIStudioThinkingLevel(t *testing.T) {	tests := []struct {		name      string		levelJSON string		want      string		unchanged bool	}{		{name: "minimal", levelJSON: `"minimal"`, want: "MINIMAL"},		{name: "low", levelJSON: `"low"`, want: "LOW"},		{name: "medium", levelJSON: `"medium"`, want: "MEDIUM"},		{name: "mixed case high", levelJSON: `"hIgH"`, want: "HIGH"},		{name: "already uppercase", levelJSON: `"HIGH"`, want: "HIGH", unchanged: true},		{name: "none", levelJSON: `"none"`, want: "none", unchanged: true},		{name: "xhigh", levelJSON: `"xhigh"`, want: "xhigh", unchanged: true},		{name: "max", levelJSON: `"max"`, want: "max", unchanged: true},		{name: "custom", levelJSON: `"custom"`, want: "custom", unchanged: true},		{name: "non-string", levelJSON: `1`, want: "1", unchanged: true},	} 	for _, test := range tests {		t.Run(test.name, func(t *testing.T) {			input := []byte(fmt.Sprintf(`{"generationConfig":{"thinkingConfig":{"thinkingLevel":%s}},"marker":true}`, test.levelJSON))			output := normalizeAIStudioThinkingLevel(input)			level := gjson.GetBytes(output, "generationConfig.thinkingConfig.thinkingLevel")			if got := level.String(); got != test.want {				t.Fatalf("thinkingLevel = %q, want %q; payload=%s", got, test.want, output)			}			if test.unchanged && string(output) != string(input) {				t.Fatalf("payload changed from %s to %s", input, output)			}		})	}} func TestAIStudioTranslateRequestPrependsLeadingUserForIssue4959ResponsesHistory(t *testing.T) {	executor := NewAIStudioExecutor(&config.Config{}, "aistudio", nil)	_, body, err := executor.translateRequest(context.Background(), cliproxyexecutor.Request{		Model:   "gemini-3.7-flash-high",		Payload: issue4959ResponsesModelFirstPayload(),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}, false)	if err != nil {		t.Fatalf("translateRequest() error = %v", err)	}	assertIssue4959LeadingUserContents(t, gjson.GetBytes(body.payload, "contents").Array())} func TestAIStudioTranslateRequestAppendsTrailingUserForTrailingModelTurn(t *testing.T) {	executor := NewAIStudioExecutor(&config.Config{}, "aistudio", nil)	_, body, err := executor.translateRequest(context.Background(), cliproxyexecutor.Request{		Model: "gemini-3.7-flash",		Payload: []byte(`{"contents":[` +			`{"role":"user","parts":[{"text":"hello"}]},` +			`{"role":"model","parts":[{"text":"answer"}]}` +			`]}`),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini}, false)	if err != nil {		t.Fatalf("translateRequest() error = %v", err)	}	contents := gjson.GetBytes(body.payload, "contents").Array()	if len(contents) != 3 || contents[0].Get("role").String() != "user" || contents[1].Get("role").String() != "model" || contents[2].Get("role").String() != "user" {		t.Fatalf("contents roles malformed: %s", body.payload)	}	if got := contents[2].Get("parts.0.text").String(); got != "" {		t.Fatalf("trailing user prompt = %q, want empty string; body=%s", got, body.payload)	}} func TestAIStudioTranslateRequestCountTokensPreservesTrailingModelTurn(t *testing.T) {	executor := NewAIStudioExecutor(&config.Config{}, "aistudio", nil)	// When action is countTokens, trailing model turn must not be modified	_, body, err := executor.translateRequest(context.Background(), cliproxyexecutor.Request{		Model: "gemini-3.7-flash",		Payload: []byte(`{"contents":[` +			`{"role":"user","parts":[{"text":"hello"}]},` +			`{"role":"model","parts":[{"text":"answer"}]}` +			`]}`),		Metadata: map[string]any{"action": "countTokens"},	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini}, false)	if err != nil {		t.Fatalf("translateRequest() error = %v", err)	}	contents := gjson.GetBytes(body.payload, "contents").Array()	if len(contents) != 2 || contents[0].Get("role").String() != "user" || contents[1].Get("role").String() != "model" {		t.Fatalf("countTokens contents roles malformed: %s", body.payload)	}} func TestAIStudioExecutorWithoutRelaySessionDoesNotMarkUpstreamAttempt(t *testing.T) {	const authID = "aistudio-not-connected"	relay := wsrelay.NewManager(wsrelay.Options{})	exec := NewAIStudioExecutor(&config.Config{}, "aistudio", relay)	auth := &cliproxyauth.Auth{ID: authID, Provider: "aistudio"}	req := cliproxyexecutor.Request{		Model:   "gemini-3.1-pro-preview",		Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`),	}	opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini} 	tests := []struct {		name string		run  func(context.Context) error	}{		{			name: "HTTP request",			run: func(ctx context.Context) error {				httpReq, errRequest := http.NewRequestWithContext(ctx, http.MethodPost, "https://example.com/generate", strings.NewReader(`{"contents":[]}`))				if errRequest != nil {					return errRequest				}				_, errRequest = exec.HttpRequest(ctx, auth, httpReq)				return errRequest			},		},		{			name: "execute",			run: func(ctx context.Context) error {				_, errExecute := exec.Execute(ctx, auth, req, opts)				return errExecute			},		},		{			name: "stream",			run: func(ctx context.Context) error {				_, errStream := exec.ExecuteStream(ctx, auth, req, opts)				return errStream			},		},		{			name: "count tokens",			run: func(ctx context.Context) error {				_, errCount := exec.CountTokens(ctx, auth, req, opts)				return errCount			},		},	}	for _, test := range tests {		t.Run(test.name, func(t *testing.T) {			ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background())			errRun := test.run(ctx)			if errRun == nil || !strings.Contains(errRun.Error(), "not connected") {				t.Fatalf("request error = %v, want provider not connected", errRun)			}			if cliproxyexecutor.UpstreamAttempted(ctx) {				t.Fatal("missing relay session was marked as an upstream attempt")			}		})	}} func TestAIStudioExecutorExecuteStartsTTFTBeforeRelayWait(t *testing.T) {	const authID = "aistudio-ttft-auth"	delay := 40 * time.Millisecond	connected := make(chan struct{})	var connectedOnce sync.Once	relay := wsrelay.NewManager(wsrelay.Options{		ProviderFactory: func(*http.Request) (string, error) {			return authID, nil		},		OnConnected: func(provider string) {			if provider == authID {				connectedOnce.Do(func() {					close(connected)				})			}		},	})	server := httptest.NewServer(relay.Handler())	defer server.Close()	defer func() {		if errStop := relay.Stop(context.Background()); errStop != nil {			t.Errorf("relay stop error = %v", errStop)		}	}() 	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + relay.Path()	conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, nil)	if errDial != nil {		t.Fatalf("dial websocket: %v", errDial)	}	defer func() {		if errClose := conn.Close(); errClose != nil {			t.Errorf("websocket close error = %v", errClose)		}	}()	select {	case <-connected:	case <-time.After(time.Second):		t.Fatal("timed out waiting for relay connection")	} 	clientDone := make(chan error, 1)	go func() {		var msg wsrelay.Message		if errReadJSON := conn.ReadJSON(&msg); errReadJSON != nil {			clientDone <- fmt.Errorf("read relay request: %w", errReadJSON)			return		}		if msg.Type != wsrelay.MessageTypeHTTPReq {			clientDone <- fmt.Errorf("relay message type = %q, want %q", msg.Type, wsrelay.MessageTypeHTTPReq)			return		}		time.Sleep(delay)		response := wsrelay.Message{			ID:   msg.ID,			Type: wsrelay.MessageTypeHTTPResp,			Payload: map[string]any{				"status":  float64(http.StatusOK),				"headers": map[string]any{"Content-Type": "application/json"},				"body":    `{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}`,			},		}		if errWriteJSON := conn.WriteJSON(response); errWriteJSON != nil {			clientDone <- fmt.Errorf("write relay response: %w", errWriteJSON)			return		}		clientDone <- nil	}() 	plugin := &captureAIStudioUsagePlugin{records: make(chan usage.Record, 16)}	usage.RegisterPlugin(plugin)	exec := NewAIStudioExecutor(&config.Config{}, "aistudio", relay)	ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background())	_, errExecute := exec.Execute(ctx, &cliproxyauth.Auth{ID: authID, Provider: "aistudio"}, cliproxyexecutor.Request{		Model:   "gemini-3.1-pro-preview",		Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`),	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatGemini})	if errExecute != nil {		t.Fatalf("Execute() error = %v", errExecute)	}	if !cliproxyexecutor.UpstreamAttempted(ctx) {		t.Fatal("relay request did not mark an upstream attempt")	}	if errClient := <-clientDone; errClient != nil {		t.Fatal(errClient)	} 	record := waitForAIStudioUsageRecord(t, plugin.records, "gemini-3.1-pro-preview")	if record.TTFT < delay {		t.Fatalf("ttft = %v, want >= %v", record.TTFT, delay)	}} type captureAIStudioUsagePlugin struct {	records chan usage.Record} func (p *captureAIStudioUsagePlugin) HandleUsage(_ context.Context, record usage.Record) {	if p == nil {		return	}	select {	case p.records <- record:	default:	}} func waitForAIStudioUsageRecord(t *testing.T, records <-chan usage.Record, model string) usage.Record {	t.Helper()	timeout := time.After(2 * time.Second)	for {		select {		case record := <-records:			if record.Provider == "aistudio" && record.Model == model {				return record			}		case <-timeout:			t.Fatalf("timed out waiting for AI Studio usage record")		}	}}