package pluginhost import ( "context" "errors" "fmt" "strconv" "sync" "sync/atomic" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi") type streamBridge struct { next atomic.Uint64 mu sync.Mutex streams map[string]*streamBridgeStream} const streamBridgeBufferSize = 16 var errStreamBridgeClosed = errors.New("stream is not open") type streamBridgeStream struct { chunks chan pluginapi.ExecutorStreamChunk emits chan streamBridgeEmit closes chan streamBridgeClose closed chan struct{} finished chan struct{} abort chan struct{} closeOnce sync.Once abortOnce sync.Once} type streamBridgeEmit struct { ctx context.Context chunk pluginapi.ExecutorStreamChunk done chan error} type streamBridgeClose struct { errorMessage string accepted chan struct{}} type rpcStreamEmitRequest struct { StreamID string `json:"stream_id"` Payload []byte `json:"payload,omitempty"` Error string `json:"error,omitempty"`} type rpcStreamCloseRequest struct { StreamID string `json:"stream_id"` Error string `json:"error,omitempty"`} func newStreamBridge() *streamBridge { return &streamBridge{streams: make(map[string]*streamBridgeStream)}} func newStreamBridgeStream() *streamBridgeStream { stream := &streamBridgeStream{ chunks: make(chan pluginapi.ExecutorStreamChunk), emits: make(chan streamBridgeEmit), closes: make(chan streamBridgeClose), closed: make(chan struct{}), finished: make(chan struct{}), abort: make(chan struct{}), } go stream.run() return stream} func (s *streamBridgeStream) run() { defer func() { s.markClosed() close(s.chunks) close(s.finished) }() queue := make([]pluginapi.ExecutorStreamChunk, 0, streamBridgeBufferSize) for { var emitC <-chan streamBridgeEmit if len(queue) < streamBridgeBufferSize { emitC = s.emits } var outputC chan pluginapi.ExecutorStreamChunk var next pluginapi.ExecutorStreamChunk if len(queue) > 0 { outputC = s.chunks next = queue[0] } select { case <-s.abort: return case request := <-s.closes: s.markClosed() close(request.accepted) if request.errorMessage != "" { queue = append(queue, pluginapi.ExecutorStreamChunk{Err: fmt.Errorf("%s", request.errorMessage)}) } for len(queue) > 0 { select { case <-s.abort: return case s.chunks <- queue[0]: queue = queue[1:] } } return case request := <-emitC: if err := request.ctx.Err(); err != nil { request.done <- err continue } queue = append(queue, request.chunk) request.done <- nil case outputC <- next: queue = queue[1:] } }} func (s *streamBridgeStream) markClosed() { if s == nil { return } s.closeOnce.Do(func() { close(s.closed) })} func (s *streamBridgeStream) abortStream() { if s == nil { return } s.abortOnce.Do(func() { s.markClosed() close(s.abort) })} func (s *streamBridgeStream) emit(ctx context.Context, chunk pluginapi.ExecutorStreamChunk) error { if s == nil { return errStreamBridgeClosed } if ctx == nil { ctx = context.Background() } request := streamBridgeEmit{ ctx: ctx, chunk: chunk, done: make(chan error, 1), } select { case <-ctx.Done(): return ctx.Err() case <-s.closed: return errStreamBridgeClosed case s.emits <- request: } return <-request.done} func (s *streamBridgeStream) close(errorMessage string) { if s == nil { return } request := streamBridgeClose{ errorMessage: errorMessage, accepted: make(chan struct{}), } select { case <-s.finished: return case s.closes <- request: } select { case <-request.accepted: case <-s.finished: }} func (b *streamBridge) open(ctx context.Context) (string, <-chan pluginapi.ExecutorStreamChunk, func()) { if b == nil { chunks := make(chan pluginapi.ExecutorStreamChunk) close(chunks) return "", chunks, func() {} } id := strconv.FormatUint(b.next.Add(1), 10) stream := newStreamBridgeStream() b.mu.Lock() b.streams[id] = stream b.mu.Unlock() cleanup := func() { b.mu.Lock() if b.streams[id] == stream { delete(b.streams, id) } b.mu.Unlock() stream.abortStream() } if ctx != nil && ctx.Done() != nil { // Abort streams canceled before ExecuteStream can install cleanupWhenStreamDone. go func() { <-ctx.Done() cleanup() }() } return id, stream.chunks, cleanup} func (b *streamBridge) emit(ctx context.Context, id string, chunk pluginapi.ExecutorStreamChunk) error { if b == nil || id == "" { return fmt.Errorf("stream id is required") } b.mu.Lock() stream := b.streams[id] b.mu.Unlock() if stream == nil { return fmt.Errorf("stream %s is not open", id) } if err := stream.emit(ctx, chunk); err != nil { if errors.Is(err, errStreamBridgeClosed) { return fmt.Errorf("stream %s is not open", id) } return err } return nil} func (b *streamBridge) close(id string, errorMessage string) { if b == nil || id == "" { return } b.mu.Lock() stream := b.streams[id] delete(b.streams, id) b.mu.Unlock() if stream == nil { return } stream.close(errorMessage)}
package pluginhost
import (
"context"
"errors"
"fmt"
"strconv"
"sync"
"sync/atomic"
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
)
type streamBridge struct {
next atomic.Uint64
mu sync.Mutex
streams map[string]*streamBridgeStream
}
const streamBridgeBufferSize = 16
var errStreamBridgeClosed = errors.New("stream is not open")
type streamBridgeStream struct {
chunks chan pluginapi.ExecutorStreamChunk
emits chan streamBridgeEmit
closes chan streamBridgeClose
closed chan struct{}
finished chan struct{}
abort chan struct{}
closeOnce sync.Once
abortOnce sync.Once
}
type streamBridgeEmit struct {
ctx context.Context
chunk pluginapi.ExecutorStreamChunk
done chan error
}
type streamBridgeClose struct {
errorMessage string
accepted chan struct{}
}
type rpcStreamEmitRequest struct {
StreamID string `json:"stream_id"`
Payload []byte `json:"payload,omitempty"`
Error string `json:"error,omitempty"`
}
type rpcStreamCloseRequest struct {
StreamID string `json:"stream_id"`
Error string `json:"error,omitempty"`
}
func newStreamBridge() *streamBridge {
return &streamBridge{streams: make(map[string]*streamBridgeStream)}
}
func newStreamBridgeStream() *streamBridgeStream {
stream := &streamBridgeStream{
chunks: make(chan pluginapi.ExecutorStreamChunk),
emits: make(chan streamBridgeEmit),
closes: make(chan streamBridgeClose),
closed: make(chan struct{}),
finished: make(chan struct{}),
abort: make(chan struct{}),
}
go stream.run()
return stream
}
func (s *streamBridgeStream) run() {
defer func() {
s.markClosed()
close(s.chunks)
close(s.finished)
}()
queue := make([]pluginapi.ExecutorStreamChunk, 0, streamBridgeBufferSize)
for {
var emitC <-chan streamBridgeEmit
if len(queue) < streamBridgeBufferSize {
emitC = s.emits
}
var outputC chan pluginapi.ExecutorStreamChunk
var next pluginapi.ExecutorStreamChunk
if len(queue) > 0 {
outputC = s.chunks
next = queue[0]
}
select {
case <-s.abort:
return
case request := <-s.closes:
s.markClosed()
close(request.accepted)
if request.errorMessage != "" {
queue = append(queue, pluginapi.ExecutorStreamChunk{Err: fmt.Errorf("%s", request.errorMessage)})
}
for len(queue) > 0 {
select {
case <-s.abort:
return
case s.chunks <- queue[0]:
queue = queue[1:]
}
}
return
case request := <-emitC:
if err := request.ctx.Err(); err != nil {
request.done <- err
continue
}
queue = append(queue, request.chunk)
request.done <- nil
case outputC <- next:
queue = queue[1:]
}
}
}
func (s *streamBridgeStream) markClosed() {
if s == nil {
return
}
s.closeOnce.Do(func() { close(s.closed) })
}
func (s *streamBridgeStream) abortStream() {
if s == nil {
return
}
s.abortOnce.Do(func() {
s.markClosed()
close(s.abort)
})
}
func (s *streamBridgeStream) emit(ctx context.Context, chunk pluginapi.ExecutorStreamChunk) error {
if s == nil {
return errStreamBridgeClosed
}
if ctx == nil {
ctx = context.Background()
}
request := streamBridgeEmit{
ctx: ctx,
chunk: chunk,
done: make(chan error, 1),
}
select {
case <-ctx.Done():
return ctx.Err()
case <-s.closed:
return errStreamBridgeClosed
case s.emits <- request:
}
return <-request.done
}
func (s *streamBridgeStream) close(errorMessage string) {
if s == nil {
return
}
request := streamBridgeClose{
errorMessage: errorMessage,
accepted: make(chan struct{}),
}
select {
case <-s.finished:
return
case s.closes <- request:
}
select {
case <-request.accepted:
case <-s.finished:
}
}
func (b *streamBridge) open(ctx context.Context) (string, <-chan pluginapi.ExecutorStreamChunk, func()) {
if b == nil {
chunks := make(chan pluginapi.ExecutorStreamChunk)
close(chunks)
return "", chunks, func() {}
}
id := strconv.FormatUint(b.next.Add(1), 10)
stream := newStreamBridgeStream()
b.mu.Lock()
b.streams[id] = stream
b.mu.Unlock()
cleanup := func() {
b.mu.Lock()
if b.streams[id] == stream {
delete(b.streams, id)
}
b.mu.Unlock()
stream.abortStream()
}
if ctx != nil && ctx.Done() != nil {
// Abort streams canceled before ExecuteStream can install cleanupWhenStreamDone.
go func() {
<-ctx.Done()
cleanup()
}()
}
return id, stream.chunks, cleanup
}
func (b *streamBridge) emit(ctx context.Context, id string, chunk pluginapi.ExecutorStreamChunk) error {
if b == nil || id == "" {
return fmt.Errorf("stream id is required")
}
b.mu.Lock()
stream := b.streams[id]
b.mu.Unlock()
if stream == nil {
return fmt.Errorf("stream %s is not open", id)
}
if err := stream.emit(ctx, chunk); err != nil {
if errors.Is(err, errStreamBridgeClosed) {
return fmt.Errorf("stream %s is not open", id)
}
return err
}
return nil
}
func (b *streamBridge) close(id string, errorMessage string) {
if b == nil || id == "" {
return
}
b.mu.Lock()
stream := b.streams[id]
delete(b.streams, id)
b.mu.Unlock()
if stream == nil {
return
}
stream.close(errorMessage)
}