From 693ce1c55aa490e1ca90601354cce6c98211147d Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Tue, 9 Jun 2026 13:39:19 +0800 Subject: [PATCH] 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. --- examples/plugin/README.md | 18 + examples/plugin/README_CN.md | 18 + examples/plugin/scheduler/README.md | 50 +++ examples/plugin/scheduler/go/go.mod | 10 + examples/plugin/scheduler/go/go.sum | 4 + examples/plugin/scheduler/go/main.go | 270 ++++++++++++++ internal/pluginhost/host.go | 1 + internal/pluginhost/rpc_client.go | 13 + internal/pluginhost/rpc_schema.go | 2 + internal/pluginhost/rpc_schema_test.go | 193 ++++++++++ internal/pluginhost/scheduler.go | 107 ++++++ internal/pluginhost/scheduler_test.go | 217 +++++++++++ internal/pluginhost/test_helpers_test.go | 22 ++ sdk/cliproxy/auth/conductor.go | 264 +++++++++++++- sdk/cliproxy/auth/scheduler_test.go | 343 ++++++++++++++++++ sdk/cliproxy/builder.go | 3 + sdk/cliproxy/service.go | 3 + sdk/cliproxy/service_plugin_scheduler_test.go | 87 +++++ sdk/pluginabi/types.go | 3 + sdk/pluginabi/types_test.go | 6 + sdk/pluginapi/types.go | 66 ++++ sdk/pluginapi/types_test.go | 70 ++++ 22 files changed, 1750 insertions(+), 20 deletions(-) create mode 100644 examples/plugin/scheduler/README.md create mode 100644 examples/plugin/scheduler/go/go.mod create mode 100644 examples/plugin/scheduler/go/go.sum create mode 100644 examples/plugin/scheduler/go/main.go create mode 100644 internal/pluginhost/scheduler.go create mode 100644 internal/pluginhost/scheduler_test.go create mode 100644 sdk/cliproxy/service_plugin_scheduler_test.go diff --git a/examples/plugin/README.md b/examples/plugin/README.md index 668ada8d7..9ee78a7a7 100644 --- a/examples/plugin/README.md +++ b/examples/plugin/README.md @@ -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 diff --git a/examples/plugin/README_CN.md b/examples/plugin/README_CN.md index a489ee022..f430aec60 100644 --- a/examples/plugin/README_CN.md +++ b/examples/plugin/README_CN.md @@ -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 diff --git a/examples/plugin/scheduler/README.md b/examples/plugin/scheduler/README.md new file mode 100644 index 000000000..2890a5034 --- /dev/null +++ b/examples/plugin/scheduler/README.md @@ -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 +``` diff --git a/examples/plugin/scheduler/go/go.mod b/examples/plugin/scheduler/go/go.mod new file mode 100644 index 000000000..99ead9836 --- /dev/null +++ b/examples/plugin/scheduler/go/go.mod @@ -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 => ../../../.. diff --git a/examples/plugin/scheduler/go/go.sum b/examples/plugin/scheduler/go/go.sum new file mode 100644 index 000000000..a62c313c5 --- /dev/null +++ b/examples/plugin/scheduler/go/go.sum @@ -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= diff --git a/examples/plugin/scheduler/go/main.go b/examples/plugin/scheduler/go/main.go new file mode 100644 index 000000000..d9190c34e --- /dev/null +++ b/examples/plugin/scheduler/go/main.go @@ -0,0 +1,270 @@ +package main + +/* +#include +#include + +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)) +} diff --git a/internal/pluginhost/host.go b/internal/pluginhost/host.go index b12bfd839..af0e8e501 100644 --- a/internal/pluginhost/host.go +++ b/internal/pluginhost/host.go @@ -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 || diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 8addde684..0d3817c28 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -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 { diff --git a/internal/pluginhost/rpc_schema.go b/internal/pluginhost/rpc_schema.go index b579354f4..61f474d44 100644 --- a/internal/pluginhost/rpc_schema.go +++ b/internal/pluginhost/rpc_schema.go @@ -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...), diff --git a/internal/pluginhost/rpc_schema_test.go b/internal/pluginhost/rpc_schema_test.go index b48e9e7c6..c0bf3dc3d 100644 --- a/internal/pluginhost/rpc_schema_test.go +++ b/internal/pluginhost/rpc_schema_test.go @@ -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) + } +} diff --git a/internal/pluginhost/scheduler.go b/internal/pluginhost/scheduler.go new file mode 100644 index 000000000..fa7489563 --- /dev/null +++ b/internal/pluginhost/scheduler.go @@ -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 + } +} diff --git a/internal/pluginhost/scheduler_test.go b/internal/pluginhost/scheduler_test.go new file mode 100644 index 000000000..374b884f3 --- /dev/null +++ b/internal/pluginhost/scheduler_test.go @@ -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 +} diff --git a/internal/pluginhost/test_helpers_test.go b/internal/pluginhost/test_helpers_test.go index 40a25500c..e7b0fbd7b 100644 --- a/internal/pluginhost/test_helpers_test.go +++ b/internal/pluginhost/test_helpers_test.go @@ -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) diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index d16c62745..db03a3516 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -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) } diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index 864fa938e..48a7673eb 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -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() diff --git a/sdk/cliproxy/builder.go b/sdk/cliproxy/builder.go index 32cad4be1..54a83c468 100644 --- a/sdk/cliproxy/builder.go +++ b/sdk/cliproxy/builder.go @@ -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, diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index 87fb18f8d..2873f0027 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -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 diff --git a/sdk/cliproxy/service_plugin_scheduler_test.go b/sdk/cliproxy/service_plugin_scheduler_test.go new file mode 100644 index 000000000..d80c75b13 --- /dev/null +++ b/sdk/cliproxy/service_plugin_scheduler_test.go @@ -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 +} diff --git a/sdk/pluginabi/types.go b/sdk/pluginabi/types.go index f6f7e9671..af80be14a 100644 --- a/sdk/pluginabi/types.go +++ b/sdk/pluginabi/types.go @@ -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" diff --git a/sdk/pluginabi/types_test.go b/sdk/pluginabi/types_test.go index 111e343ab..f95624483 100644 --- a/sdk/pluginabi/types_test.go +++ b/sdk/pluginabi/types_test.go @@ -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) + } +} diff --git a/sdk/pluginapi/types.go b/sdk/pluginapi/types.go index 308f04116..a0b749075 100644 --- a/sdk/pluginapi/types.go +++ b/sdk/pluginapi/types.go @@ -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 diff --git a/sdk/pluginapi/types_test.go b/sdk/pluginapi/types_test.go index 3f5694a3d..3f4cd1683 100644 --- a/sdk/pluginapi/types_test.go +++ b/sdk/pluginapi/types_test.go @@ -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 }