444 lines · 15.5 KB · Go
RawDownload
package executor import (	"bytes"	"context"	"io"	"mime"	"mime/multipart"	"net/http"	"net/http/httptest"	"net/textproto"	"strings"	"testing" 	"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") func TestOpenAICompatExecutorCompactPassthrough(t *testing.T) {	var gotPath string	var gotBody []byte	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		gotPath = r.URL.Path		body, _ := io.ReadAll(r.Body)		gotBody = body		w.Header().Set("Content-Type", "application/json")		_, _ = w.Write([]byte(`{"id":"resp_1","object":"response.compaction","usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}`))	}))	defer server.Close() 	executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{})	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"base_url": server.URL + "/v1",		"api_key":  "test",	}}	payload := []byte(`{"model":"gpt-5.1-codex-max","input":[{"role":"user","content":"hi"}]}`)	resp, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{		Model:   "gpt-5.1-codex-max",		Payload: payload,	}, cliproxyexecutor.Options{		SourceFormat: sdktranslator.FromString("openai-response"),		Alt:          "responses/compact",		Stream:       false,	})	if err != nil {		t.Fatalf("Execute error: %v", err)	}	if gotPath != "/v1/responses/compact" {		t.Fatalf("path = %q, want %q", gotPath, "/v1/responses/compact")	}	if !gjson.GetBytes(gotBody, "input").Exists() {		t.Fatalf("expected input in body")	}	if gjson.GetBytes(gotBody, "messages").Exists() {		t.Fatalf("unexpected messages in body")	}	if string(resp.Payload) != `{"id":"resp_1","object":"response.compaction","usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}` {		t.Fatalf("payload = %s", string(resp.Payload))	}} func TestOpenAICompatExecutorPayloadOverrideWinsOverThinkingSuffix(t *testing.T) {	var gotBody []byte	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		body, _ := io.ReadAll(r.Body)		gotBody = body		w.Header().Set("Content-Type", "application/json")		_, _ = w.Write([]byte(`{"id":"chatcmpl_1","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`))	}))	defer server.Close() 	executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{		Payload: config.PayloadConfig{			Override: []config.PayloadRule{				{					Models: []config.PayloadModelRule{						{Name: "custom-openai", Protocol: "openai"},					},					Params: map[string]any{						"reasoning_effort": "low",					},				},			},		},	})	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"base_url": server.URL + "/v1",		"api_key":  "test",	}}	payload := []byte(`{"model":"custom-openai(high)","messages":[{"role":"user","content":"hi"}]}`)	_, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{		Model:   "custom-openai(high)",		Payload: payload,	}, cliproxyexecutor.Options{		SourceFormat: sdktranslator.FromString("openai"),		Stream:       false,	})	if err != nil {		t.Fatalf("Execute error: %v", err)	}	if got := gjson.GetBytes(gotBody, "reasoning_effort").String(); got != "low" {		t.Fatalf("reasoning_effort = %q, want %q; body=%s", got, "low", string(gotBody))	}} func TestOpenAICompatExecutorImagesGenerationsPassthrough(t *testing.T) {	var gotPath string	var gotBody []byte	var gotContentType string	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		gotPath = r.URL.Path		gotContentType = r.Header.Get("Content-Type")		body, _ := io.ReadAll(r.Body)		gotBody = body		w.Header().Set("Content-Type", "application/json")		_, _ = w.Write([]byte(`{"created":123,"data":[{"b64_json":"AA=="}],"usage":{"total_tokens":1}}`))	}))	defer server.Close() 	executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{})	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"base_url": server.URL + "/v1",		"api_key":  "test",	}}	resp, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{		Model:   "upstream-image",		Payload: []byte(`{"model":"compat-image","prompt":"draw"}`),	}, cliproxyexecutor.Options{		SourceFormat: sdktranslator.FromString("openai-image"),		Stream:       false,		Headers: http.Header{			"Content-Type": []string{"application/json"},		},		Metadata: map[string]any{			cliproxyexecutor.RequestPathMetadataKey: "/v1/images/generations",		},	})	if err != nil {		t.Fatalf("Execute error: %v", err)	}	if gotPath != "/v1/images/generations" {		t.Fatalf("path = %q, want %q", gotPath, "/v1/images/generations")	}	if gotContentType != "application/json" {		t.Fatalf("content type = %q, want application/json", gotContentType)	}	if got := gjson.GetBytes(gotBody, "model").String(); got != "upstream-image" {		t.Fatalf("model = %q, want upstream-image; body=%s", got, string(gotBody))	}	if got := gjson.GetBytes(resp.Payload, "data.0.b64_json").String(); got != "AA==" {		t.Fatalf("response payload = %s", string(resp.Payload))	}} func TestOpenAICompatExecutorImagesGenerationsStreamsUpstream(t *testing.T) {	var gotPath string	var gotBody []byte	var gotAccept string	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		gotPath = r.URL.Path		gotAccept = r.Header.Get("Accept")		body, _ := io.ReadAll(r.Body)		gotBody = body		w.Header().Set("Content-Type", "text/event-stream")		_, _ = w.Write([]byte("event: image_generation.partial\ndata: {\"type\":\"image_generation.partial\"}\n\n"))		if flusher, ok := w.(http.Flusher); ok {			flusher.Flush()		}		_, _ = w.Write([]byte("data: [DONE]\n\n"))	}))	defer server.Close() 	executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{})	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"base_url": server.URL + "/v1",		"api_key":  "test",	}}	streamResult, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{		Model:   "upstream-image",		Payload: []byte(`{"model":"compat-image","prompt":"draw","stream":true}`),	}, cliproxyexecutor.Options{		SourceFormat: sdktranslator.FromString("openai-image"),		Stream:       true,		Headers: http.Header{			"Content-Type": []string{"application/json"},		},		Metadata: map[string]any{			cliproxyexecutor.RequestPathMetadataKey: "/v1/images/generations",		},	})	if err != nil {		t.Fatalf("ExecuteStream error: %v", err)	}	var streamed bytes.Buffer	for chunk := range streamResult.Chunks {		if chunk.Err != nil {			t.Fatalf("stream chunk error: %v", chunk.Err)		}		streamed.Write(chunk.Payload)	}	if gotPath != "/v1/images/generations" {		t.Fatalf("path = %q, want %q", gotPath, "/v1/images/generations")	}	if gotAccept != "text/event-stream" {		t.Fatalf("accept = %q, want text/event-stream", gotAccept)	}	if got := gjson.GetBytes(gotBody, "model").String(); got != "upstream-image" {		t.Fatalf("model = %q, want upstream-image; body=%s", got, string(gotBody))	}	if !gjson.GetBytes(gotBody, "stream").Bool() {		t.Fatalf("stream flag missing from upstream body: %s", string(gotBody))	}	if !strings.Contains(streamed.String(), "event: image_generation.partial") || !strings.Contains(streamed.String(), "data: [DONE]") {		t.Fatalf("streamed body = %q", streamed.String())	}} func TestOpenAICompatExecutorImagesEditsMultipartRewritesModel(t *testing.T) {	var body bytes.Buffer	writer := multipart.NewWriter(&body)	if errWrite := writer.WriteField("model", "compat-image"); errWrite != nil {		t.Fatalf("write model field: %v", errWrite)	}	if errWrite := writer.WriteField("prompt", "edit"); errWrite != nil {		t.Fatalf("write prompt field: %v", errWrite)	}	header := make(textproto.MIMEHeader)	header.Set("Content-Disposition", multipart.FileContentDisposition("image", "image.png"))	header.Set("Content-Type", "image/png")	part, errCreate := writer.CreatePart(header)	if errCreate != nil {		t.Fatalf("create image field: %v", errCreate)	}	if _, errWrite := part.Write([]byte("png-data")); errWrite != nil {		t.Fatalf("write image field: %v", errWrite)	}	if errClose := writer.Close(); errClose != nil {		t.Fatalf("close multipart writer: %v", errClose)	}	contentType := writer.FormDataContentType() 	var gotPath string	var gotModel string	var gotPrompt string	var gotFile string	var gotFileContentType string	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		gotPath = r.URL.Path		if errParse := r.ParseMultipartForm(32 << 20); errParse != nil {			t.Fatalf("parse multipart form: %v", errParse)		}		gotModel = r.FormValue("model")		gotPrompt = r.FormValue("prompt")		file, fileHeader, errFile := r.FormFile("image")		if errFile != nil {			t.Fatalf("read image file: %v", errFile)		}		gotFileContentType = fileHeader.Header.Get("Content-Type")		data, errRead := io.ReadAll(file)		if errClose := file.Close(); errClose != nil {			t.Fatalf("close image file: %v", errClose)		}		if errRead != nil {			t.Fatalf("read image file: %v", errRead)		}		gotFile = string(data)		w.Header().Set("Content-Type", "application/json")		_, _ = w.Write([]byte(`{"created":123,"data":[{"b64_json":"AA=="}]}`))	}))	defer server.Close() 	executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{})	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"base_url": server.URL + "/v1",		"api_key":  "test",	}}	_, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{		Model:   "upstream-image",		Payload: body.Bytes(),	}, cliproxyexecutor.Options{		SourceFormat: sdktranslator.FromString("openai-image"),		Stream:       false,		Headers: http.Header{			"Content-Type": []string{contentType},		},		Metadata: map[string]any{			cliproxyexecutor.RequestPathMetadataKey: "/v1/images/edits",		},	})	if err != nil {		t.Fatalf("Execute error: %v", err)	}	if gotPath != "/v1/images/edits" {		t.Fatalf("path = %q, want %q", gotPath, "/v1/images/edits")	}	if gotModel != "upstream-image" {		t.Fatalf("model = %q, want upstream-image", gotModel)	}	if gotPrompt != "edit" {		t.Fatalf("prompt = %q, want edit", gotPrompt)	}	if gotFile != "png-data" {		t.Fatalf("file = %q, want png-data", gotFile)	}	if gotFileContentType != "image/png" {		t.Fatalf("file content type = %q, want image/png", gotFileContentType)	}} func TestRewriteOpenAICompatImagesMultipartPayloadPreservesStreamAndFileContentType(t *testing.T) {	var body bytes.Buffer	writer := multipart.NewWriter(&body)	if errWrite := writer.WriteField("model", "compat-image"); errWrite != nil {		t.Fatalf("write model field: %v", errWrite)	}	if errWrite := writer.WriteField("stream", "false"); errWrite != nil {		t.Fatalf("write stream field: %v", errWrite)	}	header := make(textproto.MIMEHeader)	header.Set("Content-Disposition", multipart.FileContentDisposition("image", "image.webp"))	header.Set("Content-Type", "image/webp")	part, errCreate := writer.CreatePart(header)	if errCreate != nil {		t.Fatalf("create image field: %v", errCreate)	}	if _, errWrite := part.Write([]byte("webp-data")); errWrite != nil {		t.Fatalf("write image field: %v", errWrite)	}	if errClose := writer.Close(); errClose != nil {		t.Fatalf("close multipart writer: %v", errClose)	} 	out, contentType, err := prepareOpenAICompatImagesPayload(body.Bytes(), "upstream-image", writer.FormDataContentType(), true)	if err != nil {		t.Fatalf("prepareOpenAICompatImagesPayload error: %v", err)	}	mediaType, params, errParse := mime.ParseMediaType(contentType)	if errParse != nil {		t.Fatalf("parse content type: %v", errParse)	}	if mediaType != "multipart/form-data" {		t.Fatalf("media type = %q, want multipart/form-data", mediaType)	}	reader := multipart.NewReader(bytes.NewReader(out), params["boundary"])	form, errRead := reader.ReadForm(32 << 20)	if errRead != nil {		t.Fatalf("read rewritten form: %v", errRead)	}	defer func() {		if errRemove := form.RemoveAll(); errRemove != nil {			t.Fatalf("remove form files: %v", errRemove)		}	}()	if got := form.Value["model"]; len(got) != 1 || got[0] != "upstream-image" {		t.Fatalf("model values = %#v, want upstream-image", got)	}	if got := form.Value["stream"]; len(got) != 1 || got[0] != "true" {		t.Fatalf("stream values = %#v, want true", got)	}	if got := form.File["image"]; len(got) != 1 || got[0].Header.Get("Content-Type") != "image/webp" {		t.Fatalf("image headers = %#v, want image/webp", got)	}} func TestOpenAICompatExecutorStreamRejectsPlainJSONAfterBlankLines(t *testing.T) {	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		w.Header().Set("Content-Type", "text/event-stream")		_, _ = w.Write([]byte("\n\n: openrouter processing\n\nevent: error\n"))		_, _ = w.Write([]byte(`{"error":{"message":"upstream failed","type":"server_error"}}` + "\n"))	}))	defer server.Close() 	executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{})	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"base_url": server.URL + "/v1",		"api_key":  "test",	}}	result, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{		Model:   "openrouter-model",		Payload: []byte(`{"model":"openrouter-model","messages":[{"role":"user","content":"hi"}],"stream":true}`),	}, cliproxyexecutor.Options{		SourceFormat: sdktranslator.FromString("openai"),		Stream:       true,	})	if err != nil {		t.Fatalf("ExecuteStream error: %v", err)	} 	var gotErr error	for chunk := range result.Chunks {		if chunk.Err != nil {			gotErr = chunk.Err			break		}	}	if gotErr == nil {		t.Fatalf("expected plain JSON stream error")	}	if status, ok := gotErr.(interface{ StatusCode() int }); !ok || status.StatusCode() != http.StatusBadGateway {		t.Fatalf("stream error status = %v, want %d", gotErr, http.StatusBadGateway)	}	if !strings.Contains(gotErr.Error(), "upstream failed") {		t.Fatalf("stream error = %v", gotErr)	}} func TestOpenAICompatExecutorStreamSkipsKeepAliveUntilDataLine(t *testing.T) {	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		w.Header().Set("Content-Type", "text/event-stream")		_, _ = w.Write([]byte("\n\n: openrouter processing\n\nevent: ping\nid: 1\nretry: 1000\n"))		_, _ = w.Write([]byte(`data: {"id":"chatcmpl_1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":null}]}` + "\n"))	}))	defer server.Close() 	executor := NewOpenAICompatExecutor("openai-compatibility", &config.Config{})	auth := &cliproxyauth.Auth{Attributes: map[string]string{		"base_url": server.URL + "/v1",		"api_key":  "test",	}}	result, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{		Model:   "openrouter-model",		Payload: []byte(`{"model":"openrouter-model","messages":[{"role":"user","content":"hi"}],"stream":true}`),	}, cliproxyexecutor.Options{		SourceFormat: sdktranslator.FromString("openai"),		Stream:       true,	})	if err != nil {		t.Fatalf("ExecuteStream error: %v", err)	} 	var got strings.Builder	for chunk := range result.Chunks {		if chunk.Err != nil {			t.Fatalf("unexpected stream error: %v", chunk.Err)		}		got.Write(chunk.Payload)	}	if gjson.Get(got.String(), "choices.0.delta.content").String() != "hello" {		t.Fatalf("stream payload = %s", got.String())	}}