From 538e3416dbd275d8037d2a8535add733ecdfcc5b Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Fri, 12 Jun 2026 02:38:51 +0800 Subject: [PATCH] feat(plugin, api): prevent plugin recursion on host model callbacks, enable targeted interceptor skipping - Updated host model callback logic to skip originating plugin's interceptors during nested model executions. - Added `SkipInterceptorPluginID` field to plugin API structs for controlling interceptor bypass behavior. - Introduced supporting logic in host API handlers, plugin host registry, and callback contexts to identify and skip specific plugins. - Enhanced unit tests across plugin host, API handlers, and execution paths to verify interceptor skipping behavior and plugin isolation. - Revised documentation to clarify non-recursive behavior of host model callbacks and the use of `SkipInterceptorPluginID`. --- examples/plugin/README.md | 2 + examples/plugin/README_CN.md | 2 + examples/plugin/host-model-callback/README.md | 6 + .../plugin/host-model-callback/go/main.go | 6 + internal/pluginhost/abi.go | 2 +- internal/pluginhost/adapters.go | 35 +++-- internal/pluginhost/adapters_test.go | 75 ++++++++++ internal/pluginhost/callback_contexts.go | 35 ++++- internal/pluginhost/host.go | 2 +- internal/pluginhost/host_callbacks.go | 61 +++++++-- internal/pluginhost/host_callbacks_test.go | 32 +++++ internal/pluginhost/host_callbacks_unix.go | 7 +- internal/pluginhost/loader_unix.go | 8 +- internal/pluginhost/loader_unsupported.go | 4 +- internal/pluginhost/loader_windows.go | 13 +- internal/pluginhost/rpc_client.go | 2 +- internal/pluginhost/test_helpers_test.go | 6 +- sdk/api/handlers/handlers.go | 89 ++++++++---- sdk/api/handlers/model_execution.go | 44 +++--- sdk/api/handlers/model_execution_test.go | 128 ++++++++++++++++++ 20 files changed, 474 insertions(+), 85 deletions(-) diff --git a/examples/plugin/README.md b/examples/plugin/README.md index 663054a17..59bd5a434 100644 --- a/examples/plugin/README.md +++ b/examples/plugin/README.md @@ -43,6 +43,8 @@ plugins: `host-model-callback` declares the Management API capability and exposes a browser resource named `Host Model Callback`. The resource calls `host.model.execute` for non-streaming requests and `host.model.execute_stream` plus `host.model.stream_read` for streaming requests. It demonstrates explicit stream close with `host.model.stream_close` and an `implicit_close=true` option for RPC-scope host cleanup. +When the resource forwards its `host_callback_id`, CPA identifies the plugin that initiated the host model callback and skips that same plugin's interceptors for the nested execution. This makes host model callbacks non-recursive for the caller while allowing other plugins to intercept the nested request. + ```yaml plugins: configs: diff --git a/examples/plugin/README_CN.md b/examples/plugin/README_CN.md index de8507421..2fe650e02 100644 --- a/examples/plugin/README_CN.md +++ b/examples/plugin/README_CN.md @@ -43,6 +43,8 @@ plugins: `host-model-callback` 声明 Management API 能力，并暴露名为 `Host Model Callback` 的浏览器资源。该资源在非流式请求中调用 `host.model.execute`，在流式请求中调用 `host.model.execute_stream` 和 `host.model.stream_read`。它演示了通过 `host.model.stream_close` 显式关闭流，也提供 `implicit_close=true` 用于演示 RPC 作用域结束时的宿主隐式清理。 +当该资源转发自身收到的 `host_callback_id` 时，CPA 会识别发起宿主模型回调的插件，并在嵌套模型执行中跳过同一个插件的拦截器。因此宿主模型回调不会递归调用发起插件自身，但其他已启用插件仍可拦截这次嵌套请求。 + ```yaml plugins: configs: diff --git a/examples/plugin/host-model-callback/README.md b/examples/plugin/host-model-callback/README.md index a69e27e3a..f0b5c3929 100644 --- a/examples/plugin/host-model-callback/README.md +++ b/examples/plugin/host-model-callback/README.md @@ -115,6 +115,12 @@ By default, streaming mode explicitly closes the host-owned stream with `host.mo When `implicit_close=true` is set, the plugin intentionally skips the explicit close call. CPA injects `host_callback_id` into the `management.handle` request, and this example forwards that callback ID to `host.model.execute_stream` so the host can close the stream when the `management.handle` RPC callback scope returns. This mode exists only to demonstrate host cleanup behavior; normal plugin code should explicitly close streams it opens. +## Recursion Guard + +This example forwards the `host_callback_id` received from `management.handle` when it calls `host.model.execute` or `host.model.execute_stream`. CPA uses that callback scope to identify the plugin that initiated the host model callback and skips that same plugin's request, response, and stream interceptors for the nested model execution. + +Host model callbacks are therefore not recursive for the caller. Other enabled plugins can still intercept the nested request. + ## Billing and Usage The callback uses the existing CPA model executor path. Usage collection, request accounting, and billing metadata are handled by the same executor and usage reporter path as normal proxied requests. The callback layer does not bill twice and does not create an additional usage record by itself. diff --git a/examples/plugin/host-model-callback/go/main.go b/examples/plugin/host-model-callback/go/main.go index 76cb1ae3f..313611161 100644 --- a/examples/plugin/host-model-callback/go/main.go +++ b/examples/plugin/host-model-callback/go/main.go @@ -427,6 +427,9 @@ func executeOnce(opts runOptions) (pluginapi.HostModelExecutionResponse, error) if errBody != nil { return pluginapi.HostModelExecutionResponse{}, errBody } + // Forward HostCallbackID so the host skips this plugin's interceptors on the + // nested model execution. Host model callbacks do not recursively call the + // originating plugin's interceptor chain. result, errCall := callHost(pluginabi.MethodHostModelExecute, hostModelExecutionRequest{ HostModelExecutionRequest: pluginapi.HostModelExecutionRequest{ EntryProtocol: opts.EntryProtocol, @@ -456,6 +459,9 @@ func executeStream(opts runOptions) (data streamPageData) { data.Error = errBody.Error() return data } + // Forward HostCallbackID so the host skips this plugin's interceptors on the + // nested model execution. Host model callbacks do not recursively call the + // originating plugin's interceptor chain. result, errCall := callHost(pluginabi.MethodHostModelExecuteStream, hostModelExecutionRequest{ HostModelExecutionRequest: pluginapi.HostModelExecutionRequest{ EntryProtocol: opts.EntryProtocol, diff --git a/internal/pluginhost/abi.go b/internal/pluginhost/abi.go index 44d75cd52..a63694faa 100644 --- a/internal/pluginhost/abi.go +++ b/internal/pluginhost/abi.go @@ -14,5 +14,5 @@ type pluginClient interface { } type pluginLoader interface { - Open(path string, host *Host) (pluginClient, error) + Open(file pluginFile, host *Host) (pluginClient, error) } diff --git a/internal/pluginhost/adapters.go b/internal/pluginhost/adapters.go index 33ca53f34..63fb33dee 100644 --- a/internal/pluginhost/adapters.go +++ b/internal/pluginhost/adapters.go @@ -569,25 +569,34 @@ func (h *Host) callStreamChunkInterceptor(ctx context.Context, pluginID string, } func (h *Host) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + return h.InterceptRequestBeforeAuthExcept(ctx, req, "") +} + +func (h *Host) InterceptRequestBeforeAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) 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) - }) + }, skipPluginID) } 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) - }) + return h.InterceptRequestAfterAuthExcept(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 { +func (h *Host) InterceptRequestAfterAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) 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) + }, skipPluginID) +} + +func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest, method string, invoke func(pluginapi.RequestInterceptor, context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error), skipPluginID string) pluginapi.RequestInterceptResponse { current := pluginapi.RequestInterceptResponse{ Headers: cloneHeader(req.Headers), Body: bytes.Clone(req.Body), } + skipPluginID = strings.TrimSpace(skipPluginID) for _, record := range h.Snapshot().records { interceptor := record.plugin.Capabilities.RequestInterceptor - if h.isPluginFused(record.id) || interceptor == nil { + if h.isPluginFused(record.id) || interceptor == nil || record.id == skipPluginID { continue } nextReq := req @@ -607,13 +616,18 @@ func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterc } func (h *Host) InterceptResponse(ctx context.Context, req pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse { + return h.InterceptResponseExcept(ctx, req, "") +} + +func (h *Host) InterceptResponseExcept(ctx context.Context, req pluginapi.ResponseInterceptRequest, skipPluginID string) pluginapi.ResponseInterceptResponse { current := pluginapi.ResponseInterceptResponse{ Headers: cloneHeader(req.ResponseHeaders), Body: bytes.Clone(req.Body), } + skipPluginID = strings.TrimSpace(skipPluginID) for _, record := range h.Snapshot().records { interceptor := record.plugin.Capabilities.ResponseInterceptor - if h.isPluginFused(record.id) || interceptor == nil { + if h.isPluginFused(record.id) || interceptor == nil || record.id == skipPluginID { continue } nextReq := req @@ -634,13 +648,18 @@ func (h *Host) InterceptResponse(ctx context.Context, req pluginapi.ResponseInte } func (h *Host) InterceptStreamChunk(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { + return h.InterceptStreamChunkExcept(ctx, req, "") +} + +func (h *Host) InterceptStreamChunkExcept(ctx context.Context, req pluginapi.StreamChunkInterceptRequest, skipPluginID string) pluginapi.StreamChunkInterceptResponse { current := pluginapi.StreamChunkInterceptResponse{ Headers: cloneHeader(req.ResponseHeaders), Body: bytes.Clone(req.Body), } + skipPluginID = strings.TrimSpace(skipPluginID) for _, record := range h.Snapshot().records { interceptor := record.plugin.Capabilities.StreamChunkInterceptor - if h.isPluginFused(record.id) || interceptor == nil || current.DropChunk { + if h.isPluginFused(record.id) || interceptor == nil || current.DropChunk || record.id == skipPluginID { continue } nextReq := req diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index b5a5d8b3e..58aa75f3a 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -1342,6 +1342,81 @@ func TestInterceptRequestAfterAuthPassesTargetFormat(t *testing.T) { } } +func TestInterceptorsSkipExceptedPlugin(t *testing.T) { + originCalls := 0 + otherCalls := 0 + host := newHostWithRecords( + capabilityRecord{ + id: "origin", + priority: 20, + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + RequestInterceptor: requestInterceptorFunc(func(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + originCalls++ + return pluginapi.RequestInterceptResponse{Body: append(req.Body, []byte("|origin-request")...)}, nil + }), + ResponseInterceptor: responseInterceptorFunc{ + interceptResponse: func(ctx context.Context, req pluginapi.ResponseInterceptRequest) (pluginapi.ResponseInterceptResponse, error) { + originCalls++ + return pluginapi.ResponseInterceptResponse{Body: append(req.Body, []byte("|origin-response")...)}, nil + }, + }, + StreamChunkInterceptor: responseInterceptorFunc{ + interceptStreamChunk: func(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) (pluginapi.StreamChunkInterceptResponse, error) { + originCalls++ + return pluginapi.StreamChunkInterceptResponse{Body: append(req.Body, []byte("|origin-stream")...)}, nil + }, + }, + }}, + }, + capabilityRecord{ + id: "other", + priority: 10, + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + RequestInterceptor: requestInterceptorFunc(func(ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + otherCalls++ + return pluginapi.RequestInterceptResponse{Body: append(req.Body, []byte("|other-request")...)}, nil + }), + ResponseInterceptor: responseInterceptorFunc{ + interceptResponse: func(ctx context.Context, req pluginapi.ResponseInterceptRequest) (pluginapi.ResponseInterceptResponse, error) { + otherCalls++ + return pluginapi.ResponseInterceptResponse{Body: append(req.Body, []byte("|other-response")...)}, nil + }, + }, + StreamChunkInterceptor: responseInterceptorFunc{ + interceptStreamChunk: func(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) (pluginapi.StreamChunkInterceptResponse, error) { + otherCalls++ + return pluginapi.StreamChunkInterceptResponse{Body: append(req.Body, []byte("|other-stream")...)}, nil + }, + }, + }}, + }, + ) + + reqOut := host.InterceptRequestBeforeAuthExcept(context.Background(), pluginapi.RequestInterceptRequest{Body: []byte("body")}, "origin") + afterOut := host.InterceptRequestAfterAuthExcept(context.Background(), pluginapi.RequestInterceptRequest{Body: []byte("body")}, "origin") + respOut := host.InterceptResponseExcept(context.Background(), pluginapi.ResponseInterceptRequest{Body: []byte("body")}, "origin") + streamOut := host.InterceptStreamChunkExcept(context.Background(), pluginapi.StreamChunkInterceptRequest{Body: []byte("body")}, "origin") + + if originCalls != 0 { + t.Fatalf("origin plugin calls = %d, want 0", originCalls) + } + if otherCalls != 4 { + t.Fatalf("other plugin calls = %d, want 4", otherCalls) + } + if string(reqOut.Body) != "body|other-request" { + t.Fatalf("request body = %q, want body|other-request", reqOut.Body) + } + if string(afterOut.Body) != "body|other-request" { + t.Fatalf("after-auth request body = %q, want body|other-request", afterOut.Body) + } + if string(respOut.Body) != "body|other-response" { + t.Fatalf("response body = %q, want body|other-response", respOut.Body) + } + if string(streamOut.Body) != "body|other-stream" { + t.Fatalf("stream body = %q, want body|other-stream", streamOut.Body) + } +} + func TestResponseInterceptorsChainAndStreamHistory(t *testing.T) { var seenHistory [][]byte var sawSecondResponse bool diff --git a/internal/pluginhost/callback_contexts.go b/internal/pluginhost/callback_contexts.go index b87e67ed6..27c5aaded 100644 --- a/internal/pluginhost/callback_contexts.go +++ b/internal/pluginhost/callback_contexts.go @@ -3,6 +3,7 @@ package pluginhost import ( "context" "strconv" + "strings" "sync" "sync/atomic" ) @@ -14,24 +15,27 @@ type callbackContextRegistry struct { } type callbackContextEntry struct { - ctx context.Context - cleanup []func() + ctx context.Context + pluginID string + cleanup []func() } func newCallbackContextRegistry() *callbackContextRegistry { return &callbackContextRegistry{contexts: make(map[string]callbackContextEntry)} } -func (r *callbackContextRegistry) open(ctx context.Context) (string, func()) { +func (r *callbackContextRegistry) open(ctx context.Context, pluginID string) (string, func()) { if r == nil { return "", func() {} } if ctx == nil { ctx = context.Background() } + pluginID = strings.TrimSpace(pluginID) + ctx = withHostCallbackPluginID(ctx, pluginID) id := strconv.FormatUint(r.next.Add(1), 10) r.mu.Lock() - r.contexts[id] = callbackContextEntry{ctx: ctx} + r.contexts[id] = callbackContextEntry{ctx: ctx, pluginID: pluginID} r.mu.Unlock() var once sync.Once @@ -52,6 +56,16 @@ func (r *callbackContextRegistry) open(ctx context.Context) (string, func()) { } } +func (r *callbackContextRegistry) pluginID(id string) string { + if r == nil || id == "" { + return "" + } + r.mu.RLock() + entry := r.contexts[id] + r.mu.RUnlock() + return strings.TrimSpace(entry.pluginID) +} + func (r *callbackContextRegistry) addCleanup(id string, cleanup func()) bool { if r == nil || id == "" || cleanup == nil { return false @@ -87,10 +101,14 @@ func (r *callbackContextRegistry) resolve(id string, fallback context.Context) c } func (h *Host) openCallbackContext(ctx context.Context) (string, func()) { + return h.openCallbackContextForPlugin(ctx, "") +} + +func (h *Host) openCallbackContextForPlugin(ctx context.Context, pluginID string) (string, func()) { if h == nil || h.callbackContexts == nil { return "", func() {} } - return h.callbackContexts.open(ctx) + return h.callbackContexts.open(ctx, pluginID) } func (h *Host) addCallbackCleanup(id string, cleanup func()) bool { @@ -112,3 +130,10 @@ func (h *Host) resolveCallbackContext(id string, fallback context.Context) conte } return h.callbackContexts.resolve(id, fallback) } + +func (h *Host) callbackContextPluginID(id string) string { + if h == nil || h.callbackContexts == nil { + return "" + } + return h.callbackContexts.pluginID(id) +} diff --git a/internal/pluginhost/host.go b/internal/pluginhost/host.go index fefc5bd86..ffa596ad5 100644 --- a/internal/pluginhost/host.go +++ b/internal/pluginhost/host.go @@ -186,7 +186,7 @@ func (h *Host) ApplyConfig(ctx context.Context, cfg *config.Config) { } func (h *Host) loadLocked(file pluginFile) (*loadedPlugin, error) { - client, errOpen := h.loader.Open(file.Path, h) + client, errOpen := h.loader.Open(file, h) if errOpen != nil { return nil, errOpen } diff --git a/internal/pluginhost/host_callbacks.go b/internal/pluginhost/host_callbacks.go index dd12ceb30..a573fbc33 100644 --- a/internal/pluginhost/host_callbacks.go +++ b/internal/pluginhost/host_callbacks.go @@ -66,6 +66,35 @@ type rpcHostModelExecutionRequest struct { HostCallbackID string `json:"host_callback_id,omitempty"` } +type dynamicHostCallbackEntry struct { + host *Host + pluginID string +} + +type hostCallbackPluginIDKey struct{} + +func withHostCallbackPluginID(ctx context.Context, pluginID string) context.Context { + pluginID = strings.TrimSpace(pluginID) + if pluginID == "" { + if ctx == nil { + return context.Background() + } + return ctx + } + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, hostCallbackPluginIDKey{}, pluginID) +} + +func hostCallbackPluginIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + pluginID, _ := ctx.Value(hostCallbackPluginIDKey{}).(string) + return strings.TrimSpace(pluginID) +} + func (h *Host) callFromPlugin(ctx context.Context, method string, request []byte) ([]byte, error) { switch method { case pluginabi.MethodHostModelExecute: @@ -95,6 +124,13 @@ func (h *Host) callFromPlugin(ctx context.Context, method string, request []byte } } +func (h *Host) callbackCallerPluginID(ctx context.Context, callbackID string) string { + if pluginID := hostCallbackPluginIDFromContext(ctx); pluginID != "" { + return pluginID + } + return h.callbackContextPluginID(callbackID) +} + func (h *Host) callHostHTTPDo(ctx context.Context, request []byte) ([]byte, error) { httpReq, callbackID, errDecode := decodeHostHTTPRequestWithCallbackID(request) if errDecode != nil { @@ -234,8 +270,9 @@ func (h *Host) callHostModelExecute(ctx context.Context, request []byte) ([]byte if executor == nil { return nil, fmt.Errorf("host model executor is unavailable") } + skipPluginID := h.callbackCallerPluginID(ctx, req.HostCallbackID) ctx = h.resolveCallbackContext(req.HostCallbackID, ctx) - resp, errMsg := executor.ExecuteModel(ctx, modelExecutionRequestFromPlugin(req.HostModelExecutionRequest)) + resp, errMsg := executor.ExecuteModel(ctx, modelExecutionRequestFromPlugin(req.HostModelExecutionRequest, skipPluginID)) if errMsg != nil { return nil, modelExecutionError(errMsg) } @@ -258,12 +295,13 @@ func (h *Host) callHostModelExecuteStream(ctx context.Context, request []byte) ( if executor == nil { return nil, fmt.Errorf("host model executor is unavailable") } + skipPluginID := h.callbackCallerPluginID(ctx, req.HostCallbackID) ctx = h.resolveCallbackContext(req.HostCallbackID, ctx) if ctx == nil { ctx = context.Background() } streamCtx, cancel := context.WithCancel(ctx) - stream, errMsg := executor.ExecuteModelStream(streamCtx, modelExecutionRequestFromPlugin(req.HostModelExecutionRequest)) + stream, errMsg := executor.ExecuteModelStream(streamCtx, modelExecutionRequestFromPlugin(req.HostModelExecutionRequest, skipPluginID)) if errMsg != nil { cancel() return nil, modelExecutionError(errMsg) @@ -322,16 +360,17 @@ func (h *Host) callHostModelStreamClose(request []byte) ([]byte, error) { return marshalRPCResult(rpcEmptyResponse{}) } -func modelExecutionRequestFromPlugin(req pluginapi.HostModelExecutionRequest) handlers.ModelExecutionRequest { +func modelExecutionRequestFromPlugin(req pluginapi.HostModelExecutionRequest, skipPluginID string) handlers.ModelExecutionRequest { return handlers.ModelExecutionRequest{ - EntryProtocol: req.EntryProtocol, - ExitProtocol: req.ExitProtocol, - Model: req.Model, - Stream: req.Stream, - Body: append([]byte(nil), req.Body...), - Headers: cloneHeader(req.Headers), - Query: cloneValues(req.Query), - Alt: req.Alt, + EntryProtocol: req.EntryProtocol, + ExitProtocol: req.ExitProtocol, + Model: req.Model, + Stream: req.Stream, + Body: append([]byte(nil), req.Body...), + Headers: cloneHeader(req.Headers), + Query: cloneValues(req.Query), + Alt: req.Alt, + SkipInterceptorPluginID: skipPluginID, } } diff --git a/internal/pluginhost/host_callbacks_test.go b/internal/pluginhost/host_callbacks_test.go index e0ca16a41..6d9f33825 100644 --- a/internal/pluginhost/host_callbacks_test.go +++ b/internal/pluginhost/host_callbacks_test.go @@ -293,6 +293,38 @@ func TestHostModelExecuteCallback(t *testing.T) { } } +func TestHostModelExecuteCallbackCarriesCallerPluginSkipID(t *testing.T) { + host := New() + var got handlers.ModelExecutionRequest + host.SetModelExecutor(&fakeHostModelExecutor{ + executeModel: func(ctx context.Context, req handlers.ModelExecutionRequest) (handlers.ModelExecutionResponse, *interfaces.ErrorMessage) { + got = req + return handlers.ModelExecutionResponse{StatusCode: http.StatusOK, Body: []byte(`{"ok":true}`)}, nil + }, + }) + callbackID, closeCallback := host.openCallbackContextForPlugin(context.Background(), "origin-plugin") + defer closeCallback() + + rawReq, errMarshal := json.Marshal(rpcHostModelExecutionRequest{ + HostModelExecutionRequest: pluginapi.HostModelExecutionRequest{ + EntryProtocol: "openai", + ExitProtocol: "openai", + Model: "model-1", + Body: []byte(`{"request":true}`), + }, + HostCallbackID: callbackID, + }) + if errMarshal != nil { + t.Fatalf("marshal request: %v", errMarshal) + } + if _, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostModelExecute, rawReq); errCall != nil { + t.Fatalf("callFromPlugin() error = %v", errCall) + } + if got.SkipInterceptorPluginID != "origin-plugin" { + t.Fatalf("SkipInterceptorPluginID = %q, want origin-plugin", got.SkipInterceptorPluginID) + } +} + func TestHostModelStreamClosesWithCallbackScope(t *testing.T) { host := New() ctxSeen := make(chan context.Context, 1) diff --git a/internal/pluginhost/host_callbacks_unix.go b/internal/pluginhost/host_callbacks_unix.go index 1f624cd2c..b1d9af6cc 100644 --- a/internal/pluginhost/host_callbacks_unix.go +++ b/internal/pluginhost/host_callbacks_unix.go @@ -32,15 +32,16 @@ func cliproxyHostCall(hostCtx unsafe.Pointer, method *C.char, request *C.uint8_t if !okHost { return 1 } - host, okHost := rawHost.(*Host) - if !okHost || host == nil { + entry, okHost := rawHost.(dynamicHostCallbackEntry) + if !okHost || entry.host == nil { return 1 } var requestBytes []byte if request != nil && requestLen > 0 { requestBytes = C.GoBytes(unsafe.Pointer(request), C.int(requestLen)) } - resp, errCall := host.callFromPlugin(context.Background(), C.GoString(method), requestBytes) + ctx := withHostCallbackPluginID(context.Background(), entry.pluginID) + resp, errCall := entry.host.callFromPlugin(ctx, C.GoString(method), requestBytes) if errCall != nil { resp = marshalRPCError("host_call_failed", errCall.Error()) } diff --git a/internal/pluginhost/loader_unix.go b/internal/pluginhost/loader_unix.go index a44ab7e35..32261752e 100644 --- a/internal/pluginhost/loader_unix.go +++ b/internal/pluginhost/loader_unix.go @@ -108,13 +108,13 @@ func defaultPluginLoader() pluginLoader { return dynamicLibraryLoader{} } -func (dynamicLibraryLoader) Open(path string, host *Host) (pluginClient, error) { - cPath := C.CString(path) +func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, error) { + cPath := C.CString(file.Path) defer C.free(unsafe.Pointer(cPath)) handle := C.cliproxy_dlopen(cPath) if handle == nil { - return nil, fmt.Errorf("dlopen %s: %s", path, dlerrorString()) + return nil, fmt.Errorf("dlopen %s: %s", file.Path, dlerrorString()) } cSymbol := C.CString("cliproxy_plugin_init") @@ -138,7 +138,7 @@ func (dynamicLibraryLoader) Open(path string, host *Host) (pluginClient, error) } id := hostCallbackID.Add(1) *(*C.uintptr_t)(hostCtx) = C.uintptr_t(id) - hostCallbackEntries.Store(id, host) + hostCallbackEntries.Store(id, dynamicHostCallbackEntry{host: host, pluginID: file.ID}) C.cliproxy_set_host_api(hostAPI, C.uint32_t(pluginHostABIVersion), hostCtx) client := &dynamicLibraryClient{ diff --git a/internal/pluginhost/loader_unsupported.go b/internal/pluginhost/loader_unsupported.go index eb2567a2b..303d106c5 100644 --- a/internal/pluginhost/loader_unsupported.go +++ b/internal/pluginhost/loader_unsupported.go @@ -6,8 +6,8 @@ import "fmt" type unsupportedLoader struct{} -func (unsupportedLoader) Open(path string, host *Host) (pluginClient, error) { - return nil, fmt.Errorf("standard dynamic library plugin loading requires cgo on this platform: %s", path) +func (unsupportedLoader) Open(file pluginFile, host *Host) (pluginClient, error) { + return nil, fmt.Errorf("standard dynamic library plugin loading requires cgo on this platform: %s", file.Path) } func defaultPluginLoader() pluginLoader { diff --git a/internal/pluginhost/loader_windows.go b/internal/pluginhost/loader_windows.go index ff42eb62c..317860e79 100644 --- a/internal/pluginhost/loader_windows.go +++ b/internal/pluginhost/loader_windows.go @@ -52,8 +52,8 @@ func defaultPluginLoader() pluginLoader { return dynamicLibraryLoader{} } -func (dynamicLibraryLoader) Open(path string, host *Host) (pluginClient, error) { - dll, errLoad := syscall.LoadDLL(path) +func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, error) { + dll, errLoad := syscall.LoadDLL(file.Path) if errLoad != nil { return nil, errLoad } @@ -65,7 +65,7 @@ func (dynamicLibraryLoader) Open(path string, host *Host) (pluginClient, error) id := windowsHostCallbackID.Add(1) hostCtx := new(uintptr) *hostCtx = id - windowsHostCallbackEntries.Store(id, host) + windowsHostCallbackEntries.Store(id, dynamicHostCallbackEntry{host: host, pluginID: file.ID}) client := &dynamicLibraryClient{ dll: dll, hostCtx: hostCtx, @@ -165,8 +165,8 @@ func windowsHostCall(hostCtx uintptr, methodPtr uintptr, requestPtr uintptr, req if !okHost { return 1 } - host, okHost := rawHost.(*Host) - if !okHost || host == nil { + entry, okHost := rawHost.(dynamicHostCallbackEntry) + if !okHost || entry.host == nil { return 1 } var request []byte @@ -174,7 +174,8 @@ func windowsHostCall(hostCtx uintptr, methodPtr uintptr, requestPtr uintptr, req request = unsafe.Slice((*byte)(unsafe.Pointer(requestPtr)), requestLen) request = append([]byte(nil), request...) } - resp, errCall := host.callFromPlugin(context.Background(), windowsString(methodPtr), request) + ctx := withHostCallbackPluginID(context.Background(), entry.pluginID) + resp, errCall := entry.host.callFromPlugin(ctx, windowsString(methodPtr), request) if errCall != nil { resp = marshalRPCError("host_call_failed", errCall.Error()) } diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 6ef163116..1df108470 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -285,7 +285,7 @@ func (a *rpcPluginAdapter) openHostCallbackContext(ctx context.Context) (string, if a == nil || a.host == nil { return "", func() {} } - return a.host.openCallbackContext(ctx) + return a.host.openCallbackContextForPlugin(ctx, a.id) } func (a *rpcPluginAdapter) RegisterModels(ctx context.Context, req pluginapi.ModelRegistrationRequest) (pluginapi.ModelRegistrationResponse, error) { diff --git a/internal/pluginhost/test_helpers_test.go b/internal/pluginhost/test_helpers_test.go index 81289eb23..f169ad70a 100644 --- a/internal/pluginhost/test_helpers_test.go +++ b/internal/pluginhost/test_helpers_test.go @@ -22,11 +22,11 @@ func newTestSymbolLoader() *testSymbolLoader { return &testSymbolLoader{lookups: make(map[string]*testSymbolLookup)} } -func (l *testSymbolLoader) Open(path string, host *Host) (pluginClient, error) { +func (l *testSymbolLoader) Open(file pluginFile, host *Host) (pluginClient, error) { l.openCalls++ - lookup := l.lookups[pluginIDFromPath(path)] + lookup := l.lookups[file.ID] if lookup == nil { - return nil, fmt.Errorf("missing test plugin for %s", path) + return nil, fmt.Errorf("missing test plugin for %s", file.Path) } return lookup, nil } diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index 6ad218550..7842295c5 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -72,6 +72,13 @@ type PluginInterceptorHost interface { InterceptStreamChunk(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse } +type pluginInterceptorSkipHost interface { + InterceptRequestBeforeAuthExcept(context.Context, pluginapi.RequestInterceptRequest, string) pluginapi.RequestInterceptResponse + InterceptRequestAfterAuthExcept(context.Context, pluginapi.RequestInterceptRequest, string) pluginapi.RequestInterceptResponse + InterceptResponseExcept(context.Context, pluginapi.ResponseInterceptRequest, string) pluginapi.ResponseInterceptResponse + InterceptStreamChunkExcept(context.Context, pluginapi.StreamChunkInterceptRequest, string) pluginapi.StreamChunkInterceptResponse +} + type streamInterceptorDetector interface { HasStreamInterceptors() bool } @@ -659,10 +666,10 @@ func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entr ResponseFormat: sdktranslator.FromString(responseProtocol), Headers: modelExecutionHeaders(ctx, execOptions.Headers), Query: cloneURLValues(execOptions.Query), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), } opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, modelName, req, opts) + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, modelName, req, opts, execOptions.SkipInterceptorPluginID) resp, err := h.AuthManager.Execute(ctx, providers, req, opts) if err != nil { err = enrichAuthSelectionError(err, providers, normalizedModel) @@ -683,7 +690,7 @@ func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entr executedReq, executedOpts := afterAuthCapture.apply(req, opts) rawResponseHeaders := cloneHeader(resp.Headers) responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, responseProtocol, normalizedModel, modelName, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK) + body, responseHeaders := h.applyResponseInterceptors(ctx, responseProtocol, normalizedModel, modelName, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) return body, responseHeaders, nil } @@ -713,10 +720,10 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle OriginalRequest: rawJSON, SourceFormat: sdktranslator.FromString(handlerType), Headers: headersFromContext(ctx), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, ""), } opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptorsBeforeAuth(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) @@ -737,7 +744,7 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle executedReq, executedOpts := afterAuthCapture.apply(req, opts) rawResponseHeaders := cloneHeader(resp.Headers) responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, modelName, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.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 } @@ -788,10 +795,10 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context ResponseFormat: sdktranslator.FromString(responseProtocol), Headers: modelExecutionHeaders(ctx, execOptions.Headers), Query: cloneURLValues(execOptions.Query), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), } opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, modelName, req, opts) + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, modelName, req, opts, execOptions.SkipInterceptorPluginID) streamResult, err := h.AuthManager.ExecuteStream(ctx, providers, req, opts) if err != nil { err = enrichAuthSelectionError(err, providers, normalizedModel) @@ -846,7 +853,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context return } executedReq, executedOpts := executedRequest() - intercepted := interceptorHost.InterceptStreamChunk(ctx, pluginapi.StreamChunkInterceptRequest{ + intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ SourceFormat: responseProtocol, Model: normalizedModel, RequestedModel: modelName, @@ -856,7 +863,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context RequestBody: cloneBytes(executedReq.Payload), ChunkIndex: pluginapi.StreamChunkHeaderInitIndex, Metadata: executedOpts.Metadata, - }) + }, execOptions.SkipInterceptorPluginID) applyStreamHeaders(intercepted.Headers) streamHeaderInitialized = true } @@ -1001,7 +1008,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context payload := cloneBytes(chunk.Payload) if streamInterceptorsActive { executedReq, executedOpts := executedRequest() - intercepted := interceptorHost.InterceptStreamChunk(ctx, pluginapi.StreamChunkInterceptRequest{ + intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ SourceFormat: responseProtocol, Model: normalizedModel, RequestedModel: modelName, @@ -1013,7 +1020,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context HistoryChunks: cloneByteSlices(historyChunks), ChunkIndex: chunkIndex, Metadata: executedOpts.Metadata, - }) + }, execOptions.SkipInterceptorPluginID) applyStreamHeaders(intercepted.Headers) if len(intercepted.Body) > 0 { payload = cloneBytes(intercepted.Body) @@ -1400,12 +1407,48 @@ func mergeRequestInterceptorHeaders(current, updates http.Header, clear []string return out } -func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, handlerType, requestedModel string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Request, coreexecutor.Options) { +func interceptRequestBeforeAuth(ctx context.Context, host PluginInterceptorHost, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { + if skipPluginID != "" { + if skipper, ok := host.(pluginInterceptorSkipHost); ok { + return skipper.InterceptRequestBeforeAuthExcept(ctx, req, skipPluginID) + } + } + return host.InterceptRequestBeforeAuth(ctx, req) +} + +func interceptRequestAfterAuth(ctx context.Context, host PluginInterceptorHost, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { + if skipPluginID != "" { + if skipper, ok := host.(pluginInterceptorSkipHost); ok { + return skipper.InterceptRequestAfterAuthExcept(ctx, req, skipPluginID) + } + } + return host.InterceptRequestAfterAuth(ctx, req) +} + +func interceptResponse(ctx context.Context, host PluginInterceptorHost, req pluginapi.ResponseInterceptRequest, skipPluginID string) pluginapi.ResponseInterceptResponse { + if skipPluginID != "" { + if skipper, ok := host.(pluginInterceptorSkipHost); ok { + return skipper.InterceptResponseExcept(ctx, req, skipPluginID) + } + } + return host.InterceptResponse(ctx, req) +} + +func interceptStreamChunk(ctx context.Context, host PluginInterceptorHost, req pluginapi.StreamChunkInterceptRequest, skipPluginID string) pluginapi.StreamChunkInterceptResponse { + if skipPluginID != "" { + if skipper, ok := host.(pluginInterceptorSkipHost); ok { + return skipper.InterceptStreamChunkExcept(ctx, req, skipPluginID) + } + } + 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) { host := h.interceptorHost() if host == nil { return req, opts } - resp := host.InterceptRequestBeforeAuth(ctx, pluginapi.RequestInterceptRequest{ + resp := interceptRequestBeforeAuth(ctx, host, pluginapi.RequestInterceptRequest{ SourceFormat: handlerType, Model: req.Model, RequestedModel: requestedModel, @@ -1413,7 +1456,7 @@ func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, Headers: cloneHeader(opts.Headers), Body: cloneBytes(req.Payload), Metadata: opts.Metadata, - }) + }, skipPluginID) opts.Headers = finalInterceptorHeaders(opts.Headers, resp.Headers) if len(resp.Body) > 0 { req.Payload = cloneBytes(resp.Body) @@ -1422,12 +1465,12 @@ func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, return req, opts } -func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCapture) coreexecutor.RequestAfterAuthInterceptor { +func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCapture, 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) + resp := h.applyRequestInterceptorsAfterAuth(ctx, req, skipPluginID) if capture != nil { capture.record(req, resp) } @@ -1435,12 +1478,12 @@ func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCa } } -func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest) coreexecutor.RequestAfterAuthInterceptResponse { +func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest, skipPluginID string) coreexecutor.RequestAfterAuthInterceptResponse { host := h.interceptorHost() if !requestInterceptorsEnabled(host) { return coreexecutor.RequestAfterAuthInterceptResponse{} } - resp := host.InterceptRequestAfterAuth(ctx, pluginapi.RequestInterceptRequest{ + resp := interceptRequestAfterAuth(ctx, host, pluginapi.RequestInterceptRequest{ SourceFormat: req.SourceFormat.String(), ToFormat: req.ToFormat.String(), Model: req.Model, @@ -1449,7 +1492,7 @@ func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, Headers: cloneHeader(req.Headers), Body: cloneBytes(req.Body), Metadata: req.Metadata, - }) + }, skipPluginID) return coreexecutor.RequestAfterAuthInterceptResponse{ Headers: resp.Headers, Body: resp.Body, @@ -1457,12 +1500,12 @@ func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, } } -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) { +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) { host := h.interceptorHost() if host == nil { return body, responseHeaders } - resp := host.InterceptResponse(ctx, pluginapi.ResponseInterceptRequest{ + resp := interceptResponse(ctx, host, pluginapi.ResponseInterceptRequest{ SourceFormat: handlerType, Model: normalizedModel, RequestedModel: requestedModel, @@ -1474,7 +1517,7 @@ func (h *BaseAPIHandler) applyResponseInterceptors(ctx context.Context, handlerT Body: cloneBytes(body), StatusCode: statusCode, Metadata: opts.Metadata, - }) + }, skipPluginID) responseHeaders = downstreamHeadersAfterInterceptors(rawResponseHeaders, finalInterceptorHeaders(rawResponseHeaders, resp.Headers), PassthroughHeadersEnabled(h.Cfg)) if len(resp.Body) > 0 { body = cloneBytes(resp.Body) diff --git a/sdk/api/handlers/model_execution.go b/sdk/api/handlers/model_execution.go index e004fea2c..1057ea0e3 100644 --- a/sdk/api/handlers/model_execution.go +++ b/sdk/api/handlers/model_execution.go @@ -15,21 +15,23 @@ const ( ) type modelExecutionOptions struct { - Headers http.Header - Query url.Values - InternalSource bool + Headers http.Header + Query url.Values + InternalSource bool + SkipInterceptorPluginID string } // ModelExecutionRequest describes an internal model execution request. type ModelExecutionRequest struct { - EntryProtocol string - ExitProtocol string - Model string - Stream bool - Body []byte - Headers http.Header - Query url.Values - Alt string + EntryProtocol string + ExitProtocol string + Model string + Stream bool + Body []byte + Headers http.Header + Query url.Values + Alt string + SkipInterceptorPluginID string } // ModelExecutionResponse describes a non-streaming internal model execution response. @@ -71,14 +73,18 @@ func (e *ModelExecutionStreamError) Error() string { } // ExecuteModel executes an internal non-streaming model request. +// Host model callbacks are non-recursive for their caller: when +// SkipInterceptorPluginID is set, that plugin's interceptors are skipped for the +// nested model execution while other plugins may still run. func (h *BaseAPIHandler) ExecuteModel(ctx context.Context, req ModelExecutionRequest) (ModelExecutionResponse, *interfaces.ErrorMessage) { if req.Stream { return ModelExecutionResponse{}, modelExecutionModeError("ExecuteModel requires Stream=false") } body, headers, errMsg := h.executeWithAuthManagerFormats(ctx, req.EntryProtocol, req.ExitProtocol, req.Model, cloneBytes(req.Body), req.Alt, false, modelExecutionOptions{ - Headers: req.Headers, - Query: req.Query, - InternalSource: true, + Headers: req.Headers, + Query: req.Query, + InternalSource: true, + SkipInterceptorPluginID: req.SkipInterceptorPluginID, }) if errMsg != nil { return ModelExecutionResponse{}, errMsg @@ -91,14 +97,18 @@ func (h *BaseAPIHandler) ExecuteModel(ctx context.Context, req ModelExecutionReq } // ExecuteModelStream executes an internal streaming model request. +// Host model callbacks are non-recursive for their caller: when +// SkipInterceptorPluginID is set, that plugin's interceptors are skipped for the +// nested model execution while other plugins may still run. func (h *BaseAPIHandler) ExecuteModelStream(ctx context.Context, req ModelExecutionRequest) (ModelExecutionStream, *interfaces.ErrorMessage) { if !req.Stream { return ModelExecutionStream{}, modelExecutionModeError("ExecuteModelStream requires Stream=true") } dataChan, headers, errChan := h.executeStreamWithAuthManagerFormats(ctx, req.EntryProtocol, req.ExitProtocol, req.Model, cloneBytes(req.Body), req.Alt, false, modelExecutionOptions{ - Headers: req.Headers, - Query: req.Query, - InternalSource: true, + Headers: req.Headers, + Query: req.Query, + InternalSource: true, + SkipInterceptorPluginID: req.SkipInterceptorPluginID, }) chunks, errMsg := prepareModelExecutionStream(ctx, dataChan, errChan) if errMsg != nil { diff --git a/sdk/api/handlers/model_execution_test.go b/sdk/api/handlers/model_execution_test.go index 642fcf42a..37f98d10a 100644 --- a/sdk/api/handlers/model_execution_test.go +++ b/sdk/api/handlers/model_execution_test.go @@ -14,6 +14,7 @@ import ( coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" ) @@ -33,6 +34,61 @@ type modelExecutionStatusHeaderError struct { headers http.Header } +type modelExecutionSkipHost struct { + beforeSkip string + afterSkip string + respSkip string + streamSkip []string +} + +func (h *modelExecutionSkipHost) InterceptRequestBeforeAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + panic("InterceptRequestBeforeAuth called without skip") +} + +func (h *modelExecutionSkipHost) InterceptRequestAfterAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + panic("InterceptRequestAfterAuth called without skip") +} + +func (h *modelExecutionSkipHost) InterceptResponse(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse { + panic("InterceptResponse called without skip") +} + +func (h *modelExecutionSkipHost) InterceptStreamChunk(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { + panic("InterceptStreamChunk called without skip") +} + +func (h *modelExecutionSkipHost) InterceptRequestBeforeAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { + h.beforeSkip = skipPluginID + return pluginapi.RequestInterceptResponse{ + Headers: cloneHeader(req.Headers), + Body: cloneBytes(req.Body), + } +} + +func (h *modelExecutionSkipHost) InterceptRequestAfterAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { + h.afterSkip = skipPluginID + return pluginapi.RequestInterceptResponse{ + Headers: cloneHeader(req.Headers), + Body: cloneBytes(req.Body), + } +} + +func (h *modelExecutionSkipHost) InterceptResponseExcept(ctx context.Context, req pluginapi.ResponseInterceptRequest, skipPluginID string) pluginapi.ResponseInterceptResponse { + h.respSkip = skipPluginID + return pluginapi.ResponseInterceptResponse{ + Headers: cloneHeader(req.ResponseHeaders), + Body: cloneBytes(req.Body), + } +} + +func (h *modelExecutionSkipHost) InterceptStreamChunkExcept(ctx context.Context, req pluginapi.StreamChunkInterceptRequest, skipPluginID string) pluginapi.StreamChunkInterceptResponse { + h.streamSkip = append(h.streamSkip, skipPluginID) + return pluginapi.StreamChunkInterceptResponse{ + Headers: cloneHeader(req.ResponseHeaders), + Body: cloneBytes(req.Body), + } +} + func (e modelExecutionStatusHeaderError) Error() string { return e.message } @@ -195,6 +251,32 @@ func TestExecuteModelCarriesEntryAndExitProtocols(t *testing.T) { } } +func TestExecuteModelSkipsOriginatingPluginInterceptors(t *testing.T) { + model := "model-execution-skip-origin-model" + requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model)) + executor := &modelExecutionCaptureExecutor{} + handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{}) + skipHost := &modelExecutionSkipHost{} + handler.SetPluginHost(skipHost) + + resp, errMsg := handler.ExecuteModel(context.Background(), ModelExecutionRequest{ + EntryProtocol: "openai", + ExitProtocol: "openai", + Model: model, + Body: requestBody, + SkipInterceptorPluginID: "origin-plugin", + }) + if errMsg != nil { + t.Fatalf("ExecuteModel() error = %+v", errMsg) + } + if string(resp.Body) != "model-execution-ok" { + t.Fatalf("body = %q, want executor response", resp.Body) + } + if skipHost.beforeSkip != "origin-plugin" || skipHost.afterSkip != "origin-plugin" || skipHost.respSkip != "origin-plugin" { + t.Fatalf("skip ids = before:%q after:%q response:%q, want origin-plugin", skipHost.beforeSkip, skipHost.afterSkip, skipHost.respSkip) + } +} + func TestExecuteModelStream(t *testing.T) { model := "model-execution-stream-model" requestBody := []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, model)) @@ -269,6 +351,52 @@ func TestExecuteModelStream(t *testing.T) { } } +func TestExecuteModelStreamSkipsOriginatingPluginInterceptors(t *testing.T) { + model := "model-execution-stream-skip-origin-model" + requestBody := []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, model)) + executor := &modelExecutionCaptureExecutor{ + stream: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk, 1) + chunks <- coreexecutor.StreamChunk{Payload: []byte("stream-one")} + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil + }, + } + handler := newModelExecutionHandler(t, model, executor, &sdkconfig.SDKConfig{}) + skipHost := &modelExecutionSkipHost{} + handler.SetPluginHost(skipHost) + + stream, errMsg := handler.ExecuteModelStream(context.Background(), ModelExecutionRequest{ + EntryProtocol: "openai", + ExitProtocol: "openai", + Model: model, + Stream: true, + Body: requestBody, + SkipInterceptorPluginID: "origin-plugin", + }) + if errMsg != nil { + t.Fatalf("ExecuteModelStream() error = %+v", errMsg) + } + chunk, ok := <-stream.Chunks + if !ok { + t.Fatal("stream chunks closed before payload") + } + if string(chunk.Payload) != "stream-one" { + t.Fatalf("stream chunk payload = %q, want stream-one", chunk.Payload) + } + if skipHost.beforeSkip != "origin-plugin" || skipHost.afterSkip != "origin-plugin" { + t.Fatalf("request skip ids = before:%q after:%q, want origin-plugin", skipHost.beforeSkip, skipHost.afterSkip) + } + if len(skipHost.streamSkip) == 0 { + t.Fatal("stream interceptor was not called with skip") + } + for _, skipID := range skipHost.streamSkip { + if skipID != "origin-plugin" { + t.Fatalf("stream skip id = %q, want origin-plugin", skipID) + } + } +} + func TestExecuteModelStreamStartupError(t *testing.T) { model := "model-execution-stream-startup-error-model" requestBody := []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, model))