| 1 | package extension |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "io" |
| 8 | "strings" |
| 9 | "sync/atomic" |
| 10 | "testing" |
| 11 | "time" |
| 12 | ) |
| 13 | |
| 14 | // TestHandshakeHappyPath verifies the initialize exchange echoes the |
| 15 | // sidecar's declaration and the barrier opens on extension/initialized. |
| 16 | func TestHandshakeHappyPath(t *testing.T) { |
| 17 | handler := basicHandler() |
| 18 | handler.result.Subscriptions = []string{"tool.before", "session.start"} |
| 19 | handler.result.Replaces = []string{"model"} |
| 20 | handler.result.UIActions = []UIActionDecl{{ActionID: "open", Label: "Open"}} |
| 21 | var seen InitializeParams |
| 22 | handler.seen = &seen |
| 23 | host, _ := startFakeHost(t, handler, Options{}) |
| 24 | result := host.handshake(t) |
| 25 | if result.ProtocolVersion != ProtocolVersion { |
| 26 | t.Fatalf("protocolVersion = %q, want %q", result.ProtocolVersion, ProtocolVersion) |
| 27 | } |
| 28 | if result.Name != "test-ext" || result.Version != "0.1.0" { |
| 29 | t.Fatalf("identity = %q/%q, want test-ext/0.1.0", result.Name, result.Version) |
| 30 | } |
| 31 | if len(result.Subscriptions) != 2 || result.Subscriptions[0] != "tool.before" { |
| 32 | t.Fatalf("subscriptions = %v", result.Subscriptions) |
| 33 | } |
| 34 | if len(result.UIActions) != 1 || result.UIActions[0].ActionID != "open" { |
| 35 | t.Fatalf("uiActions = %v", result.UIActions) |
| 36 | } |
| 37 | if seen.Session.SessionID != "sess-1" || seen.Session.Generation != 7 || seen.Session.WorkspaceRoot != "/repo" { |
| 38 | t.Fatalf("session context = %+v", seen.Session) |
| 39 | } |
| 40 | if !seen.Capabilities.ContentRefs || seen.Capabilities.UIHost != UIHostHeadless { |
| 41 | t.Fatalf("host capabilities = %+v", seen.Capabilities) |
| 42 | } |
| 43 | if seen.Manifest.Capabilities == nil { |
| 44 | t.Fatalf("manifest expectation = %+v", seen.Manifest) |
| 45 | } |
| 46 | } |
| 47 | |
| 48 | // TestHostRequestBeforeInitializeRejected sends a non-initialize request |
| 49 | // first: it must be answered with the frozen protocol_error. |
| 50 | func TestHostRequestBeforeInitializeRejected(t *testing.T) { |
| 51 | host, _ := startFakeHost(t, basicHandler(), Options{}) |
| 52 | resp := host.request(MethodExtensionIntercept, InterceptParams{ |
| 53 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), |
| 54 | }) |
| 55 | if resp.Err == nil { |
| 56 | t.Fatal("expected a protocol error for a request before initialize") |
| 57 | } |
| 58 | if resp.Err.Code != CodeInvalidRequest { |
| 59 | t.Fatalf("code = %d, want %d", resp.Err.Code, CodeInvalidRequest) |
| 60 | } |
| 61 | data, ok := resp.Err.Data.(ProtocolErrorData) |
| 62 | if !ok || data.Reason != ErrProtocolError { |
| 63 | t.Fatalf("error data = %+v, want reason protocol_error", resp.Err.Data) |
| 64 | } |
| 65 | } |
| 66 | |
| 67 | // TestHostRequestBeforeInitializedRejected covers the barrier window: the |
| 68 | // handshake answer is out but extension/initialized has not arrived, so |
| 69 | // intercepts are still refused. |
| 70 | func TestHostRequestBeforeInitializedRejected(t *testing.T) { |
| 71 | host, _ := startFakeHost(t, basicHandler(), Options{}) |
| 72 | resp := host.request(MethodExtensionInitialize, InitializeParams{ |
| 73 | ProtocolVersion: ProtocolVersion, ProtocolID: ProtocolID, |
| 74 | Session: SessionContext{SessionID: "s", WorkspaceRoot: "/r"}, |
| 75 | Capabilities: HostCapabilities{ProtocolVersion: ProtocolVersion}, |
| 76 | }) |
| 77 | if resp.Err != nil { |
| 78 | t.Fatalf("initialize failed: %+v", resp.Err) |
| 79 | } |
| 80 | resp = host.request(MethodExtensionIntercept, InterceptParams{ |
| 81 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), |
| 82 | }) |
| 83 | if resp.Err == nil || resp.Err.Code != CodeInvalidRequest { |
| 84 | t.Fatalf("expected protocol_error before initialized, got %+v", resp.Err) |
| 85 | } |
| 86 | // Notifications before the barrier are dropped silently (no response is |
| 87 | // possible); the connection must survive and the barrier must still open. |
| 88 | host.notify(MethodExtensionEvent, EventParams{Event: EventSessionStart, Payload: json.RawMessage(`{}`)}) |
| 89 | host.notify(MethodExtensionInitialized, InitializedParams{}) |
| 90 | resp = host.request(MethodExtensionIntercept, InterceptParams{ |
| 91 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), |
| 92 | }) |
| 93 | if resp.Err != nil { |
| 94 | t.Fatalf("intercept after initialized failed: %+v", resp.Err) |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | // TestOutboundCallBeforeInitializedFails checks the sidecar side of the |
| 99 | // barrier: Extension → Host calls from the Initialize handler return |
| 100 | // ErrNotReady. |
| 101 | func TestOutboundCallBeforeInitializedFails(t *testing.T) { |
| 102 | handler := &testHandler{} |
| 103 | var readyErr error |
| 104 | handlerWithProbe := HandlerFunc(func(ctx context.Context, p InitializeParams) (*InitializeResult, error) { |
| 105 | ui := HostUI{} |
| 106 | readyErr = ui.PublishNotification(ctx, p.Session.SessionID, p.Session.Generation, "probe", |
| 107 | UINotificationPayload{Title: "hi"}) |
| 108 | return &InitializeResult{Name: "probe", Version: "1"}, nil |
| 109 | }) |
| 110 | handler.result = &InitializeResult{Name: "x", Version: "1"} |
| 111 | host, _ := startFakeHost(t, handlerWithProbe, Options{}) |
| 112 | host.handshake(t) |
| 113 | if !errors.Is(readyErr, ErrNotReady) { |
| 114 | t.Fatalf("outbound call before initialized = %v, want ErrNotReady", readyErr) |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | // TestInitializeVersionMismatch answers unsupported_version and ends Serve. |
| 119 | func TestInitializeVersionMismatch(t *testing.T) { |
| 120 | host, serveDone := startFakeHost(t, basicHandler(), Options{}) |
| 121 | resp := host.request(MethodExtensionInitialize, InitializeParams{ |
| 122 | ProtocolVersion: "2", ProtocolID: ProtocolID, |
| 123 | Session: SessionContext{SessionID: "s", WorkspaceRoot: "/r"}, |
| 124 | Capabilities: HostCapabilities{ProtocolVersion: ProtocolVersion}, |
| 125 | }) |
| 126 | if resp.Err == nil { |
| 127 | t.Fatal("expected unsupported_version") |
| 128 | } |
| 129 | data, _ := resp.Err.Data.(ProtocolErrorData) |
| 130 | if data.Reason != ErrUnsupportedVersion { |
| 131 | t.Fatalf("reason = %q, want unsupported_version", data.Reason) |
| 132 | } |
| 133 | err, ok := serveDone.wait(5 * time.Second) |
| 134 | if !ok { |
| 135 | t.Fatal("Serve did not end after a failed handshake") |
| 136 | } |
| 137 | if err == nil { |
| 138 | t.Fatal("Serve returned nil after a failed handshake") |
| 139 | } |
| 140 | } |
| 141 | |
| 142 | // TestUnknownMethod verifies the frozen unknown_method answer. |
| 143 | func TestUnknownMethod(t *testing.T) { |
| 144 | host, _ := startFakeHost(t, basicHandler(), Options{}) |
| 145 | host.handshake(t) |
| 146 | resp := host.request("extension/bogus", struct{}{}) |
| 147 | if resp.Err == nil || resp.Err.Code != CodeMethodNotFound { |
| 148 | t.Fatalf("expected -32601, got %+v", resp.Err) |
| 149 | } |
| 150 | data, _ := resp.Err.Data.(ProtocolErrorData) |
| 151 | if data.Reason != ErrUnknownMethod { |
| 152 | t.Fatalf("reason = %q, want unknown_method", data.Reason) |
| 153 | } |
| 154 | } |
| 155 | |
| 156 | // TestOversizedInboundFrame kills the connection with a frame error. |
| 157 | func TestOversizedInboundFrame(t *testing.T) { |
| 158 | host, serveDone := startFakeHost(t, basicHandler(), Options{}) |
| 159 | big := make([]byte, FrameBytes+16) |
| 160 | for i := range big { |
| 161 | big[i] = ' ' |
| 162 | } |
| 163 | copy(big, []byte(`{"jsonrpc":"2.0","id":1,"method":"extension/initialize","params":{}}`)) |
| 164 | host.writeRaw(big) |
| 165 | err, ok := serveDone.wait(5 * time.Second) |
| 166 | if !ok { |
| 167 | t.Fatal("Serve did not end on an oversized frame") |
| 168 | } |
| 169 | var tooLarge *FrameTooLargeError |
| 170 | if !errors.As(err, &tooLarge) { |
| 171 | t.Fatalf("Serve error = %v, want FrameTooLargeError", err) |
| 172 | } |
| 173 | } |
| 174 | |
| 175 | // TestStrictFrameViolations checks envelope rejection: wrong jsonrpc |
| 176 | // version, string ids, and non-object params all answer -32600, and the |
| 177 | // connection survives to complete the handshake afterwards. |
| 178 | func TestStrictFrameViolations(t *testing.T) { |
| 179 | host, _ := startFakeHost(t, basicHandler(), Options{}) |
| 180 | cases := []struct { |
| 181 | name string |
| 182 | frame string |
| 183 | wantID string // expected id of the rejection response |
| 184 | }{ |
| 185 | {"wrong jsonrpc", `{"jsonrpc":"1.0","id":1,"method":"extension/initialize","params":{}}`, "1"}, |
| 186 | {"missing jsonrpc", `{"id":2,"method":"extension/initialize","params":{}}`, "2"}, |
| 187 | {"string id", `{"jsonrpc":"2.0","id":"abc","method":"extension/initialize","params":{}}`, "null"}, |
| 188 | {"fractional id", `{"jsonrpc":"2.0","id":1.5,"method":"extension/initialize","params":{}}`, "null"}, |
| 189 | {"array params", `{"jsonrpc":"2.0","id":3,"method":"extension/initialize","params":[]}`, "3"}, |
| 190 | {"scalar params", `{"jsonrpc":"2.0","id":4,"method":"extension/initialize","params":42}`, "4"}, |
| 191 | {"result and method", `{"jsonrpc":"2.0","id":5,"method":"extension/initialize","result":{}}`, "5"}, |
| 192 | } |
| 193 | for _, tc := range cases { |
| 194 | host.writeRaw([]byte(tc.frame)) |
| 195 | stray := host.nextStray() |
| 196 | if stray.Error == nil { |
| 197 | t.Fatalf("%s: frame %s: expected an error response, got a result", tc.name, tc.frame) |
| 198 | } |
| 199 | if stray.Error.Code != CodeInvalidRequest { |
| 200 | t.Fatalf("%s: code = %d, want %d", tc.name, stray.Error.Code, CodeInvalidRequest) |
| 201 | } |
| 202 | if string(stray.ID) != tc.wantID { |
| 203 | t.Fatalf("%s: response id = %s, want %s", tc.name, stray.ID, tc.wantID) |
| 204 | } |
| 205 | } |
| 206 | // The connection survived every rejection. |
| 207 | host.handshake(t) |
| 208 | } |
| 209 | |
| 210 | // TestInterceptRouting exercises exact and wildcard routing plus the default |
| 211 | // continue answer. |
| 212 | func TestInterceptRouting(t *testing.T) { |
| 213 | var calls atomic.Int64 |
| 214 | interceptors := map[string]InterceptorFunc{ |
| 215 | "tool.before": func(_ context.Context, event string, payload json.RawMessage) (*InterceptResult, error) { |
| 216 | calls.Add(1) |
| 217 | if event != "tool.before" { |
| 218 | t.Errorf("event = %q, want tool.before", event) |
| 219 | } |
| 220 | return Block("no tools today"), nil |
| 221 | }, |
| 222 | "*": func(_ context.Context, event string, payload json.RawMessage) (*InterceptResult, error) { |
| 223 | calls.Add(1) |
| 224 | return Continue(), nil |
| 225 | }, |
| 226 | } |
| 227 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 228 | host.handshake(t) |
| 229 | |
| 230 | resp := host.request(MethodExtensionIntercept, InterceptParams{ |
| 231 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{"tool":"bash"}`), TimeoutMillis: 5000, |
| 232 | }) |
| 233 | if resp.Err != nil { |
| 234 | t.Fatalf("intercept failed: %+v", resp.Err) |
| 235 | } |
| 236 | var result InterceptResult |
| 237 | if err := json.Unmarshal(resp.Result, &result); err != nil { |
| 238 | t.Fatalf("decode intercept result: %v", err) |
| 239 | } |
| 240 | if result.Decision != DecisionBlock || result.Reason != "no tools today" { |
| 241 | t.Fatalf("result = %+v", result) |
| 242 | } |
| 243 | |
| 244 | resp = host.request(MethodExtensionIntercept, InterceptParams{ |
| 245 | Event: EventSessionStart, Seq: 2, Payload: json.RawMessage(`{}`), |
| 246 | }) |
| 247 | if err := json.Unmarshal(resp.Result, &result); err != nil { |
| 248 | t.Fatalf("decode wildcard result: %v", err) |
| 249 | } |
| 250 | if result.Decision != DecisionContinue { |
| 251 | t.Fatalf("wildcard decision = %q, want continue", result.Decision) |
| 252 | } |
| 253 | if calls.Load() != 2 { |
| 254 | t.Fatalf("interceptor calls = %d, want 2", calls.Load()) |
| 255 | } |
| 256 | } |
| 257 | |
| 258 | // TestInterceptReplaceHelper checks Replace marshaling and the wire shape. |
| 259 | func TestInterceptReplace(t *testing.T) { |
| 260 | interceptors := map[string]InterceptorFunc{ |
| 261 | "input.receive": func(_ context.Context, _ string, _ json.RawMessage) (*InterceptResult, error) { |
| 262 | return Replace(map[string]any{"text": "rewritten", "n": 2}) |
| 263 | }, |
| 264 | } |
| 265 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 266 | host.handshake(t) |
| 267 | resp := host.request(MethodExtensionIntercept, InterceptParams{ |
| 268 | Event: EventInputReceive, Seq: 1, Payload: json.RawMessage(`{"text":"original"}`), |
| 269 | }) |
| 270 | if resp.Err != nil { |
| 271 | t.Fatalf("intercept failed: %+v", resp.Err) |
| 272 | } |
| 273 | var wire struct { |
| 274 | Decision string `json:"decision"` |
| 275 | Replacement json.RawMessage `json:"replacement"` |
| 276 | } |
| 277 | if err := json.Unmarshal(resp.Result, &wire); err != nil { |
| 278 | t.Fatalf("decode: %v", err) |
| 279 | } |
| 280 | if wire.Decision != "replace" { |
| 281 | t.Fatalf("decision = %q", wire.Decision) |
| 282 | } |
| 283 | var replacement map[string]any |
| 284 | if err := json.Unmarshal(wire.Replacement, &replacement); err != nil { |
| 285 | t.Fatalf("replacement is not an object: %v", err) |
| 286 | } |
| 287 | if replacement["text"] != "rewritten" || replacement["n"] != float64(2) { |
| 288 | t.Fatalf("replacement = %v", replacement) |
| 289 | } |
| 290 | } |
| 291 | |
| 292 | // TestInterceptAllowDeny covers the permission.decision rulings. |
| 293 | func TestInterceptAllowDeny(t *testing.T) { |
| 294 | interceptors := map[string]InterceptorFunc{ |
| 295 | "permission.decision": func(_ context.Context, _ string, payload json.RawMessage) (*InterceptResult, error) { |
| 296 | if strings.Contains(string(payload), "dangerous") { |
| 297 | return Deny("too dangerous"), nil |
| 298 | } |
| 299 | return Allow(), nil |
| 300 | }, |
| 301 | } |
| 302 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 303 | host.handshake(t) |
| 304 | var result InterceptResult |
| 305 | resp := host.request(MethodExtensionIntercept, InterceptParams{ |
| 306 | Event: EventPermissionDecision, Seq: 1, Payload: json.RawMessage(`{"command":"ls"}`), |
| 307 | }) |
| 308 | if err := json.Unmarshal(resp.Result, &result); err != nil || result.Decision != DecisionAllow { |
| 309 | t.Fatalf("allow: result=%+v err=%v", result, err) |
| 310 | } |
| 311 | resp = host.request(MethodExtensionIntercept, InterceptParams{ |
| 312 | Event: EventPermissionDecision, Seq: 2, Payload: json.RawMessage(`{"command":"dangerous"}`), |
| 313 | }) |
| 314 | if err := json.Unmarshal(resp.Result, &result); err != nil || result.Decision != DecisionDeny || result.Reason != "too dangerous" { |
| 315 | t.Fatalf("deny: result=%+v err=%v", result, err) |
| 316 | } |
| 317 | } |
| 318 | |
| 319 | // TestInterceptInvalidEnvelope verifies envelope validation stays with the |
| 320 | // SDK: unknown events, bad seqs, and missing payload keys answer |
| 321 | // invalid_params without reaching the interceptor. |
| 322 | func TestInterceptInvalidEnvelope(t *testing.T) { |
| 323 | var calls atomic.Int64 |
| 324 | interceptors := map[string]InterceptorFunc{ |
| 325 | "*": func(context.Context, string, json.RawMessage) (*InterceptResult, error) { |
| 326 | calls.Add(1) |
| 327 | return Continue(), nil |
| 328 | }, |
| 329 | } |
| 330 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 331 | host.handshake(t) |
| 332 | frames := []string{ |
| 333 | `{"event":"not.an.event","seq":1,"payload":{},"timeoutMillis":0}`, |
| 334 | `{"event":"tool.before","seq":0,"payload":{},"timeoutMillis":0}`, |
| 335 | `{"event":"tool.before","seq":1,"timeoutMillis":0}`, |
| 336 | `{"event":"tool.before","seq":1,"payload":{},"timeoutMillis":-1,"extra":1}`, |
| 337 | } |
| 338 | for _, params := range frames { |
| 339 | resp := host.request(MethodExtensionIntercept, json.RawMessage(params)) |
| 340 | if resp.Err == nil || resp.Err.Code != CodeInvalidParams { |
| 341 | t.Fatalf("params %s: expected invalid_params, got %+v", params, resp.Err) |
| 342 | } |
| 343 | } |
| 344 | if calls.Load() != 0 { |
| 345 | t.Fatalf("interceptor ran %d times on invalid envelopes", calls.Load()) |
| 346 | } |
| 347 | } |
| 348 | |
| 349 | // TestInterceptHandlerPanic verifies a panicking interceptor answers the |
| 350 | // frozen internal error and the connection survives. |
| 351 | func TestInterceptHandlerPanic(t *testing.T) { |
| 352 | interceptors := map[string]InterceptorFunc{ |
| 353 | "tool.before": func(context.Context, string, json.RawMessage) (*InterceptResult, error) { |
| 354 | panic("boom") |
| 355 | }, |
| 356 | } |
| 357 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 358 | host.handshake(t) |
| 359 | resp := host.request(MethodExtensionIntercept, InterceptParams{ |
| 360 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), |
| 361 | }) |
| 362 | if resp.Err == nil || resp.Err.Code != CodeInternal { |
| 363 | t.Fatalf("expected internal error, got %+v", resp.Err) |
| 364 | } |
| 365 | data, _ := resp.Err.Data.(ProtocolErrorData) |
| 366 | if data.Reason != ErrInternal { |
| 367 | t.Fatalf("reason = %q, want internal", data.Reason) |
| 368 | } |
| 369 | resp = host.request(MethodExtensionIntercept, InterceptParams{ |
| 370 | Event: EventToolBefore, Seq: 2, Payload: json.RawMessage(`{}`), |
| 371 | }) |
| 372 | if resp.Err == nil { |
| 373 | t.Fatal("connection died after a handler panic") |
| 374 | } |
| 375 | } |
| 376 | |
| 377 | // TestInterceptDeadlineReturnsFrozenTimeout pins the equal-deadline race |
| 378 | // between the SDK callback budget and the host request budget. When the SDK |
| 379 | // timer wins, it must answer intercept_timeout rather than a generic internal |
| 380 | // error so the host observes one deterministic reason either way. |
| 381 | func TestInterceptDeadlineReturnsFrozenTimeout(t *testing.T) { |
| 382 | interceptors := map[string]InterceptorFunc{ |
| 383 | "tool.before": func(ctx context.Context, _ string, _ json.RawMessage) (*InterceptResult, error) { |
| 384 | <-ctx.Done() |
| 385 | return nil, ctx.Err() |
| 386 | }, |
| 387 | } |
| 388 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 389 | host.handshake(t) |
| 390 | resp := host.request(MethodExtensionIntercept, InterceptParams{ |
| 391 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), TimeoutMillis: 20, |
| 392 | }) |
| 393 | if resp.Err == nil || resp.Err.Code != DomainErrorCode { |
| 394 | t.Fatalf("expected a domain timeout error, got %+v", resp.Err) |
| 395 | } |
| 396 | data, _ := resp.Err.Data.(ProtocolErrorData) |
| 397 | if data.Reason != ErrInterceptTimeout { |
| 398 | t.Fatalf("reason = %q, want %q", data.Reason, ErrInterceptTimeout) |
| 399 | } |
| 400 | } |
| 401 | |
| 402 | // TestEventObservation checks extension/event reaches the observer and a |
| 403 | // panicking observer does not kill the loop. |
| 404 | func TestEventObservation(t *testing.T) { |
| 405 | seen := make(chan string, 4) |
| 406 | opts := Options{ |
| 407 | Observer: func(_ context.Context, event string, payload json.RawMessage) { |
| 408 | seen <- event + ":" + string(payload) |
| 409 | }, |
| 410 | } |
| 411 | host, _ := startFakeHost(t, basicHandler(), opts) |
| 412 | host.handshake(t) |
| 413 | host.notify(MethodExtensionEvent, EventParams{Event: EventToolAfter, Payload: json.RawMessage(`{"ok":true}`)}) |
| 414 | select { |
| 415 | case got := <-seen: |
| 416 | if got != `tool.after:{"ok":true}` { |
| 417 | t.Fatalf("observation = %q", got) |
| 418 | } |
| 419 | case <-time.After(5 * time.Second): |
| 420 | t.Fatal("observer not called") |
| 421 | } |
| 422 | // A nil observer extension must not choke on events either. |
| 423 | host2, _ := startFakeHost(t, basicHandler(), Options{}) |
| 424 | host2.handshake(t) |
| 425 | host2.notify(MethodExtensionEvent, EventParams{Event: EventToolAfter, Payload: json.RawMessage(`{}`)}) |
| 426 | } |
| 427 | |
| 428 | // TestResourcesChanged checks the resources/changed notification. |
| 429 | func TestResourcesChanged(t *testing.T) { |
| 430 | seen := make(chan []string, 1) |
| 431 | opts := Options{ResourcesChanged: func(_ context.Context, paths []string) { seen <- paths }} |
| 432 | host, _ := startFakeHost(t, basicHandler(), opts) |
| 433 | host.handshake(t) |
| 434 | host.notify(MethodExtensionResourcesChanged, ResourcesChangedParams{Paths: []string{"skills/a", "themes/b"}}) |
| 435 | select { |
| 436 | case paths := <-seen: |
| 437 | if len(paths) != 2 || paths[0] != "skills/a" { |
| 438 | t.Fatalf("paths = %v", paths) |
| 439 | } |
| 440 | case <-time.After(5 * time.Second): |
| 441 | t.Fatal("ResourcesChanged not called") |
| 442 | } |
| 443 | } |
| 444 | |
| 445 | // TestUIActionSubmit covers the Host → Extension UI calls. |
| 446 | func TestUIActionSubmit(t *testing.T) { |
| 447 | ui := UIHandler{ |
| 448 | Action: func(_ context.Context, actionID string, args map[string]string) error { |
| 449 | if actionID == "fail" { |
| 450 | return errTest |
| 451 | } |
| 452 | if actionID != "open" || args["k"] != "v" { |
| 453 | return errors.New("bad action invocation") |
| 454 | } |
| 455 | return nil |
| 456 | }, |
| 457 | Submit: func(_ context.Context, surfaceID string, values map[string]any) error { |
| 458 | if surfaceID != "form-1" || values["name"] != "reasonix" { |
| 459 | return errTest |
| 460 | } |
| 461 | return nil |
| 462 | }, |
| 463 | } |
| 464 | host, _ := startFakeHost(t, basicHandler(), Options{UI: ui}) |
| 465 | host.handshake(t) |
| 466 | |
| 467 | resp := host.request(MethodExtensionUIAction, UIActionParams{ |
| 468 | ActionID: "open", SessionID: "sess-1", Generation: 7, Args: map[string]string{"k": "v"}, |
| 469 | }) |
| 470 | var actionResult UIActionResult |
| 471 | if err := json.Unmarshal(resp.Result, &actionResult); err != nil || !actionResult.Accepted { |
| 472 | t.Fatalf("action: result=%+v err=%v respErr=%+v", actionResult, err, resp.Err) |
| 473 | } |
| 474 | resp = host.request(MethodExtensionUIAction, UIActionParams{ActionID: "fail", SessionID: "sess-1", Generation: 7}) |
| 475 | if err := json.Unmarshal(resp.Result, &actionResult); err != nil { |
| 476 | t.Fatalf("action fail decode: %v", err) |
| 477 | } |
| 478 | if actionResult.Accepted || actionResult.Message != errTest.Error() { |
| 479 | t.Fatalf("action fail result = %+v", actionResult) |
| 480 | } |
| 481 | |
| 482 | resp = host.request(MethodExtensionUISubmit, UISubmitParams{ |
| 483 | SurfaceID: "form-1", SessionID: "sess-1", Generation: 7, Values: map[string]any{"name": "reasonix"}, |
| 484 | }) |
| 485 | var submitResult UISubmitResult |
| 486 | if err := json.Unmarshal(resp.Result, &submitResult); err != nil || !submitResult.Accepted { |
| 487 | t.Fatalf("submit: result=%+v err=%v", submitResult, err) |
| 488 | } |
| 489 | |
| 490 | // Without UI configured the methods answer unknown_method. |
| 491 | host2, _ := startFakeHost(t, basicHandler(), Options{}) |
| 492 | host2.handshake(t) |
| 493 | resp = host2.request(MethodExtensionUIAction, UIActionParams{ActionID: "x", SessionID: "s", Generation: 1}) |
| 494 | if resp.Err == nil || resp.Err.Code != CodeMethodNotFound { |
| 495 | t.Fatalf("expected unknown_method without UI, got %+v", resp.Err) |
| 496 | } |
| 497 | } |
| 498 | |
| 499 | // TestShutdownSequence runs the graceful stop: the fn runs with its timeout, |
| 500 | // {accepted:true} is answered, and Serve returns nil. |
| 501 | func TestShutdownSequence(t *testing.T) { |
| 502 | var ran atomic.Bool |
| 503 | opts := Options{Shutdown: func(ctx context.Context) { |
| 504 | ran.Store(true) |
| 505 | if _, hasDeadline := ctx.Deadline(); !hasDeadline { |
| 506 | t.Errorf("shutdown ctx has no deadline despite timeoutMillis") |
| 507 | } |
| 508 | }} |
| 509 | host, serveDone := startFakeHost(t, basicHandler(), opts) |
| 510 | host.handshake(t) |
| 511 | resp := host.request(MethodExtensionShutdown, ShutdownParams{TimeoutMillis: 5000}) |
| 512 | if resp.Err != nil { |
| 513 | t.Fatalf("shutdown failed: %+v", resp.Err) |
| 514 | } |
| 515 | var result ShutdownResult |
| 516 | if err := json.Unmarshal(resp.Result, &result); err != nil || !result.Accepted { |
| 517 | t.Fatalf("shutdown result = %+v", result) |
| 518 | } |
| 519 | if !ran.Load() { |
| 520 | t.Fatal("shutdown fn did not run") |
| 521 | } |
| 522 | err, ok := serveDone.wait(5 * time.Second) |
| 523 | if !ok { |
| 524 | t.Fatal("Serve did not return after shutdown") |
| 525 | } |
| 526 | if err != nil { |
| 527 | t.Fatalf("Serve returned %v after an orderly shutdown, want nil", err) |
| 528 | } |
| 529 | } |
| 530 | |
| 531 | // TestShutdownWithoutFn still answers and exits when no Shutdown fn is set. |
| 532 | func TestShutdownWithoutFn(t *testing.T) { |
| 533 | host, serveDone := startFakeHost(t, basicHandler(), Options{}) |
| 534 | host.handshake(t) |
| 535 | resp := host.request(MethodExtensionShutdown, ShutdownParams{TimeoutMillis: 0}) |
| 536 | var result ShutdownResult |
| 537 | if err := json.Unmarshal(resp.Result, &result); err != nil || !result.Accepted { |
| 538 | t.Fatalf("shutdown result = %+v err=%v", result, resp.Err) |
| 539 | } |
| 540 | err, ok := serveDone.wait(5 * time.Second) |
| 541 | if !ok { |
| 542 | t.Fatal("Serve did not return") |
| 543 | } |
| 544 | if err != nil { |
| 545 | t.Fatalf("Serve returned %v, want nil", err) |
| 546 | } |
| 547 | } |
| 548 | |
| 549 | // TestServeContextCancel verifies canceling the Serve context tears the |
| 550 | // transport down promptly. |
| 551 | func TestServeContextCancel(t *testing.T) { |
| 552 | sdkStdinR, sdkStdinW := io.Pipe() |
| 553 | sdkStdoutR, sdkStdoutW := io.Pipe() |
| 554 | defer sdkStdoutR.Close() |
| 555 | ctx, cancel := context.WithCancel(context.Background()) |
| 556 | serveDone := make(chan error, 1) |
| 557 | go func() { |
| 558 | serveDone <- Serve(ctx, basicHandler(), Options{Stdin: sdkStdinR, Stdout: sdkStdoutW}) |
| 559 | }() |
| 560 | go io.Copy(io.Discard, sdkStdoutR) |
| 561 | time.Sleep(50 * time.Millisecond) |
| 562 | cancel() |
| 563 | select { |
| 564 | case err := <-serveDone: |
| 565 | if !errors.Is(err, context.Canceled) { |
| 566 | t.Fatalf("Serve returned %v, want context.Canceled", err) |
| 567 | } |
| 568 | case <-time.After(5 * time.Second): |
| 569 | t.Fatal("Serve did not return after ctx cancel") |
| 570 | } |
| 571 | _ = sdkStdinW.Close() |
| 572 | } |
| 573 | |
| 574 | // TestHostEOFCleanReturn covers the host closing the transport: Serve |
| 575 | // returns nil. |
| 576 | func TestHostEOFCleanReturn(t *testing.T) { |
| 577 | host, serveDone := startFakeHost(t, basicHandler(), Options{}) |
| 578 | host.handshake(t) |
| 579 | if err := host.toSDK.Close(); err != nil { |
| 580 | t.Fatalf("close host pipe: %v", err) |
| 581 | } |
| 582 | err, ok := serveDone.wait(5 * time.Second) |
| 583 | if !ok { |
| 584 | t.Fatal("Serve did not return on host EOF") |
| 585 | } |
| 586 | if err != nil { |
| 587 | t.Fatalf("Serve returned %v on host EOF, want nil", err) |
| 588 | } |
| 589 | } |
| 590 |