mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-08 17:36:47 +08:00
feat(pluginhost, scheduler): introduce Go-based plugin with scheduler capabilities
- Added a Go scheduler plugin demonstrating CLIProxyAPI capabilities, such as `plugin.register`, `plugin.reconfigure`, and `scheduler.pick`. - Implemented methods for plugin configuration, built-in scheduler delegation (`fill-first`, `round-robin`), dynamic candidate selection, and error handling. - Extended `pluginhost` with scheduler handling, candidate normalization, and fallback mechanisms. - Included examples, tests, and detailed documentation for scheduler usage and implementation.
This commit is contained in:
@@ -14,6 +14,7 @@ This directory contains standard dynamic library plugin examples for the CLIProx
|
||||
- `request-translator/`: request translation capability only.
|
||||
- `request-normalizer/`: request normalization capability only.
|
||||
- `codex-service-tier/`: Go-only request normalizer that sets Codex `gpt-5.4` requests to the priority service tier when enabled.
|
||||
- `scheduler/`: Go-only scheduler that can select a configured auth ID, delegate to a built-in scheduler, or deny picks.
|
||||
- `response-translator/`: response translation capability only.
|
||||
- `response-normalizer/`: response normalization capability only.
|
||||
- `thinking/`: thinking applier capability only.
|
||||
@@ -37,6 +38,23 @@ plugins:
|
||||
fast: false
|
||||
```
|
||||
|
||||
## Scheduler
|
||||
|
||||
`scheduler` declares the scheduler capability. It can select a configured auth ID from the candidate list, delegate to the built-in `fill-first` or `round-robin` scheduler, or reject picks when `deny` is `true`.
|
||||
|
||||
```yaml
|
||||
plugins:
|
||||
configs:
|
||||
scheduler:
|
||||
enabled: true
|
||||
priority: 1
|
||||
auth_id: ""
|
||||
delegate: ""
|
||||
deny: false
|
||||
```
|
||||
|
||||
`auth_id` selects a matching candidate when `delegate` is empty. `delegate` accepts `""`, `fill-first`, or `round-robin`; other non-empty values leave the pick unhandled. `deny` returns a scheduler error.
|
||||
|
||||
## Build All Examples
|
||||
|
||||
```bash
|
||||
|
||||
@@ -14,6 +14,7 @@
|
||||
- `request-translator/`:只演示请求转换能力。
|
||||
- `request-normalizer/`:只演示请求规整能力。
|
||||
- `codex-service-tier/`:仅 Go 实现的请求规整插件,启用后会将 Codex `gpt-5.4` 请求设置为 priority service tier。
|
||||
- `scheduler/`:仅 Go 实现的调度插件,可选择指定 auth ID、委托内置调度器或拒绝调度。
|
||||
- `response-translator/`:只演示响应转换能力。
|
||||
- `response-normalizer/`:只演示响应规整能力。
|
||||
- `thinking/`:只演示 Thinking 处理能力。
|
||||
@@ -37,6 +38,23 @@ plugins:
|
||||
fast: false
|
||||
```
|
||||
|
||||
## Scheduler
|
||||
|
||||
`scheduler` 声明调度能力。它可以从候选列表中选择配置的 auth ID,委托内置的 `fill-first` 或 `round-robin` 调度器,或在 `deny` 为 `true` 时拒绝调度。
|
||||
|
||||
```yaml
|
||||
plugins:
|
||||
configs:
|
||||
scheduler:
|
||||
enabled: true
|
||||
priority: 1
|
||||
auth_id: ""
|
||||
delegate: ""
|
||||
deny: false
|
||||
```
|
||||
|
||||
`auth_id` 会在 `delegate` 为空时选择匹配候选。`delegate` 支持 `""`、`fill-first` 和 `round-robin`;其他非空值会让本插件不处理本次调度。`deny` 会返回调度错误。
|
||||
|
||||
## 构建全部示例
|
||||
|
||||
```bash
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
# Scheduler Plugin
|
||||
|
||||
This plugin demonstrates the CLIProxyAPI C ABI scheduler capability from Go.
|
||||
|
||||
It implements:
|
||||
|
||||
- `plugin.register`
|
||||
- `plugin.reconfigure`
|
||||
- `scheduler.pick`
|
||||
|
||||
The plugin can select a configured auth ID, delegate routing to a built-in scheduler, or reject scheduler picks.
|
||||
|
||||
## Configuration
|
||||
|
||||
Add the plugin under `plugins.configs`:
|
||||
|
||||
```yaml
|
||||
plugins:
|
||||
configs:
|
||||
scheduler:
|
||||
enabled: true
|
||||
priority: 1
|
||||
auth_id: ""
|
||||
delegate: ""
|
||||
deny: false
|
||||
```
|
||||
|
||||
Fields:
|
||||
|
||||
- `auth_id`: selects this auth ID when it appears in the scheduler candidates.
|
||||
- `delegate`: delegates selection to a built-in scheduler. Supported values are `""`, `fill-first`, and `round-robin`.
|
||||
- `deny`: returns a scheduler error when set to `true`.
|
||||
|
||||
Behavior:
|
||||
|
||||
- When `deny` is `true`, the plugin returns an error envelope with code `scheduler_denied`.
|
||||
- When `delegate` is `fill-first` or `round-robin`, the plugin returns `DelegateBuiltin` and marks the pick as handled.
|
||||
- When `delegate` is any other non-empty value, the plugin leaves the pick unhandled.
|
||||
- When `delegate` is empty and `auth_id` exists in the candidates, the plugin returns that auth ID and marks the pick as handled.
|
||||
- When no rule matches, the plugin leaves the pick unhandled.
|
||||
|
||||
## Build
|
||||
|
||||
From this directory:
|
||||
|
||||
```bash
|
||||
cd go
|
||||
go build -buildmode=c-shared -o /tmp/cliproxy-scheduler-plugin.so .
|
||||
rm -f /tmp/cliproxy-scheduler-plugin.so /tmp/cliproxy-scheduler-plugin.h
|
||||
```
|
||||
@@ -0,0 +1,10 @@
|
||||
module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/scheduler/go
|
||||
|
||||
go 1.26.0
|
||||
|
||||
require (
|
||||
github.com/router-for-me/CLIProxyAPI/v7 v7.0.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
replace github.com/router-for-me/CLIProxyAPI/v7 => ../../../..
|
||||
@@ -0,0 +1,4 @@
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
@@ -0,0 +1,270 @@
|
||||
package main
|
||||
|
||||
/*
|
||||
#include <stdint.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
typedef struct {
|
||||
void* ptr;
|
||||
size_t len;
|
||||
} cliproxy_buffer;
|
||||
|
||||
typedef struct {
|
||||
uint32_t abi_version;
|
||||
void* host_ctx;
|
||||
void* call;
|
||||
void* free_buffer;
|
||||
} cliproxy_host_api;
|
||||
|
||||
typedef int (*cliproxy_plugin_call_fn)(char*, uint8_t*, size_t, cliproxy_buffer*);
|
||||
typedef void (*cliproxy_plugin_free_fn)(void*, size_t);
|
||||
typedef void (*cliproxy_plugin_shutdown_fn)(void);
|
||||
|
||||
typedef struct {
|
||||
uint32_t abi_version;
|
||||
cliproxy_plugin_call_fn call;
|
||||
cliproxy_plugin_free_fn free_buffer;
|
||||
cliproxy_plugin_shutdown_fn shutdown;
|
||||
} cliproxy_plugin_api;
|
||||
|
||||
extern int cliproxyPluginCall(char*, uint8_t*, size_t, cliproxy_buffer*);
|
||||
extern void cliproxyPluginFree(void*, size_t);
|
||||
extern void cliproxyPluginShutdown(void);
|
||||
*/
|
||||
import "C"
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"unsafe"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
var currentConfig atomic.Value
|
||||
|
||||
type envelope struct {
|
||||
OK bool `json:"ok"`
|
||||
Result json.RawMessage `json:"result,omitempty"`
|
||||
Error *envelopeError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type envelopeError struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type lifecycleRequest struct {
|
||||
ConfigYAML []byte `json:"config_yaml"`
|
||||
}
|
||||
|
||||
type pluginConfig struct {
|
||||
AuthID string `yaml:"auth_id"`
|
||||
Delegate string `yaml:"delegate"`
|
||||
Deny bool `yaml:"deny"`
|
||||
}
|
||||
|
||||
type registration struct {
|
||||
SchemaVersion uint32 `json:"schema_version"`
|
||||
Metadata pluginapi.Metadata `json:"metadata"`
|
||||
Capabilities registrationCapability `json:"capabilities"`
|
||||
}
|
||||
|
||||
type registrationCapability struct {
|
||||
Scheduler bool `json:"scheduler"`
|
||||
}
|
||||
|
||||
func main() {}
|
||||
|
||||
//export cliproxy_plugin_init
|
||||
func cliproxy_plugin_init(_ *C.cliproxy_host_api, plugin *C.cliproxy_plugin_api) C.int {
|
||||
if plugin == nil {
|
||||
return 1
|
||||
}
|
||||
plugin.abi_version = C.uint32_t(pluginabi.ABIVersion)
|
||||
plugin.call = C.cliproxy_plugin_call_fn(C.cliproxyPluginCall)
|
||||
plugin.free_buffer = C.cliproxy_plugin_free_fn(C.cliproxyPluginFree)
|
||||
plugin.shutdown = C.cliproxy_plugin_shutdown_fn(C.cliproxyPluginShutdown)
|
||||
return 0
|
||||
}
|
||||
|
||||
//export cliproxyPluginCall
|
||||
func cliproxyPluginCall(method *C.char, request *C.uint8_t, requestLen C.size_t, response *C.cliproxy_buffer) C.int {
|
||||
if response != nil {
|
||||
response.ptr = nil
|
||||
response.len = 0
|
||||
}
|
||||
if method == nil {
|
||||
writeResponse(response, errorEnvelope("invalid_method", "method is required"))
|
||||
return 1
|
||||
}
|
||||
var requestBytes []byte
|
||||
if request != nil && requestLen > 0 {
|
||||
requestBytes = C.GoBytes(unsafe.Pointer(request), C.int(requestLen))
|
||||
}
|
||||
raw, errHandle := handleMethod(C.GoString(method), requestBytes)
|
||||
if errHandle != nil {
|
||||
writeResponse(response, errorEnvelope("plugin_error", errHandle.Error()))
|
||||
return 1
|
||||
}
|
||||
writeResponse(response, raw)
|
||||
return 0
|
||||
}
|
||||
|
||||
//export cliproxyPluginFree
|
||||
func cliproxyPluginFree(ptr unsafe.Pointer, len C.size_t) {
|
||||
if ptr != nil {
|
||||
C.free(ptr)
|
||||
}
|
||||
}
|
||||
|
||||
//export cliproxyPluginShutdown
|
||||
func cliproxyPluginShutdown() {}
|
||||
|
||||
func handleMethod(method string, request []byte) ([]byte, error) {
|
||||
switch method {
|
||||
case pluginabi.MethodPluginRegister, pluginabi.MethodPluginReconfigure:
|
||||
if errConfigure := configure(request); errConfigure != nil {
|
||||
return nil, errConfigure
|
||||
}
|
||||
return okEnvelope(pluginRegistration())
|
||||
case pluginabi.MethodSchedulerPick:
|
||||
return pickAuth(request)
|
||||
default:
|
||||
return errorEnvelope("unknown_method", "unknown method: "+method), nil
|
||||
}
|
||||
}
|
||||
|
||||
func configure(raw []byte) error {
|
||||
var req lifecycleRequest
|
||||
if len(raw) > 0 {
|
||||
if errUnmarshal := json.Unmarshal(raw, &req); errUnmarshal != nil {
|
||||
return errUnmarshal
|
||||
}
|
||||
}
|
||||
|
||||
cfg := pluginConfig{}
|
||||
if len(req.ConfigYAML) > 0 {
|
||||
decoded, errDecode := decodeConfig(req.ConfigYAML)
|
||||
if errDecode != nil {
|
||||
return errDecode
|
||||
}
|
||||
cfg = decoded
|
||||
}
|
||||
cfg.AuthID = strings.TrimSpace(cfg.AuthID)
|
||||
cfg.Delegate = strings.TrimSpace(cfg.Delegate)
|
||||
currentConfig.Store(cfg)
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeConfig(raw []byte) (pluginConfig, error) {
|
||||
var cfg pluginConfig
|
||||
if errUnmarshal := yaml.Unmarshal(raw, &cfg); errUnmarshal != nil {
|
||||
return pluginConfig{}, errUnmarshal
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
func pluginRegistration() registration {
|
||||
return registration{
|
||||
SchemaVersion: pluginabi.SchemaVersion,
|
||||
Metadata: pluginapi.Metadata{
|
||||
Name: "scheduler",
|
||||
Version: "0.1.0",
|
||||
Author: "router-for-me",
|
||||
GitHubRepository: "https://github.com/router-for-me/CLIProxyAPI",
|
||||
Logo: "https://raw.githubusercontent.com/router-for-me/CLIProxyAPI/main/docs/logo.png",
|
||||
ConfigFields: []pluginapi.ConfigField{
|
||||
{
|
||||
Name: "auth_id",
|
||||
Type: pluginapi.ConfigFieldTypeString,
|
||||
Description: "Selects this auth ID when it is present in the scheduler candidates.",
|
||||
},
|
||||
{
|
||||
Name: "delegate",
|
||||
Type: pluginapi.ConfigFieldTypeEnum,
|
||||
EnumValues: []string{"", pluginapi.SchedulerBuiltinFillFirst, pluginapi.SchedulerBuiltinRoundRobin},
|
||||
Description: "Delegates selection to a built-in scheduler when set to fill-first or round-robin.",
|
||||
},
|
||||
{
|
||||
Name: "deny",
|
||||
Type: pluginapi.ConfigFieldTypeBoolean,
|
||||
Description: "Rejects scheduler picks with an explicit error when enabled.",
|
||||
},
|
||||
},
|
||||
},
|
||||
Capabilities: registrationCapability{
|
||||
Scheduler: true,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func pickAuth(raw []byte) ([]byte, error) {
|
||||
var req pluginapi.SchedulerPickRequest
|
||||
if errUnmarshal := json.Unmarshal(raw, &req); errUnmarshal != nil {
|
||||
return nil, errUnmarshal
|
||||
}
|
||||
|
||||
cfg := loadedConfig()
|
||||
if cfg.Deny {
|
||||
return errorEnvelope("scheduler_denied", "scheduler pick denied by plugin configuration"), nil
|
||||
}
|
||||
switch cfg.Delegate {
|
||||
case pluginapi.SchedulerBuiltinFillFirst, pluginapi.SchedulerBuiltinRoundRobin:
|
||||
return okEnvelope(pluginapi.SchedulerPickResponse{
|
||||
DelegateBuiltin: cfg.Delegate,
|
||||
Handled: true,
|
||||
})
|
||||
case "":
|
||||
default:
|
||||
return okEnvelope(pluginapi.SchedulerPickResponse{Handled: false})
|
||||
}
|
||||
if cfg.AuthID == "" {
|
||||
return okEnvelope(pluginapi.SchedulerPickResponse{Handled: false})
|
||||
}
|
||||
for _, candidate := range req.Candidates {
|
||||
if candidate.ID == cfg.AuthID {
|
||||
return okEnvelope(pluginapi.SchedulerPickResponse{
|
||||
AuthID: cfg.AuthID,
|
||||
Handled: true,
|
||||
})
|
||||
}
|
||||
}
|
||||
return okEnvelope(pluginapi.SchedulerPickResponse{Handled: false})
|
||||
}
|
||||
|
||||
func loadedConfig() pluginConfig {
|
||||
raw := currentConfig.Load()
|
||||
if cfg, ok := raw.(pluginConfig); ok {
|
||||
return cfg
|
||||
}
|
||||
return pluginConfig{}
|
||||
}
|
||||
|
||||
func okEnvelope(v any) ([]byte, error) {
|
||||
raw, errMarshal := json.Marshal(v)
|
||||
if errMarshal != nil {
|
||||
return nil, errMarshal
|
||||
}
|
||||
return json.Marshal(envelope{OK: true, Result: raw})
|
||||
}
|
||||
|
||||
func errorEnvelope(code, message string) []byte {
|
||||
raw, _ := json.Marshal(envelope{OK: false, Error: &envelopeError{Code: code, Message: message}})
|
||||
return raw
|
||||
}
|
||||
|
||||
func writeResponse(response *C.cliproxy_buffer, raw []byte) {
|
||||
if response == nil || len(raw) == 0 {
|
||||
return
|
||||
}
|
||||
ptr := C.CBytes(raw)
|
||||
if ptr == nil {
|
||||
return
|
||||
}
|
||||
response.ptr = ptr
|
||||
response.len = C.size_t(len(raw))
|
||||
}
|
||||
@@ -264,6 +264,7 @@ func validPlugin(plugin pluginapi.Plugin) bool {
|
||||
caps.ModelProvider != nil ||
|
||||
caps.AuthProvider != nil ||
|
||||
caps.FrontendAuthProvider != nil ||
|
||||
caps.Scheduler != nil ||
|
||||
caps.Executor != nil ||
|
||||
caps.RequestTranslator != nil ||
|
||||
caps.RequestNormalizer != nil ||
|
||||
|
||||
@@ -70,6 +70,9 @@ func registerRPCPlugin(ctx context.Context, host *Host, id string, client plugin
|
||||
if resp.Capabilities.FrontendAuthProvider {
|
||||
plugin.Capabilities.FrontendAuthProvider = rpcFrontendAuthProvider{rpcPluginAdapter: adapter}
|
||||
}
|
||||
if resp.Capabilities.Scheduler {
|
||||
plugin.Capabilities.Scheduler = adapter
|
||||
}
|
||||
if resp.Capabilities.Executor {
|
||||
plugin.Capabilities.Executor = rpcProviderExecutor{rpcPluginAdapter: adapter}
|
||||
}
|
||||
@@ -147,6 +150,12 @@ func sanitizePluginRequest(request any) any {
|
||||
case pluginapi.AuthModelRequest:
|
||||
req.HTTPClient = nil
|
||||
return req
|
||||
case pluginapi.SchedulerPickRequest:
|
||||
req.Options.Metadata = sanitizePluginMetadata(req.Options.Metadata)
|
||||
for index := range req.Candidates {
|
||||
req.Candidates[index].Metadata = sanitizePluginMetadata(req.Candidates[index].Metadata)
|
||||
}
|
||||
return req
|
||||
case pluginapi.ExecutorRequest:
|
||||
req.HTTPClient = nil
|
||||
req.Metadata = sanitizePluginMetadata(req.Metadata)
|
||||
@@ -287,6 +296,10 @@ func (a *rpcPluginAdapter) ModelsForAuth(ctx context.Context, req pluginapi.Auth
|
||||
})
|
||||
}
|
||||
|
||||
func (a *rpcPluginAdapter) Pick(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
return callPlugin[pluginapi.SchedulerPickResponse](ctx, a.client, pluginabi.MethodSchedulerPick, req)
|
||||
}
|
||||
|
||||
func callPluginIdentifier(client pluginClient, method string) string {
|
||||
resp, errCall := callPlugin[rpcIdentifierResponse](context.Background(), client, method, rpcEmptyResponse{})
|
||||
if errCall != nil {
|
||||
|
||||
@@ -23,6 +23,7 @@ type rpcCapabilities struct {
|
||||
AuthProvider bool `json:"auth_provider"`
|
||||
FrontendAuthProvider bool `json:"frontend_auth_provider"`
|
||||
FrontendAuthProviderExclusive bool `json:"frontend_auth_provider_exclusive"`
|
||||
Scheduler bool `json:"scheduler"`
|
||||
Executor bool `json:"executor"`
|
||||
ExecutorModelScope pluginapi.ExecutorModelScope `json:"executor_model_scope"`
|
||||
ExecutorInputFormats []string `json:"executor_input_formats,omitempty"`
|
||||
@@ -100,6 +101,7 @@ func rpcCapabilitiesFromPlugin(plugin pluginapi.Plugin) rpcCapabilities {
|
||||
AuthProvider: caps.AuthProvider != nil,
|
||||
FrontendAuthProvider: caps.FrontendAuthProvider != nil,
|
||||
FrontendAuthProviderExclusive: caps.FrontendAuthProvider != nil && caps.FrontendAuthProviderExclusive,
|
||||
Scheduler: caps.Scheduler != nil,
|
||||
Executor: caps.Executor != nil,
|
||||
ExecutorModelScope: normalizedExecutorModelScope(caps),
|
||||
ExecutorInputFormats: append([]string(nil), caps.ExecutorInputFormats...),
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
package pluginhost
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
)
|
||||
|
||||
@@ -38,3 +41,193 @@ func TestRPCCapabilitiesIncludeFrontendAuthProviderExclusive(t *testing.T) {
|
||||
t.Fatalf("frontend_auth_provider_exclusive = %#v, want true", decoded["frontend_auth_provider_exclusive"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestRPCCapabilitiesIncludeScheduler(t *testing.T) {
|
||||
plugin := pluginapi.Plugin{
|
||||
Capabilities: pluginapi.Capabilities{
|
||||
Scheduler: schedulerFunc(func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
return pluginapi.SchedulerPickResponse{}, nil
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
caps := rpcCapabilitiesFromPlugin(plugin)
|
||||
if !caps.Scheduler {
|
||||
t.Fatal("Scheduler = false, want true")
|
||||
}
|
||||
|
||||
raw, errMarshal := json.Marshal(caps)
|
||||
if errMarshal != nil {
|
||||
t.Fatalf("Marshal() error = %v", errMarshal)
|
||||
}
|
||||
if !json.Valid(raw) {
|
||||
t.Fatalf("marshaled capabilities are invalid JSON: %s", raw)
|
||||
}
|
||||
var decoded map[string]any
|
||||
if errUnmarshal := json.Unmarshal(raw, &decoded); errUnmarshal != nil {
|
||||
t.Fatalf("Unmarshal() error = %v", errUnmarshal)
|
||||
}
|
||||
if decoded["scheduler"] != true {
|
||||
t.Fatalf("scheduler = %#v, want true", decoded["scheduler"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestRPCSchedulerPickUsesAdapter(t *testing.T) {
|
||||
var pickCalls int
|
||||
var gotReq pluginapi.SchedulerPickRequest
|
||||
lookup := newTestSymbolLookup(&testPlugin{
|
||||
registerResult: pluginapi.Plugin{
|
||||
Metadata: pluginapi.Metadata{
|
||||
Name: "scheduler",
|
||||
Version: "1.0.0",
|
||||
Author: "test",
|
||||
GitHubRepository: "https://github.com/router-for-me/CLIProxyAPI",
|
||||
},
|
||||
Capabilities: pluginapi.Capabilities{
|
||||
Scheduler: schedulerFunc(func(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
pickCalls++
|
||||
gotReq = req
|
||||
return pluginapi.SchedulerPickResponse{
|
||||
AuthID: "auth-2",
|
||||
Handled: true,
|
||||
}, nil
|
||||
}),
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
plugin, errRegister := registerRPCPlugin(context.Background(), nil, "scheduler", lookup, pluginabi.MethodPluginRegister, nil)
|
||||
if errRegister != nil {
|
||||
t.Fatalf("registerRPCPlugin() error = %v", errRegister)
|
||||
}
|
||||
if plugin.Capabilities.Scheduler == nil {
|
||||
t.Fatal("Scheduler = nil, want adapter")
|
||||
}
|
||||
|
||||
req := pluginapi.SchedulerPickRequest{
|
||||
Provider: "openai",
|
||||
Providers: []string{"openai", "codex"},
|
||||
Model: "gpt-5.4",
|
||||
Stream: true,
|
||||
Options: pluginapi.SchedulerOptions{
|
||||
Headers: map[string][]string{"X-Test": {"one", "two"}},
|
||||
},
|
||||
Candidates: []pluginapi.SchedulerAuthCandidate{
|
||||
{
|
||||
ID: "auth-1",
|
||||
Provider: "openai",
|
||||
Priority: 10,
|
||||
Status: "ready",
|
||||
Attributes: map[string]string{"region": "us"},
|
||||
},
|
||||
{
|
||||
ID: "auth-2",
|
||||
Provider: "codex",
|
||||
Priority: 20,
|
||||
Status: "ready",
|
||||
Attributes: map[string]string{"region": "eu"},
|
||||
},
|
||||
},
|
||||
}
|
||||
resp, errPick := plugin.Capabilities.Scheduler.Pick(context.Background(), req)
|
||||
if errPick != nil {
|
||||
t.Fatalf("Scheduler.Pick() error = %v", errPick)
|
||||
}
|
||||
if resp.AuthID != "auth-2" || !resp.Handled {
|
||||
t.Fatalf("Scheduler.Pick() response = %#v, want auth-2 handled", resp)
|
||||
}
|
||||
if pickCalls != 1 {
|
||||
t.Fatalf("scheduler pick calls = %d, want 1", pickCalls)
|
||||
}
|
||||
if gotReq.Provider != req.Provider || !reflect.DeepEqual(gotReq.Providers, req.Providers) ||
|
||||
gotReq.Model != req.Model || gotReq.Stream != req.Stream {
|
||||
t.Fatalf("scheduler request main fields = %#v, want %#v", gotReq, req)
|
||||
}
|
||||
if !reflect.DeepEqual(gotReq.Options.Headers, req.Options.Headers) {
|
||||
t.Fatalf("scheduler request headers = %#v, want %#v", gotReq.Options.Headers, req.Options.Headers)
|
||||
}
|
||||
if len(gotReq.Candidates) != len(req.Candidates) {
|
||||
t.Fatalf("scheduler candidates len = %d, want %d", len(gotReq.Candidates), len(req.Candidates))
|
||||
}
|
||||
for index := range req.Candidates {
|
||||
gotCandidate := gotReq.Candidates[index]
|
||||
wantCandidate := req.Candidates[index]
|
||||
if gotCandidate.ID != wantCandidate.ID ||
|
||||
gotCandidate.Provider != wantCandidate.Provider ||
|
||||
gotCandidate.Priority != wantCandidate.Priority ||
|
||||
gotCandidate.Status != wantCandidate.Status ||
|
||||
!reflect.DeepEqual(gotCandidate.Attributes, wantCandidate.Attributes) {
|
||||
t.Fatalf("scheduler candidate[%d] = %#v, want %#v", index, gotCandidate, wantCandidate)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizePluginRequestScheduler(t *testing.T) {
|
||||
req := pluginapi.SchedulerPickRequest{
|
||||
Provider: "openai",
|
||||
Providers: []string{"openai", "codex"},
|
||||
Model: "gpt-5.4",
|
||||
Stream: true,
|
||||
Options: pluginapi.SchedulerOptions{
|
||||
Headers: map[string][]string{"X-Test": {"one", "two"}},
|
||||
Metadata: map[string]any{
|
||||
"keep": "value",
|
||||
"drop": make(chan struct{}),
|
||||
},
|
||||
},
|
||||
Candidates: []pluginapi.SchedulerAuthCandidate{
|
||||
{
|
||||
ID: "auth-1",
|
||||
Provider: "openai",
|
||||
Priority: 10,
|
||||
Status: "ready",
|
||||
Attributes: map[string]string{"region": "us"},
|
||||
Metadata: map[string]any{
|
||||
"keep": "candidate",
|
||||
"drop": make(chan struct{}),
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
raw, errMarshal := json.Marshal(sanitizePluginRequest(req))
|
||||
if errMarshal != nil {
|
||||
t.Fatalf("Marshal(sanitized scheduler request) error = %v", errMarshal)
|
||||
}
|
||||
var decoded pluginapi.SchedulerPickRequest
|
||||
if errUnmarshal := json.Unmarshal(raw, &decoded); errUnmarshal != nil {
|
||||
t.Fatalf("Unmarshal(sanitized scheduler request) error = %v", errUnmarshal)
|
||||
}
|
||||
|
||||
if decoded.Provider != req.Provider || !reflect.DeepEqual(decoded.Providers, req.Providers) ||
|
||||
decoded.Model != req.Model || decoded.Stream != req.Stream {
|
||||
t.Fatalf("scheduler request main fields = %#v, want %#v", decoded, req)
|
||||
}
|
||||
if !reflect.DeepEqual(decoded.Options.Headers, req.Options.Headers) {
|
||||
t.Fatalf("scheduler request headers = %#v, want %#v", decoded.Options.Headers, req.Options.Headers)
|
||||
}
|
||||
if decoded.Options.Metadata["keep"] != "value" {
|
||||
t.Fatalf("scheduler options metadata keep = %#v, want value", decoded.Options.Metadata["keep"])
|
||||
}
|
||||
if _, ok := decoded.Options.Metadata["drop"]; ok {
|
||||
t.Fatalf("scheduler options metadata drop survived sanitize: %#v", decoded.Options.Metadata)
|
||||
}
|
||||
if len(decoded.Candidates) != 1 {
|
||||
t.Fatalf("scheduler candidates len = %d, want 1", len(decoded.Candidates))
|
||||
}
|
||||
gotCandidate := decoded.Candidates[0]
|
||||
wantCandidate := req.Candidates[0]
|
||||
if gotCandidate.ID != wantCandidate.ID ||
|
||||
gotCandidate.Provider != wantCandidate.Provider ||
|
||||
gotCandidate.Priority != wantCandidate.Priority ||
|
||||
gotCandidate.Status != wantCandidate.Status ||
|
||||
!reflect.DeepEqual(gotCandidate.Attributes, wantCandidate.Attributes) {
|
||||
t.Fatalf("scheduler candidate = %#v, want %#v", gotCandidate, wantCandidate)
|
||||
}
|
||||
if gotCandidate.Metadata["keep"] != "candidate" {
|
||||
t.Fatalf("scheduler candidate metadata keep = %#v, want candidate", gotCandidate.Metadata["keep"])
|
||||
}
|
||||
if _, ok := gotCandidate.Metadata["drop"]; ok {
|
||||
t.Fatalf("scheduler candidate metadata drop survived sanitize: %#v", gotCandidate.Metadata)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
package pluginhost
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
func (h *Host) PickAuth(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) {
|
||||
record := h.schedulerRecord()
|
||||
if record == nil {
|
||||
return pluginapi.SchedulerPickResponse{}, false, nil
|
||||
}
|
||||
|
||||
resp, handled, errPick := h.callScheduler(ctx, *record, req)
|
||||
if errPick != nil || !handled {
|
||||
return resp, handled, errPick
|
||||
}
|
||||
if !resp.Handled {
|
||||
return pluginapi.SchedulerPickResponse{}, false, nil
|
||||
}
|
||||
|
||||
resp, valid, reason := normalizeSchedulerResponse(resp, req)
|
||||
if !valid {
|
||||
log.WithField("plugin_id", record.id).Warnf("pluginhost: scheduler returned invalid response: %s", reason)
|
||||
return pluginapi.SchedulerPickResponse{}, false, nil
|
||||
}
|
||||
return resp, true, nil
|
||||
}
|
||||
|
||||
func (h *Host) schedulerRecord() *capabilityRecord {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
for _, record := range h.Snapshot().records {
|
||||
if h.isPluginFused(record.id) || record.plugin.Capabilities.Scheduler == nil {
|
||||
continue
|
||||
}
|
||||
copyRecord := record
|
||||
return ©Record
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *Host) callScheduler(ctx context.Context, record capabilityRecord, req pluginapi.SchedulerPickRequest) (resp pluginapi.SchedulerPickResponse, handled bool, err error) {
|
||||
scheduler := record.plugin.Capabilities.Scheduler
|
||||
if h == nil || scheduler == nil || h.isPluginFused(record.id) {
|
||||
return pluginapi.SchedulerPickResponse{}, false, nil
|
||||
}
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
h.fusePlugin(record.id, "Scheduler.Pick", recovered)
|
||||
resp = pluginapi.SchedulerPickResponse{}
|
||||
handled = false
|
||||
err = nil
|
||||
}
|
||||
}()
|
||||
|
||||
req.Plugin = record.meta
|
||||
resp, errPick := scheduler.Pick(ctx, req)
|
||||
if errPick != nil {
|
||||
log.WithField("plugin_id", record.id).WithError(errPick).Warn("pluginhost: scheduler rejected auth pick")
|
||||
return pluginapi.SchedulerPickResponse{}, true, errPick
|
||||
}
|
||||
return resp, true, nil
|
||||
}
|
||||
|
||||
func normalizeSchedulerResponse(resp pluginapi.SchedulerPickResponse, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, string) {
|
||||
resp.AuthID = strings.TrimSpace(resp.AuthID)
|
||||
resp.DelegateBuiltin = strings.TrimSpace(resp.DelegateBuiltin)
|
||||
|
||||
hasAuthID := resp.AuthID != ""
|
||||
hasDelegate := resp.DelegateBuiltin != ""
|
||||
if !hasAuthID && !hasDelegate {
|
||||
return pluginapi.SchedulerPickResponse{}, false, "missing auth id or delegate"
|
||||
}
|
||||
if hasAuthID {
|
||||
if !schedulerCandidateExists(req.Candidates, resp.AuthID) {
|
||||
return pluginapi.SchedulerPickResponse{}, false, "unknown auth id"
|
||||
}
|
||||
return resp, true, ""
|
||||
}
|
||||
if !validSchedulerBuiltin(resp.DelegateBuiltin) {
|
||||
return pluginapi.SchedulerPickResponse{}, false, "unknown delegate"
|
||||
}
|
||||
return resp, true, ""
|
||||
}
|
||||
|
||||
func schedulerCandidateExists(candidates []pluginapi.SchedulerAuthCandidate, authID string) bool {
|
||||
for _, candidate := range candidates {
|
||||
if strings.TrimSpace(candidate.ID) == authID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func validSchedulerBuiltin(delegate string) bool {
|
||||
switch delegate {
|
||||
case pluginapi.SchedulerBuiltinRoundRobin, pluginapi.SchedulerBuiltinFillFirst:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
package pluginhost
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
)
|
||||
|
||||
func TestHostPickAuthUsesHighestPrioritySchedulerOnly(t *testing.T) {
|
||||
var highCalls int
|
||||
var lowCalls int
|
||||
host := newHostWithRecords(
|
||||
capabilityRecord{
|
||||
id: "low",
|
||||
priority: 1,
|
||||
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{Scheduler: schedulerFunc(func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
lowCalls++
|
||||
return pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-low"}, nil
|
||||
})}},
|
||||
},
|
||||
capabilityRecord{
|
||||
id: "high",
|
||||
priority: 10,
|
||||
meta: pluginapi.Metadata{Name: "high", Version: "1.0.0"},
|
||||
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{Scheduler: schedulerFunc(func(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
highCalls++
|
||||
if req.Plugin.Name != "high" {
|
||||
t.Fatalf("req.Plugin.Name = %q, want high", req.Plugin.Name)
|
||||
}
|
||||
return pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-high"}, nil
|
||||
})}},
|
||||
},
|
||||
)
|
||||
|
||||
resp, handled, errPick := host.PickAuth(context.Background(), schedulerRequest("auth-high", "auth-low"))
|
||||
if errPick != nil {
|
||||
t.Fatalf("PickAuth() error = %v, want nil", errPick)
|
||||
}
|
||||
if !handled {
|
||||
t.Fatal("PickAuth() handled = false, want true")
|
||||
}
|
||||
if resp.AuthID != "auth-high" {
|
||||
t.Fatalf("PickAuth() AuthID = %q, want auth-high", resp.AuthID)
|
||||
}
|
||||
if highCalls != 1 {
|
||||
t.Fatalf("high calls = %d, want 1", highCalls)
|
||||
}
|
||||
if lowCalls != 0 {
|
||||
t.Fatalf("low calls = %d, want 0", lowCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHostPickAuthReturnsSchedulerError(t *testing.T) {
|
||||
host := newHostWithRecords(capabilityRecord{
|
||||
id: "scheduler",
|
||||
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{Scheduler: schedulerFunc(func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
return pluginapi.SchedulerPickResponse{}, errors.New("tenant quota exhausted")
|
||||
})}},
|
||||
})
|
||||
|
||||
_, handled, errPick := host.PickAuth(context.Background(), schedulerRequest("auth-1"))
|
||||
if !handled {
|
||||
t.Fatal("PickAuth() handled = false, want true")
|
||||
}
|
||||
if errPick == nil || !strings.Contains(errPick.Error(), "tenant quota exhausted") {
|
||||
t.Fatalf("PickAuth() error = %v, want tenant quota exhausted", errPick)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHostPickAuthPanicFusesAndFallsBack(t *testing.T) {
|
||||
host := newHostWithRecords(capabilityRecord{
|
||||
id: "scheduler",
|
||||
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{Scheduler: schedulerFunc(func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
panic("boom")
|
||||
})}},
|
||||
})
|
||||
|
||||
_, handled, errPick := host.PickAuth(context.Background(), schedulerRequest("auth-1"))
|
||||
if handled {
|
||||
t.Fatal("PickAuth() handled = true, want false")
|
||||
}
|
||||
if errPick != nil {
|
||||
t.Fatalf("PickAuth() error = %v, want nil", errPick)
|
||||
}
|
||||
if !host.isPluginFused("scheduler") {
|
||||
t.Fatal("scheduler plugin was not fused after panic")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHostPickAuthUnhandledDoesNotCallLowerPriorityScheduler(t *testing.T) {
|
||||
var lowCalls int
|
||||
host := newHostWithRecords(
|
||||
capabilityRecord{
|
||||
id: "low",
|
||||
priority: 1,
|
||||
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{Scheduler: schedulerFunc(func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
lowCalls++
|
||||
return pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-low"}, nil
|
||||
})}},
|
||||
},
|
||||
capabilityRecord{
|
||||
id: "high",
|
||||
priority: 10,
|
||||
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{Scheduler: schedulerFunc(func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
return pluginapi.SchedulerPickResponse{Handled: false}, nil
|
||||
})}},
|
||||
},
|
||||
)
|
||||
|
||||
_, handled, errPick := host.PickAuth(context.Background(), schedulerRequest("auth-low"))
|
||||
if errPick != nil {
|
||||
t.Fatalf("PickAuth() error = %v, want nil", errPick)
|
||||
}
|
||||
if handled {
|
||||
t.Fatal("PickAuth() handled = true, want false")
|
||||
}
|
||||
if lowCalls != 0 {
|
||||
t.Fatalf("low calls = %d, want 0", lowCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHostPickAuthInvalidResponseFallsBack(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
resp pluginapi.SchedulerPickResponse
|
||||
}{
|
||||
{
|
||||
name: "unknown auth id",
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "missing"},
|
||||
},
|
||||
{
|
||||
name: "unknown delegate",
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: true, DelegateBuiltin: "unknown"},
|
||||
},
|
||||
{
|
||||
name: "handled without decision",
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: true},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
host := newHostWithRecords(capabilityRecord{
|
||||
id: "scheduler",
|
||||
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{Scheduler: schedulerFunc(func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
return tt.resp, nil
|
||||
})}},
|
||||
})
|
||||
|
||||
_, handled, errPick := host.PickAuth(context.Background(), schedulerRequest("auth-1"))
|
||||
if errPick != nil {
|
||||
t.Fatalf("PickAuth() error = %v, want nil", errPick)
|
||||
}
|
||||
if handled {
|
||||
t.Fatal("PickAuth() handled = true, want false")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHostPickAuthPrefersValidAuthIDOverInvalidDelegate(t *testing.T) {
|
||||
host := newHostWithRecords(capabilityRecord{
|
||||
id: "scheduler",
|
||||
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{Scheduler: schedulerFunc(func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
return pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-a", DelegateBuiltin: "unknown"}, nil
|
||||
})}},
|
||||
})
|
||||
|
||||
resp, handled, errPick := host.PickAuth(context.Background(), schedulerRequest("auth-a"))
|
||||
if errPick != nil {
|
||||
t.Fatalf("PickAuth() error = %v, want nil", errPick)
|
||||
}
|
||||
if !handled {
|
||||
t.Fatal("PickAuth() handled = false, want true")
|
||||
}
|
||||
if resp.AuthID != "auth-a" {
|
||||
t.Fatalf("PickAuth() AuthID = %q, want auth-a", resp.AuthID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHostPickAuthAllowsKnownBuiltinDelegates(t *testing.T) {
|
||||
for _, delegate := range []string{pluginapi.SchedulerBuiltinRoundRobin, pluginapi.SchedulerBuiltinFillFirst} {
|
||||
t.Run(delegate, func(t *testing.T) {
|
||||
host := newHostWithRecords(capabilityRecord{
|
||||
id: "scheduler",
|
||||
plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{Scheduler: schedulerFunc(func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
return pluginapi.SchedulerPickResponse{Handled: true, DelegateBuiltin: delegate}, nil
|
||||
})}},
|
||||
})
|
||||
|
||||
resp, handled, errPick := host.PickAuth(context.Background(), schedulerRequest("auth-1"))
|
||||
if errPick != nil {
|
||||
t.Fatalf("PickAuth() error = %v, want nil", errPick)
|
||||
}
|
||||
if !handled {
|
||||
t.Fatal("PickAuth() handled = false, want true")
|
||||
}
|
||||
if resp.DelegateBuiltin != delegate {
|
||||
t.Fatalf("PickAuth() DelegateBuiltin = %q, want %q", resp.DelegateBuiltin, delegate)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func schedulerRequest(ids ...string) pluginapi.SchedulerPickRequest {
|
||||
req := pluginapi.SchedulerPickRequest{
|
||||
Provider: "test",
|
||||
Model: "test-model",
|
||||
}
|
||||
for _, id := range ids {
|
||||
req.Candidates = append(req.Candidates, pluginapi.SchedulerAuthCandidate{ID: id})
|
||||
}
|
||||
return req
|
||||
}
|
||||
@@ -107,6 +107,19 @@ func (l *testSymbolLookup) Call(ctx context.Context, method string, request []by
|
||||
return nil, fmt.Errorf("missing auth provider")
|
||||
}
|
||||
return marshalRPCResult(rpcIdentifierResponse{Identifier: l.active.Capabilities.AuthProvider.Identifier()})
|
||||
case pluginabi.MethodSchedulerPick:
|
||||
if l.active.Capabilities.Scheduler == nil {
|
||||
return nil, fmt.Errorf("missing scheduler")
|
||||
}
|
||||
var req pluginapi.SchedulerPickRequest
|
||||
if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil {
|
||||
return nil, errUnmarshal
|
||||
}
|
||||
resp, errPick := l.active.Capabilities.Scheduler.Pick(ctx, req)
|
||||
if errPick != nil {
|
||||
return nil, errPick
|
||||
}
|
||||
return marshalRPCResult(resp)
|
||||
case pluginabi.MethodUsageHandle:
|
||||
if l.active.Capabilities.UsagePlugin == nil {
|
||||
return marshalRPCResult(rpcEmptyResponse{})
|
||||
@@ -225,6 +238,15 @@ func (f requestInterceptorFunc) InterceptRequest(ctx context.Context, req plugin
|
||||
return f(ctx, req)
|
||||
}
|
||||
|
||||
type schedulerFunc func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error)
|
||||
|
||||
func (f schedulerFunc) Pick(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, error) {
|
||||
if f == nil {
|
||||
return pluginapi.SchedulerPickResponse{}, fmt.Errorf("missing scheduler callback")
|
||||
}
|
||||
return f(ctx, req)
|
||||
}
|
||||
|
||||
type responseInterceptorFunc struct {
|
||||
interceptResponse func(context.Context, pluginapi.ResponseInterceptRequest) (pluginapi.ResponseInterceptResponse, error)
|
||||
interceptStreamChunk func(context.Context, pluginapi.StreamChunkInterceptRequest) (pluginapi.StreamChunkInterceptResponse, error)
|
||||
|
||||
+244
-20
@@ -24,6 +24,7 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
@@ -121,6 +122,10 @@ type Selector interface {
|
||||
Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error)
|
||||
}
|
||||
|
||||
type PluginScheduler interface {
|
||||
PickAuth(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error)
|
||||
}
|
||||
|
||||
// StoppableSelector is an optional interface for selectors that hold resources.
|
||||
// Selectors that implement this interface will have Stop called during shutdown.
|
||||
type StoppableSelector interface {
|
||||
@@ -159,6 +164,9 @@ type Manager struct {
|
||||
mu sync.RWMutex
|
||||
auths map[string]*Auth
|
||||
scheduler *authScheduler
|
||||
// pluginScheduler runs outside m.mu before falling back to native selection.
|
||||
pluginScheduler PluginScheduler
|
||||
pluginDelegateRoundRobin *RoundRobinSelector
|
||||
// homeRuntimeAuths caches auths returned by Home so websocket sessions can
|
||||
// reuse an established upstream credential without dispatching every turn.
|
||||
homeRuntimeAuths map[string]map[string]*Auth
|
||||
@@ -203,14 +211,15 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager {
|
||||
hook = NoopHook{}
|
||||
}
|
||||
manager := &Manager{
|
||||
store: store,
|
||||
executors: make(map[string]ProviderExecutor),
|
||||
selector: selector,
|
||||
hook: hook,
|
||||
auths: make(map[string]*Auth),
|
||||
homeRuntimeAuths: make(map[string]map[string]*Auth),
|
||||
providerOffsets: make(map[string]int),
|
||||
modelPoolOffsets: make(map[string]int),
|
||||
store: store,
|
||||
executors: make(map[string]ProviderExecutor),
|
||||
selector: selector,
|
||||
hook: hook,
|
||||
auths: make(map[string]*Auth),
|
||||
pluginDelegateRoundRobin: &RoundRobinSelector{},
|
||||
homeRuntimeAuths: make(map[string]map[string]*Auth),
|
||||
providerOffsets: make(map[string]int),
|
||||
modelPoolOffsets: make(map[string]int),
|
||||
}
|
||||
// atomic.Value requires non-nil initial value.
|
||||
manager.runtimeConfig.Store(&internalconfig.Config{})
|
||||
@@ -219,6 +228,28 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager {
|
||||
return manager
|
||||
}
|
||||
|
||||
func (m *Manager) SetPluginScheduler(scheduler PluginScheduler) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.pluginScheduler = scheduler
|
||||
if m.pluginDelegateRoundRobin == nil {
|
||||
m.pluginDelegateRoundRobin = &RoundRobinSelector{}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
func (m *Manager) hasPluginScheduler() bool {
|
||||
if m == nil {
|
||||
return false
|
||||
}
|
||||
m.mu.RLock()
|
||||
ok := m.pluginScheduler != nil
|
||||
m.mu.RUnlock()
|
||||
return ok
|
||||
}
|
||||
|
||||
func isBuiltInSelector(selector Selector) bool {
|
||||
switch selector.(type) {
|
||||
case *RoundRobinSelector, *FillFirstSelector:
|
||||
@@ -722,6 +753,182 @@ func selectionArgForSelector(selector Selector, routeModel string) string {
|
||||
return routeModel
|
||||
}
|
||||
|
||||
func schedulerAttributeSensitive(key string) bool {
|
||||
key = strings.ToLower(strings.TrimSpace(key))
|
||||
normalized := strings.NewReplacer("-", "_", ".", "_", " ", "_").Replace(key)
|
||||
compact := strings.NewReplacer("_", "", "-", "", ".", "", " ", "").Replace(key)
|
||||
for _, fragment := range []string{
|
||||
"api_key",
|
||||
"apikey",
|
||||
"token",
|
||||
"secret",
|
||||
"cookie",
|
||||
"credential",
|
||||
"password",
|
||||
"storage",
|
||||
"authorization",
|
||||
"auth_header",
|
||||
"proxy_url",
|
||||
} {
|
||||
if strings.Contains(key, fragment) || strings.Contains(normalized, fragment) || strings.Contains(compact, fragment) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func schedulerSafeAttributes(src map[string]string) map[string]string {
|
||||
if len(src) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]string, len(src))
|
||||
for key, value := range src {
|
||||
if schedulerAttributeSensitive(key) {
|
||||
continue
|
||||
}
|
||||
out[key] = value
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneSchedulerAnyMap(src map[string]any) map[string]any {
|
||||
if len(src) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]any, len(src))
|
||||
for key, value := range src {
|
||||
out[key] = value
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneAuthSlice(auths []*Auth) []*Auth {
|
||||
if len(auths) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*Auth, 0, len(auths))
|
||||
for _, auth := range auths {
|
||||
if auth == nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, auth.Clone())
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func schedulerAuthCandidates(auths []*Auth) []pluginapi.SchedulerAuthCandidate {
|
||||
if len(auths) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]pluginapi.SchedulerAuthCandidate, 0, len(auths))
|
||||
for _, auth := range auths {
|
||||
if auth == nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, pluginapi.SchedulerAuthCandidate{
|
||||
ID: auth.ID,
|
||||
Provider: strings.ToLower(strings.TrimSpace(auth.Provider)),
|
||||
Priority: authPriority(auth),
|
||||
Status: string(auth.Status),
|
||||
Attributes: schedulerSafeAttributes(auth.Attributes),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func schedulerProviders(provider string, providers []string) []string {
|
||||
out := make([]string, 0, len(providers)+1)
|
||||
seen := make(map[string]struct{}, len(providers)+1)
|
||||
addProvider := func(value string) {
|
||||
value = strings.ToLower(strings.TrimSpace(value))
|
||||
if value == "" || value == "mixed" {
|
||||
return
|
||||
}
|
||||
if _, ok := seen[value]; ok {
|
||||
return
|
||||
}
|
||||
seen[value] = struct{}{}
|
||||
out = append(out, value)
|
||||
}
|
||||
addProvider(provider)
|
||||
for _, value := range providers {
|
||||
addProvider(value)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func schedulerOptions(opts cliproxyexecutor.Options) pluginapi.SchedulerOptions {
|
||||
return pluginapi.SchedulerOptions{
|
||||
Headers: cloneHTTPHeader(opts.Headers),
|
||||
Metadata: cloneSchedulerAnyMap(opts.Metadata),
|
||||
}
|
||||
}
|
||||
|
||||
func pickSchedulerAuthByID(candidates []*Auth, authID string) *Auth {
|
||||
authID = strings.TrimSpace(authID)
|
||||
if authID == "" {
|
||||
return nil
|
||||
}
|
||||
for _, candidate := range candidates {
|
||||
if candidate != nil && candidate.ID == authID {
|
||||
return candidate
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) pickViaPluginScheduler(ctx context.Context, scheduler PluginScheduler, roundRobin *RoundRobinSelector, provider string, providers []string, model string, opts cliproxyexecutor.Options, candidates []*Auth) (*Auth, bool, error) {
|
||||
if scheduler == nil || len(candidates) == 0 {
|
||||
return nil, false, nil
|
||||
}
|
||||
providerKey := strings.ToLower(strings.TrimSpace(provider))
|
||||
requestProvider := providerKey
|
||||
if providerKey == "mixed" {
|
||||
requestProvider = ""
|
||||
}
|
||||
req := pluginapi.SchedulerPickRequest{
|
||||
Provider: requestProvider,
|
||||
Providers: schedulerProviders(providerKey, providers),
|
||||
Model: model,
|
||||
Stream: opts.Stream,
|
||||
Options: schedulerOptions(opts),
|
||||
Candidates: schedulerAuthCandidates(candidates),
|
||||
}
|
||||
resp, handled, errPick := scheduler.PickAuth(ctx, req)
|
||||
if errPick != nil {
|
||||
return nil, true, errPick
|
||||
}
|
||||
if !handled || !resp.Handled {
|
||||
return nil, false, nil
|
||||
}
|
||||
if selected := pickSchedulerAuthByID(candidates, resp.AuthID); selected != nil {
|
||||
return selected, true, nil
|
||||
}
|
||||
|
||||
switch strings.TrimSpace(resp.DelegateBuiltin) {
|
||||
case pluginapi.SchedulerBuiltinRoundRobin:
|
||||
if roundRobin == nil {
|
||||
roundRobin = &RoundRobinSelector{}
|
||||
}
|
||||
selected, errSelect := roundRobin.Pick(ctx, providerKey, "", opts, candidates)
|
||||
if errSelect != nil {
|
||||
return nil, true, errSelect
|
||||
}
|
||||
return selected, true, nil
|
||||
case pluginapi.SchedulerBuiltinFillFirst:
|
||||
selected, errSelect := (&FillFirstSelector{}).Pick(ctx, providerKey, "", opts, candidates)
|
||||
if errSelect != nil {
|
||||
return nil, true, errSelect
|
||||
}
|
||||
return selected, true, nil
|
||||
default:
|
||||
return nil, false, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) authSupportsRouteModel(registryRef *registry.ModelRegistry, auth *Auth, routeModel string) bool {
|
||||
if registryRef == nil || auth == nil {
|
||||
return true
|
||||
@@ -3167,6 +3374,9 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op
|
||||
disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata)
|
||||
|
||||
m.mu.RLock()
|
||||
selector := m.selector
|
||||
pluginScheduler := m.pluginScheduler
|
||||
pluginDelegateRoundRobin := m.pluginDelegateRoundRobin
|
||||
executor, okExecutor := m.executors[provider]
|
||||
if !okExecutor {
|
||||
m.mu.RUnlock()
|
||||
@@ -3209,17 +3419,23 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op
|
||||
m.mu.RUnlock()
|
||||
return nil, nil, errAvailable
|
||||
}
|
||||
selected, errPick := m.selector.Pick(ctx, provider, selectionArgForSelector(m.selector, model), opts, available)
|
||||
available = cloneAuthSlice(available)
|
||||
m.mu.RUnlock()
|
||||
|
||||
selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, pluginDelegateRoundRobin, provider, []string{provider}, model, opts, available)
|
||||
if errPick != nil {
|
||||
m.mu.RUnlock()
|
||||
return nil, nil, errPick
|
||||
}
|
||||
if !handled {
|
||||
selected, errPick = selector.Pick(ctx, provider, selectionArgForSelector(selector, model), opts, available)
|
||||
if errPick != nil {
|
||||
return nil, nil, errPick
|
||||
}
|
||||
}
|
||||
if selected == nil {
|
||||
m.mu.RUnlock()
|
||||
return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"}
|
||||
}
|
||||
authCopy := selected.Clone()
|
||||
m.mu.RUnlock()
|
||||
if !selected.indexAssigned {
|
||||
m.mu.Lock()
|
||||
if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned {
|
||||
@@ -3237,7 +3453,7 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cli
|
||||
return auth, exec, err
|
||||
}
|
||||
|
||||
if !m.useSchedulerFastPath() {
|
||||
if m.hasPluginScheduler() || !m.useSchedulerFastPath() {
|
||||
return m.pickNextLegacy(ctx, provider, model, opts, tried)
|
||||
}
|
||||
if strings.TrimSpace(model) != "" {
|
||||
@@ -3314,6 +3530,9 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m
|
||||
}
|
||||
|
||||
m.mu.RLock()
|
||||
selector := m.selector
|
||||
pluginScheduler := m.pluginScheduler
|
||||
pluginDelegateRoundRobin := m.pluginDelegateRoundRobin
|
||||
candidates := make([]*Auth, 0, len(m.auths))
|
||||
modelKey := strings.TrimSpace(model)
|
||||
// Always use base model name (without thinking suffix) for auth matching.
|
||||
@@ -3361,23 +3580,28 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m
|
||||
m.mu.RUnlock()
|
||||
return nil, nil, "", errAvailable
|
||||
}
|
||||
selected, errPick := m.selector.Pick(ctx, "mixed", selectionArgForSelector(m.selector, model), opts, available)
|
||||
available = cloneAuthSlice(available)
|
||||
m.mu.RUnlock()
|
||||
|
||||
selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, pluginDelegateRoundRobin, "mixed", providers, model, opts, available)
|
||||
if errPick != nil {
|
||||
m.mu.RUnlock()
|
||||
return nil, nil, "", errPick
|
||||
}
|
||||
if !handled {
|
||||
selected, errPick = selector.Pick(ctx, "mixed", selectionArgForSelector(selector, model), opts, available)
|
||||
if errPick != nil {
|
||||
return nil, nil, "", errPick
|
||||
}
|
||||
}
|
||||
if selected == nil {
|
||||
m.mu.RUnlock()
|
||||
return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"}
|
||||
}
|
||||
providerKey := strings.TrimSpace(strings.ToLower(selected.Provider))
|
||||
executor, okExecutor := m.executors[providerKey]
|
||||
executor, okExecutor := m.Executor(providerKey)
|
||||
if !okExecutor {
|
||||
m.mu.RUnlock()
|
||||
return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"}
|
||||
}
|
||||
authCopy := selected.Clone()
|
||||
m.mu.RUnlock()
|
||||
if !selected.indexAssigned {
|
||||
m.mu.Lock()
|
||||
if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned {
|
||||
@@ -3394,7 +3618,7 @@ func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model s
|
||||
return m.pickNextViaHome(ctx, model, opts, tried)
|
||||
}
|
||||
|
||||
if !m.useSchedulerFastPath() {
|
||||
if m.hasPluginScheduler() || !m.useSchedulerFastPath() {
|
||||
return m.pickNextMixedLegacy(ctx, providers, model, opts, tried)
|
||||
}
|
||||
|
||||
|
||||
@@ -2,12 +2,15 @@ package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
)
|
||||
|
||||
type schedulerTestExecutor struct{}
|
||||
@@ -34,6 +37,24 @@ func (schedulerTestExecutor) HttpRequest(ctx context.Context, auth *Auth, req *h
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
type fakePluginScheduler struct {
|
||||
resp pluginapi.SchedulerPickResponse
|
||||
handled bool
|
||||
err error
|
||||
calls int
|
||||
requests []pluginapi.SchedulerPickRequest
|
||||
pick func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error)
|
||||
}
|
||||
|
||||
func (s *fakePluginScheduler) PickAuth(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) {
|
||||
s.calls++
|
||||
s.requests = append(s.requests, req)
|
||||
if s.pick != nil {
|
||||
return s.pick(ctx, req)
|
||||
}
|
||||
return s.resp, s.handled, s.err
|
||||
}
|
||||
|
||||
type trackingSelector struct {
|
||||
calls int
|
||||
lastAuthID []string
|
||||
@@ -366,6 +387,328 @@ func TestManager_PickNextMixed_DisallowFreeAuthSkipsCodexFreePlan(t *testing.T)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerSelectsAuthID(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-a) error = %v", errRegister)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-b) error = %v", errRegister)
|
||||
}
|
||||
|
||||
scheduler := &fakePluginScheduler{
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-b"},
|
||||
handled: true,
|
||||
}
|
||||
manager.SetPluginScheduler(scheduler)
|
||||
|
||||
got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{Stream: true}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNext() error = %v", errPick)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("pickNext() auth = nil")
|
||||
}
|
||||
if got.ID != "auth-b" {
|
||||
t.Fatalf("pickNext() auth.ID = %q, want %q", got.ID, "auth-b")
|
||||
}
|
||||
if scheduler.calls != 1 {
|
||||
t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 1)
|
||||
}
|
||||
if len(scheduler.requests) != 1 {
|
||||
t.Fatalf("len(scheduler.requests) = %d, want %d", len(scheduler.requests), 1)
|
||||
}
|
||||
if !scheduler.requests[0].Stream {
|
||||
t.Fatalf("scheduler request Stream = false, want true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerSkippedWhenHomeEnabled(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
||||
scheduler := &fakePluginScheduler{
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-a"},
|
||||
handled: true,
|
||||
}
|
||||
manager.SetPluginScheduler(scheduler)
|
||||
|
||||
_, _, _ = manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil)
|
||||
|
||||
if scheduler.calls != 0 {
|
||||
t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 0)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerCalledOutsideManagerLock(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-a) error = %v", errRegister)
|
||||
}
|
||||
|
||||
scheduler := &fakePluginScheduler{
|
||||
handled: true,
|
||||
pick: func(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) {
|
||||
if !manager.mu.TryLock() {
|
||||
t.Fatalf("plugin scheduler called while manager lock is held")
|
||||
}
|
||||
manager.mu.Unlock()
|
||||
return pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-a"}, true, nil
|
||||
},
|
||||
}
|
||||
manager.SetPluginScheduler(scheduler)
|
||||
|
||||
got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNext() error = %v", errPick)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("pickNext() auth = nil")
|
||||
}
|
||||
if got.ID != "auth-a" {
|
||||
t.Fatalf("pickNext() auth.ID = %q, want auth-a", got.ID)
|
||||
}
|
||||
if scheduler.calls != 1 {
|
||||
t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerErrorStopsPick(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-a) error = %v", errRegister)
|
||||
}
|
||||
|
||||
scheduler := &fakePluginScheduler{
|
||||
handled: true,
|
||||
err: errors.New("tenant denied"),
|
||||
}
|
||||
manager.SetPluginScheduler(scheduler)
|
||||
|
||||
got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick == nil {
|
||||
t.Fatalf("pickNext() error = nil, want tenant denied")
|
||||
}
|
||||
if errPick.Error() != "tenant denied" {
|
||||
t.Fatalf("pickNext() error = %v, want tenant denied", errPick)
|
||||
}
|
||||
if got != nil {
|
||||
t.Fatalf("pickNext() auth = %v, want nil", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerFallsBackWhenUnhandledOrUnknown(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
resp pluginapi.SchedulerPickResponse
|
||||
handled bool
|
||||
}{
|
||||
{
|
||||
name: "unhandled",
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: false},
|
||||
handled: false,
|
||||
},
|
||||
{
|
||||
name: "unknown auth id",
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "missing"},
|
||||
handled: true,
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
manager := NewManager(nil, &FillFirstSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-b) error = %v", errRegister)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-a) error = %v", errRegister)
|
||||
}
|
||||
|
||||
scheduler := &fakePluginScheduler{resp: tc.resp, handled: tc.handled}
|
||||
manager.SetPluginScheduler(scheduler)
|
||||
|
||||
got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNext() error = %v", errPick)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("pickNext() auth = nil")
|
||||
}
|
||||
if got.ID != "auth-a" {
|
||||
t.Fatalf("pickNext() auth.ID = %q, want %q", got.ID, "auth-a")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerDelegatesBuiltin(t *testing.T) {
|
||||
t.Run("round-robin", func(t *testing.T) {
|
||||
manager := NewManager(nil, &FillFirstSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-a) error = %v", errRegister)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-b) error = %v", errRegister)
|
||||
}
|
||||
manager.SetPluginScheduler(&fakePluginScheduler{
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: true, DelegateBuiltin: pluginapi.SchedulerBuiltinRoundRobin},
|
||||
handled: true,
|
||||
})
|
||||
|
||||
gotA, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNext() first error = %v", errPick)
|
||||
}
|
||||
gotB, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNext() second error = %v", errPick)
|
||||
}
|
||||
if gotA == nil || gotB == nil {
|
||||
t.Fatalf("pickNext() auths = %v, %v; want non-nil", gotA, gotB)
|
||||
}
|
||||
if gotA.ID != "auth-a" || gotB.ID != "auth-b" {
|
||||
t.Fatalf("round-robin picks = %q, %q; want auth-a, auth-b", gotA.ID, gotB.ID)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("fill-first", func(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-b) error = %v", errRegister)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-a) error = %v", errRegister)
|
||||
}
|
||||
manager.SetPluginScheduler(&fakePluginScheduler{
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: true, DelegateBuiltin: pluginapi.SchedulerBuiltinFillFirst},
|
||||
handled: true,
|
||||
})
|
||||
|
||||
got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNext() error = %v", errPick)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("pickNext() auth = nil")
|
||||
}
|
||||
if got.ID != "auth-a" {
|
||||
t.Fatalf("fill-first pick = %q, want auth-a", got.ID)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerPickNextMixedSelectsProvider(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
manager.executors["claude"] = schedulerTestExecutor{}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-a", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(gemini-a) error = %v", errRegister)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "claude-a", Provider: "claude"}); errRegister != nil {
|
||||
t.Fatalf("Register(claude-a) error = %v", errRegister)
|
||||
}
|
||||
scheduler := &fakePluginScheduler{
|
||||
resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "claude-a"},
|
||||
handled: true,
|
||||
}
|
||||
manager.SetPluginScheduler(scheduler)
|
||||
|
||||
got, executor, provider, errPick := manager.pickNextMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNextMixed() error = %v", errPick)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("pickNextMixed() auth = nil")
|
||||
}
|
||||
if got.ID != "claude-a" {
|
||||
t.Fatalf("pickNextMixed() auth.ID = %q, want claude-a", got.ID)
|
||||
}
|
||||
if provider != "claude" {
|
||||
t.Fatalf("pickNextMixed() provider = %q, want claude", provider)
|
||||
}
|
||||
if executor == nil {
|
||||
t.Fatalf("pickNextMixed() executor = nil")
|
||||
}
|
||||
if len(scheduler.requests) != 1 {
|
||||
t.Fatalf("len(scheduler.requests) = %d, want %d", len(scheduler.requests), 1)
|
||||
}
|
||||
req := scheduler.requests[0]
|
||||
if req.Provider != "" {
|
||||
t.Fatalf("scheduler request Provider = %q, want empty for mixed provider pick", req.Provider)
|
||||
}
|
||||
if len(req.Providers) != 2 || req.Providers[0] != "gemini" || req.Providers[1] != "claude" {
|
||||
t.Fatalf("scheduler request Providers = %#v, want [gemini claude]", req.Providers)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerCandidatesAreSafeCopies(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
auth := &Auth{
|
||||
ID: "auth-a",
|
||||
Provider: "gemini",
|
||||
Status: StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"access_token": "token-value",
|
||||
"api_key": "api-key-value",
|
||||
"cookie": "cookie-value",
|
||||
"priority": "7",
|
||||
"team": "alpha",
|
||||
},
|
||||
Metadata: map[string]any{"tenant": "one"},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("Register(auth-a) error = %v", errRegister)
|
||||
}
|
||||
|
||||
scheduler := &fakePluginScheduler{
|
||||
handled: true,
|
||||
pick: func(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) {
|
||||
if len(req.Candidates) != 1 {
|
||||
t.Fatalf("len(req.Candidates) = %d, want %d", len(req.Candidates), 1)
|
||||
}
|
||||
candidate := req.Candidates[0]
|
||||
if candidate.ID != "auth-a" || candidate.Provider != "gemini" || candidate.Priority != 7 || candidate.Status != string(StatusActive) {
|
||||
t.Fatalf("scheduler candidate = %#v, want sanitized auth-a metadata", candidate)
|
||||
}
|
||||
for _, key := range []string{"access_token", "api_key", "cookie"} {
|
||||
if _, ok := candidate.Attributes[key]; ok {
|
||||
t.Fatalf("scheduler candidate Attributes contains sensitive key %q", key)
|
||||
}
|
||||
}
|
||||
if candidate.Attributes["priority"] != "7" {
|
||||
t.Fatalf("scheduler candidate priority attribute = %q, want 7", candidate.Attributes["priority"])
|
||||
}
|
||||
if len(candidate.Metadata) != 0 {
|
||||
t.Fatalf("scheduler candidate Metadata = %#v, want empty", candidate.Metadata)
|
||||
}
|
||||
candidate.Attributes["team"] = "mutated"
|
||||
req.Candidates[0] = candidate
|
||||
return pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-a"}, true, nil
|
||||
},
|
||||
}
|
||||
manager.SetPluginScheduler(scheduler)
|
||||
|
||||
if _, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil); errPick != nil {
|
||||
t.Fatalf("pickNext() error = %v", errPick)
|
||||
}
|
||||
|
||||
manager.mu.RLock()
|
||||
gotAttr := manager.auths["auth-a"].Attributes["team"]
|
||||
gotAPIKey := manager.auths["auth-a"].Attributes["api_key"]
|
||||
manager.mu.RUnlock()
|
||||
if gotAttr != "alpha" {
|
||||
t.Fatalf("manager auth attribute team = %q, want alpha", gotAttr)
|
||||
}
|
||||
if gotAPIKey != "api-key-value" {
|
||||
t.Fatalf("manager auth attribute api_key = %q, want api-key-value", gotAPIKey)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerCustomSelector_FallsBackToLegacyPath(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -266,6 +266,9 @@ func (b *Builder) Build() (*Service, error) {
|
||||
coreManager.SetRoundTripperProvider(newDefaultRoundTripperProvider())
|
||||
coreManager.SetConfig(b.cfg)
|
||||
coreManager.SetOAuthModelAlias(b.cfg.OAuthModelAlias)
|
||||
if pluginHost != nil {
|
||||
coreManager.SetPluginScheduler(pluginHost)
|
||||
}
|
||||
|
||||
service := &Service{
|
||||
cfg: b.cfg,
|
||||
|
||||
@@ -164,6 +164,9 @@ func (s *Service) syncPluginRuntimeConfig(ctx context.Context) bool {
|
||||
if s.pluginHost != nil {
|
||||
s.pluginHost.ApplyConfig(ctx, cfg)
|
||||
}
|
||||
if s.coreManager != nil {
|
||||
s.coreManager.SetPluginScheduler(s.pluginHost)
|
||||
}
|
||||
s.registerPluginAuthParser()
|
||||
if s.pluginHost == nil {
|
||||
return false
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
package cliproxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"testing"
|
||||
"unsafe"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
func TestBuilderBuildInjectsPluginHostScheduler(t *testing.T) {
|
||||
host := pluginhost.New()
|
||||
service, errBuild := NewBuilder().
|
||||
WithConfig(&config.Config{AuthDir: t.TempDir()}).
|
||||
WithConfigPath(t.TempDir() + "/config.yaml").
|
||||
WithPluginHost(host).
|
||||
Build()
|
||||
if errBuild != nil {
|
||||
t.Fatalf("Build() error = %v", errBuild)
|
||||
}
|
||||
|
||||
got := pluginSchedulerFromManager(t, service.coreManager)
|
||||
if got != host {
|
||||
t.Fatalf("plugin scheduler = %p, want host %p", got, host)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServiceSyncPluginRuntimeConfigInjectsPluginHostScheduler(t *testing.T) {
|
||||
host := pluginhost.New()
|
||||
service := &Service{
|
||||
cfg: &config.Config{},
|
||||
coreManager: coreauth.NewManager(nil, nil, nil),
|
||||
pluginHost: host,
|
||||
}
|
||||
|
||||
if ok := service.syncPluginRuntimeConfig(context.Background()); !ok {
|
||||
t.Fatal("syncPluginRuntimeConfig() = false, want true")
|
||||
}
|
||||
|
||||
got := pluginSchedulerFromManager(t, service.coreManager)
|
||||
if got != host {
|
||||
t.Fatalf("plugin scheduler = %p, want host %p", got, host)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServiceSyncPluginRuntimeConfigClearsPluginSchedulerWithoutHost(t *testing.T) {
|
||||
host := pluginhost.New()
|
||||
service := &Service{
|
||||
cfg: &config.Config{},
|
||||
coreManager: coreauth.NewManager(nil, nil, nil),
|
||||
pluginHost: host,
|
||||
}
|
||||
service.coreManager.SetPluginScheduler(host)
|
||||
service.pluginHost = nil
|
||||
|
||||
if ok := service.syncPluginRuntimeConfig(context.Background()); ok {
|
||||
t.Fatal("syncPluginRuntimeConfig() = true, want false")
|
||||
}
|
||||
|
||||
got := pluginSchedulerFromManager(t, service.coreManager)
|
||||
if got != nil {
|
||||
t.Fatalf("plugin scheduler = %p, want nil", got)
|
||||
}
|
||||
}
|
||||
|
||||
func pluginSchedulerFromManager(t *testing.T, manager *coreauth.Manager) *pluginhost.Host {
|
||||
t.Helper()
|
||||
if manager == nil {
|
||||
t.Fatal("manager = nil")
|
||||
}
|
||||
value := reflect.ValueOf(manager).Elem().FieldByName("pluginScheduler")
|
||||
if !value.IsValid() {
|
||||
t.Fatal("pluginScheduler field not found")
|
||||
}
|
||||
scheduler := reflect.NewAt(value.Type(), unsafe.Pointer(value.UnsafeAddr())).Elem().Interface()
|
||||
if scheduler == nil {
|
||||
return nil
|
||||
}
|
||||
host, ok := scheduler.(*pluginhost.Host)
|
||||
if !ok {
|
||||
t.Fatalf("pluginScheduler type = %T, want *pluginhost.Host", scheduler)
|
||||
}
|
||||
return host
|
||||
}
|
||||
@@ -25,6 +25,9 @@ const (
|
||||
MethodFrontendAuthIdentifier = "frontend_auth.identifier"
|
||||
MethodFrontendAuthAuthenticate = "frontend_auth.authenticate"
|
||||
|
||||
// MethodSchedulerPick asks a scheduler plugin to select an auth candidate.
|
||||
MethodSchedulerPick = "scheduler.pick"
|
||||
|
||||
MethodExecutorIdentifier = "executor.identifier"
|
||||
MethodExecutorExecute = "executor.execute"
|
||||
MethodExecutorExecuteStream = "executor.execute_stream"
|
||||
|
||||
@@ -49,3 +49,9 @@ func TestMethodNamesAreStable(t *testing.T) {
|
||||
t.Fatalf("MethodExecutorExecuteStream = %q", MethodExecutorExecuteStream)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerPickMethodName(t *testing.T) {
|
||||
if MethodSchedulerPick != "scheduler.pick" {
|
||||
t.Fatalf("MethodSchedulerPick = %q", MethodSchedulerPick)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -76,6 +76,8 @@ type Capabilities struct {
|
||||
FrontendAuthProvider FrontendAuthProvider
|
||||
// FrontendAuthProviderExclusive makes this frontend auth provider the only active request auth provider when selected.
|
||||
FrontendAuthProviderExclusive bool
|
||||
// Scheduler chooses an auth candidate before the built-in scheduler runs.
|
||||
Scheduler Scheduler
|
||||
// Executor sends requests to an upstream provider or local backend.
|
||||
Executor ProviderExecutor
|
||||
// ExecutorModelScope declares whether Executor serves static models, OAuth auth models, or both.
|
||||
@@ -441,6 +443,70 @@ type FrontendAuthResponse struct {
|
||||
Metadata map[string]string
|
||||
}
|
||||
|
||||
const (
|
||||
// SchedulerBuiltinRoundRobin delegates auth selection to the built-in round-robin scheduler.
|
||||
SchedulerBuiltinRoundRobin = "round-robin"
|
||||
// SchedulerBuiltinFillFirst delegates auth selection to the built-in fill-first scheduler.
|
||||
SchedulerBuiltinFillFirst = "fill-first"
|
||||
)
|
||||
|
||||
// Scheduler chooses an auth candidate before the built-in scheduler runs.
|
||||
type Scheduler interface {
|
||||
Pick(context.Context, SchedulerPickRequest) (SchedulerPickResponse, error)
|
||||
}
|
||||
|
||||
// SchedulerPickRequest describes the routing context offered to a scheduler plugin.
|
||||
type SchedulerPickRequest struct {
|
||||
// Plugin is the metadata of the plugin being executed.
|
||||
Plugin Metadata
|
||||
// Provider is the primary provider key requested by the route.
|
||||
Provider string
|
||||
// Providers contains every provider key accepted by the route.
|
||||
Providers []string
|
||||
// Model is the requested model identifier.
|
||||
Model string
|
||||
// Stream reports whether the request expects streaming output.
|
||||
Stream bool
|
||||
// Options contains request-scoped scheduler inputs.
|
||||
Options SchedulerOptions
|
||||
// Candidates contains auth records available for selection.
|
||||
Candidates []SchedulerAuthCandidate
|
||||
}
|
||||
|
||||
// SchedulerOptions carries request-scoped scheduler inputs.
|
||||
type SchedulerOptions struct {
|
||||
// Headers contains request headers relevant to scheduling.
|
||||
Headers map[string][]string
|
||||
// Metadata carries host-provided scheduler context.
|
||||
Metadata map[string]any
|
||||
}
|
||||
|
||||
// SchedulerAuthCandidate describes one auth candidate available to a scheduler.
|
||||
type SchedulerAuthCandidate struct {
|
||||
// ID identifies the auth record.
|
||||
ID string
|
||||
// Provider identifies the auth provider.
|
||||
Provider string
|
||||
// Priority is the host priority assigned to the auth record.
|
||||
Priority int
|
||||
// Status is the current host-visible auth status.
|
||||
Status string
|
||||
// Attributes contains immutable routing and provider attributes.
|
||||
Attributes map[string]string
|
||||
// Metadata contains mutable host-managed auth metadata.
|
||||
Metadata map[string]any
|
||||
}
|
||||
|
||||
// SchedulerPickResponse returns a scheduler plugin routing decision.
|
||||
type SchedulerPickResponse struct {
|
||||
// AuthID identifies the selected auth record.
|
||||
AuthID string
|
||||
// DelegateBuiltin asks the host to use a named built-in scheduler.
|
||||
DelegateBuiltin string
|
||||
// Handled reports whether the plugin made a scheduling decision.
|
||||
Handled bool
|
||||
}
|
||||
|
||||
// ProviderExecutor handles model execution, streaming, HTTP bridging, and token counting.
|
||||
type ProviderExecutor interface {
|
||||
Identifier() string
|
||||
|
||||
@@ -13,6 +13,7 @@ var _ ModelRegistrar = (*compileTimePlugin)(nil)
|
||||
var _ ModelProvider = (*compileTimePlugin)(nil)
|
||||
var _ AuthProvider = (*compileTimePlugin)(nil)
|
||||
var _ FrontendAuthProvider = (*compileTimePlugin)(nil)
|
||||
var _ Scheduler = (*compileTimePlugin)(nil)
|
||||
var _ ProviderExecutor = (*compileTimePlugin)(nil)
|
||||
var _ HostHTTPClient = (*compileTimePlugin)(nil)
|
||||
var _ RequestTranslator = (*compileTimePlugin)(nil)
|
||||
@@ -113,6 +114,71 @@ func TestHostInjectedHTTPClientIsNotEncodedInPluginJSON(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerTypesExposeRoutingFields(t *testing.T) {
|
||||
request := SchedulerPickRequest{
|
||||
Plugin: Metadata{Name: "scheduler-plugin"},
|
||||
Provider: "openai",
|
||||
Providers: []string{"openai", "gemini"},
|
||||
Model: "gpt-test",
|
||||
Stream: true,
|
||||
Options: SchedulerOptions{
|
||||
Headers: map[string][]string{"X-Test": []string{"1"}},
|
||||
Metadata: map[string]any{"tenant": "demo"},
|
||||
},
|
||||
Candidates: []SchedulerAuthCandidate{{
|
||||
ID: "auth-1",
|
||||
Provider: "openai",
|
||||
Priority: 10,
|
||||
Status: "ready",
|
||||
Attributes: map[string]string{"region": "us"},
|
||||
Metadata: map[string]any{"load": float64(0.5)},
|
||||
}},
|
||||
}
|
||||
response := SchedulerPickResponse{
|
||||
AuthID: request.Candidates[0].ID,
|
||||
DelegateBuiltin: SchedulerBuiltinRoundRobin,
|
||||
Handled: true,
|
||||
}
|
||||
|
||||
if request.Plugin.Name != "scheduler-plugin" {
|
||||
t.Fatalf("Plugin.Name = %q", request.Plugin.Name)
|
||||
}
|
||||
if request.Provider != "openai" {
|
||||
t.Fatalf("Provider = %q", request.Provider)
|
||||
}
|
||||
if len(request.Providers) != 2 || request.Providers[1] != "gemini" {
|
||||
t.Fatalf("Providers = %#v", request.Providers)
|
||||
}
|
||||
if request.Model != "gpt-test" {
|
||||
t.Fatalf("Model = %q", request.Model)
|
||||
}
|
||||
if !request.Stream {
|
||||
t.Fatalf("Stream = %v", request.Stream)
|
||||
}
|
||||
if got := request.Options.Headers["X-Test"]; len(got) != 1 || got[0] != "1" {
|
||||
t.Fatalf("Options.Headers = %#v", request.Options.Headers)
|
||||
}
|
||||
if request.Options.Metadata["tenant"] != "demo" {
|
||||
t.Fatalf("Options.Metadata = %#v", request.Options.Metadata)
|
||||
}
|
||||
if len(request.Candidates) != 1 {
|
||||
t.Fatalf("Candidates = %#v", request.Candidates)
|
||||
}
|
||||
candidate := request.Candidates[0]
|
||||
if candidate.ID != "auth-1" || candidate.Provider != "openai" || candidate.Priority != 10 || candidate.Status != "ready" {
|
||||
t.Fatalf("Candidate = %#v", candidate)
|
||||
}
|
||||
if candidate.Attributes["region"] != "us" {
|
||||
t.Fatalf("Candidate.Attributes = %#v", candidate.Attributes)
|
||||
}
|
||||
if candidate.Metadata["load"] != float64(0.5) {
|
||||
t.Fatalf("Candidate.Metadata = %#v", candidate.Metadata)
|
||||
}
|
||||
if response.AuthID != "auth-1" || response.DelegateBuiltin != SchedulerBuiltinRoundRobin || !response.Handled {
|
||||
t.Fatalf("SchedulerPickResponse = %#v", response)
|
||||
}
|
||||
}
|
||||
|
||||
func (compileTimePlugin) RegisterModels(context.Context, ModelRegistrationRequest) (ModelRegistrationResponse, error) {
|
||||
return ModelRegistrationResponse{}, nil
|
||||
}
|
||||
@@ -147,6 +213,10 @@ func (compileTimePlugin) Authenticate(context.Context, FrontendAuthRequest) (Fro
|
||||
return FrontendAuthResponse{}, nil
|
||||
}
|
||||
|
||||
func (compileTimePlugin) Pick(context.Context, SchedulerPickRequest) (SchedulerPickResponse, error) {
|
||||
return SchedulerPickResponse{}, nil
|
||||
}
|
||||
|
||||
func (compileTimePlugin) Execute(context.Context, ExecutorRequest) (ExecutorResponse, error) {
|
||||
return ExecutorResponse{}, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user