Add projects

backend/internal/runtime/executor/aistudio_executor_test.go

@@ -0,0 +1,170 @@
+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 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 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)
+	_, errExecute := exec.Execute(context.Background(), &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 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")
+		}
+	}
+}