| 1 | package extension |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "sync" |
| 8 | "sync/atomic" |
| 9 | "testing" |
| 10 | ) |
| 11 | |
| 12 | // scriptProvider is a test Provider with scripted catalog and streams. |
| 13 | type scriptProvider struct { |
| 14 | catalog []ProviderDescriptor |
| 15 | catalogErr error |
| 16 | |
| 17 | streamErr error |
| 18 | // makeChannel builds the chunk channel for one Stream call; the test owns |
| 19 | // the channel lifecycle. |
| 20 | makeChannel func(req StreamRequest) <-chan StreamChunk |
| 21 | |
| 22 | mu sync.Mutex |
| 23 | requests []StreamRequest |
| 24 | } |
| 25 | |
| 26 | func (p *scriptProvider) Catalog(context.Context) ([]ProviderDescriptor, error) { |
| 27 | if p.catalogErr != nil { |
| 28 | return nil, p.catalogErr |
| 29 | } |
| 30 | return p.catalog, nil |
| 31 | } |
| 32 | |
| 33 | func (p *scriptProvider) Stream(_ context.Context, req StreamRequest) (<-chan StreamChunk, error) { |
| 34 | p.mu.Lock() |
| 35 | p.requests = append(p.requests, req) |
| 36 | p.mu.Unlock() |
| 37 | if p.streamErr != nil { |
| 38 | return nil, p.streamErr |
| 39 | } |
| 40 | return p.makeChannel(req), nil |
| 41 | } |
| 42 | |
| 43 | func providerHandler() *testHandler { |
| 44 | return &testHandler{result: &InitializeResult{ |
| 45 | Name: "provider-ext", Version: "1.0.0", |
| 46 | Providers: []ProviderDescriptor{{Ref: "plugin/provider-ext/echo", DisplayName: "Echo", Model: "echo-1"}}, |
| 47 | }} |
| 48 | } |
| 49 | |
| 50 | func openStreamRequest(streamID string) StreamOpenParams { |
| 51 | return StreamOpenParams{ |
| 52 | StreamID: streamID, |
| 53 | ProviderRef: "plugin/provider-ext/echo", |
| 54 | Model: "echo-1", |
| 55 | Request: ProviderRequest{ |
| 56 | Messages: []ProviderMessage{{Role: ProviderRoleUser, Content: "hi"}}, |
| 57 | Tools: []ProviderToolSchema{}, |
| 58 | }, |
| 59 | SeqBase: 1, |
| 60 | } |
| 61 | } |
| 62 | |
| 63 | // TestProviderCatalog serves extension/provider/catalog. |
| 64 | func TestProviderCatalog(t *testing.T) { |
| 65 | provider := &scriptProvider{catalog: []ProviderDescriptor{ |
| 66 | {Ref: "plugin/provider-ext/echo", DisplayName: "Echo", Model: "echo-1", ContextWindow: 8192, Tools: true}, |
| 67 | }} |
| 68 | host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider}) |
| 69 | host.handshake(t) |
| 70 | resp := host.request(MethodExtensionProviderCatalog, ProviderCatalogParams{}) |
| 71 | if resp.Err != nil { |
| 72 | t.Fatalf("catalog failed: %+v", resp.Err) |
| 73 | } |
| 74 | var result ProviderCatalogResult |
| 75 | if err := json.Unmarshal(resp.Result, &result); err != nil { |
| 76 | t.Fatalf("decode catalog: %v", err) |
| 77 | } |
| 78 | if len(result.Providers) != 1 || result.Providers[0].Ref != "plugin/provider-ext/echo" || !result.Providers[0].Tools { |
| 79 | t.Fatalf("catalog = %+v", result.Providers) |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | // TestProviderCatalogNil ensures the array shape survives an empty catalog: |
| 84 | // the wire requires "providers":[], never null. |
| 85 | func TestProviderCatalogNil(t *testing.T) { |
| 86 | provider := &scriptProvider{} |
| 87 | host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider}) |
| 88 | host.handshake(t) |
| 89 | resp := host.request(MethodExtensionProviderCatalog, ProviderCatalogParams{}) |
| 90 | var raw struct { |
| 91 | Providers json.RawMessage `json:"providers"` |
| 92 | } |
| 93 | if err := json.Unmarshal(resp.Result, &raw); err != nil { |
| 94 | t.Fatalf("decode: %v", err) |
| 95 | } |
| 96 | if string(raw.Providers) != "[]" { |
| 97 | t.Fatalf("providers = %s, want []", raw.Providers) |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | // TestProviderCatalogWithoutProvider answers unknown_method when no Provider |
| 102 | // is configured. |
| 103 | func TestProviderCatalogWithoutProvider(t *testing.T) { |
| 104 | host, _ := startFakeHost(t, basicHandler(), Options{}) |
| 105 | host.handshake(t) |
| 106 | resp := host.request(MethodExtensionProviderCatalog, ProviderCatalogParams{}) |
| 107 | if resp.Err == nil || resp.Err.Code != CodeMethodNotFound { |
| 108 | t.Fatalf("expected unknown_method, got %+v", resp.Err) |
| 109 | } |
| 110 | } |
| 111 | |
| 112 | // TestProviderStreamPump verifies contiguous 1-based seqs and the terminal |
| 113 | // stream/end lastSeq. |
| 114 | func TestProviderStreamPump(t *testing.T) { |
| 115 | chunks := make(chan StreamChunk, 4) |
| 116 | chunks <- TextChunk("Hello") |
| 117 | chunks <- ReasoningChunk("thinking", "sig-1") |
| 118 | chunks <- UsageChunk(ProviderUsage{PromptTokens: 3, CompletionTokens: 2, TotalTokens: 5, FinishReason: "stop"}) |
| 119 | close(chunks) |
| 120 | provider := &scriptProvider{makeChannel: func(StreamRequest) <-chan StreamChunk { return chunks }} |
| 121 | host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider}) |
| 122 | host.handshake(t) |
| 123 | |
| 124 | resp := host.request(MethodExtensionProviderStreamOpen, openStreamRequest("stream-1")) |
| 125 | if resp.Err != nil { |
| 126 | t.Fatalf("stream open failed: %+v", resp.Err) |
| 127 | } |
| 128 | var opened StreamOpenResult |
| 129 | if err := json.Unmarshal(resp.Result, &opened); err != nil || !opened.Accepted { |
| 130 | t.Fatalf("open result = %+v", opened) |
| 131 | } |
| 132 | |
| 133 | endParams := host.waitStreamEnd() |
| 134 | sentChunks, ends := host.streamNotifications() |
| 135 | if len(ends) != 1 { |
| 136 | t.Fatalf("stream/end count = %d, want exactly 1", len(ends)) |
| 137 | } |
| 138 | if endParams.StreamID != "stream-1" || endParams.LastSeq != 3 || endParams.Error != "" || endParams.Interrupted { |
| 139 | t.Fatalf("end = %+v, want lastSeq 3 clean", endParams) |
| 140 | } |
| 141 | for i, chunk := range sentChunks { |
| 142 | if chunk.Seq != int64(i+1) { |
| 143 | t.Fatalf("chunk %d seq = %d, want contiguous 1-based", i, chunk.Seq) |
| 144 | } |
| 145 | if chunk.StreamID != "stream-1" { |
| 146 | t.Fatalf("chunk %d streamId = %q", i, chunk.StreamID) |
| 147 | } |
| 148 | } |
| 149 | if sentChunks[0].Chunk.Type != ChunkText || sentChunks[0].Chunk.Text != "Hello" { |
| 150 | t.Fatalf("chunk 0 = %+v", sentChunks[0].Chunk) |
| 151 | } |
| 152 | if sentChunks[1].Chunk.Type != ChunkReasoning || sentChunks[1].Chunk.Signature != "sig-1" { |
| 153 | t.Fatalf("chunk 1 = %+v", sentChunks[1].Chunk) |
| 154 | } |
| 155 | if sentChunks[2].Chunk.Usage == nil || sentChunks[2].Chunk.Usage.TotalTokens != 5 { |
| 156 | t.Fatalf("chunk 2 = %+v", sentChunks[2].Chunk) |
| 157 | } |
| 158 | } |
| 159 | |
| 160 | // TestProviderStreamCancel asserts a processed cancel stops chunk production: |
| 161 | // no chunk may be sent after the cancel response, and the stream ends |
| 162 | // interrupted. |
| 163 | func TestProviderStreamCancel(t *testing.T) { |
| 164 | chunks := make(chan StreamChunk) // unbuffered: every send is visible |
| 165 | provider := &scriptProvider{makeChannel: func(StreamRequest) <-chan StreamChunk { return chunks }} |
| 166 | host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider}) |
| 167 | host.handshake(t) |
| 168 | resp := host.request(MethodExtensionProviderStreamOpen, openStreamRequest("stream-c")) |
| 169 | var opened StreamOpenResult |
| 170 | if err := json.Unmarshal(resp.Result, &opened); err != nil || !opened.Accepted { |
| 171 | t.Fatalf("open result = %+v", opened) |
| 172 | } |
| 173 | |
| 174 | // Feed one chunk, wait for it on the wire. |
| 175 | go func() { chunks <- TextChunk("one") }() |
| 176 | first := host.nextNotification(MethodExtensionProviderStreamChunk) |
| 177 | |
| 178 | // Cancel; the response means the SDK processed it. |
| 179 | resp = host.request(MethodExtensionProviderStreamCancel, StreamCancelParams{StreamID: "stream-c"}) |
| 180 | var cancelled StreamCancelResult |
| 181 | if err := json.Unmarshal(resp.Result, &cancelled); err != nil || !cancelled.Cancelled { |
| 182 | t.Fatalf("cancel result = %+v respErr=%+v", cancelled, resp.Err) |
| 183 | } |
| 184 | |
| 185 | // Keep producing: none of these may reach the wire. |
| 186 | go func() { |
| 187 | for i := 0; i < 5; i++ { |
| 188 | chunks <- TextChunk("late") |
| 189 | } |
| 190 | }() |
| 191 | endParams := host.waitStreamEnd() |
| 192 | if !endParams.Interrupted || endParams.LastSeq != 1 { |
| 193 | t.Fatalf("end = %+v, want interrupted lastSeq 1", endParams) |
| 194 | } |
| 195 | sentChunks, _ := host.streamNotifications() |
| 196 | for _, chunk := range sentChunks { |
| 197 | if chunk.Seq > 1 { |
| 198 | t.Fatalf("chunk seq %d sent after the cancel was processed", chunk.Seq) |
| 199 | } |
| 200 | } |
| 201 | var firstParams StreamChunkParams |
| 202 | if err := json.Unmarshal(first.Params, &firstParams); err != nil || firstParams.Seq != 1 { |
| 203 | t.Fatalf("first chunk = %+v", firstParams) |
| 204 | } |
| 205 | } |
| 206 | |
| 207 | // TestProviderStreamErrorChunk maps a provider error chunk to stream/end's |
| 208 | // error field without forwarding the chunk. |
| 209 | func TestProviderStreamErrorChunk(t *testing.T) { |
| 210 | chunks := make(chan StreamChunk, 2) |
| 211 | chunks <- TextChunk("partial") |
| 212 | chunks <- ErrorChunk("provider upstream unavailable") |
| 213 | close(chunks) |
| 214 | provider := &scriptProvider{makeChannel: func(StreamRequest) <-chan StreamChunk { return chunks }} |
| 215 | host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider}) |
| 216 | host.handshake(t) |
| 217 | host.request(MethodExtensionProviderStreamOpen, openStreamRequest("stream-e")) |
| 218 | |
| 219 | endParams := host.waitStreamEnd() |
| 220 | if endParams.Error != "provider upstream unavailable" { |
| 221 | t.Fatalf("end.error = %q", endParams.Error) |
| 222 | } |
| 223 | if endParams.LastSeq != 1 || endParams.Interrupted { |
| 224 | t.Fatalf("end = %+v, want lastSeq 1 not interrupted", endParams) |
| 225 | } |
| 226 | sentChunks, _ := host.streamNotifications() |
| 227 | if len(sentChunks) != 1 || sentChunks[0].Chunk.Type != ChunkText { |
| 228 | t.Fatalf("chunks = %+v, want only the text chunk forwarded", sentChunks) |
| 229 | } |
| 230 | } |
| 231 | |
| 232 | // TestProviderStreamOpenError answers provider_failed when Stream refuses to |
| 233 | // open. |
| 234 | func TestProviderStreamOpenError(t *testing.T) { |
| 235 | provider := &scriptProvider{streamErr: errors.New("quota exhausted")} |
| 236 | host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider}) |
| 237 | host.handshake(t) |
| 238 | resp := host.request(MethodExtensionProviderStreamOpen, openStreamRequest("stream-f")) |
| 239 | if resp.Err == nil { |
| 240 | t.Fatal("expected provider_failed") |
| 241 | } |
| 242 | data, _ := resp.Err.Data.(ProtocolErrorData) |
| 243 | if data.Reason != ErrProviderFailed { |
| 244 | t.Fatalf("reason = %q, want provider_failed", data.Reason) |
| 245 | } |
| 246 | } |
| 247 | |
| 248 | // TestProviderStreamOpenInvalidEnvelope rejects malformed opens before they |
| 249 | // reach the Provider. |
| 250 | func TestProviderStreamOpenInvalidEnvelope(t *testing.T) { |
| 251 | var calls atomic.Int64 |
| 252 | provider := &scriptProvider{makeChannel: func(StreamRequest) <-chan StreamChunk { |
| 253 | calls.Add(1) |
| 254 | return make(chan StreamChunk) |
| 255 | }} |
| 256 | host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider}) |
| 257 | host.handshake(t) |
| 258 | frames := []string{ |
| 259 | `{"providerRef":"x","request":{"messages":[],"tools":[]},"seqBase":1,"streamId":""}`, |
| 260 | `{"providerRef":"x","request":{"messages":null,"tools":[]},"seqBase":1,"streamId":"s"}`, |
| 261 | `{"providerRef":"x","request":{"messages":[],"tools":[{"name":"t","parameters":[1]}]},"seqBase":1,"streamId":"s"}`, |
| 262 | `{"providerRef":"x","request":{"messages":[],"tools":[]},"seqBase":-1,"streamId":"s"}`, |
| 263 | } |
| 264 | for _, params := range frames { |
| 265 | resp := host.request(MethodExtensionProviderStreamOpen, json.RawMessage(params)) |
| 266 | if resp.Err == nil || resp.Err.Code != CodeInvalidParams { |
| 267 | t.Fatalf("params %s: expected invalid_params, got %+v", params, resp.Err) |
| 268 | } |
| 269 | } |
| 270 | if calls.Load() != 0 { |
| 271 | t.Fatalf("Stream called %d times on invalid envelopes", calls.Load()) |
| 272 | } |
| 273 | } |
| 274 | |
| 275 | // TestStreamRequestPassedThrough checks the helper-level StreamRequest maps |
| 276 | // the wire params faithfully. |
| 277 | func TestStreamRequestPassedThrough(t *testing.T) { |
| 278 | chunks := make(chan StreamChunk) |
| 279 | close(chunks) |
| 280 | provider := &scriptProvider{makeChannel: func(StreamRequest) <-chan StreamChunk { return chunks }} |
| 281 | host, _ := startFakeHost(t, providerHandler(), Options{Provider: provider}) |
| 282 | host.handshake(t) |
| 283 | open := openStreamRequest("stream-req") |
| 284 | open.Effort = "high" |
| 285 | open.Request.MaxTokens = 128 |
| 286 | temp := 0.5 |
| 287 | open.Request.Temperature = &temp |
| 288 | host.request(MethodExtensionProviderStreamOpen, open) |
| 289 | host.waitStreamEnd() |
| 290 | provider.mu.Lock() |
| 291 | defer provider.mu.Unlock() |
| 292 | if len(provider.requests) != 1 { |
| 293 | t.Fatalf("Stream calls = %d", len(provider.requests)) |
| 294 | } |
| 295 | req := provider.requests[0] |
| 296 | if req.StreamID != "stream-req" || req.ProviderRef != "plugin/provider-ext/echo" || req.Model != "echo-1" || req.Effort != "high" { |
| 297 | t.Fatalf("request = %+v", req) |
| 298 | } |
| 299 | if req.Request.MaxTokens != 128 || req.Request.Temperature == nil || *req.Request.Temperature != 0.5 { |
| 300 | t.Fatalf("provider request = %+v", req.Request) |
| 301 | } |
| 302 | if len(req.Request.Messages) != 1 || req.Request.Messages[0].Content != "hi" { |
| 303 | t.Fatalf("messages = %+v", req.Request.Messages) |
| 304 | } |
| 305 | } |
| 306 |