fix(sdk): reject negative provider token usage

sdk/go/provider_usage_lifetime_test.go

@@ -0,0 +1,163 @@
+package extension
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+	"time"
+)
+
+type usageLifetimeRun struct {
+	ctx  context.Context
+	done chan struct{}
+}
+
+type usageLifetimeProvider struct {
+	invalid StreamChunk
+	stop    chan struct{}
+	runs    chan usageLifetimeRun
+}
+
+func (*usageLifetimeProvider) Catalog(context.Context) ([]ProviderDescriptor, error) {
+	return []ProviderDescriptor{}, nil
+}
+
+func (p *usageLifetimeProvider) Stream(ctx context.Context, req StreamRequest) (<-chan StreamChunk, error) {
+	chunks := make(chan StreamChunk)
+	run := usageLifetimeRun{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 TestProviderNegativeUsageStopsUnbufferedProducer(t *testing.T) {
+	for _, tc := range []struct {
+		name  string
+		chunk StreamChunk
+	}{
+		{name: "prompt", chunk: UsageChunk(ProviderUsage{PromptTokens: -1})},
+		{name: "completion", chunk: UsageChunk(ProviderUsage{CompletionTokens: -1})},
+		{name: "total", chunk: UsageChunk(ProviderUsage{TotalTokens: -1})},
+		{name: "cache hit", chunk: UsageChunk(ProviderUsage{CacheHitTokens: -1})},
+		{name: "cache miss", chunk: UsageChunk(ProviderUsage{CacheMissTokens: -1})},
+		{name: "reasoning", chunk: UsageChunk(ProviderUsage{ReasoningTokens: -1})},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			provider := &usageLifetimeProvider{invalid: tc.chunk, stop: make(chan struct{}), runs: make(chan usageLifetimeRun, 3)}
+			var runs []usageLifetimeRun
+			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_usage_test.go

@@ -0,0 +1,70 @@
+package extension
+
+import (
+	"encoding/json"
+	"reflect"
+	"testing"
+)
+
+func TestProviderStreamRejectsNegativeUsage(t *testing.T) {
+	cases := []struct {
+		name    string
+		chunk   StreamChunk
+		invalid bool
+	}{
+		{"prompt", UsageChunk(ProviderUsage{PromptTokens: -1}), true},
+		{"completion", UsageChunk(ProviderUsage{CompletionTokens: -1}), true},
+		{"total", UsageChunk(ProviderUsage{TotalTokens: -1}), true},
+		{"cache hit", UsageChunk(ProviderUsage{CacheHitTokens: -1}), true},
+		{"cache miss", UsageChunk(ProviderUsage{CacheMissTokens: -1}), true},
+		{"reasoning", UsageChunk(ProviderUsage{ReasoningTokens: -1}), true},
+		{"text with negative usage", StreamChunk{Type: ChunkText, Text: "bad", Usage: &ProviderUsage{PromptTokens: -1}}, true},
+		{"nil usage", StreamChunk{Type: ChunkUsage}, true},
+		{"zero usage", UsageChunk(ProviderUsage{}), false},
+		{"positive usage", UsageChunk(ProviderUsage{PromptTokens: 5, CompletionTokens: 7, TotalTokens: 12, CacheHitTokens: 2, CacheMissTokens: 3, ReasoningTokens: 4, FinishReason: "stop"}), false},
+		{"text with usage", StreamChunk{Type: ChunkText, Text: "ok", Usage: &ProviderUsage{TotalTokens: 1}}, false},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			provider := &scriptProvider{makeChannel: func(StreamRequest) <-chan StreamChunk {
+				chunks := make(chan StreamChunk, 4)
+				chunks <- TextChunk("prefix")
+				chunks <- tc.chunk
+				chunks <- TextChunk("trailer")
+				chunks <- DoneChunk()
+				close(chunks)
+				return chunks
+			}}
+			host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider})
+			host.handshake(t)
+			opened := host.request(MethodExtensionProviderStreamOpen, openStreamRequest("usage"))
+			if opened.Err != nil {
+				t.Fatalf("stream open: %+v", opened.Err)
+			}
+			end := host.waitStreamEnd()
+			chunks, ends := host.streamNotifications()
+			if len(ends) != 1 || end.StreamID != "usage" || end.Interrupted || len(chunks) == 0 || chunks[0].Chunk.Text != "prefix" {
+				t.Fatalf("chunks=%+v ends=%+v", chunks, ends)
+			}
+			if tc.invalid {
+				if len(chunks) != 1 || end.LastSeq != 1 || end.Error != "the extension provider produced an invalid chunk" {
+					t.Fatalf("invalid usage delivered: chunks=%+v end=%+v", chunks, end)
+				}
+			} else {
+				if len(chunks) != 4 || end.LastSeq != 4 || end.Error != "" || !reflect.DeepEqual(chunks[1].Chunk, tc.chunk) || chunks[2].Chunk.Text != "trailer" || chunks[3].Chunk.Type != ChunkDone {
+					t.Fatalf("valid usage changed: chunks=%+v end=%+v", chunks, end)
+				}
+			}
+			for i, chunk := range chunks {
+				if chunk.Seq != int64(i+1) || chunk.StreamID != "usage" {
+					t.Fatalf("chunk %d=%+v", i, chunk)
+				}
+			}
+			catalog := host.request(MethodExtensionProviderCatalog, ProviderCatalogParams{})
+			var result ProviderCatalogResult
+			if catalog.Err != nil || json.Unmarshal(catalog.Result, &result) != nil || result.Providers == nil {
+				t.Fatalf("same connection catalog: %+v", catalog)
+			}
+		})
+	}
+}

sdk/go/types_ext.go

@@ -107,6 +107,13 @@ func (chunk ProviderChunk) Validate() error {
 	if chunk.Type == ChunkUsage && chunk.Usage == nil {
 		return validationError("usage chunks require usage")
 	}
+	if usage := chunk.Usage; usage != nil {
+		for _, tokens := range []int{usage.PromptTokens, usage.CompletionTokens, usage.TotalTokens, usage.CacheHitTokens, usage.CacheMissTokens, usage.ReasoningTokens} {
+			if tokens < 0 {
+				return validationError("usage token counts must be non-negative")
+			}
+		}
+	}
 	return nil
 }