feat: publish sanitized custom.97 source

2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
@@ -1,6 +1,20 @@
 package plugins_test

-import "testing"
+import (
+	"net/http"
+	"testing"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/model"
+	"github.com/QuantumNous/new-api/pkg/jsplugin"
+	builtinplugins "github.com/QuantumNous/new-api/plugins"
+	"github.com/QuantumNous/new-api/relay/channel"
+	taskplugin "github.com/QuantumNous/new-api/relay/channel/task/jsplugin"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)

 func TestGoogleResponsesProtocol(t *testing.T) {
 	testVideoResponsesProtocol(t, videoResponsesTestCase{
@@ -28,3 +42,127 @@ func TestGoogleResponsesProtocol(t *testing.T) {
 		wantVendorName: "gemini",
 	})
 }
+
+// A New API channel reaches Veo through another gateway's /google routes. The
+// gateway rebuilds the request this driver sent, names the operation after its
+// own task so it is polled there, and serves the video file only its Gemini
+// key can read.
+func TestGoogleThroughNewAPIGateway(t *testing.T) {
+	source, err := builtinplugins.Source("google")
+	require.NoError(t, err)
+	plugin, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: "google"})
+	require.NoError(t, err)
+	call := func(export string, members []string, args ...any) map[string]any {
+		t.Helper()
+		value, err := plugin.Engine.CallPath(t.Context(), export, members, args...)
+		require.NoError(t, err)
+		return decodePluginValue(t, value)
+	}
+	const veo = "veo-3.1-fast-generate-preview"
+	gateway := map[string]any{"kind": "new_api"}
+	submitParams := map[string]any{"model": veo + ":predictLongRunning"}
+
+	for _, tc := range []struct {
+		name        string
+		requestBody map[string]any
+	}{
+		{"text with parameters", map[string]any{"model": veo, "prompt": "p", "duration": 6, "metadata": map[string]any{"resolution": "1080p", "aspectRatio": "9:16", "negativePrompt": "rain"}}},
+		{"first frame", map[string]any{"model": veo, "prompt": "p", "images": []any{"data:image/png;base64,aGVsbG8="}, "size": "1280x720"}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			downstream := map[string]any{"requestBody": tc.requestBody, "model": veo, "upstreamModel": veo,
+				"baseUrl": "https://gateway.example", "apiKey": "sk-gateway", "upstream": gateway}
+			sent := call("buildSubmitRequest", nil, downstream)
+			assert.Equal(t, "https://gateway.example/google/v1beta/models/"+veo+":predictLongRunning", sent["url"])
+			assert.Equal(t, map[string]any{"Accept": "application/json", "Content-Type": "application/json", "Authorization": "Bearer sk-gateway"}, sent["headers"])
+
+			intent := call("native", []string{"predictLongRunning"}, map[string]any{"params": submitParams, "body": map[string]any{"kind": "json", "value": sent["body"]}})
+			assert.Equal(t, veo, intent["model"])
+			assert.Equal(t, sent["action"], intent["action"])
+			upstream := map[string]any{"requestBody": intent["requestBody"], "model": veo, "upstreamModel": veo,
+				"baseUrl": "https://generativelanguage.example", "apiKey": "vendor-key", "upstream": map[string]any{"kind": "vendor"}}
+			forwarded := call("buildSubmitRequest", nil, upstream)
+			assert.Equal(t, "https://generativelanguage.example/v1beta/models/"+veo+":predictLongRunning", forwarded["url"])
+			assert.Equal(t, "vendor-key", forwarded["headers"].(map[string]any)["x-goog-api-key"])
+			assert.Equal(t, sent["body"], forwarded["body"])
+			assert.Equal(t, call("extractUsage", nil, downstream), call("extractUsage", nil, upstream))
+		})
+	}
+
+	created := call("native", []string{"operation"}, map[string]any{"params": submitParams},
+		map[string]any{"task_id": "task_pub", "status": "SUBMITTED", "data": map[string]any{"name": "models/" + veo + "/operations/op-private"}})
+	assert.Equal(t, map[string]any{"name": "models/" + veo + "/operations/task_pub"}, created)
+	upstreamTaskID := call("parseSubmitResponse", nil, map[string]any{}, map[string]any{"body": created})["taskId"]
+	query := call("buildQueryRequest", nil, map[string]any{"taskId": upstreamTaskID, "baseUrl": "https://gateway.example", "apiKey": "sk-gateway", "upstream": gateway})
+	assert.Equal(t, "https://gateway.example/google/v1beta/models/"+veo+"/operations/task_pub", query["url"])
+
+	adaptor := taskplugin.New(plugin)
+	adaptor.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{
+		ChannelType: constant.ChannelTypeNewAPI, ApiKey: "sk-gateway", ChannelBaseUrl: "https://gateway.example",
+	}})
+	queryParams := map[string]any{"model": veo, "task_id": "task_pub"}
+	for _, tc := range []struct {
+		name       string
+		task       map[string]any
+		wantStatus string
+		wantReason string
+	}{
+		{
+			name: "running", wantStatus: "IN_PROGRESS",
+			task: map[string]any{"status": "IN_PROGRESS", "data": map[string]any{"name": "models/" + veo + "/operations/op-private"}},
+		},
+		{
+			name: "succeeded", wantStatus: "SUCCESS",
+			task: map[string]any{"status": "SUCCESS", "data": map[string]any{"name": "models/" + veo + "/operations/op-private", "done": true,
+				"response": map[string]any{"generateVideoResponse": map[string]any{"generatedVideos": []any{
+					map[string]any{"video": map[string]any{"uri": "https://generativelanguage.example/v1beta/files/f1:download?alt=media"}},
+				}}}}},
+		},
+		{
+			name: "failed by the gateway poller", wantStatus: "FAILURE", wantReason: "task timed out",
+			task: map[string]any{"status": "FAILURE", "fail_reason": "task timed out", "data": map[string]any{"name": "models/" + veo + "/operations/op-private"}},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			tc.task["task_id"] = "task_pub"
+			rendered := call("native", []string{"operation"}, map[string]any{"params": queryParams}, tc.task)
+			assert.Equal(t, "models/"+veo+"/operations/task_pub", rendered["name"])
+			result := call("parseTaskResult", nil, map[string]any{}, rendered)
+			assert.Equal(t, tc.wantStatus, result["status"])
+			assert.Equal(t, tc.wantReason, common.Interface2String(result["reason"]))
+			if tc.wantStatus != "SUCCESS" {
+				return
+			}
+			data, err := common.Marshal(rendered)
+			require.NoError(t, err)
+			task := &model.Task{TaskID: "task_down", Status: model.TaskStatusSuccess, Data: data, PrivateData: model.TaskPrivateData{UpstreamTaskID: upstreamTaskID.(string)}}
+			artifacts, err := adaptor.ListArtifacts(task)
+			require.NoError(t, err)
+			assert.Equal(t, []channel.TaskArtifact{{Key: "video", Type: "video"}}, artifacts)
+			descriptor, err := adaptor.BuildContentRequest(t.Context(), task, "video", channel.TaskArtifactClientRequest{Method: http.MethodGet})
+			require.NoError(t, err)
+			assert.Equal(t, &channel.TaskContentRequest{
+				URL:     "https://gateway.example/v1/tasks/task_pub/artifacts/video/content",
+				Method:  http.MethodGet,
+				Headers: map[string]string{"Authorization": "Bearer sk-gateway"},
+			}, descriptor)
+		})
+	}
+
+	t.Run("native bodies are checked before billing", func(t *testing.T) {
+		for _, tc := range []struct {
+			model   string
+			body    map[string]any
+			wantErr string
+		}{
+			{veo + ":generateVideos", map[string]any{"instances": []any{map[string]any{"prompt": "p"}}}, "models/{model}:predictLongRunning"},
+			{veo + ":predictLongRunning", map[string]any{"instances": []any{map[string]any{"prompt": "p", "lastFrame": map[string]any{}}}}, "instances[0].lastFrame is not supported"},
+			{veo + ":predictLongRunning", map[string]any{"instances": []any{map[string]any{"prompt": "p"}}, "parameters": map[string]any{"durationSeconds": 5}}, "durationSeconds must be one of 4, 6, or 8"},
+		} {
+			_, err := plugin.Engine.CallPath(t.Context(), "native", []string{"predictLongRunning"}, map[string]any{
+				"params": map[string]any{"model": tc.model}, "body": map[string]any{"kind": "json", "value": tc.body},
+			})
+			require.ErrorContains(t, err, tc.wantErr)
+		}
+	})
+}