mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-11 06:30:21 +00:00
feat(task): replace built-in task adaptors with a sandboxed JS plugin system (#7076)
This commit is contained in:
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,45 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
vertexcore "github.com/QuantumNous/new-api/relay/channel/vertex"
|
||||
)
|
||||
|
||||
type cachedAuth struct {
|
||||
header, projectID string
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
var pluginAuthCache sync.Map
|
||||
var acquireAccessToken = vertexcore.AcquireAccessToken
|
||||
|
||||
func resolveAuth(meta pluginruntime.AuthMeta, apiKey, proxy string) (map[string]any, error) {
|
||||
typeName := strings.TrimSpace(meta.Type)
|
||||
if typeName == "" || typeName == "none" || typeName == "api_key" {
|
||||
return map[string]any{"authHeader": apiKey}, nil
|
||||
}
|
||||
cacheKey := apiKey + "\x00" + proxy
|
||||
if value, ok := pluginAuthCache.Load(cacheKey); ok {
|
||||
entry := value.(cachedAuth)
|
||||
if time.Now().Before(entry.expiresAt) {
|
||||
return map[string]any{"authHeader": entry.header, "projectId": entry.projectID}, nil
|
||||
}
|
||||
}
|
||||
var credentials vertexcore.Credentials
|
||||
if err := common.Unmarshal([]byte(apiKey), &credentials); err != nil {
|
||||
return nil, fmt.Errorf("decode oauth2_jwt credentials: %w", err)
|
||||
}
|
||||
token, err := acquireAccessToken(credentials, proxy)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
entry := cachedAuth{header: "Bearer " + token, projectID: credentials.ProjectID, expiresAt: time.Now().Add(25 * time.Minute)}
|
||||
pluginAuthCache.Store(cacheKey, entry)
|
||||
return map[string]any{"authHeader": entry.header, "projectId": entry.projectID}, nil
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
vertexcore "github.com/QuantumNous/new-api/relay/channel/vertex"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOAuth2JWTAuthCachesAndRefreshes(t *testing.T) {
|
||||
pluginAuthCache = sync.Map{}
|
||||
original := acquireAccessToken
|
||||
t.Cleanup(func() { acquireAccessToken = original; pluginAuthCache = sync.Map{} })
|
||||
calls := 0
|
||||
acquireAccessToken = func(_ vertexcore.Credentials, _ string) (string, error) {
|
||||
calls++
|
||||
return fmt.Sprintf("token-%d", calls), nil
|
||||
}
|
||||
credentials, err := common.Marshal(vertexcore.Credentials{ProjectID: "project", ClientEmail: "a@example.com", PrivateKey: "secret"})
|
||||
require.NoError(t, err)
|
||||
meta := pluginruntime.AuthMeta{Type: "oauth2_jwt"}
|
||||
first, err := resolveAuth(meta, string(credentials), "")
|
||||
require.NoError(t, err)
|
||||
second, err := resolveAuth(meta, string(credentials), "")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Bearer token-1", first["authHeader"])
|
||||
assert.Equal(t, first, second)
|
||||
assert.Equal(t, 1, calls)
|
||||
pluginAuthCache.Store(string(credentials)+"\x00", cachedAuth{expiresAt: time.Now().Add(-time.Second)})
|
||||
refreshed, err := resolveAuth(meta, string(credentials), "")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Bearer token-2", refreshed["authHeader"])
|
||||
assert.Equal(t, 2, calls)
|
||||
}
|
||||
|
||||
func TestOAuth2JWTContextDoesNotExposeServiceAccountKey(t *testing.T) {
|
||||
pluginAuthCache = sync.Map{}
|
||||
original := acquireAccessToken
|
||||
t.Cleanup(func() { acquireAccessToken = original; pluginAuthCache = sync.Map{} })
|
||||
acquireAccessToken = func(_ vertexcore.Credentials, _ string) (string, error) {
|
||||
return "access-token", nil
|
||||
}
|
||||
credentials, err := common.Marshal(vertexcore.Credentials{ProjectID: "project", ClientEmail: "a@example.com", PrivateKey: "secret"})
|
||||
require.NoError(t, err)
|
||||
source := `
|
||||
export const meta = {apiVersion:1,key:"oauth",name:"OAuth",version:"1.0.0",author:{name:"Test"},models:["m"],fetchMode:"per_task",auth:{type:"oauth2_jwt"}};
|
||||
export function buildSubmitRequest(ctx) {
|
||||
if (ctx.apiKey !== undefined) throw new Error("raw key exposed");
|
||||
return {url:ctx.baseUrl+"/submit",headers:{Authorization:ctx.authHeader}};
|
||||
}
|
||||
export function parseSubmitResponse(){return {taskId:"1"}}
|
||||
export function buildQueryRequest(){return {url:"https://example.com"}}
|
||||
export function parseTaskResult(){return {status:"SUCCESS"}}
|
||||
`
|
||||
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
adaptor := New(plugin)
|
||||
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://provider.example", ApiKey: string(credentials)}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
||||
adaptor.Init(info)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
|
||||
c.Set("task_request", relaycommon.TaskSubmitReq{Prompt: "p"})
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
req := httptest.NewRequest(http.MethodPost, "https://provider.example/submit", nil)
|
||||
require.NoError(t, adaptor.BuildRequestHeader(c, req, info))
|
||||
assert.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
|
||||
}
|
||||
Reference in New Issue
Block a user