diff --git a/examples/plugin/README.md b/examples/plugin/README.md index 849305612..2e7b2de0c 100644 --- a/examples/plugin/README.md +++ b/examples/plugin/README.md @@ -4,8 +4,7 @@ This directory contains standard dynamic library plugin examples for the CLIProx ## Layout -- `simple/`- : Go-only plugin resource that calls host auth file callbacks (, , , ). -- : full provider-native skeleton that declares every supported capability. +- `simple/`: full provider-native skeleton that declares every supported capability. - `model/`: model capability only. - `auth/`: auth provider capability only. - `frontend-auth/`: frontend auth provider capability only. @@ -15,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.5` requests to the priority service tier when enabled. +- `request-lifecycle/`: Go-only request admission example with concurrency control, active HTTP termination, and terminal callbacks. - `scheduler/`: Go-only scheduler that can select a configured auth ID, delegate to a built-in scheduler, or deny picks. - `claude-web-search-router/`: ModelRouter + executor for Claude Code built-in `web_search` (antigravity / codex / xai / Tavily). See `claude-web-search-router/README.md`. - `response-translator/`: response translation capability only. @@ -42,7 +42,21 @@ plugins: fast: false ``` +## Request Lifecycle +`request-lifecycle` combines `request_interceptor` with `request_lifecycle_plugin`. It acquires a concurrency slot before auth selection, can return a custom `403` or `429` response without contacting an upstream model, and releases admitted slots from `request.complete` on success, failure, rejection, or cancellation. + +```yaml +plugins: + configs: + request-lifecycle: + enabled: true + priority: 100 + max_concurrency: 2 + reject_keyword: "blocked" +``` + +See `request-lifecycle/README.md` for build instructions and lifecycle semantics. ## Host Auth Files Callback diff --git a/examples/plugin/README_CN.md b/examples/plugin/README_CN.md index b1987e7c6..a9d2e316b 100644 --- a/examples/plugin/README_CN.md +++ b/examples/plugin/README_CN.md @@ -1,5 +1,4 @@ -- :仅 Go 实现的插件资源,演示 host 凭证文件回调(、、、)。 -- # 标准动态库插件示例 +# 标准动态库插件示例 本目录包含 CLIProxyAPI C ABI 的标准动态库插件示例。 @@ -15,6 +14,7 @@ - `request-translator/`:只演示请求转换能力。 - `request-normalizer/`:只演示请求规整能力。 - `codex-service-tier/`:仅 Go 实现的请求规整插件,启用后会将 Codex `gpt-5.5` 请求设置为 priority service tier。 +- `request-lifecycle/`:仅 Go 实现的请求生命周期插件,演示并发控制、主动终止 HTTP 请求和终态回调。 - `scheduler/`:仅 Go 实现的调度插件,可选择指定 auth ID、委托内置调度器或拒绝调度。 - `response-translator/`:只演示响应转换能力。 - `response-normalizer/`:只演示响应规整能力。 @@ -41,7 +41,21 @@ plugins: fast: false ``` +## 请求生命周期 +`request-lifecycle` 同时声明 `request_interceptor` 和 `request_lifecycle_plugin`。它会在认证选择前占用并发槽位,可以直接返回自定义 `403` 或 `429` 响应而不请求上游模型,并在成功、失败、拒绝或取消时通过 `request.complete` 释放已接入请求的槽位。 + +```yaml +plugins: + configs: + request-lifecycle: + enabled: true + priority: 100 + max_concurrency: 2 + reject_keyword: "blocked" +``` + +构建方式和生命周期语义详见 `request-lifecycle/README.md`。 ## Host Auth Files 回调 diff --git a/examples/plugin/request-lifecycle/README.md b/examples/plugin/request-lifecycle/README.md new file mode 100644 index 000000000..2f8a1aa4d --- /dev/null +++ b/examples/plugin/request-lifecycle/README.md @@ -0,0 +1,61 @@ +# Request Lifecycle Plugin + +This Go dynamic-library plugin demonstrates request admission, active termination, and exactly-once terminal lifecycle handling. It requires a host that supports plugin RPC schema version 2 or newer. + +It declares two optional capabilities: + +- `request_interceptor`: acquires a concurrency slot in `request.intercept_before` and can terminate the request before any upstream executor runs. +- `request_lifecycle_plugin`: releases the slot in `request.complete` for successful, failed, rejected, and canceled requests. + +The host passes the same `RequestID` to request interception, response interception, stream interception, and the terminal `RequestCompletion` event. + +## Behavior + +- Allows at most `max_concurrency` requests in flight. +- Returns a custom `429` JSON response with `Retry-After: 1` when the limit is reached. +- Returns a custom `403` JSON response when the raw request body contains `reject_keyword`. +- Does not send terminated requests to an upstream model. +- Releases only request IDs that were previously admitted, so rejected requests and duplicate terminal events do not underflow the counter. + +## Configuration + +```yaml +plugins: + enabled: true + configs: + request-lifecycle: + enabled: true + priority: 100 + max_concurrency: 2 + reject_keyword: "blocked" +``` + +Set `reject_keyword` to an empty string to disable keyword rejection. + +## Build + +From the repository root on macOS: + +```bash +mkdir -p plugins/darwin/$(go env GOARCH) +go build -buildmode=c-shared \ + -o plugins/darwin/$(go env GOARCH)/request-lifecycle.dylib \ + ./examples/plugin/request-lifecycle/go +rm -f plugins/darwin/$(go env GOARCH)/request-lifecycle.h +``` + +Use `.so` on Linux or FreeBSD and `.dll` on Windows. + +The output filename is the plugin ID, so the example artifact must be named `request-lifecycle` for the configuration above. + +## Relevant RPC Methods + +```text +plugin.register +plugin.reconfigure +request.intercept_before +request.intercept_after +request.complete +``` + +`request.complete` is an observational callback. The host schedules it asynchronously so a blocked plugin cannot delay response delivery, logs callback errors, and uses a context detached from downstream cancellation so a canceled request can still release its slot. diff --git a/examples/plugin/request-lifecycle/go/go.mod b/examples/plugin/request-lifecycle/go/go.mod new file mode 100644 index 000000000..420d628e8 --- /dev/null +++ b/examples/plugin/request-lifecycle/go/go.mod @@ -0,0 +1,10 @@ +module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/request-lifecycle/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/request-lifecycle/go/go.sum b/examples/plugin/request-lifecycle/go/go.sum new file mode 100644 index 000000000..a62c313c5 --- /dev/null +++ b/examples/plugin/request-lifecycle/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/request-lifecycle/go/main.go b/examples/plugin/request-lifecycle/go/main.go new file mode 100644 index 000000000..318accd07 --- /dev/null +++ b/examples/plugin/request-lifecycle/go/main.go @@ -0,0 +1,308 @@ +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" + "fmt" + "net/http" + "strings" + "sync" + "unsafe" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + "gopkg.in/yaml.v3" +) + +var state = pluginState{ + config: pluginConfig{MaxConcurrency: 2, RejectKeyword: "blocked"}, + active: make(map[string]struct{}), +} + +type pluginState struct { + mu sync.Mutex + config pluginConfig + active map[string]struct{} +} + +type pluginConfig struct { + MaxConcurrency int `yaml:"max_concurrency"` + RejectKeyword string `yaml:"reject_keyword"` +} + +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"` + SchemaVersion uint32 `json:"schema_version"` +} + +type registration struct { + SchemaVersion uint32 `json:"schema_version"` + Metadata pluginapi.Metadata `json:"metadata"` + Capabilities registrationCapability `json:"capabilities"` +} + +type registrationCapability struct { + RequestInterceptor bool `json:"request_interceptor"` + RequestLifecyclePlugin bool `json:"request_lifecycle_plugin"` +} + +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) + } + _ = len +} + +//export cliproxyPluginShutdown +func cliproxyPluginShutdown() { + state.mu.Lock() + defer state.mu.Unlock() + state.active = make(map[string]struct{}) +} + +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.MethodRequestInterceptBefore: + return interceptBeforeAuth(request) + case pluginabi.MethodRequestInterceptAfter: + return passThroughRequest(request) + case pluginabi.MethodRequestComplete: + return completeRequest(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 + } + } + if req.SchemaVersion < 2 { + return fmt.Errorf("request lifecycle plugin requires host schema version 2 or newer") + } + cfg := pluginConfig{MaxConcurrency: 2, RejectKeyword: "blocked"} + if len(req.ConfigYAML) > 0 { + if errUnmarshal := yaml.Unmarshal(req.ConfigYAML, &cfg); errUnmarshal != nil { + return errUnmarshal + } + } + if cfg.MaxConcurrency < 1 { + return fmt.Errorf("max_concurrency must be greater than zero") + } + cfg.RejectKeyword = strings.TrimSpace(cfg.RejectKeyword) + state.mu.Lock() + defer state.mu.Unlock() + state.config = cfg + return nil +} + +func pluginRegistration() registration { + return registration{ + SchemaVersion: pluginabi.SchemaVersion, + Metadata: pluginapi.Metadata{ + Name: "request-lifecycle", + 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: "max_concurrency", + Type: pluginapi.ConfigFieldTypeInteger, + Description: "Maximum number of intercepted requests allowed in flight.", + }, + { + Name: "reject_keyword", + Type: pluginapi.ConfigFieldTypeString, + Description: "Terminates requests whose raw JSON body contains this keyword.", + }, + }, + }, + Capabilities: registrationCapability{ + RequestInterceptor: true, + RequestLifecyclePlugin: true, + }, + } +} + +func interceptBeforeAuth(raw []byte) ([]byte, error) { + var req pluginapi.RequestInterceptRequest + if errUnmarshal := json.Unmarshal(raw, &req); errUnmarshal != nil { + return nil, errUnmarshal + } + if req.RequestID == "" { + return nil, fmt.Errorf("request ID is required") + } + + state.mu.Lock() + defer state.mu.Unlock() + if _, exists := state.active[req.RequestID]; exists { + return okEnvelope(pluginapi.RequestInterceptResponse{Headers: req.Headers, Body: req.Body}) + } + if state.config.RejectKeyword != "" && strings.Contains(string(req.Body), state.config.RejectKeyword) { + return terminatedResponse(http.StatusForbidden, "request blocked by plugin policy", nil) + } + if len(state.active) >= state.config.MaxConcurrency { + return terminatedResponse(http.StatusTooManyRequests, "plugin concurrency limit reached", http.Header{"Retry-After": {"1"}}) + } + state.active[req.RequestID] = struct{}{} + return okEnvelope(pluginapi.RequestInterceptResponse{Headers: req.Headers, Body: req.Body}) +} + +func passThroughRequest(raw []byte) ([]byte, error) { + var req pluginapi.RequestInterceptRequest + if errUnmarshal := json.Unmarshal(raw, &req); errUnmarshal != nil { + return nil, errUnmarshal + } + return okEnvelope(pluginapi.RequestInterceptResponse{Headers: req.Headers, Body: req.Body}) +} + +func terminatedResponse(statusCode int, message string, headers http.Header) ([]byte, error) { + body, errMarshal := json.Marshal(map[string]any{ + "error": map[string]any{ + "type": "plugin_request_rejected", + "message": message, + }, + }) + if errMarshal != nil { + return nil, errMarshal + } + if headers == nil { + headers = make(http.Header) + } + headers.Set("Content-Type", "application/json") + return okEnvelope(pluginapi.RequestInterceptResponse{ + Terminate: true, + StatusCode: statusCode, + ResponseHeaders: headers, + ResponseBody: body, + }) +} + +func completeRequest(raw []byte) ([]byte, error) { + var completion pluginapi.RequestCompletion + if errUnmarshal := json.Unmarshal(raw, &completion); errUnmarshal != nil { + return nil, errUnmarshal + } + state.mu.Lock() + defer state.mu.Unlock() + delete(state.active, completion.RequestID) + return okEnvelope(struct{}{}) +} + +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, errMarshal := json.Marshal(envelope{OK: false, Error: &envelopeError{Code: code, Message: message}}) + if errMarshal != nil { + return []byte(`{"ok":false,"error":{"code":"plugin_error","message":"encode error"}}`) + } + 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/examples/plugin/request-lifecycle/go/main_test.go b/examples/plugin/request-lifecycle/go/main_test.go new file mode 100644 index 000000000..948e4d1e3 --- /dev/null +++ b/examples/plugin/request-lifecycle/go/main_test.go @@ -0,0 +1,89 @@ +package main + +import ( + "encoding/json" + "net/http" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func TestConfigureRejectsLegacyHostSchema(t *testing.T) { + raw, errMarshal := json.Marshal(lifecycleRequest{SchemaVersion: 1}) + if errMarshal != nil { + t.Fatalf("marshal lifecycle request: %v", errMarshal) + } + if errConfigure := configure(raw); errConfigure == nil { + t.Fatal("configure() error = nil for schema version 1") + } +} + +func TestConcurrencySlotReleasedByCompletion(t *testing.T) { + resetState(pluginConfig{MaxConcurrency: 1}) + first := interceptForTest(t, pluginapi.RequestInterceptRequest{RequestID: "first", Body: []byte(`{"model":"test"}`)}) + if first.Terminate { + t.Fatalf("first request was terminated: %#v", first) + } + second := interceptForTest(t, pluginapi.RequestInterceptRequest{RequestID: "second", Body: []byte(`{"model":"test"}`)}) + if !second.Terminate || second.StatusCode != http.StatusTooManyRequests { + t.Fatalf("second response = %#v", second) + } + + completionRaw, errMarshal := json.Marshal(pluginapi.RequestCompletion{RequestID: "first", Outcome: pluginapi.RequestCompletionSucceeded}) + if errMarshal != nil { + t.Fatalf("marshal completion: %v", errMarshal) + } + completeRaw, errComplete := completeRequest(completionRaw) + if errComplete != nil { + t.Fatalf("completeRequest() error = %v", errComplete) + } + if len(completeRaw) == 0 { + t.Fatal("completeRequest() response is empty") + } + third := interceptForTest(t, pluginapi.RequestInterceptRequest{RequestID: "third", Body: []byte(`{"model":"test"}`)}) + if third.Terminate { + t.Fatalf("third request was terminated after release: %#v", third) + } +} + +func TestPolicyTerminationReturnsCustomResponse(t *testing.T) { + resetState(pluginConfig{MaxConcurrency: 1, RejectKeyword: "blocked"}) + response := interceptForTest(t, pluginapi.RequestInterceptRequest{RequestID: "blocked", Body: []byte(`{"prompt":"blocked"}`)}) + if !response.Terminate || response.StatusCode != http.StatusForbidden { + t.Fatalf("response = %#v", response) + } + if response.ResponseHeaders.Get("Content-Type") != "application/json" { + t.Fatalf("response headers = %#v", response.ResponseHeaders) + } + if len(response.ResponseBody) == 0 { + t.Fatal("response body is empty") + } +} + +func resetState(cfg pluginConfig) { + state.mu.Lock() + defer state.mu.Unlock() + state.config = cfg + state.active = make(map[string]struct{}) +} + +func interceptForTest(t *testing.T, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + t.Helper() + raw, errMarshal := json.Marshal(req) + if errMarshal != nil { + t.Fatalf("marshal request: %v", errMarshal) + } + rawEnvelope, errIntercept := interceptBeforeAuth(raw) + if errIntercept != nil { + t.Fatalf("interceptBeforeAuth() error = %v", errIntercept) + } + var env envelope + if errUnmarshal := json.Unmarshal(rawEnvelope, &env); errUnmarshal != nil { + t.Fatalf("unmarshal envelope: %v", errUnmarshal) + } + var response pluginapi.RequestInterceptResponse + if errUnmarshal := json.Unmarshal(env.Result, &response); errUnmarshal != nil { + t.Fatalf("unmarshal response: %v", errUnmarshal) + } + return response +} diff --git a/internal/interfaces/error_message.go b/internal/interfaces/error_message.go index eecdc9cbe..93fa3acbe 100644 --- a/internal/interfaces/error_message.go +++ b/internal/interfaces/error_message.go @@ -15,6 +15,15 @@ type ErrorMessage struct { // Error is the underlying error that occurred. Error error - // Addon contains additional headers to be added to the response. + // Addon contains upstream headers that may be passed through when enabled. Addon http.Header + + // DirectResponse reports that Body and Headers were explicitly supplied by a trusted in-process component. + DirectResponse bool + + // Body contains a preformatted downstream response when DirectResponse is true. + Body []byte + + // Headers contains downstream response headers when DirectResponse is true. + Headers http.Header } diff --git a/internal/pluginhost/adapters_interceptors.go b/internal/pluginhost/adapters_interceptors.go index 228784044..91d0227f5 100644 --- a/internal/pluginhost/adapters_interceptors.go +++ b/internal/pluginhost/adapters_interceptors.go @@ -113,11 +113,54 @@ func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterc if len(resp.Body) > 0 { current.Body = bytes.Clone(resp.Body) } + if resp.Terminate { + current.Terminate = true + current.StatusCode = resp.StatusCode + current.ResponseHeaders = cloneHeader(resp.ResponseHeaders) + current.ResponseBody = bytes.Clone(resp.ResponseBody) + break + } } } return current } +// CompleteRequest schedules terminal notifications without blocking response delivery. +func (h *Host) CompleteRequest(ctx context.Context, completion pluginapi.RequestCompletion) { + h.CompleteRequestExcept(ctx, completion, "") +} + +// CompleteRequestExcept notifies lifecycle plugins except the plugin that initiated a nested host execution. +func (h *Host) CompleteRequestExcept(ctx context.Context, completion pluginapi.RequestCompletion, skipPluginID string) { + if h == nil { + return + } + if ctx == nil { + ctx = context.Background() + } else { + ctx = context.WithoutCancel(ctx) + } + skipPluginID = strings.TrimSpace(skipPluginID) + for _, record := range h.activeRecords() { + plugin := record.plugin.Capabilities.RequestLifecyclePlugin + if h.isPluginFused(record.id) || plugin == nil || record.id == skipPluginID || !h.recordCurrent(record) { + continue + } + next := completion + next.Metadata = cloneInterceptorMetadata(completion.Metadata) + go func(record capabilityRecord, plugin pluginapi.RequestLifecyclePlugin, completion pluginapi.RequestCompletion) { + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, "RequestLifecyclePlugin.HandleRequestComplete", recovered) + } + }() + if errComplete := plugin.HandleRequestComplete(ctx, completion); errComplete != nil { + log.Warnf("pluginhost: request lifecycle plugin %s failed: %v", record.id, errComplete) + } + }(record, plugin, next) + } +} + func (h *Host) InterceptResponse(ctx context.Context, req pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse { return h.InterceptResponseExcept(ctx, req, "") } diff --git a/internal/pluginhost/host.go b/internal/pluginhost/host.go index c58e36c78..0fc56cf16 100644 --- a/internal/pluginhost/host.go +++ b/internal/pluginhost/host.go @@ -856,6 +856,7 @@ func validPlugin(plugin pluginapi.Plugin) bool { caps.RequestTranslator != nil || caps.RequestNormalizer != nil || caps.RequestInterceptor != nil || + caps.RequestLifecyclePlugin != nil || caps.ResponseTranslator != nil || caps.ResponseBeforeTranslator != nil || caps.ResponseAfterTranslator != nil || diff --git a/internal/pluginhost/request_lifecycle_test.go b/internal/pluginhost/request_lifecycle_test.go new file mode 100644 index 000000000..e7dea2c04 --- /dev/null +++ b/internal/pluginhost/request_lifecycle_test.go @@ -0,0 +1,164 @@ +package pluginhost + +import ( + "context" + "encoding/json" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func TestRequestInterceptorTerminationStopsChain(t *testing.T) { + lowCalls := 0 + host := newHostWithRecords( + capabilityRecord{ + id: "high", + priority: 20, + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + RequestInterceptor: requestInterceptorFunc(func(context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return pluginapi.RequestInterceptResponse{ + Terminate: true, + StatusCode: http.StatusForbidden, + ResponseHeaders: http.Header{"Content-Type": {"application/json"}}, + ResponseBody: []byte(`{"error":"blocked"}`), + }, nil + }), + }}, + }, + capabilityRecord{ + id: "low", + priority: 10, + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + RequestInterceptor: requestInterceptorFunc(func(context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + lowCalls++ + return pluginapi.RequestInterceptResponse{}, nil + }), + }}, + }, + ) + + response := host.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{RequestID: "request-1"}) + if !response.Terminate || response.StatusCode != http.StatusForbidden { + t.Fatalf("termination response = %#v", response) + } + if response.ResponseHeaders.Get("Content-Type") != "application/json" || string(response.ResponseBody) != `{"error":"blocked"}` { + t.Fatalf("termination payload = %#v", response) + } + if lowCalls != 0 { + t.Fatalf("lower-priority interceptor calls = %d, want 0", lowCalls) + } +} + +func TestCompleteRequestUsesUncancelledContextAndClonesMetadata(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + originalNested := map[string]any{"value": "original"} + var got pluginapi.RequestCompletion + var callbackContextError error + done := make(chan struct{}) + host := newHostWithRecords(capabilityRecord{ + id: "lifecycle", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + RequestLifecyclePlugin: requestLifecyclePluginFunc(func(callbackCtx context.Context, completion pluginapi.RequestCompletion) { + callbackContextError = callbackCtx.Err() + got = completion + completion.Metadata["nested"].(map[string]any)["value"] = "mutated" + close(done) + }), + }}, + }) + + host.CompleteRequest(ctx, pluginapi.RequestCompletion{ + RequestID: "request-1", + Outcome: pluginapi.RequestCompletionCanceled, + StartedAt: time.Now().Add(-time.Second), + CompletedAt: time.Now(), + Metadata: map[string]any{"nested": originalNested}, + }) + <-done + + if callbackContextError != nil { + t.Fatalf("callback context error = %v", callbackContextError) + } + if got.RequestID != "request-1" || got.Outcome != pluginapi.RequestCompletionCanceled { + t.Fatalf("completion = %#v", got) + } + if originalNested["value"] != "original" { + t.Fatalf("input metadata was mutated: %#v", originalNested) + } +} + +func TestCompleteRequestDoesNotWaitForBlockingPlugin(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + host := newHostWithRecords(capabilityRecord{ + id: "blocking-lifecycle", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + RequestLifecyclePlugin: requestLifecyclePluginFunc(func(context.Context, pluginapi.RequestCompletion) { + close(started) + <-release + }), + }}, + }) + + returned := make(chan struct{}) + go func() { + host.CompleteRequest(context.Background(), pluginapi.RequestCompletion{RequestID: "request-blocking"}) + close(returned) + }() + select { + case <-returned: + case <-time.After(time.Second): + t.Fatal("CompleteRequest blocked on lifecycle plugin") + } + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("lifecycle plugin was not invoked") + } + close(release) +} + +func TestRPCCapabilitiesAndAdapterIncludeRequestLifecycle(t *testing.T) { + var got pluginapi.RequestCompletion + plugin := validTestPlugin("request-lifecycle") + plugin.Capabilities.RequestLifecyclePlugin = requestLifecyclePluginFunc(func(_ context.Context, completion pluginapi.RequestCompletion) { + got = completion + }) + caps := rpcCapabilitiesFromPlugin(plugin) + if !caps.RequestLifecyclePlugin { + t.Fatal("RequestLifecyclePlugin = false, want true") + } + rawCaps, errMarshal := json.Marshal(caps) + if errMarshal != nil { + t.Fatalf("Marshal() error = %v", errMarshal) + } + var decoded map[string]any + if errUnmarshal := json.Unmarshal(rawCaps, &decoded); errUnmarshal != nil { + t.Fatalf("Unmarshal() error = %v", errUnmarshal) + } + if decoded["request_lifecycle_plugin"] != true { + t.Fatalf("request_lifecycle_plugin = %#v", decoded["request_lifecycle_plugin"]) + } + + lookup := newTestSymbolLookup(&testPlugin{registerResult: plugin}) + registered, errRegister := registerRPCPlugin(context.Background(), nil, "request-lifecycle", lookup, pluginabi.MethodPluginRegister, nil) + if errRegister != nil { + t.Fatalf("registerRPCPlugin() error = %v", errRegister) + } + if registered.Capabilities.RequestLifecyclePlugin == nil { + t.Fatal("RequestLifecyclePlugin = nil, want RPC adapter") + } + if errComplete := registered.Capabilities.RequestLifecyclePlugin.HandleRequestComplete(context.Background(), pluginapi.RequestCompletion{ + RequestID: "request-rpc", + Outcome: pluginapi.RequestCompletionSucceeded, + }); errComplete != nil { + t.Fatalf("HandleRequestComplete() error = %v", errComplete) + } + if got.RequestID != "request-rpc" || got.Outcome != pluginapi.RequestCompletionSucceeded { + t.Fatalf("RPC completion = %#v", got) + } +} diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 10f767a5a..01319fd78 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -107,6 +107,9 @@ func registerRPCPlugin(ctx context.Context, host *Host, id string, client plugin if resp.Capabilities.RequestInterceptor { plugin.Capabilities.RequestInterceptor = adapter } + if resp.Capabilities.RequestLifecyclePlugin { + plugin.Capabilities.RequestLifecyclePlugin = adapter + } if resp.Capabilities.ResponseTranslator { plugin.Capabilities.ResponseTranslator = adapter } @@ -191,6 +194,9 @@ func sanitizePluginRequest(request any) any { case pluginapi.RequestInterceptRequest: req.Metadata = sanitizePluginMetadata(req.Metadata) return req + case pluginapi.RequestCompletion: + req.Metadata = sanitizePluginMetadata(req.Metadata) + return req case pluginapi.ResponseInterceptRequest: req.Metadata = sanitizePluginMetadata(req.Metadata) return req @@ -203,6 +209,9 @@ func sanitizePluginRequest(request any) any { case rpcModelRouteRequest: req.Metadata = sanitizePluginMetadata(req.Metadata) return req + case rpcRequestCompletion: + req.Metadata = sanitizePluginMetadata(req.Metadata) + return req case rpcResponseInterceptRequest: req.Metadata = sanitizePluginMetadata(req.Metadata) return req @@ -476,6 +485,16 @@ func (a *rpcPluginAdapter) InterceptRequestAfterAuth(ctx context.Context, req pl }) } +func (a *rpcPluginAdapter) HandleRequestComplete(ctx context.Context, completion pluginapi.RequestCompletion) error { + callbackID, closeCallback := a.openHostCallbackContext(ctx) + defer closeCallback() + _, errCall := callPlugin[rpcEmptyResponse](ctx, a.client, pluginabi.MethodRequestComplete, rpcRequestCompletion{ + RequestCompletion: completion, + HostCallbackID: callbackID, + }) + return errCall +} + func (a *rpcPluginAdapter) TranslateResponse(ctx context.Context, req pluginapi.ResponseTransformRequest) (pluginapi.PayloadResponse, error) { return callPlugin[pluginapi.PayloadResponse](ctx, a.client, pluginabi.MethodResponseTranslate, req) } diff --git a/internal/pluginhost/rpc_schema.go b/internal/pluginhost/rpc_schema.go index b88711009..306d9166a 100644 --- a/internal/pluginhost/rpc_schema.go +++ b/internal/pluginhost/rpc_schema.go @@ -33,6 +33,7 @@ type rpcCapabilities struct { RequestTranslator bool `json:"request_translator"` RequestNormalizer bool `json:"request_normalizer"` RequestInterceptor bool `json:"request_interceptor"` + RequestLifecyclePlugin bool `json:"request_lifecycle_plugin"` ResponseTranslator bool `json:"response_translator"` ResponseBeforeTranslator bool `json:"response_before_translator"` ResponseAfterTranslator bool `json:"response_after_translator"` @@ -94,6 +95,11 @@ type rpcModelRouteRequest struct { HostCallbackID string `json:"host_callback_id,omitempty"` } +type rpcRequestCompletion struct { + pluginapi.RequestCompletion + HostCallbackID string `json:"host_callback_id,omitempty"` +} + type rpcResponseInterceptRequest struct { pluginapi.ResponseInterceptRequest HostCallbackID string `json:"host_callback_id,omitempty"` @@ -138,6 +144,7 @@ func rpcCapabilitiesFromPlugin(plugin pluginapi.Plugin) rpcCapabilities { RequestTranslator: caps.RequestTranslator != nil, RequestNormalizer: caps.RequestNormalizer != nil, RequestInterceptor: caps.RequestInterceptor != nil, + RequestLifecyclePlugin: caps.RequestLifecyclePlugin != nil, ResponseTranslator: caps.ResponseTranslator != nil, ResponseBeforeTranslator: caps.ResponseBeforeTranslator != nil, ResponseAfterTranslator: caps.ResponseAfterTranslator != nil, diff --git a/internal/pluginhost/test_helpers_test.go b/internal/pluginhost/test_helpers_test.go index c3deb906f..46146ad6f 100644 --- a/internal/pluginhost/test_helpers_test.go +++ b/internal/pluginhost/test_helpers_test.go @@ -94,6 +94,18 @@ func (l *testSymbolLookup) Call(ctx context.Context, method string, request []by return nil, errIntercept } return marshalRPCResult(resp) + case pluginabi.MethodRequestComplete: + if l.active.Capabilities.RequestLifecyclePlugin == nil { + return nil, fmt.Errorf("missing request lifecycle plugin") + } + var req pluginapi.RequestCompletion + if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { + return nil, errUnmarshal + } + if errComplete := l.active.Capabilities.RequestLifecyclePlugin.HandleRequestComplete(ctx, req); errComplete != nil { + return nil, errComplete + } + return marshalRPCResult(rpcEmptyResponse{}) case pluginabi.MethodResponseInterceptAfter: if l.active.Capabilities.ResponseInterceptor == nil { return nil, fmt.Errorf("missing response interceptor") @@ -245,6 +257,13 @@ type testUsageCapability struct{} func (testUsageCapability) HandleUsage(ctx context.Context, record pluginapi.UsageRecord) {} +type requestLifecyclePluginFunc func(context.Context, pluginapi.RequestCompletion) + +func (f requestLifecyclePluginFunc) HandleRequestComplete(ctx context.Context, completion pluginapi.RequestCompletion) error { + f(ctx, completion) + return nil +} + type testThinkingCapability struct { provider string } diff --git a/sdk/api/handlers/claude/code_handlers.go b/sdk/api/handlers/claude/code_handlers.go index 8fd42a6ff..e988d7d56 100644 --- a/sdk/api/handlers/claude/code_handlers.go +++ b/sdk/api/handlers/claude/code_handlers.go @@ -378,6 +378,25 @@ func (h *ClaudeCodeAPIHandler) WriteErrorResponse(c *gin.Context, msg *interface if msg != nil && msg.StatusCode > 0 { status = msg.StatusCode } + if msg != nil && msg.DirectResponse { + for key, values := range handlers.FilterUpstreamHeaders(msg.Headers) { + if len(values) == 0 || handlers.IsCPAReservedResponseHeader(key) { + continue + } + c.Writer.Header().Del(key) + for _, value := range values { + c.Writer.Header().Add(key, value) + } + } + body := bytes.Clone(msg.Body) + appendClaudeAPIResponse(c, body) + if !c.Writer.Written() && c.Writer.Header().Get("Content-Type") == "" { + c.Writer.Header().Set("Content-Type", "application/json") + } + c.Status(status) + _, _ = c.Writer.Write(body) + return + } if msg != nil && msg.Addon != nil && handlers.PassthroughHeadersEnabled(h.Cfg) { for key, values := range msg.Addon { if len(values) == 0 || handlers.IsCPAReservedResponseHeader(key) { diff --git a/sdk/api/handlers/handlers_error_response_test.go b/sdk/api/handlers/handlers_error_response_test.go index 15ce8279c..561d98d62 100644 --- a/sdk/api/handlers/handlers_error_response_test.go +++ b/sdk/api/handlers/handlers_error_response_test.go @@ -42,6 +42,47 @@ func TestWriteErrorResponse_AddonHeadersDisabledByDefault(t *testing.T) { } } +func TestWriteErrorResponseDirectResponse(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + c.Writer.Header().Set("X-Cpa-Trace-Id", "local-trace") + c.Writer.Header().Set("Access-Control-Allow-Origin", "https://trusted.example") + + handler := NewBaseAPIHandlers(nil, nil) + handler.WriteErrorResponse(c, &interfaces.ErrorMessage{ + StatusCode: http.StatusForbidden, + DirectResponse: true, + Body: []byte(`{"error":"blocked"}`), + Headers: http.Header{ + "Content-Type": {"application/problem+json"}, + "X-Plugin-Policy": {"blocked"}, + "X-Cpa-Trace-Id": {"plugin-trace"}, + "Access-Control-Allow-Origin": {"https://untrusted.example"}, + }, + }) + + if recorder.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusForbidden) + } + if got := recorder.Body.String(); got != `{"error":"blocked"}` { + t.Fatalf("body = %q", got) + } + if got := recorder.Header().Get("Content-Type"); got != "application/problem+json" { + t.Fatalf("Content-Type = %q", got) + } + if got := recorder.Header().Get("X-Plugin-Policy"); got != "blocked" { + t.Fatalf("X-Plugin-Policy = %q", got) + } + if got := recorder.Header().Get("X-Cpa-Trace-Id"); got != "local-trace" { + t.Fatalf("X-Cpa-Trace-Id = %q, want local value", got) + } + if got := recorder.Header().Get("Access-Control-Allow-Origin"); got != "https://trusted.example" { + t.Fatalf("Access-Control-Allow-Origin = %q, want trusted origin", got) + } +} + func TestInternalConcurrencyBusyWritesRetryAfterWithoutPassthrough(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() diff --git a/sdk/api/handlers/handlers_errors.go b/sdk/api/handlers/handlers_errors.go index 960442bd4..1f555e223 100644 --- a/sdk/api/handlers/handlers_errors.go +++ b/sdk/api/handlers/handlers_errors.go @@ -79,6 +79,10 @@ func (h *BaseAPIHandler) WriteErrorResponse(c *gin.Context, msg *interfaces.Erro if msg != nil && msg.StatusCode > 0 { status = msg.StatusCode } + if msg != nil && msg.DirectResponse { + writeDirectErrorResponse(c, status, msg) + return + } if msg != nil && msg.Error != nil { for _, value := range coreauth.SafeResponseHeaders(msg.Error).Values("Retry-After") { c.Writer.Header().Add("Retry-After", value) @@ -128,6 +132,25 @@ func (h *BaseAPIHandler) WriteErrorResponse(c *gin.Context, msg *interfaces.Erro _, _ = c.Writer.Write(body) } +func writeDirectErrorResponse(c *gin.Context, status int, msg *interfaces.ErrorMessage) { + for key, values := range FilterUpstreamHeaders(msg.Headers) { + if len(values) == 0 || IsCPAReservedResponseHeader(key) { + continue + } + c.Writer.Header().Del(key) + for _, value := range values { + c.Writer.Header().Add(key, value) + } + } + body := bytes.Clone(msg.Body) + appendAPIResponse(c, body) + if !c.Writer.Written() && c.Writer.Header().Get("Content-Type") == "" { + c.Writer.Header().Set("Content-Type", "application/json") + } + c.Status(status) + _, _ = c.Writer.Write(body) +} + func (h *BaseAPIHandler) LoggingAPIResponseError(ctx context.Context, err *interfaces.ErrorMessage) { if h.Cfg.RequestLog { if ginContext, ok := ctx.Value("gin").(*gin.Context); ok { diff --git a/sdk/api/handlers/handlers_execution.go b/sdk/api/handlers/handlers_execution.go index 18508b781..7c25ab146 100644 --- a/sdk/api/handlers/handlers_execution.go +++ b/sdk/api/handlers/handlers_execution.go @@ -1,11 +1,13 @@ package handlers import ( + "errors" "fmt" "net/http" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "golang.org/x/net/context" ) @@ -67,6 +69,7 @@ func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entr Payload: payload, } afterAuthCapture := &requestAfterAuthCapture{} + lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, normalizedModel, originalRequestedModel, false, reqMeta, execOptions.SkipInterceptorPluginID) opts := coreexecutor.Options{ Stream: false, Alt: alt, @@ -75,31 +78,27 @@ func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entr ResponseFormat: sdktranslator.FromString(responseProtocol), Headers: modelExecutionHeaders(ctx, execOptions.Headers), Query: modelExecutionQuery(ctx, execOptions.Query), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID), } opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + var interceptErr *interfaces.ErrorMessage + req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) + if interceptErr != nil { + lifecycle.completeError(ctx, interceptErr) + return nil, nil, interceptErr + } resp, err := h.AuthManager.Execute(ctx, providers, req, opts) if err != nil { err = enrichAuthSelectionError(err, providers, normalizedModel) - status := http.StatusInternalServerError - if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { - if code := se.StatusCode(); code > 0 { - status = code - } - } - var addon http.Header - if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { - if hdr := he.Headers(); hdr != nil { - addon = hdr.Clone() - } - } - return nil, nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} + errMsg := executionErrorMessage(err) + lifecycle.completeError(ctx, errMsg) + return nil, nil, errMsg } executedReq, executedOpts := afterAuthCapture.apply(req, opts) rawResponseHeaders := cloneHeader(resp.Headers) responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, responseProtocol, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + body, responseHeaders := h.applyResponseInterceptors(ctx, lifecycle.requestID(), responseProtocol, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + lifecycle.complete(pluginapi.RequestCompletionSucceeded, http.StatusOK, nil) return body, responseHeaders, nil } @@ -135,6 +134,7 @@ func (h *BaseAPIHandler) executeCountWithAuthManager(ctx context.Context, handle Payload: payload, } afterAuthCapture := &requestAfterAuthCapture{} + lifecycle := h.newRequestLifecycleTracker(ctx, handlerType, normalizedModel, originalRequestedModel, false, reqMeta, execOptions.SkipInterceptorPluginID) opts := coreexecutor.Options{ Stream: false, Alt: alt, @@ -142,31 +142,27 @@ func (h *BaseAPIHandler) executeCountWithAuthManager(ctx context.Context, handle SourceFormat: sdktranslator.FromString(handlerType), Headers: modelExecutionHeaders(ctx, execOptions.Headers), Query: modelExecutionQuery(ctx, execOptions.Query), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID), } opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + var interceptErr *interfaces.ErrorMessage + req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) + if interceptErr != nil { + lifecycle.completeError(ctx, interceptErr) + return nil, nil, interceptErr + } resp, err := h.AuthManager.ExecuteCount(ctx, providers, req, opts) if err != nil { err = enrichAuthSelectionError(err, providers, normalizedModel) - status := http.StatusInternalServerError - if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { - if code := se.StatusCode(); code > 0 { - status = code - } - } - var addon http.Header - if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { - if hdr := he.Headers(); hdr != nil { - addon = hdr.Clone() - } - } - return nil, nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} + errMsg := executionErrorMessage(err) + lifecycle.completeError(ctx, errMsg) + return nil, nil, errMsg } executedReq, executedOpts := afterAuthCapture.apply(req, opts) rawResponseHeaders := cloneHeader(resp.Headers) responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + body, responseHeaders := h.applyResponseInterceptors(ctx, lifecycle.requestID(), handlerType, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + lifecycle.complete(pluginapi.RequestCompletionSucceeded, http.StatusOK, nil) return body, responseHeaders, nil } @@ -179,15 +175,28 @@ func (h *BaseAPIHandler) executeWithPluginExecutor(ctx context.Context, entryPro return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} } req, opts := h.pluginExecutorRequest(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, false, execOptions) - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - req, opts = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, modelName, originalRequestedModel, false, opts.Metadata, execOptions.SkipInterceptorPluginID) + var interceptErr *interfaces.ErrorMessage + req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) + if interceptErr != nil { + lifecycle.completeError(ctx, interceptErr) + return nil, nil, interceptErr + } + req, opts, interceptErr = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) + if interceptErr != nil { + lifecycle.completeError(ctx, interceptErr) + return nil, nil, interceptErr + } resp, errExecute := host.ExecutePluginExecutor(ctx, executorPluginID, req, opts) if errExecute != nil { - return nil, nil, executionErrorMessage(errExecute) + errMsg := executionErrorMessage(errExecute) + lifecycle.completeError(ctx, errMsg) + return nil, nil, errMsg } rawResponseHeaders := cloneHeader(resp.Headers) responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, responseProtocol, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + body, responseHeaders := h.applyResponseInterceptors(ctx, lifecycle.requestID(), responseProtocol, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + lifecycle.complete(pluginapi.RequestCompletionSucceeded, http.StatusOK, nil) return body, responseHeaders, nil } @@ -200,15 +209,28 @@ func (h *BaseAPIHandler) countWithPluginExecutor(ctx context.Context, handlerTyp return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} } req, opts := h.pluginExecutorRequest(ctx, handlerType, handlerType, modelName, originalRequestedModel, rawJSON, alt, false, execOptions) - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - req, opts = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, handlerType, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + lifecycle := h.newRequestLifecycleTracker(ctx, handlerType, modelName, originalRequestedModel, false, opts.Metadata, execOptions.SkipInterceptorPluginID) + var interceptErr *interfaces.ErrorMessage + req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) + if interceptErr != nil { + lifecycle.completeError(ctx, interceptErr) + return nil, nil, interceptErr + } + req, opts, interceptErr = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, handlerType, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) + if interceptErr != nil { + lifecycle.completeError(ctx, interceptErr) + return nil, nil, interceptErr + } resp, errCount := host.CountPluginExecutor(ctx, executorPluginID, req, opts) if errCount != nil { - return nil, nil, executionErrorMessage(errCount) + errMsg := executionErrorMessage(errCount) + lifecycle.completeError(ctx, errMsg) + return nil, nil, errMsg } rawResponseHeaders := cloneHeader(resp.Headers) responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + body, responseHeaders := h.applyResponseInterceptors(ctx, lifecycle.requestID(), handlerType, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + lifecycle.complete(pluginapi.RequestCompletionSucceeded, http.StatusOK, nil) return body, responseHeaders, nil } @@ -238,9 +260,9 @@ func (h *BaseAPIHandler) pluginExecutorRequest(ctx context.Context, entryProtoco return req, opts } -func (h *BaseAPIHandler) applyRequestInterceptorsAfterPluginExecutorRoute(ctx context.Context, host PluginExecutorHost, executorPluginID, entryProtocol, originalRequestedModel string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options) { +func (h *BaseAPIHandler) applyRequestInterceptorsAfterPluginExecutorRoute(ctx context.Context, host PluginExecutorHost, executorPluginID, entryProtocol, originalRequestedModel, requestID string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options, *interfaces.ErrorMessage) { if !requestInterceptorsEnabled(h.interceptorHost()) { - return req, opts + return req, opts, nil } toFormat := sdktranslator.FromString(entryProtocol) if resolver, ok := host.(pluginExecutorFormatResolver); ok && resolver != nil { @@ -257,16 +279,29 @@ func (h *BaseAPIHandler) applyRequestInterceptorsAfterPluginExecutorRoute(ctx co Headers: cloneHeader(opts.Headers), Body: cloneBytes(req.Payload), Metadata: opts.Metadata, - }, skipPluginID) + }, requestID, skipPluginID) opts.Headers = mergeRequestInterceptorHeaders(opts.Headers, resp.Headers, resp.ClearHeaders) if len(resp.Body) > 0 { req.Payload = cloneBytes(resp.Body) opts.OriginalRequest = cloneBytes(resp.Body) } - return req, opts + if resp.Terminate { + return req, opts, directTerminationError(resp.StatusCode, resp.ResponseHeaders, resp.ResponseBody) + } + return req, opts, nil } func executionErrorMessage(err error) *interfaces.ErrorMessage { + var terminated *coreexecutor.RequestTerminatedError + if errors.As(err, &terminated) && terminated != nil { + return &interfaces.ErrorMessage{ + StatusCode: normalizedTerminationStatus(terminated.StatusCode()), + Error: err, + DirectResponse: true, + Body: terminated.ResponseBody(), + Headers: terminated.ResponseHeaders(), + } + } status := http.StatusInternalServerError if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { if code := se.StatusCode(); code > 0 { diff --git a/sdk/api/handlers/handlers_interceptors.go b/sdk/api/handlers/handlers_interceptors.go index eb90d7289..2c9cb282c 100644 --- a/sdk/api/handlers/handlers_interceptors.go +++ b/sdk/api/handlers/handlers_interceptors.go @@ -3,7 +3,11 @@ package handlers import ( "net/http" "sync" + "time" + "github.com/google/uuid" + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" "golang.org/x/net/context" @@ -32,6 +36,112 @@ type requestInterceptorDetector interface { HasRequestInterceptors() bool } +type requestLifecycleHost interface { + CompleteRequest(context.Context, pluginapi.RequestCompletion) +} + +type requestLifecycleSkipHost interface { + CompleteRequestExcept(context.Context, pluginapi.RequestCompletion, string) +} + +type requestLifecycleTracker struct { + once sync.Once + ctx context.Context + host PluginInterceptorHost + skipPluginID string + completion pluginapi.RequestCompletion +} + +func (h *BaseAPIHandler) newRequestLifecycleTracker(ctx context.Context, sourceFormat, model, requestedModel string, stream bool, metadata map[string]any, skipPluginID string) *requestLifecycleTracker { + requestID := uuid.NewString() + traceID := logging.GetRequestID(ctx) + return &requestLifecycleTracker{ + ctx: ctx, + host: h.interceptorHost(), + skipPluginID: skipPluginID, + completion: pluginapi.RequestCompletion{ + RequestID: requestID, + TraceID: traceID, + SourceFormat: sourceFormat, + Model: model, + RequestedModel: requestedModel, + Stream: stream, + StartedAt: time.Now(), + Metadata: metadata, + }, + } +} + +func (t *requestLifecycleTracker) requestID() string { + if t == nil { + return "" + } + return t.completion.RequestID +} + +func (t *requestLifecycleTracker) complete(outcome pluginapi.RequestCompletionOutcome, statusCode int, err error) { + if t == nil { + return + } + t.once.Do(func() { + completion := t.completion + completion.Outcome = outcome + completion.StatusCode = statusCode + completion.CompletedAt = time.Now() + if err != nil { + completion.Error = err.Error() + } + if t.skipPluginID != "" { + if host, ok := t.host.(requestLifecycleSkipHost); ok { + host.CompleteRequestExcept(t.ctx, completion, t.skipPluginID) + return + } + } + if host, ok := t.host.(requestLifecycleHost); ok { + host.CompleteRequest(t.ctx, completion) + } + }) +} + +func (t *requestLifecycleTracker) completeError(ctx context.Context, msg *interfaces.ErrorMessage) { + outcome := pluginapi.RequestCompletionFailed + if msg != nil && msg.DirectResponse { + outcome = pluginapi.RequestCompletionRejected + } else if ctx != nil && ctx.Err() != nil { + outcome = pluginapi.RequestCompletionCanceled + } + statusCode := 0 + var err error + if msg != nil { + statusCode = msg.StatusCode + err = msg.Error + } + if outcome == pluginapi.RequestCompletionCanceled { + statusCode = 0 + } + t.complete(outcome, statusCode, err) +} + +func normalizedTerminationStatus(statusCode int) int { + if statusCode < http.StatusOK || statusCode > 599 { + return http.StatusForbidden + } + return statusCode +} + +func requestTerminationError(resp pluginapi.RequestInterceptResponse) *interfaces.ErrorMessage { + return directTerminationError(resp.StatusCode, resp.ResponseHeaders, resp.ResponseBody) +} + +func directTerminationError(statusCode int, headers http.Header, body []byte) *interfaces.ErrorMessage { + return &interfaces.ErrorMessage{ + StatusCode: normalizedTerminationStatus(statusCode), + DirectResponse: true, + Body: cloneBytes(body), + Headers: cloneHeader(headers), + } +} + func cloneHeader(src http.Header) http.Header { if src == nil { return nil @@ -294,12 +404,14 @@ func interceptStreamChunk(ctx context.Context, host PluginInterceptorHost, req p return host.InterceptStreamChunk(ctx, req) } -func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, handlerType, requestedModel string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options) { +func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, handlerType, requestedModel, requestID string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options, *interfaces.ErrorMessage) { host := h.interceptorHost() if host == nil { - return req, opts + return req, opts, nil } resp := interceptRequestBeforeAuth(ctx, host, pluginapi.RequestInterceptRequest{ + RequestID: requestID, + TraceID: logging.GetRequestID(ctx), SourceFormat: handlerType, Model: req.Model, RequestedModel: requestedModel, @@ -313,15 +425,18 @@ func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, req.Payload = cloneBytes(resp.Body) opts.OriginalRequest = cloneBytes(resp.Body) } - return req, opts + if resp.Terminate { + return req, opts, requestTerminationError(resp) + } + return req, opts, nil } -func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCapture, skipPluginID string) coreexecutor.RequestAfterAuthInterceptor { +func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCapture, requestID, skipPluginID string) coreexecutor.RequestAfterAuthInterceptor { if !requestInterceptorsEnabled(h.interceptorHost()) { return nil } return func(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest) coreexecutor.RequestAfterAuthInterceptResponse { - resp := h.applyRequestInterceptorsAfterAuth(ctx, req, skipPluginID) + resp := h.applyRequestInterceptorsAfterAuth(ctx, req, requestID, skipPluginID) if capture != nil { capture.record(req, resp) } @@ -329,12 +444,14 @@ func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCa } } -func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest, skipPluginID string) coreexecutor.RequestAfterAuthInterceptResponse { +func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest, requestID, skipPluginID string) coreexecutor.RequestAfterAuthInterceptResponse { host := h.interceptorHost() if !requestInterceptorsEnabled(host) { return coreexecutor.RequestAfterAuthInterceptResponse{} } resp := interceptRequestAfterAuth(ctx, host, pluginapi.RequestInterceptRequest{ + RequestID: requestID, + TraceID: logging.GetRequestID(ctx), SourceFormat: req.SourceFormat.String(), ToFormat: req.ToFormat.String(), Model: req.Model, @@ -345,18 +462,23 @@ func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, Metadata: req.Metadata, }, skipPluginID) return coreexecutor.RequestAfterAuthInterceptResponse{ - Headers: resp.Headers, - Body: resp.Body, - ClearHeaders: resp.ClearHeaders, + Headers: resp.Headers, + Body: resp.Body, + ClearHeaders: resp.ClearHeaders, + Terminate: resp.Terminate, + StatusCode: normalizedTerminationStatus(resp.StatusCode), + ResponseHeaders: resp.ResponseHeaders, + ResponseBody: resp.ResponseBody, } } -func (h *BaseAPIHandler) applyResponseInterceptors(ctx context.Context, handlerType, normalizedModel, requestedModel string, opts coreexecutor.Options, rawResponseHeaders, responseHeaders http.Header, originalRequest, requestBody, body []byte, statusCode int, skipPluginID string) ([]byte, http.Header) { +func (h *BaseAPIHandler) applyResponseInterceptors(ctx context.Context, requestID, handlerType, normalizedModel, requestedModel string, opts coreexecutor.Options, rawResponseHeaders, responseHeaders http.Header, originalRequest, requestBody, body []byte, statusCode int, skipPluginID string) ([]byte, http.Header) { host := h.interceptorHost() if host == nil { return body, responseHeaders } resp := interceptResponse(ctx, host, pluginapi.ResponseInterceptRequest{ + RequestID: requestID, SourceFormat: handlerType, Model: normalizedModel, RequestedModel: requestedModel, diff --git a/sdk/api/handlers/handlers_interceptors_test.go b/sdk/api/handlers/handlers_interceptors_test.go index bdd1de070..738b1a776 100644 --- a/sdk/api/handlers/handlers_interceptors_test.go +++ b/sdk/api/handlers/handlers_interceptors_test.go @@ -8,9 +8,11 @@ import ( "net/url" "sync" "testing" + "time" "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" @@ -23,6 +25,7 @@ type handlerInterceptorTestHost struct { interceptRequestAfterAuth func(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse interceptResponse func(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse interceptStreamChunk func(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse + completeRequest func(context.Context, pluginapi.RequestCompletion) } type handlerInterceptorNoStreamTestHost struct { @@ -73,6 +76,12 @@ func (h *handlerInterceptorTestHost) InterceptStreamChunk(ctx context.Context, r } } +func (h *handlerInterceptorTestHost) CompleteRequest(ctx context.Context, completion pluginapi.RequestCompletion) { + if h != nil && h.completeRequest != nil { + h.completeRequest(ctx, completion) + } +} + type interceptorCaptureExecutor struct { provider string @@ -199,6 +208,297 @@ func contextWithQuery(query url.Values) context.Context { return context.WithValue(context.Background(), "gin", c) } +func TestRequestLifecycleTrackerUsesUniqueExecutionIDs(t *testing.T) { + handler := NewBaseAPIHandlers(nil, nil) + ctx := logging.WithRequestID(context.Background(), "trace-1") + first := handler.newRequestLifecycleTracker(ctx, "openai", "model", "model", false, nil, "") + second := handler.newRequestLifecycleTracker(ctx, "openai", "model", "model", false, nil, "") + if first.requestID() == "" || second.requestID() == "" || first.requestID() == second.requestID() { + t.Fatalf("lifecycle request IDs = %q and %q", first.requestID(), second.requestID()) + } + if first.completion.TraceID != "trace-1" || second.completion.TraceID != "trace-1" { + t.Fatalf("trace IDs = %q and %q", first.completion.TraceID, second.completion.TraceID) + } +} + +func TestHandlerRequestInterceptorTerminatesBeforeAuth(t *testing.T) { + model := "handler-interceptor-terminate-before-auth" + executor := &interceptorCaptureExecutor{} + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) + var requestID string + var completion pluginapi.RequestCompletion + handler.SetPluginHost(&handlerInterceptorTestHost{ + interceptRequestBeforeAuth: func(_ context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + requestID = req.RequestID + return pluginapi.RequestInterceptResponse{ + Terminate: true, + StatusCode: http.StatusForbidden, + ResponseHeaders: http.Header{"Content-Type": {"application/json"}, "X-Policy": {"blocked"}}, + ResponseBody: []byte(`{"error":"blocked"}`), + } + }, + completeRequest: func(_ context.Context, got pluginapi.RequestCompletion) { + completion = got + }, + }) + + body, headers, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", model, []byte(`{"model":"`+model+`"}`), "") + if body != nil || headers != nil { + t.Fatalf("terminated response body = %q, headers = %#v", body, headers) + } + if errMsg == nil || !errMsg.DirectResponse || errMsg.StatusCode != http.StatusForbidden { + t.Fatalf("termination error = %#v", errMsg) + } + if string(errMsg.Body) != `{"error":"blocked"}` || errMsg.Headers.Get("X-Policy") != "blocked" { + t.Fatalf("termination response = body %q, headers %#v", errMsg.Body, errMsg.Headers) + } + if requestID == "" || completion.RequestID != requestID { + t.Fatalf("request IDs = start %q, completion %q", requestID, completion.RequestID) + } + if completion.Outcome != pluginapi.RequestCompletionRejected || completion.StatusCode != http.StatusForbidden { + t.Fatalf("completion = %#v", completion) + } + capturedReq, _ := executor.captured() + if capturedReq.Model != "" { + t.Fatalf("executor received terminated request: %#v", capturedReq) + } +} + +func TestHandlerRequestInterceptorTerminatesAfterAuth(t *testing.T) { + model := "handler-interceptor-terminate-after-auth" + executor := &interceptorCaptureExecutor{} + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) + var beforeRequestID string + var afterRequestID string + var afterCalls int + var completion pluginapi.RequestCompletion + handler.SetPluginHost(&handlerInterceptorTestHost{ + interceptRequestBeforeAuth: func(_ context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + beforeRequestID = req.RequestID + return pluginapi.RequestInterceptResponse{Headers: req.Headers, Body: req.Body} + }, + interceptRequestAfterAuth: func(_ context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + afterCalls++ + afterRequestID = req.RequestID + return pluginapi.RequestInterceptResponse{ + Terminate: true, + StatusCode: http.StatusTooManyRequests, + ResponseHeaders: http.Header{"Retry-After": {"3"}}, + ResponseBody: []byte(`{"error":"busy"}`), + } + }, + completeRequest: func(_ context.Context, got pluginapi.RequestCompletion) { + completion = got + }, + }) + + _, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", model, []byte(`{"model":"`+model+`"}`), "") + if errMsg == nil || !errMsg.DirectResponse || errMsg.StatusCode != http.StatusTooManyRequests { + t.Fatalf("termination error = %#v", errMsg) + } + if beforeRequestID == "" || afterRequestID != beforeRequestID || completion.RequestID != beforeRequestID { + t.Fatalf("request IDs = before %q, after %q, completion %q", beforeRequestID, afterRequestID, completion.RequestID) + } + if completion.Outcome != pluginapi.RequestCompletionRejected { + t.Fatalf("completion outcome = %q", completion.Outcome) + } + if afterCalls != 1 { + t.Fatalf("after-auth interceptor calls = %d, want 1", afterCalls) + } + capturedReq, _ := executor.captured() + if capturedReq.Model != "" { + t.Fatalf("executor received terminated request: %#v", capturedReq) + } +} + +func TestHandlerAfterAuthTerminationSkipsCountAndStreamExecutors(t *testing.T) { + for _, operation := range []string{"count", "stream"} { + t.Run(operation, func(t *testing.T) { + model := "handler-interceptor-terminate-" + operation + executor := &interceptorCaptureExecutor{} + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) + afterCalls := 0 + handler.SetPluginHost(&handlerInterceptorTestHost{ + interceptRequestAfterAuth: func(_ context.Context, _ pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + afterCalls++ + return pluginapi.RequestInterceptResponse{ + Terminate: true, + StatusCode: http.StatusForbidden, + ResponseBody: []byte(`{"error":"blocked"}`), + } + }, + }) + + var errMsg *interfaces.ErrorMessage + if operation == "count" { + _, _, errMsg = handler.ExecuteCountWithAuthManager(context.Background(), "openai", model, []byte(`{"model":"`+model+`"}`), "") + } else { + dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", model, []byte(`{"model":"`+model+`","stream":true}`), "") + if dataChan != nil { + t.Fatal("terminated stream returned a data channel") + } + errMsg = <-errChan + } + if errMsg == nil || !errMsg.DirectResponse || errMsg.StatusCode != http.StatusForbidden { + t.Fatalf("termination error = %#v", errMsg) + } + if afterCalls != 1 { + t.Fatalf("after-auth interceptor calls = %d, want 1", afterCalls) + } + capturedReq, _ := executor.captured() + if capturedReq.Model != "" { + t.Fatalf("executor received terminated request: %#v", capturedReq) + } + }) + } +} + +func TestHandlerLifecycleCompletesSuccessfulRequestOnce(t *testing.T) { + model := "handler-interceptor-lifecycle-success" + executor := &interceptorCaptureExecutor{} + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) + var requestID string + var completionCount int + var completion pluginapi.RequestCompletion + handler.SetPluginHost(&handlerInterceptorTestHost{ + interceptRequestBeforeAuth: func(_ context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + requestID = req.RequestID + return pluginapi.RequestInterceptResponse{Headers: req.Headers, Body: req.Body} + }, + interceptResponse: func(_ context.Context, req pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse { + if req.RequestID != requestID { + t.Fatalf("response request ID = %q, want %q", req.RequestID, requestID) + } + return pluginapi.ResponseInterceptResponse{Headers: req.ResponseHeaders, Body: req.Body} + }, + completeRequest: func(_ context.Context, got pluginapi.RequestCompletion) { + completionCount++ + completion = got + }, + }) + + body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", model, []byte(`{"model":"`+model+`"}`), "") + if errMsg != nil || string(body) != "ok" { + t.Fatalf("ExecuteWithAuthManager() body = %q, error = %#v", body, errMsg) + } + if completionCount != 1 || completion.Outcome != pluginapi.RequestCompletionSucceeded || completion.RequestID != requestID { + t.Fatalf("completion count = %d, completion = %#v", completionCount, completion) + } + if completion.StartedAt.IsZero() || completion.CompletedAt.Before(completion.StartedAt) { + t.Fatalf("completion timestamps = %#v", completion) + } +} + +func TestHandlerLifecycleCompletesFailedRequest(t *testing.T) { + model := "handler-interceptor-lifecycle-failed" + executor := &interceptorCaptureExecutor{ + execute: func(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, fmt.Errorf("upstream failed") + }, + } + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) + var completion pluginapi.RequestCompletion + handler.SetPluginHost(&handlerInterceptorTestHost{ + completeRequest: func(_ context.Context, got pluginapi.RequestCompletion) { + completion = got + }, + }) + + _, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", model, []byte(`{"model":"`+model+`"}`), "") + if errMsg == nil { + t.Fatal("ExecuteWithAuthManager() error = nil") + } + if completion.Outcome != pluginapi.RequestCompletionFailed || completion.Error == "" { + t.Fatalf("completion = %#v", completion) + } +} + +func TestHandlerLifecycleCompletesSuccessfulStreamOnce(t *testing.T) { + model := "handler-interceptor-lifecycle-stream" + executor := &interceptorCaptureExecutor{ + stream: func(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk, 1) + chunks <- coreexecutor.StreamChunk{Payload: []byte(`{"chunk":true}`)} + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil + }, + } + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) + completions := make(chan pluginapi.RequestCompletion, 2) + handler.SetPluginHost(&handlerInterceptorTestHost{ + completeRequest: func(_ context.Context, completion pluginapi.RequestCompletion) { + completions <- completion + }, + }) + + dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", model, []byte(`{"model":"`+model+`","stream":true}`), "") + for dataChan != nil || errChan != nil { + select { + case _, ok := <-dataChan: + if !ok { + dataChan = nil + } + case errMsg, ok := <-errChan: + if ok && errMsg != nil { + t.Fatalf("stream error = %#v", errMsg) + } + if !ok { + errChan = nil + } + } + } + select { + case completion := <-completions: + if completion.Outcome != pluginapi.RequestCompletionSucceeded || !completion.Stream || completion.RequestID == "" { + t.Fatalf("stream completion = %#v", completion) + } + case <-time.After(time.Second): + t.Fatal("missing stream completion") + } + select { + case duplicate := <-completions: + t.Fatalf("duplicate stream completion = %#v", duplicate) + default: + } +} + +func TestHandlerLifecycleCompletesCanceledStream(t *testing.T) { + model := "handler-interceptor-lifecycle-canceled-stream" + chunks := make(chan coreexecutor.StreamChunk, 1) + chunks <- coreexecutor.StreamChunk{Payload: []byte(`{"chunk":true}`)} + executor := &interceptorCaptureExecutor{ + stream: func(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { + return &coreexecutor.StreamResult{Chunks: chunks}, nil + }, + } + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) + completions := make(chan pluginapi.RequestCompletion, 1) + handler.SetPluginHost(&handlerInterceptorTestHost{ + completeRequest: func(_ context.Context, completion pluginapi.RequestCompletion) { + completions <- completion + }, + }) + ctx, cancel := context.WithCancel(context.Background()) + dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", model, []byte(`{"model":"`+model+`","stream":true}`), "") + cancel() + for dataChan != nil || errChan != nil { + select { + case _, ok := <-dataChan: + if !ok { + dataChan = nil + } + case _, ok := <-errChan: + if !ok { + errChan = nil + } + } + } + completion := <-completions + if completion.Outcome != pluginapi.RequestCompletionCanceled || completion.StatusCode != 0 { + t.Fatalf("completion = %#v", completion) + } +} + func TestHandlerRequestInterceptorRewritesExecutorRequest(t *testing.T) { model := "handler-interceptor-request-model" executor := &interceptorCaptureExecutor{} diff --git a/sdk/api/handlers/handlers_stream.go b/sdk/api/handlers/handlers_stream.go index 51caffaee..d0861764b 100644 --- a/sdk/api/handlers/handlers_stream.go +++ b/sdk/api/handlers/handlers_stream.go @@ -40,18 +40,38 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt return nil, nil, errChan } req, opts := h.pluginExecutorRequest(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, true, execOptions) - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - req, opts = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, modelName, originalRequestedModel, true, opts.Metadata, execOptions.SkipInterceptorPluginID) + var interceptErr *interfaces.ErrorMessage + req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) + if interceptErr != nil { + lifecycle.completeError(ctx, interceptErr) + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- interceptErr + close(errChan) + return nil, nil, errChan + } + req, opts, interceptErr = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) + if interceptErr != nil { + lifecycle.completeError(ctx, interceptErr) + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- interceptErr + close(errChan) + return nil, nil, errChan + } streamResult, errStream := host.ExecutePluginExecutorStream(ctx, executorPluginID, req, opts) if errStream != nil { + errMsg := executionErrorMessage(errStream) + lifecycle.completeError(ctx, errMsg) errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- executionErrorMessage(errStream) + errChan <- errMsg close(errChan) return nil, nil, errChan } if streamResult == nil { + errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor returned nil stream")} + lifecycle.completeError(ctx, errMsg) errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor returned nil stream")} + errChan <- errMsg close(errChan) return nil, nil, errChan } @@ -66,6 +86,7 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt } if streamInterceptorsActive { intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ + RequestID: lifecycle.requestID(), SourceFormat: responseProtocol, Model: modelName, RequestedModel: originalRequestedModel, @@ -96,6 +117,12 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt chunks = closed } go func() { + completionOutcome := pluginapi.RequestCompletionSucceeded + completionStatus := http.StatusOK + var completionErr error + defer func() { + lifecycle.complete(completionOutcome, completionStatus, completionErr) + }() defer close(dataChan) defer close(errChan) chunkIndex := 0 @@ -103,15 +130,29 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt for { chunk, ok, canceled := nextStreamChunk(ctx, nil, nil, chunks) if canceled { + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + if ctx != nil { + completionErr = ctx.Err() + } return } if !ok { return } if chunk.Err != nil { + errMsg := executionErrorMessage(chunk.Err) + completionOutcome = pluginapi.RequestCompletionFailed + completionStatus = errMsg.StatusCode + completionErr = chunk.Err select { - case errChan <- executionErrorMessage(chunk.Err): + case errChan <- errMsg: case <-done: + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + if ctx != nil { + completionErr = ctx.Err() + } } return } @@ -121,6 +162,7 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt payload := cloneBytes(chunk.Payload) if streamInterceptorsActive { intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ + RequestID: lifecycle.requestID(), SourceFormat: responseProtocol, Model: modelName, RequestedModel: originalRequestedModel, @@ -146,9 +188,17 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt } if responseProtocol == "openai-response" { if errValidate := validateSSEDataJSON(payload); errValidate != nil { + completionOutcome = pluginapi.RequestCompletionFailed + completionStatus = http.StatusBadGateway + completionErr = errValidate select { case errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate}: case <-done: + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + if ctx != nil { + completionErr = ctx.Err() + } } return } @@ -159,6 +209,11 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt historyChunks = appendStreamInterceptorHistory(historyChunks, payload) } case <-done: + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + if ctx != nil { + completionErr = ctx.Err() + } return } } @@ -210,6 +265,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context Payload: payload, } afterAuthCapture := &requestAfterAuthCapture{} + lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, normalizedModel, originalRequestedModel, true, reqMeta, execOptions.SkipInterceptorPluginID) opts := coreexecutor.Options{ Stream: true, Alt: alt, @@ -218,33 +274,33 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context ResponseFormat: sdktranslator.FromString(responseProtocol), Headers: modelExecutionHeaders(ctx, execOptions.Headers), Query: modelExecutionQuery(ctx, execOptions.Query), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID), } opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + var interceptErr *interfaces.ErrorMessage + req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) + if interceptErr != nil { + lifecycle.completeError(ctx, interceptErr) + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- interceptErr + close(errChan) + return nil, nil, errChan + } streamResult, err := h.AuthManager.ExecuteStream(ctx, providers, req, opts) if err != nil { err = enrichAuthSelectionError(err, providers, normalizedModel) + errMsg := executionErrorMessage(err) + lifecycle.completeError(ctx, errMsg) errChan := make(chan *interfaces.ErrorMessage, 1) - status := http.StatusInternalServerError - if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { - if code := se.StatusCode(); code > 0 { - status = code - } - } - var addon http.Header - if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { - if hdr := he.Headers(); hdr != nil { - addon = hdr.Clone() - } - } - errChan <- &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} + errChan <- errMsg close(errChan) return nil, nil, errChan } if streamResult == nil { + errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("auth manager returned nil stream")} + lifecycle.completeError(ctx, errMsg) errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("auth manager returned nil stream")} + errChan <- errMsg close(errChan) return nil, nil, errChan } @@ -278,6 +334,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context } executedReq, executedOpts := executedRequest() intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ + RequestID: lifecycle.requestID(), SourceFormat: responseProtocol, Model: normalizedModel, RequestedModel: originalRequestedModel, @@ -298,6 +355,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context if streamInterceptorsActive { executedReq, executedOpts := executedRequest() intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ + RequestID: lifecycle.requestID(), SourceFormat: responseProtocol, Model: normalizedModel, RequestedModel: originalRequestedModel, @@ -434,9 +492,20 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context errChan := make(chan *interfaces.ErrorMessage, 1) go func() { + completionOutcome := pluginapi.RequestCompletionSucceeded + completionStatus := http.StatusOK + var completionErr error + defer func() { + lifecycle.complete(completionOutcome, completionStatus, completionErr) + }() defer close(dataChan) defer close(errChan) if streamCanceledBeforeRead { + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + if ctx != nil { + completionErr = ctx.Err() + } return } @@ -467,7 +536,17 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context } if bootstrapErr != nil { - _ = sendErr(bootstrapErr) + completionOutcome = pluginapi.RequestCompletionFailed + if bootstrapErr.DirectResponse { + completionOutcome = pluginapi.RequestCompletionRejected + } + completionStatus = bootstrapErr.StatusCode + completionErr = bootstrapErr.Error + if !sendErr(bootstrapErr) && ctx != nil && ctx.Err() != nil { + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + completionErr = ctx.Err() + } return } @@ -475,6 +554,11 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context historyChunks := bootstrapHistoryChunks if bootstrapPayload != nil { if okSendData := sendData(bootstrapPayload); !okSendData { + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + if ctx != nil { + completionErr = ctx.Err() + } return } if streamInterceptorsActive { @@ -483,11 +567,27 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context } for { chunk, ok, canceled := nextStreamChunk(ctx, nil, &streamClosedBeforeRead, chunks) - if canceled || !ok { + if canceled { + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + if ctx != nil { + completionErr = ctx.Err() + } + return + } + if !ok { return } if chunk.Err != nil { - _ = sendErr(executionErrorMessage(chunk.Err)) + errMsg := executionErrorMessage(chunk.Err) + completionOutcome = pluginapi.RequestCompletionFailed + completionStatus = errMsg.StatusCode + completionErr = chunk.Err + if !sendErr(errMsg) && ctx != nil && ctx.Err() != nil { + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + completionErr = ctx.Err() + } return } if len(chunk.Payload) == 0 { @@ -495,13 +595,25 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context } payload, deliverable, errMsg := transformStreamPayload(chunk.Payload, &chunkIndex, historyChunks) if errMsg != nil { - _ = sendErr(errMsg) + completionOutcome = pluginapi.RequestCompletionFailed + completionStatus = errMsg.StatusCode + completionErr = errMsg.Error + if !sendErr(errMsg) && ctx != nil && ctx.Err() != nil { + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + completionErr = ctx.Err() + } return } if !deliverable { continue } if okSendData := sendData(payload); !okSendData { + completionOutcome = pluginapi.RequestCompletionCanceled + completionStatus = 0 + if ctx != nil { + completionErr = ctx.Err() + } return } if streamInterceptorsActive { diff --git a/sdk/api/handlers/header_filter.go b/sdk/api/handlers/header_filter.go index 4e2617ecd..e72467459 100644 --- a/sdk/api/handlers/header_filter.go +++ b/sdk/api/handlers/header_filter.go @@ -37,8 +37,13 @@ var hopByHopHeaders = map[string]struct{}{ } var cpaReservedResponseHeaders = map[string]struct{}{ - "Access-Control-Expose-Headers": {}, - "X-Cpa-Trace-Id": {}, + "Access-Control-Allow-Credentials": {}, + "Access-Control-Allow-Headers": {}, + "Access-Control-Allow-Methods": {}, + "Access-Control-Allow-Origin": {}, + "Access-Control-Expose-Headers": {}, + "Access-Control-Max-Age": {}, + "X-Cpa-Trace-Id": {}, } // IsCPAReservedResponseHeader reports whether a downstream response header is managed by CPA. diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 9f31ab793..7958604d6 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -41,6 +41,9 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye if errExec == nil { return resp, nil } + if isRequestTerminatedError(errExec) { + return cliproxyexecutor.Response{}, errExec + } lastErr = errExec wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) if !shouldRetry { @@ -83,6 +86,9 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip if errExec == nil { return resp, nil } + if isRequestTerminatedError(errExec) { + return cliproxyexecutor.Response{}, errExec + } lastErr = errExec wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) if !shouldRetry { @@ -121,6 +127,9 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli if errStream == nil { return result, nil } + if isRequestTerminatedError(errStream) { + return nil, errStream + } lastErr = errStream wait, shouldRetry := m.shouldRetryAfterError(errStream, attempt, normalized, retryModel, maxWait) if !shouldRetry { @@ -151,9 +160,14 @@ type requestToFormatResolver interface { RequestToFormat(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) sdktranslator.Format } -func applyRequestAfterAuthInterceptor(ctx context.Context, executor ProviderExecutor, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, requestedModel string) (cliproxyexecutor.Request, cliproxyexecutor.Options) { +func isRequestTerminatedError(err error) bool { + var terminated *cliproxyexecutor.RequestTerminatedError + return errors.As(err, &terminated) && terminated != nil +} + +func applyRequestAfterAuthInterceptor(ctx context.Context, executor ProviderExecutor, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, requestedModel string) (cliproxyexecutor.Request, cliproxyexecutor.Options, error) { if opts.RequestAfterAuthInterceptor == nil { - return req, opts + return req, opts, nil } toFormat := requestToFormat(provider, executor, req, opts) resp := opts.RequestAfterAuthInterceptor(ctx, cliproxyexecutor.RequestAfterAuthInterceptRequest{ @@ -171,7 +185,14 @@ func applyRequestAfterAuthInterceptor(ctx context.Context, executor ProviderExec req.Payload = bytes.Clone(resp.Body) opts.OriginalRequest = bytes.Clone(resp.Body) } - return req, opts + if resp.Terminate { + return req, opts, &cliproxyexecutor.RequestTerminatedError{ + HTTPStatus: resp.StatusCode, + Header: cloneRequestHeaders(resp.ResponseHeaders), + Body: bytes.Clone(resp.ResponseBody), + } + } + return req, opts, nil } func requestToFormat(provider string, executor ProviderExecutor, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) sdktranslator.Format { @@ -304,7 +325,11 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req execReq.Model = executionModel } execOpts := opts - execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + var errIntercept error + execReq, execOpts, errIntercept = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + if errIntercept != nil { + return cliproxyexecutor.Response{}, errIntercept + } resp, errExec := executor.Execute(execCtx, auth, execReq, execOpts) if errExec != nil { if errCtx := execCtx.Err(); errCtx != nil { @@ -417,7 +442,11 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, execReq.Model = executionModel } execOpts := opts - execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + var errIntercept error + execReq, execOpts, errIntercept = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + if errIntercept != nil { + return cliproxyexecutor.Response{}, errIntercept + } resp, errExec := executor.CountTokens(execCtx, auth, execReq, execOpts) if errExec != nil { if errCtx := execCtx.Err(); errCtx != nil { diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index 92ef35084..45ede889b 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -982,6 +982,9 @@ func hasAntigravityProvider(providers []string) bool { } func shouldAttemptAntigravityCreditsFallback(m *Manager, lastErr error, providers []string) bool { + if isRequestTerminatedError(lastErr) { + return false + } status := statusCodeFromError(lastErr) log.WithFields(log.Fields{ "lastErr": errorString(lastErr), diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 8b655cbbc..c16b909e8 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -90,7 +90,13 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr } execOpts := opts execOpts.ExecutionLifecycle = selection - execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, selection.Executor, selection.Provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + var errIntercept error + execReq, execOpts, errIntercept = applyRequestAfterAuthInterceptor(execCtx, selection.Executor, selection.Provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + if errIntercept != nil { + releaseAttempt() + selection.End("request_intercepted") + return cliproxyexecutor.Response{}, errIntercept + } if errCtx := execCtx.Err(); errCtx != nil { releaseAttempt() selection.End("attempt_canceled") diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index b81920002..551dc3ef6 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -195,7 +195,11 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi execReq.Model = executionModel } execOpts := opts - execReq, execOpts = applyRequestAfterAuthInterceptor(ctx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + var errIntercept error + execReq, execOpts, errIntercept = applyRequestAfterAuthInterceptor(ctx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + if errIntercept != nil { + return nil, errIntercept + } if errCtx := ctx.Err(); errCtx != nil { return nil, errCtx } diff --git a/sdk/cliproxy/auth/request_termination_test.go b/sdk/cliproxy/auth/request_termination_test.go new file mode 100644 index 000000000..3ebd92e3b --- /dev/null +++ b/sdk/cliproxy/auth/request_termination_test.go @@ -0,0 +1,18 @@ +package auth + +import ( + "net/http" + "testing" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestRequestTerminatedErrorSkipsCreditsFallback(t *testing.T) { + errTerminated := &cliproxyexecutor.RequestTerminatedError{HTTPStatus: http.StatusTooManyRequests} + if !isRequestTerminatedError(errTerminated) { + t.Fatal("isRequestTerminatedError() = false") + } + if shouldAttemptAntigravityCreditsFallback(&Manager{}, errTerminated, []string{"antigravity"}) { + t.Fatal("terminated request must not use Antigravity credits fallback") + } +} diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index d300b2df5..85e5fdc23 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -93,6 +93,49 @@ type RequestAfterAuthInterceptResponse struct { Body []byte // ClearHeaders explicitly removes current request headers before Headers is applied. ClearHeaders []string + // Terminate prevents the selected executor from receiving the request. + Terminate bool + // StatusCode is the downstream HTTP status used when Terminate is true. + StatusCode int + // ResponseHeaders contains downstream response headers used when Terminate is true. + ResponseHeaders http.Header + // ResponseBody contains the downstream response body used when Terminate is true. + ResponseBody []byte +} + +// RequestTerminatedError carries a plugin-defined downstream response without executing upstream. +type RequestTerminatedError struct { + HTTPStatus int + Header http.Header + Body []byte +} + +func (e *RequestTerminatedError) Error() string { + return "request terminated by plugin" +} + +// StatusCode returns the plugin-defined downstream HTTP status. +func (e *RequestTerminatedError) StatusCode() int { + if e == nil { + return 0 + } + return e.HTTPStatus +} + +// ResponseHeaders returns a copy of the plugin-defined downstream headers. +func (e *RequestTerminatedError) ResponseHeaders() http.Header { + if e == nil { + return nil + } + return e.Header.Clone() +} + +// ResponseBody returns a copy of the plugin-defined downstream body. +func (e *RequestTerminatedError) ResponseBody() []byte { + if e == nil { + return nil + } + return append([]byte(nil), e.Body...) } // Options controls execution behavior for both streaming and non-streaming calls. diff --git a/sdk/pluginabi/types.go b/sdk/pluginabi/types.go index 5db85b0d6..916eb4eb0 100644 --- a/sdk/pluginabi/types.go +++ b/sdk/pluginabi/types.go @@ -6,9 +6,8 @@ const ( // ABIVersion tracks the native C ABI shape (native plugin exports). ABIVersion uint32 = 1 // SchemaVersion tracks the RPC JSON contract exchanged at plugin.register. - // Increment only for breaking RPC changes. New capabilities such as ModelRouter - // are gated by capability flags and method names while the version stays at 1. - SchemaVersion uint32 = 1 + // Version 2 adds request lifecycle completion and active request termination. + SchemaVersion uint32 = 2 ) const ( @@ -44,6 +43,7 @@ const ( MethodRequestNormalize = "request.normalize" MethodRequestInterceptBefore = "request.intercept_before" MethodRequestInterceptAfter = "request.intercept_after" + MethodRequestComplete = "request.complete" MethodResponseTranslate = "response.translate" MethodResponseNormalizeBefore = "response.normalize_before" diff --git a/sdk/pluginabi/types_test.go b/sdk/pluginabi/types_test.go index 3863d1ffc..c594546d7 100644 --- a/sdk/pluginabi/types_test.go +++ b/sdk/pluginabi/types_test.go @@ -27,6 +27,9 @@ func TestEnvelopeRoundTrip(t *testing.T) { } func TestMethodNamesAreStable(t *testing.T) { + if SchemaVersion != 2 { + t.Fatalf("SchemaVersion = %d, want 2", SchemaVersion) + } if MethodPluginRegister != "plugin.register" { t.Fatalf("MethodPluginRegister = %q", MethodPluginRegister) } @@ -36,6 +39,9 @@ func TestMethodNamesAreStable(t *testing.T) { if MethodRequestInterceptAfter != "request.intercept_after" { t.Fatalf("MethodRequestInterceptAfter = %q", MethodRequestInterceptAfter) } + if MethodRequestComplete != "request.complete" { + t.Fatalf("MethodRequestComplete = %q", MethodRequestComplete) + } if MethodResponseInterceptAfter != "response.intercept_after" { t.Fatalf("MethodResponseInterceptAfter = %q", MethodResponseInterceptAfter) } diff --git a/sdk/pluginapi/types.go b/sdk/pluginapi/types.go index ae59a2104..b1edde46b 100644 --- a/sdk/pluginapi/types.go +++ b/sdk/pluginapi/types.go @@ -103,6 +103,8 @@ type Capabilities struct { ResponseAfterTranslator ResponseNormalizer // RequestInterceptor rewrites execution requests before and after credential selection. RequestInterceptor RequestInterceptor + // RequestLifecyclePlugin asynchronously receives one terminal event for each request that reached request interception. + RequestLifecyclePlugin RequestLifecyclePlugin // ResponseInterceptor rewrites successful non-streaming HTTP execution responses before downstream delivery. ResponseInterceptor ResponseInterceptor // StreamChunkInterceptor rewrites successful HTTP stream chunks before downstream delivery. @@ -929,6 +931,11 @@ type RequestInterceptor interface { InterceptRequestAfterAuth(context.Context, RequestInterceptRequest) (RequestInterceptResponse, error) } +// RequestLifecyclePlugin receives asynchronous terminal events after execution finishes, fails, is rejected, or is canceled. +type RequestLifecyclePlugin interface { + HandleRequestComplete(context.Context, RequestCompletion) error +} + // ResponseInterceptor rewrites successful non-streaming execution responses before downstream delivery. type ResponseInterceptor interface { InterceptResponse(context.Context, ResponseInterceptRequest) (ResponseInterceptResponse, error) @@ -976,6 +983,10 @@ type ResponseTransformRequest struct { // RequestInterceptRequest describes a request about to be executed upstream. type RequestInterceptRequest struct { + // RequestID uniquely identifies one model execution and correlates it with RequestCompletion. + RequestID string + // TraceID identifies the parent inbound HTTP request when available. + TraceID string // SourceFormat is the original client protocol format. SourceFormat string // ToFormat is the selected upstream protocol format. It is empty before credential selection. @@ -1002,10 +1013,49 @@ type RequestInterceptResponse struct { Body []byte // ClearHeaders explicitly removes current request headers before Headers is applied. ClearHeaders []string + // Terminate stops the interceptor chain and prevents the request from reaching an upstream executor. + Terminate bool + // StatusCode is the downstream HTTP status used when Terminate is true. Invalid values default to 403. + StatusCode int + // ResponseHeaders contains downstream response headers used when Terminate is true. + ResponseHeaders http.Header + // ResponseBody contains the downstream response body used when Terminate is true. + ResponseBody []byte +} + +// RequestCompletionOutcome identifies how an intercepted request ended. +type RequestCompletionOutcome string + +const ( + // RequestCompletionSucceeded means the request completed successfully. + RequestCompletionSucceeded RequestCompletionOutcome = "succeeded" + // RequestCompletionFailed means model execution failed. + RequestCompletionFailed RequestCompletionOutcome = "failed" + // RequestCompletionRejected means a request interceptor terminated the request before execution. + RequestCompletionRejected RequestCompletionOutcome = "rejected" + // RequestCompletionCanceled means the request context was canceled or the downstream client disconnected. + RequestCompletionCanceled RequestCompletionOutcome = "canceled" +) + +// RequestCompletion describes the terminal state of an intercepted request. +type RequestCompletion struct { + RequestID string + TraceID string + SourceFormat string + Model string + RequestedModel string + Stream bool + Outcome RequestCompletionOutcome + StatusCode int + Error string + StartedAt time.Time + CompletedAt time.Time + Metadata map[string]any } // ResponseInterceptRequest describes a successful non-streaming response. type ResponseInterceptRequest struct { + RequestID string SourceFormat string Model string RequestedModel string @@ -1031,6 +1081,7 @@ type ResponseInterceptResponse struct { // StreamChunkInterceptRequest describes a successful stream chunk before downstream delivery. type StreamChunkInterceptRequest struct { + RequestID string SourceFormat string Model string RequestedModel string diff --git a/sdk/pluginapi/types_test.go b/sdk/pluginapi/types_test.go index de0d5c4e1..0cbd10edd 100644 --- a/sdk/pluginapi/types_test.go +++ b/sdk/pluginapi/types_test.go @@ -24,6 +24,7 @@ var _ RequestNormalizer = (*compileTimePlugin)(nil) var _ ResponseTranslator = (*compileTimePlugin)(nil) var _ ResponseNormalizer = (*compileTimePlugin)(nil) var _ RequestInterceptor = (*compileTimePlugin)(nil) +var _ RequestLifecyclePlugin = (*compileTimePlugin)(nil) var _ ResponseInterceptor = (*compileTimePlugin)(nil) var _ StreamChunkInterceptor = (*compileTimePlugin)(nil) var _ ThinkingApplier = (*compileTimePlugin)(nil) @@ -518,6 +519,8 @@ func (compileTimePlugin) InterceptRequestAfterAuth(context.Context, RequestInter return RequestInterceptResponse{}, nil } +func (compileTimePlugin) HandleRequestComplete(context.Context, RequestCompletion) error { return nil } + func (compileTimePlugin) InterceptResponse(context.Context, ResponseInterceptRequest) (ResponseInterceptResponse, error) { return ResponseInterceptResponse{}, nil }