diff --git a/internal/pluginhost/stream_bridge.go b/internal/pluginhost/stream_bridge.go index e240b062..9002e282 100644 --- a/internal/pluginhost/stream_bridge.go +++ b/internal/pluginhost/stream_bridge.go @@ -158,12 +158,7 @@ func (s *streamBridgeStream) emit(ctx context.Context, chunk pluginapi.ExecutorS return errStreamBridgeClosed case s.emits <- request: } - select { - case err := <-request.done: - return err - case <-ctx.Done(): - return ctx.Err() - } + return <-request.done } func (s *streamBridgeStream) close(errorMessage string) { @@ -205,6 +200,7 @@ func (b *streamBridge) open(ctx context.Context) (string, <-chan pluginapi.Execu stream.abortStream() } if ctx != nil && ctx.Done() != nil { + // Abort streams canceled before ExecuteStream can install cleanupWhenStreamDone. go func() { <-ctx.Done() cleanup() diff --git a/internal/pluginhost/stream_bridge_test.go b/internal/pluginhost/stream_bridge_test.go index 4892dce1..8cdb1a6a 100644 --- a/internal/pluginhost/stream_bridge_test.go +++ b/internal/pluginhost/stream_bridge_test.go @@ -71,6 +71,25 @@ func TestStreamBridgeCloseUnblocksPendingEmit(t *testing.T) { } } +func TestStreamBridgeEmitUsesAcceptedPumpResultAfterContextCancellation(t *testing.T) { + for range 1000 { + ctx, cancel := context.WithCancel(context.Background()) + stream := &streamBridgeStream{ + emits: make(chan streamBridgeEmit), + closed: make(chan struct{}), + } + go func() { + request := <-stream.emits + cancel() + request.done <- nil + }() + + if err := stream.emit(ctx, pluginapi.ExecutorStreamChunk{Payload: []byte("accepted")}); err != nil { + t.Fatalf("accepted emit returned error: %v", err) + } + } +} + func TestStreamBridgeAbortClosesSaturatedStreamWithoutConsumer(t *testing.T) { bridge := newStreamBridge() streamID, chunks, cleanup := bridge.open(context.Background())