fix(sdk): reject unregistered provider chunk types

sdk/go/provider_chunk_lifetime_test.go

@@ -0,0 +1,159 @@
+package extension
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+	"time"
+)
+
+type chunkLifetimeRun struct {
+	ctx  context.Context
+	done chan struct{}
+}
+
+type chunkLifetimeProvider struct {
+	invalid StreamChunk
+	stop    chan struct{}
+	runs    chan chunkLifetimeRun
+}
+
+func (*chunkLifetimeProvider) Catalog(context.Context) ([]ProviderDescriptor, error) {
+	return []ProviderDescriptor{}, nil
+}
+
+func (p *chunkLifetimeProvider) Stream(ctx context.Context, req StreamRequest) (<-chan StreamChunk, error) {
+	chunks := make(chan StreamChunk)
+	run := chunkLifetimeRun{ctx: ctx, done: make(chan struct{})}
+	p.runs <- run
+	go func() {
+		defer close(run.done)
+		defer close(chunks)
+		sequence := []StreamChunk{TextChunk("prefix")}
+		if req.StreamID == "reject" {
+			sequence = append(sequence, p.invalid, TextChunk("trailer"))
+		}
+		for _, chunk := range sequence {
+			select {
+			case chunks <- chunk:
+			case <-ctx.Done():
+				return
+			case <-p.stop:
+				return
+			}
+		}
+		if req.StreamID == "cancel" {
+			select {
+			case <-ctx.Done():
+			case <-p.stop:
+			}
+		}
+	}()
+	return chunks, nil
+}
+
+func TestProviderInvalidChunkStopsUnbufferedProducer(t *testing.T) {
+	for _, tc := range []struct {
+		name  string
+		chunk StreamChunk
+	}{
+		{name: "zero"},
+		{name: "unknown", chunk: StreamChunk{Type: "unknown"}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			provider := &chunkLifetimeProvider{invalid: tc.chunk, stop: make(chan struct{}), runs: make(chan chunkLifetimeRun, 3)}
+			var runs []chunkLifetimeRun
+			t.Cleanup(func() {
+				close(provider.stop)
+				for _, run := range runs {
+					select {
+					case <-run.done:
+					case <-time.After(time.Second):
+						t.Error("test producer did not stop during cleanup")
+					}
+				}
+			})
+			host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider})
+			host.handshake(t)
+			waitEnd := func(id string) StreamEndParams {
+				for deadline := time.Now().Add(time.Second); time.Now().Before(deadline); {
+					_, ends := host.streamNotifications()
+					for _, end := range ends {
+						if end.StreamID == id {
+							return end
+						}
+					}
+					time.Sleep(5 * time.Millisecond)
+				}
+				t.Fatalf("stream %s did not end", id)
+				return StreamEndParams{}
+			}
+			open := openStreamRequest("reject")
+			open.SeqBase = 7
+			resp := host.request(MethodExtensionProviderStreamOpen, open)
+			if resp.Err != nil {
+				t.Fatalf("open: %+v", resp.Err)
+			}
+			run := <-provider.runs
+			runs = append(runs, run)
+			end := waitEnd("reject")
+			sent, ends := host.streamNotifications()
+			if len(sent) != 1 || len(ends) != 1 || sent[0].StreamID != "reject" || sent[0].Seq != 7 || sent[0].Chunk.Text != "prefix" || end.StreamID != "reject" || end.LastSeq != 7 || end.Error == "" || end.Interrupted {
+				t.Fatalf("rejected stream: sent=%+v ends=%+v", sent, ends)
+			}
+			select {
+			case <-run.done:
+				if !errors.Is(run.ctx.Err(), context.Canceled) {
+					t.Fatalf("producer context = %v", run.ctx.Err())
+				}
+			case <-time.After(time.Second):
+				t.Fatal("unbuffered producer remains blocked after chunk rejection")
+			}
+			resp = host.request(MethodExtensionProviderStreamOpen, openStreamRequest("reused"))
+			if resp.Err != nil {
+				t.Fatalf("subsequent stream: %+v", resp.Err)
+			}
+			runs = append(runs, <-provider.runs)
+			end = waitEnd("reused")
+			if end.StreamID != "reused" || end.LastSeq != 1 || end.Error != "" || end.Interrupted {
+				t.Fatalf("subsequent stream end = %+v", end)
+			}
+			resp = host.request(MethodExtensionProviderStreamOpen, openStreamRequest("cancel"))
+			if resp.Err != nil {
+				t.Fatalf("cancellable stream: %+v", resp.Err)
+			}
+			run = <-provider.runs
+			runs = append(runs, run)
+			deadline := time.Now().Add(time.Second)
+			for {
+				sent, _ = host.streamNotifications()
+				if len(sent) == 3 && sent[2].StreamID == "cancel" {
+					break
+				}
+				if time.Now().After(deadline) {
+					t.Fatal("cancellable stream did not deliver its prefix")
+				}
+				time.Sleep(5 * time.Millisecond)
+			}
+			resp = host.request(MethodExtensionProviderStreamCancel, StreamCancelParams{StreamID: "cancel"})
+			var cancelled StreamCancelResult
+			if resp.Err != nil || json.Unmarshal(resp.Result, &cancelled) != nil || !cancelled.Cancelled {
+				t.Fatalf("cancel: %+v", resp)
+			}
+			end = waitEnd("cancel")
+			if end.StreamID != "cancel" || end.LastSeq != 1 || !end.Interrupted || end.Error != "" {
+				t.Fatalf("cancel end = %+v", end)
+			}
+			select {
+			case <-run.done:
+			case <-time.After(time.Second):
+				t.Fatal("explicit cancellation did not release producer")
+			}
+			sent, ends = host.streamNotifications()
+			if len(sent) != 3 || len(ends) != 3 {
+				t.Fatalf("notifications: sent=%+v ends=%+v", sent, ends)
+			}
+		})
+	}
+}

