diff --git a/sdk/go/provider_open.go b/sdk/go/provider_open.go new file mode 100644 index 0000000000..76e70838bf --- /dev/null +++ b/sdk/go/provider_open.go @@ -0,0 +1,77 @@ +package extension + +import ( + "context" + "encoding/json" + "errors" + "strings" +) + +func (s *server) handleStreamOpen(ctx context.Context, raw json.RawMessage) (result any, err error) { + if s.opts.Provider == nil { + return nil, MustProtocolError(ErrUnknownMethod) + } + var p StreamOpenParams + if err := strictDecode(raw, &p); err != nil { + return nil, MustProtocolError(ErrInvalidParams) + } + if p.SeqBase < 0 { + return nil, MustProtocolError(ErrInvalidParams) + } + if err := p.Validate(); err != nil { + return nil, MustProtocolError(ErrInvalidParams) + } + streamCtx, cancel := context.WithCancel(ctx) + handle := &streamHandle{cancel: cancel, done: make(chan struct{})} + s.streamsMu.Lock() + if _, exists := s.streams[p.StreamID]; exists { + s.streamsMu.Unlock() + cancel() + return nil, &ProtocolError{Reason: ErrProtocolError, Message: "duplicate stream id " + p.StreamID} + } + s.streams[p.StreamID] = handle + s.streamsMu.Unlock() + defer func() { + if result != nil { + return + } + cancel() + s.streamsMu.Lock() + delete(s.streams, p.StreamID) + s.streamsMu.Unlock() + close(handle.done) + }() + chunks, err := s.opts.Provider.Stream(streamCtx, StreamRequest{ + StreamID: p.StreamID, + ProviderRef: p.ProviderRef, + Model: p.Model, + Effort: p.Effort, + Request: p.Request, + }) + if err != nil { + s.log.Printf("extension: provider stream %q failed to open: %v", p.StreamID, err) + return nil, MustProtocolError(ErrProviderFailed) + } + if chunks == nil { + return nil, errors.New("extension: provider returned a nil chunk channel") + } + return deferredResult{ + result: StreamOpenResult{Accepted: true}, + after: func() { go s.pumpStream(streamCtx, p.StreamID, p.SeqBase, chunks, handle) }, + }, nil +} + +func (s *server) handleStreamCancel(_ context.Context, raw json.RawMessage) (any, error) { + var p StreamCancelParams + if err := strictDecode(raw, &p); err != nil || strings.TrimSpace(p.StreamID) == "" { + return nil, MustProtocolError(ErrInvalidParams) + } + s.streamsMu.Lock() + handle := s.streams[p.StreamID] + s.streamsMu.Unlock() + if handle == nil { + return StreamCancelResult{Cancelled: false}, nil + } + handle.cancel() + return StreamCancelResult{Cancelled: true}, nil +} diff --git a/sdk/go/provider_open_cancel_test.go b/sdk/go/provider_open_cancel_test.go new file mode 100644 index 0000000000..80c49905b2 --- /dev/null +++ b/sdk/go/provider_open_cancel_test.go @@ -0,0 +1,135 @@ +package extension + +import ( + "context" + "encoding/json" + "errors" + "sync/atomic" + "testing" + "time" +) + +type openingProvider struct { + entered chan context.Context + release chan struct{} + channelOnStop bool +} + +func (*openingProvider) Catalog(context.Context) ([]ProviderDescriptor, error) { + return nil, nil +} + +func (p *openingProvider) Stream(ctx context.Context, _ StreamRequest) (<-chan StreamChunk, error) { + p.entered <- ctx + select { + case <-ctx.Done(): + if !p.channelOnStop { + return nil, ctx.Err() + } + case <-p.release: + } + return make(chan StreamChunk), nil +} + +func TestProviderCancelDuringStreamOpen(t *testing.T) { + for _, channelOnStop := range []bool{false, true} { + name := "error" + if channelOnStop { + name = "channel" + } + t.Run(name, func(t *testing.T) { + p := &openingProvider{entered: make(chan context.Context, 1), release: make(chan struct{}), channelOnStop: channelOnStop} + host, _ := startFakeHost(t, providerHandler(), Options{Provider: p}) + host.handshake(t) + _, opened := host.startRequest(MethodExtensionProviderStreamOpen, openStreamRequest("opening")) + var streamCtx context.Context + select { + case streamCtx = <-p.entered: + case <-time.After(time.Second): + t.Fatal("provider did not receive the stream open") + } + resp := host.request(MethodExtensionProviderStreamCancel, StreamCancelParams{StreamID: "opening"}) + var cancelled StreamCancelResult + if err := json.Unmarshal(resp.Result, &cancelled); err != nil || !cancelled.Cancelled { + t.Errorf("cancel during open = %+v, error=%+v, decode=%v", cancelled, resp.Err, err) + } + select { + case <-streamCtx.Done(): + case <-time.After(200 * time.Millisecond): + t.Error("processed cancel did not cancel the opening provider context") + } + close(p.release) + select { + case resp = <-opened: + case <-time.After(time.Second): + t.Fatal("stream open did not finish") + } + if !channelOnStop && streamCtx.Err() != nil { + if resp.Err == nil { + t.Fatal("provider's cancelled open should return its existing error response") + } + host.request(MethodExtensionProviderCatalog, ProviderCatalogParams{}) + return + } + if resp.Err != nil { + t.Fatalf("channel open failed: %+v", resp.Err) + } + if streamCtx.Err() == nil { + host.request(MethodExtensionProviderStreamCancel, StreamCancelParams{StreamID: "opening"}) + } + end := host.waitStreamEnd() + chunks, ends := host.streamNotifications() + if !end.Interrupted || end.LastSeq != 0 || len(chunks) != 0 || len(ends) != 1 { + t.Fatalf("cancelled open end=%+v chunks=%d ends=%d", end, len(chunks), len(ends)) + } + }) + } +} + +type streamOpenFunc func(context.Context, StreamRequest) (<-chan StreamChunk, error) + +func (streamOpenFunc) Catalog(context.Context) ([]ProviderDescriptor, error) { return nil, nil } + +func (f streamOpenFunc) Stream(ctx context.Context, req StreamRequest) (<-chan StreamChunk, error) { + return f(ctx, req) +} + +func TestFailedProviderOpenReleasesStreamID(t *testing.T) { + for _, failure := range []string{"error", "nil channel", "panic"} { + t.Run(failure, func(t *testing.T) { + var calls atomic.Int64 + p := streamOpenFunc(func(context.Context, StreamRequest) (<-chan StreamChunk, error) { + if calls.Add(1) == 1 { + switch failure { + case "error": + return nil, errors.New("fixture open failed") + case "nil channel": + return nil, nil + case "panic": + panic("fixture open panic") + } + } + return make(chan StreamChunk), nil + }) + host, _ := startFakeHost(t, providerHandler(), Options{Provider: p}) + host.handshake(t) + resp := host.request(MethodExtensionProviderStreamOpen, openStreamRequest("retry")) + if resp.Err == nil { + t.Fatal("expected the first open to fail") + } + resp = host.request(MethodExtensionProviderStreamCancel, StreamCancelParams{StreamID: "retry"}) + var cancelled StreamCancelResult + if err := json.Unmarshal(resp.Result, &cancelled); err != nil || cancelled.Cancelled { + t.Fatalf("failed stream still registered: %+v, decode=%v", cancelled, err) + } + resp = host.request(MethodExtensionProviderStreamOpen, openStreamRequest("retry")) + if resp.Err != nil { + t.Fatalf("failed stream id not released: %+v", resp.Err) + } + host.request(MethodExtensionProviderStreamCancel, StreamCancelParams{StreamID: "retry"}) + if end := host.waitStreamEnd(); !end.Interrupted { + t.Fatalf("retry end=%+v", end) + } + }) + } +} diff --git a/sdk/go/sdk.go b/sdk/go/sdk.go index a76a590579..87a45018f4 100644 --- a/sdk/go/sdk.go +++ b/sdk/go/sdk.go @@ -538,67 +538,6 @@ func (s *server) handleProviderCatalog(ctx context.Context, raw json.RawMessage) return ProviderCatalogResult{Providers: providers}, nil } -func (s *server) handleStreamOpen(ctx context.Context, raw json.RawMessage) (any, error) { - if s.opts.Provider == nil { - return nil, MustProtocolError(ErrUnknownMethod) - } - var p StreamOpenParams - if err := strictDecode(raw, &p); err != nil { - return nil, MustProtocolError(ErrInvalidParams) - } - if p.SeqBase < 0 { - return nil, MustProtocolError(ErrInvalidParams) - } - if err := p.Validate(); err != nil { - return nil, MustProtocolError(ErrInvalidParams) - } - streamCtx, cancel := context.WithCancel(ctx) - chunks, err := s.opts.Provider.Stream(streamCtx, StreamRequest{ - StreamID: p.StreamID, - ProviderRef: p.ProviderRef, - Model: p.Model, - Effort: p.Effort, - Request: p.Request, - }) - if err != nil { - cancel() - s.log.Printf("extension: provider stream %q failed to open: %v", p.StreamID, err) - return nil, MustProtocolError(ErrProviderFailed) - } - if chunks == nil { - cancel() - return nil, errors.New("extension: provider returned a nil chunk channel") - } - handle := &streamHandle{cancel: cancel, done: make(chan struct{})} - s.streamsMu.Lock() - if _, exists := s.streams[p.StreamID]; exists { - s.streamsMu.Unlock() - cancel() - return nil, &ProtocolError{Reason: ErrProtocolError, Message: "duplicate stream id " + p.StreamID} - } - s.streams[p.StreamID] = handle - s.streamsMu.Unlock() - return deferredResult{ - result: StreamOpenResult{Accepted: true}, - after: func() { go s.pumpStream(streamCtx, p.StreamID, p.SeqBase, chunks, handle) }, - }, nil -} - -func (s *server) handleStreamCancel(_ context.Context, raw json.RawMessage) (any, error) { - var p StreamCancelParams - if err := strictDecode(raw, &p); err != nil || strings.TrimSpace(p.StreamID) == "" { - return nil, MustProtocolError(ErrInvalidParams) - } - s.streamsMu.Lock() - handle := s.streams[p.StreamID] - s.streamsMu.Unlock() - if handle == nil { - return StreamCancelResult{Cancelled: false}, nil - } - handle.cancel() - return StreamCancelResult{Cancelled: true}, nil -} - func (s *server) sendStreamEnd(end *StreamEndParams) { if err := s.conn.notify(MethodExtensionProviderStreamEnd, *end); err != nil { s.log.Printf("extension: provider stream %q could not deliver stream end: %v", end.StreamID, err)