diff --git a/sdk/go/provider_usage_lifetime_test.go b/sdk/go/provider_usage_lifetime_test.go new file mode 100644 index 0000000000..bd96ff4af5 --- /dev/null +++ b/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) + } + }) + } +} diff --git a/sdk/go/provider_usage_test.go b/sdk/go/provider_usage_test.go new file mode 100644 index 0000000000..303ed0fded --- /dev/null +++ b/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) + } + }) + } +} diff --git a/sdk/go/types_ext.go b/sdk/go/types_ext.go index f5d95632da..c82d8962d1 100644 --- a/sdk/go/types_ext.go +++ b/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 }