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:
Luis Pater
2026-06-09 13:57:36 +08:00
parent a2db6a6108
commit 693ce1c55a
22 changed files with 1750 additions and 20 deletions
+18
View File
@@ -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
+18
View File
@@ -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
+50
View File
@@ -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
```
+10
View File
@@ -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 => ../../../..
+4
View File
@@ -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=
+270
View File
@@ -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))
}
+1
View File
@@ -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 ||
+13
View File
@@ -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 {
+2
View File
@@ -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...),
+193
View File
@@ -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)
}
}
+107
View File
@@ -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 &copyRecord
}
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
}
}
+217
View File
@@ -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
}
+22
View File
@@ -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
View File
@@ -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)
}
+343
View File
@@ -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()
+3
View File
@@ -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,
+3
View File
@@ -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
}
+3
View File
@@ -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"
+6
View File
@@ -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)
}
}
+66
View File
@@ -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
+70
View File
@@ -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
}