From 1ca048abdc6af78c1d8ae3381ce5bf380976e2e7 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Wed, 10 Jun 2026 20:58:59 +0800 Subject: [PATCH] feat(auth, interceptor, jshandler): add post-auth request interceptors and enhance format handling - Introduced `applyRequestAfterAuthInterceptor` to modify requests after credential selection and before executor translation. - Added `InterceptRequestAfterAuth` method across plugin adapters with corresponding tests for context validation. - Enhanced format resolution logic (`requestToFormat`) to support additional providers and formats. - Updated JavaScript handler to include a new `on_after_auth_request` hook for post-auth request handling. - Refactored interceptor methods for clarity and better encapsulation of request/response lifecycles. --- examples/plugin/jshandler/README.md | 30 ++- examples/plugin/jshandler/abi.go | 9 +- examples/plugin/jshandler/go.mod | 2 + examples/plugin/jshandler/interceptor.go | 34 +-- examples/plugin/jshandler/interceptor_test.go | 50 ++++- .../jshandler/scripts/copilot_handler.js | 8 + internal/pluginhost/adapters.go | 53 ++++- internal/pluginhost/adapters_test.go | 67 +++++- internal/pluginhost/host_test.go | 30 ++- internal/pluginhost/rpc_client.go | 11 +- internal/pluginhost/test_helpers_test.go | 24 ++- sdk/api/handlers/handlers.go | 194 +++++++++++++++--- .../handlers/handlers_interceptors_test.go | 144 ++++++++++++- sdk/cliproxy/auth/conductor.go | 106 +++++++++- sdk/cliproxy/executor/types.go | 36 ++++ sdk/pluginabi/types.go | 1 + sdk/pluginabi/types_test.go | 3 + sdk/pluginapi/types.go | 28 ++- sdk/pluginapi/types_test.go | 6 +- 19 files changed, 744 insertions(+), 92 deletions(-) diff --git a/examples/plugin/jshandler/README.md b/examples/plugin/jshandler/README.md index e9b5aca4..e69264da 100644 --- a/examples/plugin/jshandler/README.md +++ b/examples/plugin/jshandler/README.md @@ -4,7 +4,7 @@ A CLIProxyAPI plugin that executes external JavaScript scripts to intercept and ## Features -- **Request Interception** (`on_before_request`): Modify request payloads and headers before upstream delivery. +- **Request Interception** (`on_before_request`, `on_after_auth_request`): Modify request payloads and headers before and after credential selection. - **Response Interception** (`on_after_nonstream_response`): Modify non-streaming response bodies and headers. - **Stream Chunk Interception** (`on_after_stream_response`): Modify individual streaming chunks with read-only `history_chunks` context. - **Hot Reload**: Scripts are automatically reloaded when modified on disk. @@ -40,7 +40,7 @@ Scripts can export these global functions: ### `on_before_request(ctx)` -Called before the request is sent upstream. +Called before credential selection. At this point the target upstream protocol is not selected yet. **ctx structure:** ```javascript @@ -50,7 +50,31 @@ Called before the request is sent upstream. "headers": {}, // Request headers "url": "", "model": "gpt-4", - "protocol": "openai" + "protocol": "openai", + "source_format": "openai", + "sourceFormat": "openai", + "to_format": "", + "toFormat": "" +} +``` + +### `on_after_auth_request(ctx)` + +Called after credential selection and before request translation, request normalization, and built-in payload configuration. + +**ctx structure:** +```javascript +{ + "id": "request-id", + "body": "...", // Request body string + "headers": {}, // Request headers + "url": "", + "model": "gpt-4", + "protocol": "openai", // Same as source_format + "source_format": "openai", + "sourceFormat": "openai", + "to_format": "codex", + "toFormat": "codex" } ``` diff --git a/examples/plugin/jshandler/abi.go b/examples/plugin/jshandler/abi.go index 59c30c88..39f506a3 100644 --- a/examples/plugin/jshandler/abi.go +++ b/examples/plugin/jshandler/abi.go @@ -218,7 +218,14 @@ func handleJSHandlerABIMethod(ctx context.Context, method string, request []byte if errDecode := json.Unmarshal(request, &req); errDecode != nil { return nil, errDecode } - resp, errCall := p.interceptRequest(ctx, req.RequestInterceptRequest, req.HostCallbackID) + resp, errCall := p.interceptRequest(ctx, req.RequestInterceptRequest, "on_before_request", req.HostCallbackID) + return abiOKEnvelopeWithError(resp, errCall) + case pluginabi.MethodRequestInterceptAfter: + var req abiRequestInterceptRequest + if errDecode := json.Unmarshal(request, &req); errDecode != nil { + return nil, errDecode + } + resp, errCall := p.interceptRequest(ctx, req.RequestInterceptRequest, "on_after_auth_request", req.HostCallbackID) return abiOKEnvelopeWithError(resp, errCall) case pluginabi.MethodResponseInterceptAfter: var req abiResponseInterceptRequest diff --git a/examples/plugin/jshandler/go.mod b/examples/plugin/jshandler/go.mod index 1cd5f78d..33f4c6bc 100644 --- a/examples/plugin/jshandler/go.mod +++ b/examples/plugin/jshandler/go.mod @@ -16,3 +16,5 @@ require ( golang.org/x/sys v0.38.0 // indirect golang.org/x/text v0.31.0 // indirect ) + +replace github.com/router-for-me/CLIProxyAPI/v7 => ../../.. diff --git a/examples/plugin/jshandler/interceptor.go b/examples/plugin/jshandler/interceptor.go index d866b7cf..3a33a418 100644 --- a/examples/plugin/jshandler/interceptor.go +++ b/examples/plugin/jshandler/interceptor.go @@ -44,11 +44,15 @@ func (p *jsHandlerPlugin) allScriptPaths() []string { return paths } -func (p *jsHandlerPlugin) InterceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { - return p.interceptRequest(ctx, req, "") +func (p *jsHandlerPlugin) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return p.interceptRequest(ctx, req, "on_before_request", "") } -func (p *jsHandlerPlugin) interceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest, hostCallbackID string) (pluginapi.RequestInterceptResponse, error) { +func (p *jsHandlerPlugin) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return p.interceptRequest(ctx, req, "on_after_auth_request", "") +} + +func (p *jsHandlerPlugin) interceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest, hookName, hostCallbackID string) (pluginapi.RequestInterceptResponse, error) { resp := pluginapi.RequestInterceptResponse{} scriptPaths := p.allScriptPaths() if len(scriptPaths) == 0 { @@ -64,7 +68,7 @@ func (p *jsHandlerPlugin) interceptRequest(ctx context.Context, req pluginapi.Re if scriptPath == "" { continue } - processed, cleared, errJS := p.applyJSBeforeRequest(scriptPath, []byte(body), req.Model, req.SourceFormat, headers, hostCallbackID) + processed, cleared, errJS := p.applyJSRequestHook(scriptPath, hookName, []byte(body), req.Model, req.SourceFormat, req.ToFormat, headers, hostCallbackID) if errJS != nil { log.Warnf("failed to execute JS request interceptor [%s]: %v", scriptPath, errJS) continue @@ -197,7 +201,7 @@ func (p *jsHandlerPlugin) interceptStreamChunk(ctx context.Context, req pluginap return resp, nil } -func (p *jsHandlerPlugin) applyJSBeforeRequest(scriptPath string, payloadBytes []byte, model, protocol string, headers http.Header, hostCallbackID string) ([]byte, []string, error) { +func (p *jsHandlerPlugin) applyJSRequestHook(scriptPath, hookName string, payloadBytes []byte, model, sourceFormat, toFormat string, headers http.Header, hostCallbackID string) ([]byte, []string, error) { program, err := getJSProgram(scriptPath) if err != nil { return nil, nil, err @@ -211,20 +215,24 @@ func (p *jsHandlerPlugin) applyJSBeforeRequest(scriptPath string, payloadBytes [ headersMap := headerToAnyMap(headers) jsCtx := map[string]any{ - "id": generateRequestID(), - "body": string(payloadBytes), - "headers": headersMap, - "url": "", - "model": model, - "protocol": protocol, + "id": generateRequestID(), + "body": string(payloadBytes), + "headers": headersMap, + "url": "", + "model": model, + "protocol": sourceFormat, + "source_format": sourceFormat, + "to_format": toFormat, + "sourceFormat": sourceFormat, + "toFormat": toFormat, } - jsVal, errCall := engine.callFunction("on_before_request", p.cfg.Timeout, jsCtx) + jsVal, errCall := engine.callFunction(hookName, p.cfg.Timeout, jsCtx) if errCall != nil { if errors.Is(errCall, ErrFunctionNotFound) { return payloadBytes, nil, nil } - return nil, nil, fmt.Errorf("on_before_request failed for %s: %w", scriptPath, errCall) + return nil, nil, fmt.Errorf("%s failed for %s: %w", hookName, scriptPath, errCall) } if jsVal == nil || goja.IsUndefined(jsVal) || goja.IsNull(jsVal) { diff --git a/examples/plugin/jshandler/interceptor_test.go b/examples/plugin/jshandler/interceptor_test.go index 37447362..cc1810a1 100644 --- a/examples/plugin/jshandler/interceptor_test.go +++ b/examples/plugin/jshandler/interceptor_test.go @@ -25,16 +25,18 @@ function on_before_request(ctx) { plugin := &jsHandlerPlugin{cfg: defaultJSHandlerConfig()} headers := http.Header{"X-Plugin": []string{"original"}} - processed, _, errApply := plugin.applyJSBeforeRequest( + processed, _, errApply := plugin.applyJSRequestHook( scriptPath, + "on_before_request", []byte(`{"messages":[{"role":"user","content":"contains sensitive_word"}]}`), "gpt-test", "openai", + "", headers, "", ) if errApply != nil { - t.Fatalf("applyJSBeforeRequest() error = %v", errApply) + t.Fatalf("applyJSRequestHook() error = %v", errApply) } if body := string(processed); !strings.Contains(body, "safe_word") || strings.Contains(body, "sensitive_word") { t.Fatalf("processed body = %q, want sensitive word rewritten", body) @@ -44,6 +46,50 @@ function on_before_request(ctx) { } } +func TestApplyJSAfterAuthRequestReceivesFormats(t *testing.T) { + scriptPath := filepath.Join(t.TempDir(), "after_auth.js") + script := ` +function on_after_auth_request(ctx) { + if (ctx.source_format !== "openai" || ctx.to_format !== "codex") { + throw new Error("unexpected formats: " + ctx.source_format + " -> " + ctx.to_format); + } + if (ctx.sourceFormat !== "openai" || ctx.toFormat !== "codex") { + throw new Error("unexpected camel formats: " + ctx.sourceFormat + " -> " + ctx.toFormat); + } + var req = JSON.parse(ctx.body); + req.after_auth = ctx.source_format + "_to_" + ctx.to_format; + ctx.headers["X-Protocol"] = req.after_auth; + ctx.body = JSON.stringify(req); + return ctx; +} +` + if errWrite := os.WriteFile(scriptPath, []byte(script), 0600); errWrite != nil { + t.Fatalf("os.WriteFile() error = %v", errWrite) + } + + plugin := &jsHandlerPlugin{cfg: defaultJSHandlerConfig()} + headers := http.Header{} + processed, _, errApply := plugin.applyJSRequestHook( + scriptPath, + "on_after_auth_request", + []byte(`{"model":"gpt-test"}`), + "gpt-test", + "openai", + "codex", + headers, + "", + ) + if errApply != nil { + t.Fatalf("applyJSRequestHook() error = %v", errApply) + } + if body := string(processed); !strings.Contains(body, `"after_auth":"openai_to_codex"`) { + t.Fatalf("processed body = %q, want after_auth marker", body) + } + if got := headers.Get("X-Protocol"); got != "openai_to_codex" { + t.Fatalf("header X-Protocol = %q, want openai_to_codex", got) + } +} + func TestApplyJSAfterResponseUsesFrozenNativeHistoryChunks(t *testing.T) { scriptPath := filepath.Join(t.TempDir(), "stream.js") script := ` diff --git a/examples/plugin/jshandler/scripts/copilot_handler.js b/examples/plugin/jshandler/scripts/copilot_handler.js index 6d50fff2..81831630 100644 --- a/examples/plugin/jshandler/scripts/copilot_handler.js +++ b/examples/plugin/jshandler/scripts/copilot_handler.js @@ -17,6 +17,14 @@ function on_before_request(ctx) { return ctx; } +function on_after_auth_request(ctx) { + console.log("[" + ctx.id + "] Selected request protocol: " + ctx.source_format + " -> " + ctx.to_format); + if (ctx.source_format === "openai" && ctx.to_format === "codex") { + ctx.headers["X-JS-Handler-Protocol"] = "openai-to-codex"; + } + return ctx; +} + function parse_stream_chunk(chunk) { var leading = ""; var payload = chunk.trim(); diff --git a/internal/pluginhost/adapters.go b/internal/pluginhost/adapters.go index 3c564546..5be00358 100644 --- a/internal/pluginhost/adapters.go +++ b/internal/pluginhost/adapters.go @@ -511,18 +511,18 @@ func (h *Host) callModelsForAuth(ctx context.Context, record capabilityRecord, p }) } -func (h *Host) callRequestInterceptor(ctx context.Context, pluginID string, interceptor pluginapi.RequestInterceptor, req pluginapi.RequestInterceptRequest) (out pluginapi.RequestInterceptResponse, ok bool) { - if h == nil || interceptor == nil || h.isPluginFused(pluginID) { +func (h *Host) callRequestInterceptor(ctx context.Context, pluginID, method string, call func(context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error), req pluginapi.RequestInterceptRequest) (out pluginapi.RequestInterceptResponse, ok bool) { + if h == nil || call == nil || h.isPluginFused(pluginID) { return pluginapi.RequestInterceptResponse{}, false } defer func() { if recovered := recover(); recovered != nil { - h.fusePlugin(pluginID, "RequestInterceptor.InterceptRequest", recovered) + h.fusePlugin(pluginID, method, recovered) out = pluginapi.RequestInterceptResponse{} ok = false } }() - resp, errIntercept := interceptor.InterceptRequest(ctx, req) + resp, errIntercept := call(ctx, req) if errIntercept != nil { log.Warnf("pluginhost: request interceptor %s failed: %v", pluginID, errIntercept) return pluginapi.RequestInterceptResponse{}, false @@ -568,7 +568,19 @@ func (h *Host) callStreamChunkInterceptor(ctx context.Context, pluginID string, return resp, true } -func (h *Host) InterceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { +func (h *Host) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + return h.interceptRequest(ctx, req, "RequestInterceptor.InterceptRequestBeforeAuth", func(interceptor pluginapi.RequestInterceptor, ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return interceptor.InterceptRequestBeforeAuth(ctx, req) + }) +} + +func (h *Host) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + return h.interceptRequest(ctx, req, "RequestInterceptor.InterceptRequestAfterAuth", func(interceptor pluginapi.RequestInterceptor, ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return interceptor.InterceptRequestAfterAuth(ctx, req) + }) +} + +func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest, method string, invoke func(pluginapi.RequestInterceptor, context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error)) pluginapi.RequestInterceptResponse { current := pluginapi.RequestInterceptResponse{ Headers: cloneHeader(req.Headers), Body: bytes.Clone(req.Body), @@ -582,7 +594,9 @@ func (h *Host) InterceptRequest(ctx context.Context, req pluginapi.RequestInterc nextReq.Headers = cloneHeader(current.Headers) nextReq.Body = bytes.Clone(current.Body) nextReq.Metadata = cloneInterceptorMetadata(req.Metadata) - if resp, ok := h.callRequestInterceptor(ctx, record.id, interceptor, nextReq); ok { + if resp, ok := h.callRequestInterceptor(ctx, record.id, method, func(callCtx context.Context, callReq pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return invoke(interceptor, callCtx, callReq) + }, nextReq); ok { current.Headers = mergeHeaders(current.Headers, resp.Headers, resp.ClearHeaders) if len(resp.Body) > 0 { current.Body = bytes.Clone(resp.Body) @@ -665,6 +679,21 @@ func (h *Host) HasStreamInterceptors() bool { return false } +func (h *Host) HasRequestInterceptors() bool { + if h == nil { + return false + } + for _, record := range h.Snapshot().records { + if h.isPluginFused(record.id) { + continue + } + if record.plugin.Capabilities.RequestInterceptor != nil { + return true + } + } + return false +} + func (h *Host) commitModelClients(snap *Snapshot, modelRegistry modelRegistry, registrations []modelClientRegistration, nextClients map[string]struct{}, nextProviders map[string]string, nextModelRegistrations map[string]pluginModelRegistration) { if h == nil || modelRegistry == nil { return @@ -1311,6 +1340,18 @@ func (a *executorAdapter) prepareExecutorCall(req coreexecutor.Request, opts cor }, nil } +func (a *executorAdapter) RequestToFormat(req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format { + if a == nil { + return "" + } + requestedFormat := executorRequestedFormat(req, opts) + inputFormat, errInput := a.selectExecutorInputFormat(requestedFormat) + if errInput != nil { + return "" + } + return inputFormat +} + func executorRequestedFormat(req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format { if opts.SourceFormat != "" { return normalizeExecutorFormatName(opts.SourceFormat.String()) diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index e7ae6d56..e718b57d 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -1271,7 +1271,7 @@ func TestInterceptRequestChainsByPriorityAndHeaders(t *testing.T) { ) headers := http.Header{"X-Remove": []string{"yes"}} - got := host.InterceptRequest(context.Background(), pluginapi.RequestInterceptRequest{ + got := host.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{ SourceFormat: "openai", Model: "normalized", RequestedModel: "requested", @@ -1291,6 +1291,31 @@ func TestInterceptRequestChainsByPriorityAndHeaders(t *testing.T) { } } +func TestInterceptRequestAfterAuthPassesTargetFormat(t *testing.T) { + host := newHostWithRecords(capabilityRecord{ + id: "after", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + RequestInterceptor: requestInterceptorFunc(func(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + if req.SourceFormat != "openai" || req.ToFormat != "codex" { + t.Fatalf("request formats = %q -> %q, want openai -> codex", req.SourceFormat, req.ToFormat) + } + return pluginapi.RequestInterceptResponse{Body: append(req.Body, []byte("|after")...)}, nil + }), + }}, + }) + + got := host.InterceptRequestAfterAuth(context.Background(), pluginapi.RequestInterceptRequest{ + SourceFormat: "openai", + ToFormat: "codex", + Model: "gpt-5.4", + Body: []byte("body"), + }) + + if string(got.Body) != "body|after" { + t.Fatalf("body = %q, want body|after", got.Body) + } +} + func TestResponseInterceptorsChainAndStreamHistory(t *testing.T) { var seenHistory [][]byte var sawSecondResponse bool @@ -1435,7 +1460,7 @@ func TestInterceptorsSkipErrorsAndFusePanics(t *testing.T) { }, ) - got := host.InterceptRequest(context.Background(), pluginapi.RequestInterceptRequest{Body: []byte("body")}) + got := host.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{Body: []byte("body")}) if string(got.Body) != "body|success" { t.Fatalf("body = %q, want body|success", got.Body) } @@ -1548,6 +1573,40 @@ func TestHasStreamInterceptorsReflectsActiveStreamInterceptors(t *testing.T) { } } +func TestHasRequestInterceptorsReflectsActiveRequestInterceptors(t *testing.T) { + responseOnly := newHostWithRecords(capabilityRecord{ + id: "response", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + ResponseInterceptor: responseInterceptorFunc{ + interceptResponse: func(ctx context.Context, req pluginapi.ResponseInterceptRequest) (pluginapi.ResponseInterceptResponse, error) { + return pluginapi.ResponseInterceptResponse{Body: req.Body}, nil + }, + }, + }}, + }) + if responseOnly.HasRequestInterceptors() { + t.Fatal("HasRequestInterceptors() = true, want false for response-only plugins") + } + + requestHost := newHostWithRecords(capabilityRecord{ + id: "request", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + RequestInterceptor: requestInterceptorFunc(func(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return pluginapi.RequestInterceptResponse{Body: req.Body}, nil + }), + }}, + }) + if !requestHost.HasRequestInterceptors() { + t.Fatal("HasRequestInterceptors() = false, want true for request interceptors") + } + requestHost.mu.Lock() + requestHost.fused["request"] = "test fused" + requestHost.mu.Unlock() + if requestHost.HasRequestInterceptors() { + t.Fatal("HasRequestInterceptors() = true, want false after request plugin is fused") + } +} + func TestInterceptorsDoNotMutateInputs(t *testing.T) { t.Run("request", func(t *testing.T) { headers := http.Header{"X-Request": []string{"input"}} @@ -1587,7 +1646,7 @@ func TestInterceptorsDoNotMutateInputs(t *testing.T) { }}, }) - got := host.InterceptRequest(context.Background(), pluginapi.RequestInterceptRequest{ + got := host.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{ Headers: headers, Body: body, Metadata: metadata, @@ -1844,7 +1903,7 @@ func TestInterceptorsDoNotMutateInputs(t *testing.T) { }}, }) - _ = host.InterceptRequest(context.Background(), pluginapi.RequestInterceptRequest{Metadata: metadata}) + _ = host.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{Metadata: metadata}) if structValue.Value != "original" || structValue.Items[0] != "original" { t.Fatalf("struct pointer metadata mutated: %#v", structValue) diff --git a/internal/pluginhost/host_test.go b/internal/pluginhost/host_test.go index 90d2a761..72dc9362 100644 --- a/internal/pluginhost/host_test.go +++ b/internal/pluginhost/host_test.go @@ -176,12 +176,12 @@ func TestHostApplyConfigDispatchesInterceptorRPCMethods(t *testing.T) { } caps := h.Snapshot().records[0].plugin.Capabilities - reqResp, errReq := caps.RequestInterceptor.InterceptRequest(context.Background(), pluginapi.RequestInterceptRequest{Body: []byte("request")}) + reqResp, errReq := caps.RequestInterceptor.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{Body: []byte("request")}) if errReq != nil { - t.Fatalf("InterceptRequest() error = %v", errReq) + t.Fatalf("InterceptRequestBeforeAuth() error = %v", errReq) } if got := string(reqResp.Body); got != "request|rpc" { - t.Fatalf("InterceptRequest() body = %q, want request|rpc", got) + t.Fatalf("InterceptRequestBeforeAuth() body = %q, want request|rpc", got) } respResp, errResp := caps.ResponseInterceptor.InterceptResponse(context.Background(), pluginapi.ResponseInterceptRequest{Body: []byte("response")}) @@ -202,8 +202,11 @@ func TestHostApplyConfigDispatchesInterceptorRPCMethods(t *testing.T) { } func TestInterceptorHelpersReturnErrorsWhenCallbackMissing(t *testing.T) { - if _, errReq := (requestInterceptorFunc(nil)).InterceptRequest(context.Background(), pluginapi.RequestInterceptRequest{}); errReq == nil { - t.Fatal("InterceptRequest() error = nil, want missing request interceptor callback") + if _, errReq := (requestInterceptorFunc(nil)).InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{}); errReq == nil { + t.Fatal("InterceptRequestBeforeAuth() error = nil, want missing request interceptor callback") + } + if _, errReq := (requestInterceptorFunc(nil)).InterceptRequestAfterAuth(context.Background(), pluginapi.RequestInterceptRequest{}); errReq == nil { + t.Fatal("InterceptRequestAfterAuth() error = nil, want missing request interceptor callback") } if _, errResp := (responseInterceptorFunc{interceptResponse: nil}).InterceptResponse(context.Background(), pluginapi.ResponseInterceptRequest{}); errResp == nil { t.Fatal("InterceptResponse() error = nil, want missing response interceptor callback") @@ -220,15 +223,26 @@ func TestRPCInterceptorsIncludeHostCallbackID(t *testing.T) { client: client, } - if _, errReq := adapter.InterceptRequest(context.Background(), pluginapi.RequestInterceptRequest{Body: []byte("request")}); errReq != nil { - t.Fatalf("InterceptRequest() error = %v", errReq) + if _, errReq := adapter.InterceptRequestBeforeAuth(context.Background(), pluginapi.RequestInterceptRequest{Body: []byte("request")}); errReq != nil { + t.Fatalf("InterceptRequestBeforeAuth() error = %v", errReq) } var req rpcRequestInterceptRequest if errDecode := json.Unmarshal(client.requests[pluginabi.MethodRequestInterceptBefore], &req); errDecode != nil { t.Fatalf("decode request interceptor request: %v", errDecode) } if req.HostCallbackID == "" { - t.Fatal("request interceptor host_callback_id is empty") + t.Fatal("request interceptor before-auth host_callback_id is empty") + } + + if _, errReq := adapter.InterceptRequestAfterAuth(context.Background(), pluginapi.RequestInterceptRequest{Body: []byte("request")}); errReq != nil { + t.Fatalf("InterceptRequestAfterAuth() error = %v", errReq) + } + var reqAfter rpcRequestInterceptRequest + if errDecode := json.Unmarshal(client.requests[pluginabi.MethodRequestInterceptAfter], &reqAfter); errDecode != nil { + t.Fatalf("decode after-auth request interceptor request: %v", errDecode) + } + if reqAfter.HostCallbackID == "" { + t.Fatal("request interceptor after-auth host_callback_id is empty") } if _, errResp := adapter.InterceptResponse(context.Background(), pluginapi.ResponseInterceptRequest{Body: []byte("response")}); errResp != nil { diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 4519beff..ff69bb20 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -432,7 +432,7 @@ func (a *rpcPluginAdapter) NormalizeRequest(ctx context.Context, req pluginapi.R return callPlugin[pluginapi.PayloadResponse](ctx, a.client, pluginabi.MethodRequestNormalize, req) } -func (a *rpcPluginAdapter) InterceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { +func (a *rpcPluginAdapter) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { callbackID, closeCallback := a.openHostCallbackContext(ctx) defer closeCallback() return callPlugin[pluginapi.RequestInterceptResponse](ctx, a.client, pluginabi.MethodRequestInterceptBefore, rpcRequestInterceptRequest{ @@ -441,6 +441,15 @@ func (a *rpcPluginAdapter) InterceptRequest(ctx context.Context, req pluginapi.R }) } +func (a *rpcPluginAdapter) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + callbackID, closeCallback := a.openHostCallbackContext(ctx) + defer closeCallback() + return callPlugin[pluginapi.RequestInterceptResponse](ctx, a.client, pluginabi.MethodRequestInterceptAfter, rpcRequestInterceptRequest{ + RequestInterceptRequest: req, + HostCallbackID: callbackID, + }) +} + 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/test_helpers_test.go b/internal/pluginhost/test_helpers_test.go index e7b0fbd7..81289eb2 100644 --- a/internal/pluginhost/test_helpers_test.go +++ b/internal/pluginhost/test_helpers_test.go @@ -71,7 +71,20 @@ func (l *testSymbolLookup) Call(ctx context.Context, method string, request []by if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { return nil, errUnmarshal } - resp, errIntercept := l.active.Capabilities.RequestInterceptor.InterceptRequest(ctx, req) + resp, errIntercept := l.active.Capabilities.RequestInterceptor.InterceptRequestBeforeAuth(ctx, req) + if errIntercept != nil { + return nil, errIntercept + } + return marshalRPCResult(resp) + case pluginabi.MethodRequestInterceptAfter: + if l.active.Capabilities.RequestInterceptor == nil { + return nil, fmt.Errorf("missing request interceptor") + } + var req pluginapi.RequestInterceptRequest + if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { + return nil, errUnmarshal + } + resp, errIntercept := l.active.Capabilities.RequestInterceptor.InterceptRequestAfterAuth(ctx, req) if errIntercept != nil { return nil, errIntercept } @@ -231,7 +244,14 @@ func (c testThinkingCapability) ApplyThinking(ctx context.Context, req pluginapi type requestInterceptorFunc func(context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) -func (f requestInterceptorFunc) InterceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { +func (f requestInterceptorFunc) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + if f == nil { + return pluginapi.RequestInterceptResponse{}, fmt.Errorf("missing request interceptor callback") + } + return f(ctx, req) +} + +func (f requestInterceptorFunc) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { if f == nil { return pluginapi.RequestInterceptResponse{}, fmt.Errorf("missing request interceptor callback") } diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index d30b01ec..42756dc0 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -66,7 +66,8 @@ type disallowFreeAuthContextKey struct{} // PluginInterceptorHost applies plugin interceptors around handler execution. type PluginInterceptorHost interface { - InterceptRequest(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse + InterceptRequestBeforeAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse + InterceptRequestAfterAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse InterceptResponse(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse InterceptStreamChunk(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse } @@ -75,6 +76,10 @@ type streamInterceptorDetector interface { HasStreamInterceptors() bool } +type requestInterceptorDetector interface { + HasRequestInterceptors() bool +} + // WithPinnedAuthID returns a child context that requests execution on a specific auth ID. func WithPinnedAuthID(ctx context.Context, authID string) context.Context { authID = strings.TrimSpace(authID) @@ -639,15 +644,17 @@ func (h *BaseAPIHandler) executeWithAuthManager(ctx context.Context, handlerType Model: normalizedModel, Payload: payload, } + afterAuthCapture := &requestAfterAuthCapture{} opts := coreexecutor.Options{ - Stream: false, - Alt: alt, - OriginalRequest: rawJSON, - SourceFormat: sdktranslator.FromString(handlerType), - Headers: headersFromContext(ctx), + Stream: false, + Alt: alt, + OriginalRequest: rawJSON, + SourceFormat: sdktranslator.FromString(handlerType), + Headers: headersFromContext(ctx), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture), } opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptors(ctx, handlerType, modelName, req, opts) + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, modelName, req, opts) resp, err := h.AuthManager.Execute(ctx, providers, req, opts) if err != nil { err = enrichAuthSelectionError(err, providers, normalizedModel) @@ -665,9 +672,10 @@ func (h *BaseAPIHandler) executeWithAuthManager(ctx context.Context, handlerType } return nil, nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} } + executedReq, executedOpts := afterAuthCapture.apply(req, opts) rawResponseHeaders := cloneHeader(resp.Headers) responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, modelName, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK) + body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, modelName, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK) return body, responseHeaders, nil } @@ -690,15 +698,17 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle Model: normalizedModel, Payload: payload, } + afterAuthCapture := &requestAfterAuthCapture{} opts := coreexecutor.Options{ - Stream: false, - Alt: alt, - OriginalRequest: rawJSON, - SourceFormat: sdktranslator.FromString(handlerType), - Headers: headersFromContext(ctx), + Stream: false, + Alt: alt, + OriginalRequest: rawJSON, + SourceFormat: sdktranslator.FromString(handlerType), + Headers: headersFromContext(ctx), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture), } opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptors(ctx, handlerType, modelName, req, opts) + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, modelName, req, opts) resp, err := h.AuthManager.ExecuteCount(ctx, providers, req, opts) if err != nil { err = enrichAuthSelectionError(err, providers, normalizedModel) @@ -716,9 +726,10 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle } return nil, nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} } + executedReq, executedOpts := afterAuthCapture.apply(req, opts) rawResponseHeaders := cloneHeader(resp.Headers) responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, modelName, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK) + body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, modelName, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK) return body, responseHeaders, nil } @@ -754,15 +765,17 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl Model: normalizedModel, Payload: payload, } + afterAuthCapture := &requestAfterAuthCapture{} opts := coreexecutor.Options{ - Stream: true, - Alt: alt, - OriginalRequest: rawJSON, - SourceFormat: sdktranslator.FromString(handlerType), - Headers: headersFromContext(ctx), + Stream: true, + Alt: alt, + OriginalRequest: rawJSON, + SourceFormat: sdktranslator.FromString(handlerType), + Headers: headersFromContext(ctx), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture), } opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptors(ctx, handlerType, modelName, req, opts) + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, modelName, req, opts) streamResult, err := h.AuthManager.ExecuteStream(ctx, providers, req, opts) if err != nil { err = enrichAuthSelectionError(err, providers, normalizedModel) @@ -783,6 +796,9 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl close(errChan) return nil, nil, errChan } + executedRequest := func() (coreexecutor.Request, coreexecutor.Options) { + return afterAuthCapture.apply(req, opts) + } passthroughHeadersEnabled := PassthroughHeadersEnabled(h.Cfg) interceptorHost := h.interceptorHost() streamInterceptorsActive := streamInterceptorsEnabled(interceptorHost) @@ -813,16 +829,17 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl if !streamInterceptorsActive || streamHeaderInitialized { return } + executedReq, executedOpts := executedRequest() intercepted := interceptorHost.InterceptStreamChunk(ctx, pluginapi.StreamChunkInterceptRequest{ SourceFormat: handlerType, Model: normalizedModel, RequestedModel: modelName, - RequestHeaders: cloneHeader(opts.Headers), + RequestHeaders: cloneHeader(executedOpts.Headers), ResponseHeaders: cloneHeader(rawStreamHeaders), - OriginalRequest: cloneBytes(opts.OriginalRequest), - RequestBody: cloneBytes(req.Payload), + OriginalRequest: cloneBytes(executedOpts.OriginalRequest), + RequestBody: cloneBytes(executedReq.Payload), ChunkIndex: pluginapi.StreamChunkHeaderInitIndex, - Metadata: opts.Metadata, + Metadata: executedOpts.Metadata, }) applyStreamHeaders(intercepted.Headers) streamHeaderInitialized = true @@ -967,18 +984,19 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl applyStreamHeaderInit() payload := cloneBytes(chunk.Payload) if streamInterceptorsActive { + executedReq, executedOpts := executedRequest() intercepted := interceptorHost.InterceptStreamChunk(ctx, pluginapi.StreamChunkInterceptRequest{ SourceFormat: handlerType, Model: normalizedModel, RequestedModel: modelName, - RequestHeaders: cloneHeader(opts.Headers), + RequestHeaders: cloneHeader(executedOpts.Headers), ResponseHeaders: cloneHeader(rawStreamHeaders), - OriginalRequest: cloneBytes(opts.OriginalRequest), - RequestBody: cloneBytes(req.Payload), + OriginalRequest: cloneBytes(executedOpts.OriginalRequest), + RequestBody: cloneBytes(executedReq.Payload), Body: payload, HistoryChunks: cloneByteSlices(historyChunks), ChunkIndex: chunkIndex, - Metadata: opts.Metadata, + Metadata: executedOpts.Metadata, }) applyStreamHeaders(intercepted.Headers) if len(intercepted.Body) > 0 { @@ -1287,12 +1305,91 @@ func streamInterceptorsEnabled(host PluginInterceptorHost) bool { return true } -func (h *BaseAPIHandler) applyRequestInterceptors(ctx context.Context, handlerType, requestedModel string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Request, coreexecutor.Options) { +func requestInterceptorsEnabled(host PluginInterceptorHost) bool { + if host == nil { + return false + } + if detector, ok := host.(requestInterceptorDetector); ok { + return detector.HasRequestInterceptors() + } + return true +} + +type requestAfterAuthCapture struct { + mu sync.Mutex + set bool + headers http.Header + body []byte + originalRequest []byte + originalRequestReplaced bool +} + +func (c *requestAfterAuthCapture) record(req coreexecutor.RequestAfterAuthInterceptRequest, resp coreexecutor.RequestAfterAuthInterceptResponse) { + if c == nil { + return + } + headers := mergeRequestInterceptorHeaders(req.Headers, resp.Headers, resp.ClearHeaders) + body := cloneBytes(req.Body) + var originalRequest []byte + originalRequestReplaced := false + if len(resp.Body) > 0 { + body = cloneBytes(resp.Body) + originalRequest = cloneBytes(resp.Body) + originalRequestReplaced = true + } + + c.mu.Lock() + defer c.mu.Unlock() + c.set = true + c.headers = headers + c.body = body + c.originalRequest = originalRequest + c.originalRequestReplaced = originalRequestReplaced +} + +func (c *requestAfterAuthCapture) apply(req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Request, coreexecutor.Options) { + if c == nil { + return req, opts + } + c.mu.Lock() + defer c.mu.Unlock() + if !c.set { + return req, opts + } + req.Payload = cloneBytes(c.body) + opts.Headers = cloneHeader(c.headers) + if c.originalRequestReplaced { + opts.OriginalRequest = cloneBytes(c.originalRequest) + } + return req, opts +} + +func mergeRequestInterceptorHeaders(current, updates http.Header, clear []string) http.Header { + if updates == nil && len(clear) == 0 { + return cloneHeader(current) + } + out := cloneHeader(current) + if out == nil && (len(updates) > 0 || len(clear) > 0) { + out = make(http.Header) + } + for _, key := range clear { + out.Del(key) + } + for key, values := range updates { + out.Del(key) + for _, value := range values { + out.Add(key, value) + } + } + return out +} + +func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, handlerType, requestedModel string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Request, coreexecutor.Options) { host := h.interceptorHost() if host == nil { return req, opts } - resp := host.InterceptRequest(ctx, pluginapi.RequestInterceptRequest{ + resp := host.InterceptRequestBeforeAuth(ctx, pluginapi.RequestInterceptRequest{ SourceFormat: handlerType, Model: req.Model, RequestedModel: requestedModel, @@ -1309,6 +1406,41 @@ func (h *BaseAPIHandler) applyRequestInterceptors(ctx context.Context, handlerTy return req, opts } +func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCapture) coreexecutor.RequestAfterAuthInterceptor { + if !requestInterceptorsEnabled(h.interceptorHost()) { + return nil + } + return func(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest) coreexecutor.RequestAfterAuthInterceptResponse { + resp := h.applyRequestInterceptorsAfterAuth(ctx, req) + if capture != nil { + capture.record(req, resp) + } + return resp + } +} + +func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest) coreexecutor.RequestAfterAuthInterceptResponse { + host := h.interceptorHost() + if !requestInterceptorsEnabled(host) { + return coreexecutor.RequestAfterAuthInterceptResponse{} + } + resp := host.InterceptRequestAfterAuth(ctx, pluginapi.RequestInterceptRequest{ + SourceFormat: req.SourceFormat.String(), + ToFormat: req.ToFormat.String(), + Model: req.Model, + RequestedModel: req.RequestedModel, + Stream: req.Stream, + Headers: cloneHeader(req.Headers), + Body: cloneBytes(req.Body), + Metadata: req.Metadata, + }) + return coreexecutor.RequestAfterAuthInterceptResponse{ + Headers: resp.Headers, + Body: resp.Body, + ClearHeaders: resp.ClearHeaders, + } +} + func (h *BaseAPIHandler) applyResponseInterceptors(ctx context.Context, handlerType, normalizedModel, requestedModel string, opts coreexecutor.Options, rawResponseHeaders, responseHeaders http.Header, originalRequest, requestBody, body []byte, statusCode int) ([]byte, http.Header) { host := h.interceptorHost() if host == nil { diff --git a/sdk/api/handlers/handlers_interceptors_test.go b/sdk/api/handlers/handlers_interceptors_test.go index 5a8280d2..9f9b5552 100644 --- a/sdk/api/handlers/handlers_interceptors_test.go +++ b/sdk/api/handlers/handlers_interceptors_test.go @@ -18,9 +18,10 @@ import ( ) type handlerInterceptorTestHost struct { - interceptRequest func(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse - interceptResponse func(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse - interceptStreamChunk func(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse + interceptRequestBeforeAuth func(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse + interceptRequestAfterAuth func(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse + interceptResponse func(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse + interceptStreamChunk func(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse } type handlerInterceptorNoStreamTestHost struct { @@ -31,9 +32,19 @@ func (h *handlerInterceptorNoStreamTestHost) HasStreamInterceptors() bool { return false } -func (h *handlerInterceptorTestHost) InterceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { - if h != nil && h.interceptRequest != nil { - return h.interceptRequest(ctx, req) +func (h *handlerInterceptorTestHost) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + if h != nil && h.interceptRequestBeforeAuth != nil { + return h.interceptRequestBeforeAuth(ctx, req) + } + return pluginapi.RequestInterceptResponse{ + Headers: cloneHeader(req.Headers), + Body: cloneBytes(req.Body), + } +} + +func (h *handlerInterceptorTestHost) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + if h != nil && h.interceptRequestAfterAuth != nil { + return h.interceptRequestAfterAuth(ctx, req) } return pluginapi.RequestInterceptResponse{ Headers: cloneHeader(req.Headers), @@ -177,7 +188,7 @@ func TestHandlerRequestInterceptorRewritesExecutorRequest(t *testing.T) { executor := &interceptorCaptureExecutor{} handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) handler.SetPluginHost(&handlerInterceptorTestHost{ - interceptRequest: func(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + interceptRequestBeforeAuth: func(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { if req.SourceFormat != "openai" || req.Model != model || req.RequestedModel != model { t.Fatalf("unexpected request context: %#v", req) } @@ -233,7 +244,7 @@ func TestHandlerRequestInterceptorEmptyBodyKeepsOriginalPayload(t *testing.T) { executor := &interceptorCaptureExecutor{} handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) handler.SetPluginHost(&handlerInterceptorTestHost{ - interceptRequest: func(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + interceptRequestBeforeAuth: func(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { return pluginapi.RequestInterceptResponse{ Headers: http.Header{"X-Plugin": []string{"empty-body"}}, Body: []byte{}, @@ -258,6 +269,89 @@ func TestHandlerRequestInterceptorEmptyBodyKeepsOriginalPayload(t *testing.T) { } } +func TestHandlerRequestInterceptorAfterAuthRewritesExecutorRequest(t *testing.T) { + model := "handler-interceptor-after-auth-model" + executor := &interceptorCaptureExecutor{} + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{}) + var calls []string + var responseChecked bool + handler.SetPluginHost(&handlerInterceptorTestHost{ + interceptRequestBeforeAuth: func(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + calls = append(calls, "before") + headers := cloneHeader(req.Headers) + if headers == nil { + headers = http.Header{} + } + headers.Set("X-Stage", "before") + return pluginapi.RequestInterceptResponse{ + Headers: headers, + Body: []byte(`{"stage":"before"}`), + } + }, + interceptRequestAfterAuth: func(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + calls = append(calls, "after") + if req.SourceFormat != "openai" || req.ToFormat != "codex" { + t.Fatalf("request formats = %q -> %q, want openai -> codex", req.SourceFormat, req.ToFormat) + } + if req.Model != model || req.RequestedModel != model { + t.Fatalf("request models = %q/%q, want %q/%q", req.Model, req.RequestedModel, model, model) + } + if string(req.Body) != `{"stage":"before"}` { + t.Fatalf("after-auth body = %q, want before-auth rewrite", req.Body) + } + headers := cloneHeader(req.Headers) + if headers == nil { + headers = http.Header{} + } + headers.Set("X-Stage", "after") + return pluginapi.RequestInterceptResponse{ + Headers: headers, + Body: []byte(`{"stage":"after"}`), + } + }, + interceptResponse: func(ctx context.Context, req pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse { + responseChecked = true + if req.RequestHeaders.Get("X-Stage") != "after" { + t.Fatalf("response request headers = %#v, want after-auth header", req.RequestHeaders) + } + if string(req.OriginalRequest) != `{"stage":"after"}` { + t.Fatalf("response original request = %q, want after-auth body", req.OriginalRequest) + } + if string(req.RequestBody) != `{"stage":"after"}` { + t.Fatalf("response request body = %q, want after-auth body", req.RequestBody) + } + return pluginapi.ResponseInterceptResponse{ + Headers: cloneHeader(req.ResponseHeaders), + Body: cloneBytes(req.Body), + } + }, + }) + + body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", model, []byte(fmt.Sprintf(`{"model":%q}`, model)), "") + if errMsg != nil { + t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg) + } + if string(body) != "ok" { + t.Fatalf("body = %q, want ok", body) + } + if fmt.Sprint(calls) != "[before after]" { + t.Fatalf("interceptor calls = %v, want [before after]", calls) + } + gotReq, gotOpts := executor.captured() + if string(gotReq.Payload) != `{"stage":"after"}` { + t.Fatalf("executor payload = %q, want after-auth body", gotReq.Payload) + } + if string(gotOpts.OriginalRequest) != `{"stage":"after"}` { + t.Fatalf("executor original request = %q, want after-auth body", gotOpts.OriginalRequest) + } + if gotOpts.Headers.Get("X-Stage") != "after" { + t.Fatalf("executor headers = %#v, want after-auth header", gotOpts.Headers) + } + if !responseChecked { + t.Fatal("response interceptor was not called") + } +} + func TestHandlerResponseInterceptorRewritesSuccessfulNonStreamResponse(t *testing.T) { model := "handler-interceptor-response-model" executor := &interceptorCaptureExecutor{ @@ -465,8 +559,42 @@ func TestHandlerStreamInterceptorRewritesAndDropsChunks(t *testing.T) { handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{PassthroughHeaders: true}) var streamCalls int handler.SetPluginHost(&handlerInterceptorTestHost{ + interceptRequestBeforeAuth: func(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + headers := cloneHeader(req.Headers) + if headers == nil { + headers = http.Header{} + } + headers.Set("X-Stage", "before") + return pluginapi.RequestInterceptResponse{ + Headers: headers, + Body: []byte(`{"stage":"before-stream"}`), + } + }, + interceptRequestAfterAuth: func(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + if string(req.Body) != `{"stage":"before-stream"}` { + t.Fatalf("after-auth stream body = %q, want before-auth rewrite", req.Body) + } + headers := cloneHeader(req.Headers) + if headers == nil { + headers = http.Header{} + } + headers.Set("X-Stage", "after") + return pluginapi.RequestInterceptResponse{ + Headers: headers, + Body: []byte(`{"stage":"after-stream"}`), + } + }, interceptStreamChunk: func(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { streamCalls++ + if req.RequestHeaders.Get("X-Stage") != "after" { + t.Fatalf("stream request headers = %#v, want after-auth header", req.RequestHeaders) + } + if string(req.OriginalRequest) != `{"stage":"after-stream"}` { + t.Fatalf("stream original request = %q, want after-auth body", req.OriginalRequest) + } + if string(req.RequestBody) != `{"stage":"after-stream"}` { + t.Fatalf("stream request body = %q, want after-auth body", req.RequestBody) + } if req.ChunkIndex == pluginapi.StreamChunkHeaderInitIndex { headers := cloneHeader(req.ResponseHeaders) headers.Set("X-Stream", "plugin") diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 61afc583..08b81dad 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -25,6 +25,7 @@ import ( cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" log "github.com/sirupsen/logrus" "github.com/tidwall/sjson" ) @@ -1139,7 +1140,9 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) execReq := req execReq.Model = execModel - streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, opts) + execOpts := opts + execReq, execOpts = applyRequestAfterAuthInterceptor(ctx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, execOpts) if errStream != nil { if errCtx := ctx.Err(); errCtx != nil { return nil, errCtx @@ -1654,6 +1657,99 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli return nil, &Error{Code: "auth_not_found", Message: "no auth available"} } +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) { + if opts.RequestAfterAuthInterceptor == nil { + return req, opts + } + toFormat := requestToFormat(provider, executor, req, opts) + resp := opts.RequestAfterAuthInterceptor(ctx, cliproxyexecutor.RequestAfterAuthInterceptRequest{ + SourceFormat: opts.SourceFormat, + ToFormat: toFormat, + Model: req.Model, + RequestedModel: requestedModel, + Stream: opts.Stream, + Headers: cloneRequestHeaders(opts.Headers), + Body: bytes.Clone(req.Payload), + Metadata: opts.Metadata, + }) + opts.Headers = mergeRequestHeaders(opts.Headers, resp.Headers, resp.ClearHeaders) + if len(resp.Body) > 0 { + req.Payload = bytes.Clone(resp.Body) + opts.OriginalRequest = bytes.Clone(resp.Body) + } + return req, opts +} + +func requestToFormat(provider string, executor ProviderExecutor, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) sdktranslator.Format { + resolver, ok := executor.(requestToFormatResolver) + if ok && resolver != nil { + formatRequestTo := resolver.RequestToFormat(req, opts) + if formatRequestTo != "" { + return formatRequestTo + } + } + source := opts.SourceFormat.String() + if source == "openai-image" || source == "openai-video" { + return opts.SourceFormat + } + if opts.Alt == "responses/compact" && !opts.Stream { + return sdktranslator.FormatOpenAIResponse + } + switch strings.ToLower(strings.TrimSpace(provider)) { + case "codex": + return sdktranslator.FormatCodex + case "xai": + return sdktranslator.FormatCodex + case "claude": + return sdktranslator.FormatClaude + case "gemini", "vertex", "aistudio": + return sdktranslator.FormatGemini + case "gemini-cli": + return sdktranslator.FormatGeminiCLI + case "kimi": + return sdktranslator.FormatOpenAI + case "antigravity": + return sdktranslator.FormatAntigravity + default: + return sdktranslator.FormatOpenAI + } +} + +func cloneRequestHeaders(src http.Header) http.Header { + if src == nil { + return nil + } + dst := make(http.Header, len(src)) + for key, values := range src { + dst[key] = append([]string(nil), values...) + } + return dst +} + +func mergeRequestHeaders(current, updates http.Header, clear []string) http.Header { + if updates == nil && len(clear) == 0 { + return current + } + out := cloneRequestHeaders(current) + if out == nil && (len(updates) > 0 || len(clear) > 0) { + out = make(http.Header) + } + for _, key := range clear { + out.Del(key) + } + for key, values := range updates { + out.Del(key) + for _, value := range values { + out.Add(key, value) + } + } + return out +} + func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) { if len(providers) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} @@ -1717,7 +1813,9 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel - resp, errExec := executor.Execute(execCtx, auth, execReq, opts) + execOpts := opts + execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + resp, errExec := executor.Execute(execCtx, auth, execReq, execOpts) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} if errExec != nil { if errCtx := execCtx.Err(); errCtx != nil { @@ -1816,7 +1914,9 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) execReq := req execReq.Model = upstreamModel - resp, errExec := executor.CountTokens(execCtx, auth, execReq, opts) + execOpts := opts + execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + resp, errExec := executor.CountTokens(execCtx, auth, execReq, execOpts) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} if errExec != nil { if errCtx := execCtx.Err(); errCtx != nil { diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index 8f0fc567..9f5c4a45 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -1,6 +1,7 @@ package executor import ( + "context" "net/http" "net/url" @@ -46,6 +47,39 @@ type Request struct { Metadata map[string]any } +// RequestAfterAuthInterceptor rewrites a request after credential selection and before executor translation. +type RequestAfterAuthInterceptor func(context.Context, RequestAfterAuthInterceptRequest) RequestAfterAuthInterceptResponse + +// RequestAfterAuthInterceptRequest describes a selected-auth request before executor translation. +type RequestAfterAuthInterceptRequest struct { + // SourceFormat is the original client protocol format. + SourceFormat sdktranslator.Format + // ToFormat is the selected upstream protocol format. + ToFormat sdktranslator.Format + // Model is the selected upstream model for this attempt. + Model string + // RequestedModel is the client-requested model before alias/model-pool rewriting. + RequestedModel string + // Stream reports whether the request expects streaming output. + Stream bool + // Headers contains the current upstream request headers. + Headers http.Header + // Body contains the current request payload. + Body []byte + // Metadata is a best-effort cloned context snapshot. Treat it as read-only and JSON-like. + Metadata map[string]any +} + +// RequestAfterAuthInterceptResponse returns selected-auth request modifications. +type RequestAfterAuthInterceptResponse struct { + // Headers replaces matching current request headers and preserves headers not mentioned here. + Headers http.Header + // Body replaces the current request body only when non-empty. + Body []byte + // ClearHeaders explicitly removes current request headers before Headers is applied. + ClearHeaders []string +} + // Options controls execution behavior for both streaming and non-streaming calls. type Options struct { // Stream toggles streaming mode. @@ -62,6 +96,8 @@ type Options struct { SourceFormat sdktranslator.Format // Metadata carries extra execution hints shared across selection and executors. Metadata map[string]any + // RequestAfterAuthInterceptor runs after credential selection and before executor translation. + RequestAfterAuthInterceptor RequestAfterAuthInterceptor } // Response wraps either a full provider response or metadata for streaming flows. diff --git a/sdk/pluginabi/types.go b/sdk/pluginabi/types.go index af80be14..69852234 100644 --- a/sdk/pluginabi/types.go +++ b/sdk/pluginabi/types.go @@ -37,6 +37,7 @@ const ( MethodRequestTranslate = "request.translate" MethodRequestNormalize = "request.normalize" MethodRequestInterceptBefore = "request.intercept_before" + MethodRequestInterceptAfter = "request.intercept_after" MethodResponseTranslate = "response.translate" MethodResponseNormalizeBefore = "response.normalize_before" diff --git a/sdk/pluginabi/types_test.go b/sdk/pluginabi/types_test.go index f9562448..7b6ff7da 100644 --- a/sdk/pluginabi/types_test.go +++ b/sdk/pluginabi/types_test.go @@ -33,6 +33,9 @@ func TestMethodNamesAreStable(t *testing.T) { if MethodRequestInterceptBefore != "request.intercept_before" { t.Fatalf("MethodRequestInterceptBefore = %q", MethodRequestInterceptBefore) } + if MethodRequestInterceptAfter != "request.intercept_after" { + t.Fatalf("MethodRequestInterceptAfter = %q", MethodRequestInterceptAfter) + } if MethodResponseInterceptAfter != "response.intercept_after" { t.Fatalf("MethodResponseInterceptAfter = %q", MethodResponseInterceptAfter) } diff --git a/sdk/pluginapi/types.go b/sdk/pluginapi/types.go index 6e6e36f8..7ec03c4d 100644 --- a/sdk/pluginapi/types.go +++ b/sdk/pluginapi/types.go @@ -97,7 +97,7 @@ type Capabilities struct { ResponseBeforeTranslator ResponseNormalizer // ResponseAfterTranslator normalizes translated responses before delivery. ResponseAfterTranslator ResponseNormalizer - // RequestInterceptor rewrites execution requests before they reach the upstream executor. + // RequestInterceptor rewrites execution requests before and after credential selection. RequestInterceptor RequestInterceptor // ResponseInterceptor rewrites successful non-streaming HTTP execution responses before downstream delivery. ResponseInterceptor ResponseInterceptor @@ -680,9 +680,10 @@ type ResponseNormalizer interface { NormalizeResponse(context.Context, ResponseTransformRequest) (PayloadResponse, error) } -// RequestInterceptor rewrites execution requests before they reach the upstream executor. +// RequestInterceptor rewrites execution requests before and after credential selection. type RequestInterceptor interface { - InterceptRequest(context.Context, RequestInterceptRequest) (RequestInterceptResponse, error) + InterceptRequestBeforeAuth(context.Context, RequestInterceptRequest) (RequestInterceptResponse, error) + InterceptRequestAfterAuth(context.Context, RequestInterceptRequest) (RequestInterceptResponse, error) } // ResponseInterceptor rewrites successful non-streaming execution responses before downstream delivery. @@ -732,13 +733,22 @@ type ResponseTransformRequest struct { // RequestInterceptRequest describes a request about to be executed upstream. type RequestInterceptRequest struct { - SourceFormat string - Model string + // SourceFormat is the original client protocol format. + SourceFormat string + // ToFormat is the selected upstream protocol format. It is empty before credential selection. + ToFormat string + // Model is the current execution model. After credential selection this is the selected upstream model. + Model string + // RequestedModel is the client-requested model before alias/model-pool rewriting. RequestedModel string - Stream bool - Headers http.Header - Body []byte - Metadata map[string]any + // Stream reports whether the request expects streaming output. + Stream bool + // Headers contains the current upstream request headers. + Headers http.Header + // Body contains the current request payload. + Body []byte + // Metadata is a best-effort cloned context snapshot. Treat it as read-only and JSON-like. + Metadata map[string]any } // RequestInterceptResponse returns request modifications. diff --git a/sdk/pluginapi/types_test.go b/sdk/pluginapi/types_test.go index 497ef30e..18725755 100644 --- a/sdk/pluginapi/types_test.go +++ b/sdk/pluginapi/types_test.go @@ -256,7 +256,11 @@ func (compileTimePlugin) NormalizeResponse(context.Context, ResponseTransformReq return PayloadResponse{}, nil } -func (compileTimePlugin) InterceptRequest(context.Context, RequestInterceptRequest) (RequestInterceptResponse, error) { +func (compileTimePlugin) InterceptRequestBeforeAuth(context.Context, RequestInterceptRequest) (RequestInterceptResponse, error) { + return RequestInterceptResponse{}, nil +} + +func (compileTimePlugin) InterceptRequestAfterAuth(context.Context, RequestInterceptRequest) (RequestInterceptResponse, error) { return RequestInterceptResponse{}, nil }