fix(sdk): cancel provider contexts when streams finish

sdk/go/provider_context_test.go

@@ -0,0 +1,76 @@
+package extension
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+	"time"
+)
+
+type contextProvider struct {
+	*scriptProvider
+	started chan context.Context
+}
+
+func (p *contextProvider) Stream(ctx context.Context, req StreamRequest) (<-chan StreamChunk, error) {
+	p.started <- ctx
+	return p.scriptProvider.Stream(ctx, req)
+}
+
+func TestProviderStreamContextEndsWithStream(t *testing.T) {
+	for _, mode := range []string{"clean", "error", "cancel"} {
+		t.Run(mode, func(t *testing.T) {
+			chunks := make(chan StreamChunk, 2)
+			chunks <- TextChunk("partial")
+			switch mode {
+			case "clean":
+				close(chunks)
+			case "error":
+				chunks <- ErrorChunk("upstream unavailable")
+			}
+			provider := &contextProvider{
+				scriptProvider: &scriptProvider{makeChannel: func(StreamRequest) <-chan StreamChunk { return chunks }},
+				started:        make(chan context.Context, 1),
+			}
+			host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider})
+			host.handshake(t)
+			resp := host.request(MethodExtensionProviderStreamOpen, openStreamRequest("lifetime"))
+			if resp.Err != nil {
+				t.Fatalf("open: %+v", resp.Err)
+			}
+			ctx := <-provider.started
+			if mode == "cancel" {
+				host.nextNotification(MethodExtensionProviderStreamChunk)
+				resp = host.request(MethodExtensionProviderStreamCancel, StreamCancelParams{StreamID: "lifetime"})
+				if resp.Err != nil {
+					t.Fatalf("cancel: %+v", resp.Err)
+				}
+			}
+			end := host.waitStreamEnd()
+			if end.LastSeq != 1 || end.Interrupted != (mode == "cancel") || (end.Error != "") != (mode == "error") {
+				t.Fatalf("end = %+v", end)
+			}
+			select {
+			case <-ctx.Done():
+				if !errors.Is(ctx.Err(), context.Canceled) {
+					t.Fatalf("stream context error = %v", ctx.Err())
+				}
+			case <-time.After(time.Second):
+				t.Fatal("provider context remains active after stream/end")
+			}
+			resp = host.request(MethodExtensionProviderCatalog, ProviderCatalogParams{})
+			if resp.Err != nil {
+				t.Fatalf("connection did not survive stream completion: %+v", resp.Err)
+			}
+			_, ends := host.streamNotifications()
+			if len(ends) != 1 {
+				t.Fatalf("stream/end count = %d", len(ends))
+			}
+			var catalog ProviderCatalogResult
+			if err := json.Unmarshal(resp.Result, &catalog); err != nil || catalog.Providers == nil || len(catalog.Providers) != 0 {
+				t.Fatalf("catalog = %+v, err = %v", catalog, err)
+			}
+		})
+	}
+}

sdk/go/provider_stream.go

