Merge pull request #11844 from KHG420/codex/studio-sdk-tool-call-identifiers

fix(sdk): validate provider tool call identifiers

sdk/go/provider_tool_call_test.go

@@ -0,0 +1,70 @@
+package extension
+
+import (
+	"encoding/json"
+	"reflect"
+	"testing"
+)
+
+func TestProviderStreamValidatesToolCallIdentifiers(t *testing.T) {
+	cases := []struct {
+		name    string
+		chunk   StreamChunk
+		invalid bool
+	}{
+		{"missing id", StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{Name: "lookup"}}, true},
+		{"blank id", StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{ID: " \t\n", Name: "lookup"}}, true},
+		{"missing name", StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{ID: "call-1"}}, true},
+		{"blank name", StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{ID: "call-1", Name: " \t\n"}}, true},
+		{"empty call", StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{}}, true},
+		{"text with invalid call", StreamChunk{Type: ChunkText, Text: "bad", ToolCall: &ProviderToolCall{ID: "call-1"}}, true},
+		{"start with invalid call", StreamChunk{Type: ChunkToolCallStart, ToolCall: &ProviderToolCall{Name: "lookup"}}, true},
+		{"delta with invalid call", StreamChunk{Type: ChunkToolCallDelta, ToolCall: &ProviderToolCall{ID: "call-1"}}, true},
+		{"nil call", StreamChunk{Type: ChunkToolCall}, false},
+		{"valid call", StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{ID: "call-1", Name: "lookup", Arguments: `{"query":"reasonix"}`, ThoughtSignature: "sig"}}, false},
+		{"unmodified identifiers", StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{ID: " 调用-1 ", Name: " 查询 ", Arguments: "partial"}}, 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("tool-call"))
+			if opened.Err != nil {
+				t.Fatalf("stream open: %+v", opened.Err)
+			}
+			end := host.waitStreamEnd()
+			chunks, ends := host.streamNotifications()
+			if len(ends) != 1 || end.StreamID != "tool-call" || 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 tool call 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 tool call changed: chunks=%+v end=%+v", chunks, end)
+				}
+			}
+			for i, chunk := range chunks {
+				if chunk.Seq != int64(i+1) || chunk.StreamID != "tool-call" {
+					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/provider_tool_id_lifetime_test.go

@@ -0,0 +1,161 @@
+package extension
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+	"time"
+)
+
+type toolIDLifetimeRun struct {
+	ctx  context.Context
+	done chan struct{}
+}
+
+type toolIDLifetimeProvider struct {
+	invalid StreamChunk
+	stop    chan struct{}
+	runs    chan toolIDLifetimeRun
+}
+
+func (*toolIDLifetimeProvider) Catalog(context.Context) ([]ProviderDescriptor, error) {
+	return []ProviderDescriptor{}, nil
+}
+
+func (p *toolIDLifetimeProvider) Stream(ctx context.Context, req StreamRequest) (<-chan StreamChunk, error) {
+	chunks := make(chan StreamChunk)
+	run := toolIDLifetimeRun{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 TestProviderInvalidToolIDStopsUnbufferedProducer(t *testing.T) {
+	for _, tc := range []struct {
+		name  string
+		chunk StreamChunk
+	}{
+		{name: "start empty id", chunk: StreamChunk{Type: ChunkToolCallStart, ToolCall: &ProviderToolCall{ID: "", Name: "lookup"}}},
+		{name: "start empty name", chunk: StreamChunk{Type: ChunkToolCallStart, ToolCall: &ProviderToolCall{ID: "c", Name: ""}}},
+		{name: "call blank id", chunk: StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{ID: "  ", Name: "lookup"}}},
+		{name: "call blank name", chunk: StreamChunk{Type: ChunkToolCall, ToolCall: &ProviderToolCall{ID: "c", Name: "  "}}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			provider := &toolIDLifetimeProvider{invalid: tc.chunk, stop: make(chan struct{}), runs: make(chan toolIDLifetimeRun, 3)}
+			var runs []toolIDLifetimeRun
+			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/types_ext.go

@@ -103,6 +103,9 @@ func (chunk ProviderChunk) Validate() error {
 	if chunk.ArgChars < 0 {
 		return validationError("argChars must be non-negative")
 	}
+	if call := chunk.ToolCall; call != nil && (strings.TrimSpace(call.ID) == "" || strings.TrimSpace(call.Name) == "") {
+		return validationError("tool call id and name must be non-empty")
+	}
 	if chunk.Type == ChunkError && chunk.Error == nil {
 		return validationError("error chunks require error")
 	}