From 9f940f162fbce4c2babcf35c208309b2452b3d8e Mon Sep 17 00:00:00 2001 From: sususu98 <33882693+sususu98@users.noreply.github.com> Date: Tue, 16 Jun 2026 17:31:11 +0800 Subject: [PATCH] fix(pluginhost): keep stream callbacks alive until stream close Keep RPC streaming executor callback scopes alive until async streams close, detach nested host.model.execute_stream contexts from request cancellation, and clean up the stream bridge on stream completion. --- internal/pluginhost/host_callbacks.go | 77 ----------- .../pluginhost/host_model_stream_callbacks.go | 87 ++++++++++++ .../host_model_stream_callbacks_test.go | 76 +++++++++++ internal/pluginhost/rpc_client.go | 29 ---- internal/pluginhost/rpc_client_stream.go | 80 +++++++++++ internal/pluginhost/rpc_client_stream_test.go | 127 ++++++++++++++++++ 6 files changed, 370 insertions(+), 106 deletions(-) create mode 100644 internal/pluginhost/host_model_stream_callbacks.go create mode 100644 internal/pluginhost/host_model_stream_callbacks_test.go create mode 100644 internal/pluginhost/rpc_client_stream.go create mode 100644 internal/pluginhost/rpc_client_stream_test.go diff --git a/internal/pluginhost/host_callbacks.go b/internal/pluginhost/host_callbacks.go index 615c7dc4a..f44874968 100644 --- a/internal/pluginhost/host_callbacks.go +++ b/internal/pluginhost/host_callbacks.go @@ -291,83 +291,6 @@ func (h *Host) callHostModelExecute(ctx context.Context, request []byte) ([]byte }) } -func (h *Host) callHostModelExecuteStream(ctx context.Context, request []byte) ([]byte, error) { - var req rpcHostModelExecutionRequest - if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { - return nil, fmt.Errorf("decode host model execution stream request: %w", errUnmarshal) - } - if !req.Stream { - return nil, fmt.Errorf("host.model.execute_stream requires stream=true") - } - executor := h.currentModelExecutor() - 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, skipPluginID)) - if errMsg != nil { - cancel() - return nil, modelExecutionError(errMsg) - } - streamID := "" - if h != nil && h.modelStreams != nil { - streamID = h.modelStreams.open(req.HostCallbackID, stream.Chunks, cancel) - } - if streamID == "" { - cancel() - return nil, fmt.Errorf("host model stream bridge is unavailable") - } - if req.HostCallbackID != "" { - h.addCallbackCleanup(req.HostCallbackID, func() { - h.modelStreams.close(streamID) - }) - } - return marshalRPCResult(pluginapi.HostModelStreamResponse{ - StatusCode: stream.StatusCode, - Headers: cloneHeader(stream.Headers), - StreamID: streamID, - }) -} - -func (h *Host) callHostModelStreamRead(ctx context.Context, request []byte) ([]byte, error) { - var req pluginapi.HostModelStreamReadRequest - if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { - return nil, fmt.Errorf("decode host model stream read request: %w", errUnmarshal) - } - if h == nil || h.modelStreams == nil { - return nil, fmt.Errorf("host model stream bridge is unavailable") - } - chunk, done, errRead := h.modelStreams.read(ctx, req.StreamID) - if errRead != nil { - return nil, errRead - } - resp := pluginapi.HostModelStreamReadResponse{ - Payload: append([]byte(nil), chunk.Payload...), - Done: done, - } - if chunk.Err != nil { - resp.Error = chunk.Err.Error() - resp.Done = true - } - return marshalRPCResult(resp) -} - -func (h *Host) callHostModelStreamClose(request []byte) ([]byte, error) { - var req pluginapi.HostModelStreamCloseRequest - if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { - return nil, fmt.Errorf("decode host model stream close request: %w", errUnmarshal) - } - if h != nil && h.modelStreams != nil { - h.modelStreams.close(req.StreamID) - } - return marshalRPCResult(rpcEmptyResponse{}) -} - func modelExecutionRequestFromPlugin(req pluginapi.HostModelExecutionRequest, skipPluginID string) handlers.ModelExecutionRequest { return handlers.ModelExecutionRequest{ EntryProtocol: req.EntryProtocol, diff --git a/internal/pluginhost/host_model_stream_callbacks.go b/internal/pluginhost/host_model_stream_callbacks.go new file mode 100644 index 000000000..be65e5fab --- /dev/null +++ b/internal/pluginhost/host_model_stream_callbacks.go @@ -0,0 +1,87 @@ +package pluginhost + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func (h *Host) callHostModelExecuteStream(ctx context.Context, request []byte) ([]byte, error) { + var req rpcHostModelExecutionRequest + if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { + return nil, fmt.Errorf("decode host model execution stream request: %w", errUnmarshal) + } + if !req.Stream { + return nil, fmt.Errorf("host.model.execute_stream requires stream=true") + } + executor := h.currentModelExecutor() + if executor == nil { + return nil, fmt.Errorf("host model executor is unavailable") + } + skipPluginID := h.callbackCallerPluginID(ctx, req.HostCallbackID) + callbackCtx := h.resolveCallbackContext(req.HostCallbackID, ctx) + if callbackCtx == nil { + callbackCtx = context.Background() + } + // Detach request cancellation while preserving callback values; callback cleanup owns the model stream lifetime. + streamCtx, cancel := context.WithCancel(context.WithoutCancel(callbackCtx)) + stream, errMsg := executor.ExecuteModelStream(streamCtx, modelExecutionRequestFromPlugin(req.HostModelExecutionRequest, skipPluginID)) + if errMsg != nil { + cancel() + return nil, modelExecutionError(errMsg) + } + streamID := "" + if h.modelStreams != nil { + streamID = h.modelStreams.open(req.HostCallbackID, stream.Chunks, cancel) + } + if streamID == "" { + cancel() + return nil, fmt.Errorf("host model stream bridge is unavailable") + } + if req.HostCallbackID != "" { + h.addCallbackCleanup(req.HostCallbackID, func() { + h.modelStreams.close(streamID) + }) + } + return marshalRPCResult(pluginapi.HostModelStreamResponse{ + StatusCode: stream.StatusCode, + Headers: cloneHeader(stream.Headers), + StreamID: streamID, + }) +} + +func (h *Host) callHostModelStreamRead(ctx context.Context, request []byte) ([]byte, error) { + var req pluginapi.HostModelStreamReadRequest + if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { + return nil, fmt.Errorf("decode host model stream read request: %w", errUnmarshal) + } + if h == nil || h.modelStreams == nil { + return nil, fmt.Errorf("host model stream bridge is unavailable") + } + chunk, done, errRead := h.modelStreams.read(ctx, req.StreamID) + if errRead != nil { + return nil, errRead + } + resp := pluginapi.HostModelStreamReadResponse{ + Payload: append([]byte(nil), chunk.Payload...), + Done: done, + } + if chunk.Err != nil { + resp.Error = chunk.Err.Error() + resp.Done = true + } + return marshalRPCResult(resp) +} + +func (h *Host) callHostModelStreamClose(request []byte) ([]byte, error) { + var req pluginapi.HostModelStreamCloseRequest + if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { + return nil, fmt.Errorf("decode host model stream close request: %w", errUnmarshal) + } + if h != nil && h.modelStreams != nil { + h.modelStreams.close(req.StreamID) + } + return marshalRPCResult(rpcEmptyResponse{}) +} diff --git a/internal/pluginhost/host_model_stream_callbacks_test.go b/internal/pluginhost/host_model_stream_callbacks_test.go new file mode 100644 index 000000000..bc8f29283 --- /dev/null +++ b/internal/pluginhost/host_model_stream_callbacks_test.go @@ -0,0 +1,76 @@ +package pluginhost + +import ( + "context" + "encoding/json" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func TestHostModelExecuteStreamDetachesFromCallbackParentCancel(t *testing.T) { + host := New() + ctxSeen := make(chan context.Context, 1) + host.SetModelExecutor(&fakeHostModelExecutor{ + executeModelStream: func(ctx context.Context, req handlers.ModelExecutionRequest) (handlers.ModelExecutionStream, *interfaces.ErrorMessage) { + ctxSeen <- ctx + return handlers.ModelExecutionStream{ + StatusCode: http.StatusOK, + Chunks: make(chan handlers.ModelExecutionChunk), + }, nil + }, + }) + parentCtx, cancelParent := context.WithCancel(context.Background()) + callbackID, closeCallback := host.openCallbackContext(parentCtx) + defer closeCallback() + + rawReq, errMarshal := json.Marshal(rpcHostModelExecutionRequest{ + HostModelExecutionRequest: pluginapi.HostModelExecutionRequest{ + EntryProtocol: "openai", + ExitProtocol: "openai", + Model: "model-1", + Stream: true, + Body: []byte(`{"stream":true}`), + }, + HostCallbackID: callbackID, + }) + if errMarshal != nil { + t.Fatalf("marshal request: %v", errMarshal) + } + rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostModelExecuteStream, rawReq) + if errCall != nil { + t.Fatalf("callFromPlugin() error = %v", errCall) + } + resp, errDecode := decodeRPCEnvelope[pluginapi.HostModelStreamResponse](rawResp) + if errDecode != nil { + t.Fatalf("decode response: %v", errDecode) + } + if resp.StreamID == "" { + t.Fatalf("stream id is empty: %#v", resp) + } + + var streamCtx context.Context + select { + case streamCtx = <-ctxSeen: + case <-time.After(time.Second): + t.Fatal("model executor was not called") + } + cancelParent() + select { + case <-streamCtx.Done(): + t.Fatal("stream context was canceled by callback parent context") + default: + } + + closeCallback() + select { + case <-streamCtx.Done(): + case <-time.After(time.Second): + t.Fatal("stream context was not canceled after callback scope closed") + } +} diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 1df108470..e4b1fb70a 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -377,35 +377,6 @@ func (a *rpcPluginAdapter) Execute(ctx context.Context, req pluginapi.ExecutorRe }) } -func (a *rpcPluginAdapter) ExecuteStream(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { - if a == nil || a.host == nil || a.host.streams == nil { - return pluginapi.ExecutorStreamResponse{}, fmt.Errorf("plugin stream bridge is unavailable") - } - streamID, chunks, cleanup := a.host.streams.open(ctx) - callbackID, closeCallback := a.openHostCallbackContext(ctx) - defer closeCallback() - rpcReq := rpcExecutorRequest{ - ExecutorRequest: req, - StreamID: streamID, - HostCallbackID: callbackID, - } - resp, errCall := callPlugin[rpcExecutorStreamResponse](ctx, a.client, pluginabi.MethodExecutorExecuteStream, rpcReq) - if errCall != nil { - cleanup() - return pluginapi.ExecutorStreamResponse{}, errCall - } - if len(resp.Chunks) > 0 { - cleanup() - out := make(chan pluginapi.ExecutorStreamChunk, len(resp.Chunks)) - for _, chunk := range resp.Chunks { - out <- chunk - } - close(out) - return pluginapi.ExecutorStreamResponse{Headers: resp.Headers, Chunks: out}, nil - } - return pluginapi.ExecutorStreamResponse{Headers: resp.Headers, Chunks: chunks}, nil -} - func (a *rpcPluginAdapter) CountTokens(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { callbackID, closeCallback := a.openHostCallbackContext(ctx) defer closeCallback() diff --git a/internal/pluginhost/rpc_client_stream.go b/internal/pluginhost/rpc_client_stream.go new file mode 100644 index 000000000..87939146a --- /dev/null +++ b/internal/pluginhost/rpc_client_stream.go @@ -0,0 +1,80 @@ +package pluginhost + +import ( + "context" + "fmt" + "sync" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func (a *rpcPluginAdapter) ExecuteStream(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorStreamResponse, error) { + if a == nil || a.host == nil || a.host.streams == nil { + return pluginapi.ExecutorStreamResponse{}, fmt.Errorf("plugin stream bridge is unavailable") + } + streamID, chunks, cleanupStream := a.host.streams.open(ctx) + callbackID, closeCallback := a.openHostCallbackContext(ctx) + cleanup := combinedCleanup(cleanupStream, closeCallback) + rpcReq := rpcExecutorRequest{ + ExecutorRequest: req, + StreamID: streamID, + HostCallbackID: callbackID, + } + resp, errCall := callPlugin[rpcExecutorStreamResponse](ctx, a.client, pluginabi.MethodExecutorExecuteStream, rpcReq) + if errCall != nil { + cleanup() + return pluginapi.ExecutorStreamResponse{}, errCall + } + if len(resp.Chunks) > 0 { + cleanup() + out := make(chan pluginapi.ExecutorStreamChunk, len(resp.Chunks)) + for _, chunk := range resp.Chunks { + out <- chunk + } + close(out) + return pluginapi.ExecutorStreamResponse{Headers: resp.Headers, Chunks: out}, nil + } + // Async streaming plugins can return before they finish emitting chunks, so keep callbacks alive until the stream ends. + return pluginapi.ExecutorStreamResponse{ + Headers: resp.Headers, + Chunks: cleanupWhenStreamDone(ctx, chunks, cleanup), + }, nil +} + +func combinedCleanup(cleanups ...func()) func() { + var once sync.Once + return func() { + once.Do(func() { + for _, cleanup := range cleanups { + if cleanup != nil { + cleanup() + } + } + }) + } +} + +func cleanupWhenStreamDone(ctx context.Context, chunks <-chan pluginapi.ExecutorStreamChunk, cleanup func()) <-chan pluginapi.ExecutorStreamChunk { + out := make(chan pluginapi.ExecutorStreamChunk) + go func() { + defer func() { + if cleanup != nil { + cleanup() + } + close(out) + }() + var done <-chan struct{} + if ctx != nil { + done = ctx.Done() + } + for chunk := range chunks { + select { + case out <- chunk: + case <-done: + return + } + } + }() + return out +} diff --git a/internal/pluginhost/rpc_client_stream_test.go b/internal/pluginhost/rpc_client_stream_test.go new file mode 100644 index 000000000..6e293a248 --- /dev/null +++ b/internal/pluginhost/rpc_client_stream_test.go @@ -0,0 +1,127 @@ +package pluginhost + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func TestRPCExecuteStreamKeepsHostCallbackScopeUntilStreamCloses(t *testing.T) { + host := New() + client := newStreamCallbackPluginClient() + adapter := &rpcPluginAdapter{ + id: "stream-plugin", + host: host, + client: client, + } + + stream, errStream := adapter.ExecuteStream(context.Background(), pluginapi.ExecutorRequest{Stream: true}) + if errStream != nil { + t.Fatalf("ExecuteStream() error = %v", errStream) + } + waitForStreamCallbackPlugin(t, client) + if client.callbackID == "" { + t.Fatal("host callback id is empty") + } + if !callbackContextExists(host, client.callbackID) { + t.Fatal("host callback scope closed before plugin stream closed") + } + + closeReq, errMarshal := json.Marshal(rpcStreamCloseRequest{StreamID: client.streamID}) + if errMarshal != nil { + t.Fatalf("marshal close request: %v", errMarshal) + } + if _, errClose := host.callFromPlugin(context.Background(), pluginabi.MethodHostStreamClose, closeReq); errClose != nil { + t.Fatalf("close stream: %v", errClose) + } + for range stream.Chunks { + } + + if callbackContextExists(host, client.callbackID) { + t.Fatal("host callback scope remained open after plugin stream closed") + } +} + +func TestRPCExecuteStreamClosesHostCallbackScopeOnContextCancelWhileChunkPending(t *testing.T) { + host := New() + client := newStreamCallbackPluginClient() + adapter := &rpcPluginAdapter{ + id: "stream-plugin", + host: host, + client: client, + } + ctx, cancel := context.WithCancel(context.Background()) + stream, errStream := adapter.ExecuteStream(ctx, pluginapi.ExecutorRequest{Stream: true}) + if errStream != nil { + t.Fatalf("ExecuteStream() error = %v", errStream) + } + waitForStreamCallbackPlugin(t, client) + + emitReq, errMarshal := json.Marshal(rpcStreamEmitRequest{StreamID: client.streamID, Payload: []byte("pending")}) + if errMarshal != nil { + t.Fatalf("marshal emit request: %v", errMarshal) + } + if _, errEmit := host.callFromPlugin(context.Background(), pluginabi.MethodHostStreamEmit, emitReq); errEmit != nil { + t.Fatalf("emit stream: %v", errEmit) + } + cancel() + for range stream.Chunks { + } + + if callbackContextExists(host, client.callbackID) { + t.Fatal("host callback scope remained open after context cancel") + } +} + +func callbackContextExists(host *Host, callbackID string) bool { + if host == nil || host.callbackContexts == nil { + return false + } + host.callbackContexts.mu.RLock() + _, exists := host.callbackContexts.contexts[callbackID] + host.callbackContexts.mu.RUnlock() + return exists +} + +type streamCallbackPluginClient struct { + called chan struct{} + streamID string + callbackID string +} + +func newStreamCallbackPluginClient() *streamCallbackPluginClient { + return &streamCallbackPluginClient{called: make(chan struct{})} +} + +func (c *streamCallbackPluginClient) Call(ctx context.Context, method string, request []byte) ([]byte, error) { + if method != pluginabi.MethodExecutorExecuteStream { + return nil, fmt.Errorf("method = %s, want %s", method, pluginabi.MethodExecutorExecuteStream) + } + var req rpcExecutorRequest + if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { + return nil, fmt.Errorf("decode executor stream request: %w", errUnmarshal) + } + c.streamID = req.StreamID + c.callbackID = req.HostCallbackID + close(c.called) + return marshalRPCResult(rpcExecutorStreamResponse{ + Headers: http.Header{"Content-Type": []string{"text/event-stream"}}, + }) +} + +func (c *streamCallbackPluginClient) Shutdown() {} + +func waitForStreamCallbackPlugin(t *testing.T, client *streamCallbackPluginClient) { + t.Helper() + select { + case <-client.called: + case <-time.After(time.Second): + t.Fatal("plugin stream method was not called") + } +}