feat(task): replace built-in task adaptors with a sandboxed JS plugin system (#7076)

2026-09-08 10:46:58 +00:00 · 2026-08-29 18:51:57 +08:00
parent 7037ac15bd
commit eb48396d5f
336 changed files with 52333 additions and 6369 deletions
@@ -0,0 +1,372 @@
+package jsplugin
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"strings"
+	"testing"
+	"time"
+	"unicode/utf8"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEngineCallsESMExportWithInjectedUtils(t *testing.T) {
+	t.Parallel()
+	logs := make([]string, 0, 1)
+	engine, err := Compile(`
+export function sign(ctx) {
+  console.log("called", ctx.name);
+  return {
+    now: utils.unixNow(),
+    digest: utils.hmacSHA256(ctx.message, ctx.secret),
+    encoded: utils.base64(ctx.message),
+  };
+}
+export const meta = { apiVersion: 1, key: "mock" };
+`, Options{
+		Key: "mock", Version: "1.0.0",
+		Now: func() time.Time { return time.Unix(1234, 0) },
+		Log: func(message string) { logs = append(logs, message) },
+	})
+	require.NoError(t, err)
+
+	result, err := engine.Call(context.Background(), "sign", map[string]any{
+		"name": "fixture", "message": "hello", "secret": "secret",
+	})
+	require.NoError(t, err)
+	assert.Equal(t, map[string]any{
+		"now": int64(1234), "digest": "88aab3ede8d3adf94d26ab90d3bafd4a2083070c3bcce9c014ee04a443847c0b", "encoded": "aGVsbG8=",
+	}, result)
+	assert.Equal(t, []string{"[plugin:mock@1.0.0] called fixture"}, logs)
+
+	meta, err := engine.Export(context.Background(), "meta")
+	require.NoError(t, err)
+	assert.Equal(t, map[string]any{"apiVersion": int64(1), "key": "mock"}, meta)
+}
+
+func TestEngineConsoleLogUsesDebugLoggerAndRequestContext(t *testing.T) {
+	previousDebug := common.DebugEnabled
+	common.DebugEnabled = false
+	t.Cleanup(func() { common.DebugEnabled = previousDebug })
+
+	var output bytes.Buffer
+	common.LogWriterMu.Lock()
+	previousWriter := gin.DefaultErrorWriter
+	gin.DefaultErrorWriter = &output
+	common.LogWriterMu.Unlock()
+	t.Cleanup(func() {
+		common.LogWriterMu.Lock()
+		gin.DefaultErrorWriter = previousWriter
+		common.LogWriterMu.Unlock()
+	})
+
+	plugin, err := CompilePlugin(`
+export const meta = {
+  apiVersion: 1,
+  key: "console-debug",
+  name: "Console debug",
+  version: "1.2.3",
+  author: {name: "Test"},
+  models: ["debug-model"],
+  fetchMode: "per_task",
+};
+export function run(label) {
+  console.log("checkpoint", label);
+  return true;
+}
+export function buildSubmitRequest() { return {url: "https://example.com"}; }
+export function parseSubmitResponse() { return {taskId: "one"}; }
+export function buildQueryRequest() { return {url: "https://example.com"}; }
+export function parseTaskResult() { return {status: "SUCCESS"}; }
+`, Options{})
+	require.NoError(t, err)
+	engine := plugin.Engine
+
+	disabledContext := context.WithValue(context.Background(), common.RequestIdKey, "plugin-console-disabled")
+	_, err = engine.Call(disabledContext, "run", "disabled")
+	require.NoError(t, err)
+	assert.Empty(t, output.String())
+
+	common.DebugEnabled = true
+	contextA := context.WithValue(context.Background(), common.RequestIdKey, "plugin-console-request-a")
+	_, err = engine.Call(contextA, "run", "context-a")
+	require.NoError(t, err)
+	logA := output.String()
+	output.Reset()
+
+	contextB := context.WithValue(context.Background(), common.RequestIdKey, "plugin-console-request-b")
+	_, err = engine.Call(contextB, "run", "context-b")
+	require.NoError(t, err)
+	logB := output.String()
+	output.Reset()
+
+	_, err = engine.Call(context.Background(), "run", "background")
+	require.NoError(t, err)
+	logBackground := output.String()
+
+	assert.Contains(t, logA, "plugin-console-request-a")
+	assert.NotContains(t, logA, "plugin-console-request-b")
+	assert.Contains(t, logA, "task_plugin subsystem=runtime event=console")
+	assert.Contains(t, logA, "[plugin:console-debug@1.2.3] checkpoint context-a")
+	assert.NotContains(t, logA, "disabled")
+
+	assert.Contains(t, logB, "plugin-console-request-b")
+	assert.NotContains(t, logB, "plugin-console-request-a")
+	assert.Contains(t, logB, "[plugin:console-debug@1.2.3] checkpoint context-b")
+
+	assert.Contains(t, logBackground, "| SYSTEM |")
+	assert.NotContains(t, logBackground, "plugin-console-request-a")
+	assert.NotContains(t, logBackground, "plugin-console-request-b")
+	assert.Contains(t, logBackground, "[plugin:console-debug@1.2.3] checkpoint background")
+}
+
+func TestCompileRejectsAsynchronousAndImportedPlugins(t *testing.T) {
+	t.Parallel()
+	for name, source := range map[string]string{
+		"async":           `export async function run() {}`,
+		"static import":   `import value from "dependency"; export function run() { return value; }`,
+		"dynamic import":  `export function run() { return import("dependency"); }`,
+		"top-level await": `const value = await work(); export function run() { return value; }`,
+	} {
+		t.Run(name, func(t *testing.T) {
+			_, err := Compile(source, Options{Key: "invalid"})
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), "unsupported plugin syntax")
+		})
+	}
+
+	_, err := Compile(`export function run() { return "import async await"; }`, Options{Key: "valid"})
+	require.NoError(t, err)
+}
+
+func TestCompileIgnoresSourceMapDirectives(t *testing.T) {
+	t.Parallel()
+	// A sourceMappingURL comment must stay inert. Sobek's default loader
+	// os.ReadFiles the referenced server path during Compile and turns any
+	// load failure into a compile error, so an unresolvable path compiling
+	// cleanly proves the loader is disabled.
+	engine, err := Compile("export function run() { return 1; }\n//# sourceMappingURL=/nonexistent/leak-probe.map\n", Options{Key: "sourcemap"})
+	require.NoError(t, err)
+
+	value, err := engine.Call(context.Background(), "run")
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), value)
+}
+
+func TestEngineInterruptsLongRunningHook(t *testing.T) {
+	t.Parallel()
+	engine, err := Compile(`export function run() { while (true) {} }`, Options{
+		Key: "loop", Version: "1", Timeout: 20 * time.Millisecond,
+	})
+	require.NoError(t, err)
+
+	_, err = engine.Call(context.Background(), "run")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "timed out")
+	var hookErr *HookError
+	assert.False(t, errors.As(err, &hookErr), "timeouts must not be HookError")
+}
+
+func TestEngineHookErrorExtractsSanitizedJSMessage(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name        string
+		source      string
+		wantMessage string
+		wantLen     int
+	}{
+		{
+			name:        "Error object",
+			source:      `export function run() { throw new Error("model is required"); }`,
+			wantMessage: "model is required",
+		},
+		{
+			name:        "raw string throw",
+			source:      `export function run() { throw "raw string"; }`,
+			wantMessage: "raw string",
+		},
+		{
+			name:        "truncates to 512 runes",
+			source:      `export function run() { throw new Error("x".repeat(2000)); }`,
+			wantLen:     512,
+			wantMessage: strings.Repeat("x", 512),
+		},
+		{
+			name:        "scrubs control characters",
+			source:      "export function run() { throw new Error(\"line1\\nline2\\x1b[31mred\"); }",
+			wantMessage: "line1 line2 [31mred",
+		},
+		{
+			name:        "throwing message getter falls back without crashing",
+			source:      `export function run() { throw {get message() { throw {get message() { return "deep"; }}; }}; }`,
+			wantMessage: "plugin hook failed",
+		},
+	}
+	for _, testCase := range tests {
+		t.Run(testCase.name, func(t *testing.T) {
+			t.Parallel()
+			engine, err := Compile(testCase.source, Options{Key: "diag", Version: "1.0.0"})
+			require.NoError(t, err)
+
+			_, err = engine.Call(context.Background(), "run")
+			require.Error(t, err)
+
+			var hookErr *HookError
+			require.True(t, errors.As(err, &hookErr))
+			assert.Equal(t, "run", hookErr.Hook)
+			assert.Equal(t, testCase.wantMessage, hookErr.Message)
+			if testCase.wantLen > 0 {
+				assert.Equal(t, testCase.wantLen, utf8.RuneCountInString(hookErr.Message))
+			}
+			assert.Contains(t, hookErr.Error(), "plugin diag@1.0.0")
+			assert.NotContains(t, hookErr.Message, "Error:")
+			assert.NotContains(t, hookErr.Message, "plugin diag@")
+		})
+	}
+}
+
+func TestEngineReportsProtocolAdmissionTimeoutSeparately(t *testing.T) {
+	engine, err := Compile(`
+export const protocols = {
+	responses: {renderEvents: function() { return {events: [], done: false}; }},
+};
+`, Options{Key: "admission", Version: "1", Concurrency: 1})
+	require.NoError(t, err)
+
+	engine.semaphore <- struct{}{}
+	_, err = engine.CallPathWithAdmissionTimeout(
+		context.Background(),
+		time.Nanosecond,
+		"protocols",
+		[]string{"responses", "renderEvents"},
+	)
+	<-engine.semaphore
+
+	require.ErrorIs(t, err, ErrCallAdmissionTimeout)
+	result, err := engine.CallPathWithAdmissionTimeout(
+		context.Background(),
+		time.Second,
+		"protocols",
+		[]string{"responses", "renderEvents"},
+	)
+	require.NoError(t, err)
+	assert.Equal(t, map[string]any{"events": []any{}, "done": false}, result)
+}
+
+func TestEngineExportInterruptsLongRunningGetter(t *testing.T) {
+	t.Parallel()
+	engine, err := Compile(`
+export const meta = {
+	apiVersion: 1,
+	get name() { while (true) {} },
+};
+`, Options{Key: "meta-loop", Version: "1.0.0", Timeout: 20 * time.Millisecond})
+	require.NoError(t, err)
+
+	_, err = engine.Export(context.Background(), "meta")
+	require.ErrorContains(t, err, "export meta interrupted")
+}
+
+func TestEngineExportReturnsThrownGetterError(t *testing.T) {
+	t.Parallel()
+	engine, err := Compile(`
+export const meta = {
+	apiVersion: 1,
+	get name() { throw new Error("getter failed"); },
+};
+`, Options{Key: "meta-throw", Version: "1.0.0"})
+	require.NoError(t, err)
+
+	_, err = engine.Export(context.Background(), "meta")
+	require.ErrorContains(t, err, "export meta failed")
+	assert.Contains(t, err.Error(), "getter failed")
+}
+
+func TestEngineNestedHooksRequireOwnProperties(t *testing.T) {
+	t.Parallel()
+	engine, err := Compile(`
+const inheritedRenderers = {
+	inherited: function(value) { return value; },
+	constructor: function(value) { return value; },
+	toString: function(value) { return value; },
+	["__proto__"]: function(value) { return value; },
+};
+export const renderers = Object.create(inheritedRenderers);
+renderers.own = function(value) { return {id: value.id}; };
+
+const inheritedProtocol = {
+	renderFinal: function(value) { return value; },
+};
+export const protocols = {
+	responses: Object.create(inheritedProtocol),
+};
+`, Options{Key: "own-hooks", Version: "1.0.0"})
+	require.NoError(t, err)
+
+	found, err := engine.HasCallablePath(context.Background(), "renderers", "own")
+	require.NoError(t, err)
+	assert.True(t, found)
+	result, err := engine.CallMember(context.Background(), "renderers", "own", map[string]any{"id": "task-1"})
+	require.NoError(t, err)
+	assert.Equal(t, map[string]any{"id": "task-1"}, result)
+
+	for _, member := range []string{"inherited", "constructor", "toString", "__proto__"} {
+		t.Run(member, func(t *testing.T) {
+			found, err := engine.HasCallablePath(context.Background(), "renderers", member)
+			require.NoError(t, err)
+			assert.False(t, found)
+
+			_, err = engine.CallMember(context.Background(), "renderers", member)
+			require.ErrorContains(t, err, "not found")
+		})
+	}
+
+	found, err = engine.HasCallablePath(context.Background(), "protocols", "responses", "renderFinal")
+	require.NoError(t, err)
+	assert.False(t, found)
+	_, err = engine.CallPath(context.Background(), "protocols", []string{"responses", "renderFinal"})
+	require.ErrorContains(t, err, "not found")
+}
+
+func TestCompileInterruptsLongRunningInitialization(t *testing.T) {
+	t.Parallel()
+	_, err := Compile(`while (true) {}; export function run() {}`, Options{
+		Key: "loop", Version: "1", Timeout: 20 * time.Millisecond,
+	})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "initialization timed out")
+}
+
+func TestValidateRequestURL(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name         string
+		requestURL   string
+		baseURL      string
+		allowedHosts []string
+		wantError    string
+	}{
+		{name: "same host", requestURL: "https://api.example.com/v1/task", baseURL: "https://api.example.com/v1"},
+		{name: "default port", requestURL: "https://api.example.com:443/v1/task", baseURL: "https://api.example.com"},
+		{name: "approved host", requestURL: "https://upload.example.com/task", baseURL: "https://api.example.com", allowedHosts: []string{"upload.example.com"}},
+		{name: "subdomain is not implicit", requestURL: "https://evil.api.example.com/task", baseURL: "https://api.example.com", wantError: "not allowed"},
+		{name: "userinfo trick", requestURL: "https://api.example.com@evil.example/task", baseURL: "https://api.example.com", wantError: "not allowed"},
+		{name: "relative URL", requestURL: "/v1/task", baseURL: "https://api.example.com", wantError: "absolute"},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			err := ValidateRequestURL(test.requestURL, test.baseURL, test.allowedHosts)
+			if test.wantError == "" {
+				require.NoError(t, err)
+				return
+			}
+			require.Error(t, err)
+			assert.True(t, strings.Contains(err.Error(), test.wantError), err.Error())
+		})
+	}
+}