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

@@ -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,

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{})
+}

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")
+	}
+}

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()

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
+}

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")
+	}
+}