@@ -0,0 +1,74 @@
+package extension
+
+import (
+	"context"
+	"strings"
+)
+
+// pumpStream forwards one provider channel onto the wire: chunks become
+// stream/chunk notifications with contiguous 1-based seqs (from SeqBase),
+// and exactly one stream/end closes the stream — clean on channel close,
+// with error on an error chunk, interrupted on cancel. A cancel processed by
+// the SDK is never trailed by another chunk.
+func (s *server) pumpStream(ctx context.Context, streamID string, seqBase int, chunks <-chan StreamChunk, handle *streamHandle) {
+	defer close(handle.done)
+	defer func() {
+		handle.cancel()
+		s.streamsMu.Lock()
+		delete(s.streams, streamID)
+		s.streamsMu.Unlock()
+	}()
+	seq := int64(seqBase)
+	if seq < 1 {
+		seq = 1
+	}
+	var lastSeq int64
+	end := StreamEndParams{StreamID: streamID}
+	for {
+		// A cancel must never be trailed by one more chunk, so check before
+		// every receive and again before every send.
+		select {
+		case <-ctx.Done():
+			end.LastSeq, end.Interrupted = lastSeq, true
+			s.sendStreamEnd(&end)
+			return
+		default:
+		}
+		select {
+		case <-ctx.Done():
+			end.LastSeq, end.Interrupted = lastSeq, true
+			s.sendStreamEnd(&end)
+			return
+		case chunk, ok := <-chunks:
+			if !ok {
+				end.LastSeq = lastSeq
+				s.sendStreamEnd(&end)
+				return
+			}
+			if chunk.Type == ChunkError {
+				end.LastSeq = lastSeq
+				end.Error = frozenErrorSpecs[ErrProviderFailed].Message
+				if chunk.Error != nil && strings.TrimSpace(chunk.Error.Message) != "" {
+					end.Error = chunk.Error.Message
+				}
+				s.sendStreamEnd(&end)
+				return
+			}
+			if err := chunk.Validate(); err != nil {
+				s.log.Printf("extension: provider stream %q produced an invalid chunk: %v", streamID, err)
+				end.LastSeq = lastSeq
+				end.Error = "the extension provider produced an invalid chunk"
+				s.sendStreamEnd(&end)
+				return
+			}
+			if err := s.conn.notify(MethodExtensionProviderStreamChunk, StreamChunkParams{
+				StreamID: streamID, Seq: seq, Chunk: chunk,
+			}); err != nil {
+				s.log.Printf("extension: provider stream %q could not deliver chunk %d: %v", streamID, seq, err)
+				return
+			}
+			lastSeq = seq
+			seq++
+		}
+	}
+}

sdk/go/sdk.go

@@ -641,73 +641,6 @@ func (s *server) handleStreamCancel(_ context.Context, raw json.RawMessage) (any
 	return StreamCancelResult{Cancelled: true}, nil
 }
 
-// pumpStream forwards one provider channel onto the wire: chunks become
-// stream/chunk notifications with contiguous 1-based seqs (from SeqBase),
-// and exactly one stream/end closes the stream — clean on channel close,
-// with error on an error chunk, interrupted on cancel. A cancel processed by
-// the SDK is never trailed by another chunk.
-func (s *server) pumpStream(ctx context.Context, streamID string, seqBase int, chunks <-chan StreamChunk, handle *streamHandle) {
-	defer close(handle.done)
-	defer func() {
-		s.streamsMu.Lock()
-		delete(s.streams, streamID)
-		s.streamsMu.Unlock()
-	}()
-	seq := int64(seqBase)
-	if seq < 1 {
-		seq = 1
-	}
-	var lastSeq int64
-	end := StreamEndParams{StreamID: streamID}
-	for {
-		// A cancel must never be trailed by one more chunk, so check before
-		// every receive and again before every send.
-		select {
-		case <-ctx.Done():
-			end.LastSeq, end.Interrupted = lastSeq, true
-			s.sendStreamEnd(&end)
-			return
-		default:
-		}
-		select {
-		case <-ctx.Done():
-			end.LastSeq, end.Interrupted = lastSeq, true
-			s.sendStreamEnd(&end)
-			return
-		case chunk, ok := <-chunks:
-			if !ok {
-				end.LastSeq = lastSeq
-				s.sendStreamEnd(&end)
-				return
-			}
-			if chunk.Type == ChunkError {
-				end.LastSeq = lastSeq
-				end.Error = frozenErrorSpecs[ErrProviderFailed].Message
-				if chunk.Error != nil && strings.TrimSpace(chunk.Error.Message) != "" {
-					end.Error = chunk.Error.Message
-				}
-				s.sendStreamEnd(&end)
-				return
-			}
-			if err := chunk.Validate(); err != nil {
-				s.log.Printf("extension: provider stream %q produced an invalid chunk: %v", streamID, err)
-				end.LastSeq = lastSeq
-				end.Error = "the extension provider produced an invalid chunk"
-				s.sendStreamEnd(&end)
-				return
-			}
-			if err := s.conn.notify(MethodExtensionProviderStreamChunk, StreamChunkParams{
-				StreamID: streamID, Seq: seq, Chunk: chunk,
-			}); err != nil {
-				s.log.Printf("extension: provider stream %q could not deliver chunk %d: %v", streamID, seq, err)
-				return
-			}
-			lastSeq = seq
-			seq++
-		}
-	}
-}
-
 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)