sdk/go/provider_test.go

@@ -372,3 +372,70 @@ func TestProviderStreamResponseFormat(t *testing.T) {
 		})
 	}
 }
+
+func TestProviderStreamChunkTypes(t *testing.T) {
+	for _, tc := range []struct {
+		name  string
+		chunk StreamChunk
+		valid bool
+	}{
+		{name: "zero"},
+		{name: "unknown", chunk: StreamChunk{Type: "unknown"}},
+		{name: "text", chunk: TextChunk("text"), valid: true},
+		{name: "reasoning", chunk: ReasoningChunk("thinking", "sig"), valid: true},
+		{name: "tool start", chunk: StreamChunk{Type: ChunkToolCallStart, ToolCall: &ProviderToolCall{ID: "c", Name: "lookup", Arguments: ""}}, valid: true},
+		{name: "tool delta", chunk: StreamChunk{Type: ChunkToolCallDelta, Text: `{"x":1}`, ArgChars: 7}, valid: true},
+		{name: "tool call", chunk: StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{ID: "c", Name: "lookup", Arguments: `{"x":1}`}}, valid: true},
+		{name: "usage", chunk: UsageChunk(ProviderUsage{PromptTokens: 1, CompletionTokens: 2, TotalTokens: 3}), valid: true},
+		{name: "done", chunk: DoneChunk(), valid: true},
+		{name: "error", chunk: ErrorChunk("upstream offline"), valid: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			chunks := make(chan StreamChunk, 3)
+			chunks <- TextChunk("prefix")
+			chunks <- tc.chunk
+			if !tc.valid {
+				chunks <- TextChunk("after invalid chunk")
+			}
+			close(chunks)
+			provider := &scriptProvider{makeChannel: func(StreamRequest) <-chan StreamChunk { return chunks }}
+			host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider})
+			host.handshake(t)
+			resp := host.request(MethodExtensionProviderStreamOpen, openStreamRequest("stream-types"))
+			var opened StreamOpenResult
+			if resp.Err != nil || json.Unmarshal(resp.Result, &opened) != nil || !opened.Accepted {
+				t.Fatalf("stream open = %+v", resp)
+			}
+			end := host.waitStreamEnd()
+			sent, ends := host.streamNotifications()
+			wantChunks := 2
+			if !tc.valid || tc.chunk.Type == ChunkError {
+				wantChunks = 1
+			}
+			if len(sent) != wantChunks || len(ends) != 1 || end.StreamID != "stream-types" || end.LastSeq != int64(wantChunks) || end.Interrupted {
+				t.Fatalf("chunks = %+v, ends = %+v", sent, ends)
+			}
+			if sent[0].Chunk.Type != ChunkText || sent[0].Chunk.Text != "prefix" {
+				t.Fatalf("delivered prefix = %+v", sent[0])
+			}
+			for i, chunk := range sent {
+				if chunk.StreamID != "stream-types" || chunk.Seq != int64(i+1) {
+					t.Fatalf("chunk %d = %+v", i, chunk)
+				}
+			}
+			if !tc.valid {
+				if end.Error == "" {
+					t.Fatal("invalid chunk ended successfully")
+				}
+				return
+			}
+			if tc.chunk.Type == ChunkError {
+				if end.Error != "upstream offline" {
+					t.Fatalf("provider error = %q", end.Error)
+				}
+			} else if end.Error != "" || sent[1].Chunk.Type != tc.chunk.Type {
+				t.Fatalf("valid chunk = %+v, end = %+v", sent[1], end)
+			}
+		})
+	}
+}

sdk/go/types_ext.go

@@ -95,6 +95,11 @@ func (request ProviderRequest) Validate() error {
 
 // Validate enforces chunk invariants the tags cannot express.
 func (chunk ProviderChunk) Validate() error {
+	switch chunk.Type {
+	case ChunkText, ChunkReasoning, ChunkToolCallStart, ChunkToolCallDelta, ChunkToolCall, ChunkUsage, ChunkDone, ChunkError:
+	default:
+		return validationError("chunk type must be a registered provider chunk type")
+	}
 	if chunk.ArgChars < 0 {
 		return validationError("argChars must be non-negative")
 	}