| 1 | package dispatch |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "fmt" |
| 8 | "reflect" |
| 9 | "strings" |
| 10 | "sync" |
| 11 | "testing" |
| 12 | |
| 13 | "reasonix/internal/extension" |
| 14 | "reasonix/internal/extension/protocol" |
| 15 | "reasonix/internal/extension/sidecar" |
| 16 | ) |
| 17 | |
| 18 | // The production Client interface exists so the real sidecar client drops in |
| 19 | // without an adapter; pin that here so a signature drift fails the build. |
| 20 | var _ Client = (*sidecar.Client)(nil) |
| 21 | |
| 22 | // testSecret is a credential shape secrets.RedactCredentials reliably masks |
| 23 | // (it appears in internal/secrets' own tests). |
| 24 | const testSecret = "sk-real-secret-value-123456" |
| 25 | |
| 26 | func interceptor(pluginID string, point extension.InterceptorPoint, priority int) extension.Contribution { |
| 27 | return extension.Contribution{ |
| 28 | Kind: extension.KindInterceptor, |
| 29 | ID: string(point), |
| 30 | Source: extension.ContributionSource{Scope: extension.ScopePlugin, PluginID: pluginID, Origin: "extension-runtime"}, |
| 31 | Priority: priority, |
| 32 | } |
| 33 | } |
| 34 | |
| 35 | func owner(pluginID string) extension.ContributionSource { |
| 36 | return extension.ContributionSource{Scope: extension.ScopePlugin, PluginID: pluginID, Origin: "extension-runtime"} |
| 37 | } |
| 38 | |
| 39 | // buildDispatcher wires a dispatcher over the given fakes; a plugin absent |
| 40 | // from fakes resolves to a nil (untyped) client, mirroring the documented |
| 41 | // adapter contract. |
| 42 | func buildDispatcher(chain map[extension.InterceptorPoint][]extension.Contribution, replacements map[extension.Slot]extension.ContributionSource, fakes map[string]*fakeClient, required map[string]bool, warns *warnRecorder) *Dispatcher { |
| 43 | clients := func(pluginID string) Client { |
| 44 | if client := fakes[pluginID]; client != nil { |
| 45 | return client |
| 46 | } |
| 47 | return nil |
| 48 | } |
| 49 | return New(chain, replacements, clients, required, Options{Warn: warns.warn}) |
| 50 | } |
| 51 | |
| 52 | func timeoutError(pluginID string, point extension.InterceptorPoint) error { |
| 53 | return &protocol.ProtocolError{ |
| 54 | Reason: protocol.ErrInterceptTimeout, |
| 55 | Message: fmt.Sprintf("extension %s did not answer %s within 5s", pluginID, point), |
| 56 | } |
| 57 | } |
| 58 | |
| 59 | // pointCase describes one intercept point for the dispatch matrix. |
| 60 | type pointCase struct { |
| 61 | point extension.InterceptorPoint |
| 62 | sample func() any |
| 63 | replaceJSON string |
| 64 | checkReplaced func(t *testing.T, payload any) |
| 65 | violateJSON string |
| 66 | } |
| 67 | |
| 68 | func userMessage(content string) protocol.ProviderMessage { |
| 69 | return protocol.ProviderMessage{Role: protocol.ProviderRoleUser, Content: content} |
| 70 | } |
| 71 | |
| 72 | func pointCases() []pointCase { |
| 73 | cases := []pointCase{ |
| 74 | { |
| 75 | point: extension.PointInputReceive, |
| 76 | sample: func() any { return &InputPayload{Text: "hello"} }, |
| 77 | replaceJSON: `{"text":"rewritten"}`, |
| 78 | checkReplaced: func(t *testing.T, payload any) { |
| 79 | t.Helper() |
| 80 | if got := payload.(*InputPayload).Text; got != "rewritten" { |
| 81 | t.Fatalf("Text = %q, want %q", got, "rewritten") |
| 82 | } |
| 83 | }, |
| 84 | violateJSON: `{"text":""}`, |
| 85 | }, |
| 86 | { |
| 87 | point: extension.PointAgentBeforeStart, |
| 88 | sample: func() any { return &AgentStartPayload{Model: "openai/gpt-5", ToolCount: 3, SessionID: "s1"} }, |
| 89 | replaceJSON: `{"model":"other/model","toolCount":7,"sessionId":"s1"}`, |
| 90 | checkReplaced: func(t *testing.T, payload any) { |
| 91 | t.Helper() |
| 92 | got := payload.(*AgentStartPayload) |
| 93 | if got.Model != "other/model" || got.ToolCount != 7 { |
| 94 | t.Fatalf("payload = %+v, want model other/model with 7 tools", got) |
| 95 | } |
| 96 | }, |
| 97 | violateJSON: `{"model":"m"}`, |
| 98 | }, |
| 99 | { |
| 100 | point: extension.PointSystemPromptBuild, |
| 101 | sample: func() any { return &SystemPromptPayload{Prompt: "base prompt", WorkspaceRoot: "/ws"} }, |
| 102 | replaceJSON: `{"prompt":"owned prompt","workspaceRoot":"/ws"}`, |
| 103 | checkReplaced: func(t *testing.T, payload any) { |
| 104 | t.Helper() |
| 105 | if got := payload.(*SystemPromptPayload).Prompt; got != "owned prompt" { |
| 106 | t.Fatalf("Prompt = %q, want %q", got, "owned prompt") |
| 107 | } |
| 108 | }, |
| 109 | violateJSON: `{"prompt":"x"}`, |
| 110 | }, |
| 111 | { |
| 112 | point: extension.PointContextPrepare, |
| 113 | sample: func() any { return &ContextPayload{Messages: []protocol.ProviderMessage{userMessage("hi")}} }, |
| 114 | replaceJSON: `{"messages":[{"role":"user","content":"replaced"}]}`, |
| 115 | checkReplaced: func(t *testing.T, payload any) { |
| 116 | t.Helper() |
| 117 | got := payload.(*ContextPayload) |
| 118 | if len(got.Messages) != 1 || got.Messages[0].Content != "replaced" { |
| 119 | t.Fatalf("Messages = %+v, want one replaced message", got.Messages) |
| 120 | } |
| 121 | }, |
| 122 | violateJSON: `{}`, |
| 123 | }, |
| 124 | { |
| 125 | point: extension.PointProviderRequest, |
| 126 | sample: func() any { |
| 127 | return &ProviderRequestPayload{Request: protocol.ProviderRequest{ |
| 128 | Messages: []protocol.ProviderMessage{userMessage("q")}, |
| 129 | Tools: []protocol.ProviderToolSchema{}, |
| 130 | }} |
| 131 | }, |
| 132 | replaceJSON: `{"request":{"messages":[{"role":"user","content":"q2"}],"tools":[],"maxTokens":99}}`, |
| 133 | checkReplaced: func(t *testing.T, payload any) { |
| 134 | t.Helper() |
| 135 | got := payload.(*ProviderRequestPayload) |
| 136 | if got.Request.MaxTokens != 99 || got.Request.Messages[0].Content != "q2" { |
| 137 | t.Fatalf("Request = %+v, want maxTokens 99 and replaced message", got.Request) |
| 138 | } |
| 139 | }, |
| 140 | // tool parameters must be a JSON object, not an array. |
| 141 | violateJSON: `{"request":{"messages":[],"tools":[{"name":"t","parameters":[1]}]}}`, |
| 142 | }, |
| 143 | { |
| 144 | point: extension.PointProviderResponse, |
| 145 | sample: func() any { |
| 146 | return &ProviderResponsePayload{Text: "answer", Usage: &protocol.ProviderUsage{PromptTokens: 1, TotalTokens: 2}} |
| 147 | }, |
| 148 | replaceJSON: `{"text":"changed","calls":[{"id":"c1","name":"bash","arguments":"{}"}]}`, |
| 149 | checkReplaced: func(t *testing.T, payload any) { |
| 150 | t.Helper() |
| 151 | got := payload.(*ProviderResponsePayload) |
| 152 | if got.Text != "changed" || len(got.Calls) != 1 { |
| 153 | t.Fatalf("payload = %+v, want changed text with one call", got) |
| 154 | } |
| 155 | // Whole-value assignment: fields absent from the replacement |
| 156 | // must not leak the previous value through. |
| 157 | if got.Usage != nil { |
| 158 | t.Fatalf("Usage = %+v, want nil (replacement omitted it)", got.Usage) |
| 159 | } |
| 160 | }, |
| 161 | violateJSON: `{"calls":[{"id":"","name":"x"}]}`, |
| 162 | }, |
| 163 | { |
| 164 | point: extension.PointToolBefore, |
| 165 | sample: func() any { return &ToolBeforePayload{Name: "bash", Arguments: `{"cmd":"ls"}`} }, |
| 166 | replaceJSON: `{"name":"bash","arguments":"{\"cmd\":\"pwd\"}"}`, |
| 167 | checkReplaced: func(t *testing.T, payload any) { |
| 168 | t.Helper() |
| 169 | if got := payload.(*ToolBeforePayload).Arguments; !strings.Contains(got, "pwd") { |
| 170 | t.Fatalf("Arguments = %q, want a pwd command", got) |
| 171 | } |
| 172 | }, |
| 173 | violateJSON: `{"name":"bash","arguments":"not json"}`, |
| 174 | }, |
| 175 | { |
| 176 | point: extension.PointToolAfter, |
| 177 | sample: func() any { return &ToolAfterPayload{Name: "bash", Arguments: `{"cmd":"ls"}`, Result: "out"} }, |
| 178 | replaceJSON: `{"name":"bash","result":"new out"}`, |
| 179 | checkReplaced: func(t *testing.T, payload any) { |
| 180 | t.Helper() |
| 181 | got := payload.(*ToolAfterPayload) |
| 182 | if got.Result != "new out" || got.Arguments != "" { |
| 183 | t.Fatalf("payload = %+v, want new result and cleared arguments", got) |
| 184 | } |
| 185 | }, |
| 186 | violateJSON: `{}`, |
| 187 | }, |
| 188 | { |
| 189 | point: extension.PointPermissionDecision, |
| 190 | sample: func() any { |
| 191 | return &PermissionPayload{Name: "bash", Arguments: `{"cmd":"rm -rf x"}`, HostDecision: "deny"} |
| 192 | }, |
| 193 | replaceJSON: `{"name":"bash","arguments":"{\"cmd\":\"ls\"}","hostDecision":"deny"}`, |
| 194 | checkReplaced: func(t *testing.T, payload any) { |
| 195 | t.Helper() |
| 196 | if got := payload.(*PermissionPayload).Arguments; !strings.Contains(got, "ls") { |
| 197 | t.Fatalf("Arguments = %q, want an ls command", got) |
| 198 | } |
| 199 | }, |
| 200 | violateJSON: `{"name":"bash","hostDecision":"maybe"}`, |
| 201 | }, |
| 202 | { |
| 203 | point: extension.PointCompactionPrepare, |
| 204 | sample: func() any { |
| 205 | return &CompactionPreparePayload{Messages: []protocol.ProviderMessage{userMessage("m")}, Guidance: "g"} |
| 206 | }, |
| 207 | replaceJSON: `{"messages":[],"guidance":"new guidance"}`, |
| 208 | checkReplaced: func(t *testing.T, payload any) { |
| 209 | t.Helper() |
| 210 | got := payload.(*CompactionPreparePayload) |
| 211 | if got.Guidance != "new guidance" || got.Messages == nil || len(got.Messages) != 0 { |
| 212 | t.Fatalf("payload = %+v, want new guidance with an empty non-nil messages array", got) |
| 213 | } |
| 214 | }, |
| 215 | violateJSON: `{}`, |
| 216 | }, |
| 217 | { |
| 218 | point: extension.PointCompactionComplete, |
| 219 | sample: func() any { return &CompactionCompletePayload{Summary: "summary"} }, |
| 220 | replaceJSON: `{"summary":"new summary"}`, |
| 221 | checkReplaced: func(t *testing.T, payload any) { |
| 222 | t.Helper() |
| 223 | if got := payload.(*CompactionCompletePayload).Summary; got != "new summary" { |
| 224 | t.Fatalf("Summary = %q, want %q", got, "new summary") |
| 225 | } |
| 226 | }, |
| 227 | violateJSON: `{}`, |
| 228 | }, |
| 229 | { |
| 230 | point: extension.PointFrontendEvent, |
| 231 | sample: func() any { return &FrontendEventPayload{Kind: "notice", Text: "t", Detail: "d"} }, |
| 232 | replaceJSON: `{"kind":"notice","text":"replaced text"}`, |
| 233 | checkReplaced: func(t *testing.T, payload any) { |
| 234 | t.Helper() |
| 235 | if got := payload.(*FrontendEventPayload).Text; got != "replaced text" { |
| 236 | t.Fatalf("Text = %q, want %q", got, "replaced text") |
| 237 | } |
| 238 | }, |
| 239 | violateJSON: `{}`, |
| 240 | }, |
| 241 | } |
| 242 | for _, phase := range []string{PhaseStart, PhaseEnd, PhaseLoad, PhaseSave, PhaseRotate} { |
| 243 | point := extension.InterceptorPoint("session." + phase) |
| 244 | cases = append(cases, pointCase{ |
| 245 | point: point, |
| 246 | sample: func() any { return &SessionPayload{SessionPath: "/tmp/s.json", Phase: phase} }, |
| 247 | replaceJSON: fmt.Sprintf(`{"sessionPath":"/tmp/other.json","phase":%q}`, phase), |
| 248 | checkReplaced: func(t *testing.T, payload any) { |
| 249 | t.Helper() |
| 250 | got := payload.(*SessionPayload) |
| 251 | if got.SessionPath != "/tmp/other.json" || got.Phase != phase { |
| 252 | t.Fatalf("payload = %+v, want replaced path at phase %q", got, phase) |
| 253 | } |
| 254 | }, |
| 255 | violateJSON: fmt.Sprintf(`{"sessionPath":"/x","phase":%q}`, "bogus"), |
| 256 | }) |
| 257 | } |
| 258 | return cases |
| 259 | } |
| 260 | |
| 261 | // TestInterceptMatrixContinue verifies all 17 points: a continue ruling |
| 262 | // passes the payload through unchanged. |
| 263 | func TestInterceptMatrixContinue(t *testing.T) { |
| 264 | for _, tc := range pointCases() { |
| 265 | t.Run(string(tc.point), func(t *testing.T) { |
| 266 | fake := &fakeClient{} |
| 267 | warns := &warnRecorder{} |
| 268 | d := buildDispatcher( |
| 269 | map[extension.InterceptorPoint][]extension.Contribution{tc.point: {interceptor("p1", tc.point, 0)}}, |
| 270 | nil, map[string]*fakeClient{"p1": fake}, nil, warns) |
| 271 | payload := tc.sample() |
| 272 | result, err := d.Intercept(context.Background(), tc.point, payload) |
| 273 | if err != nil { |
| 274 | t.Fatalf("Intercept: %v", err) |
| 275 | } |
| 276 | if result.Blocked || result.Permission != nil || len(result.Applied) != 0 { |
| 277 | t.Fatalf("result = %+v, want a clean pass-through", result) |
| 278 | } |
| 279 | if !reflect.DeepEqual(payload, tc.sample()) { |
| 280 | t.Fatalf("payload = %+v, want unchanged %+v", payload, tc.sample()) |
| 281 | } |
| 282 | if fake.interceptCount() != 1 { |
| 283 | t.Fatalf("intercept calls = %d, want 1", fake.interceptCount()) |
| 284 | } |
| 285 | if warns.count() != 0 { |
| 286 | t.Fatalf("warns = %v, want none", warns.msgs) |
| 287 | } |
| 288 | }) |
| 289 | } |
| 290 | } |
| 291 | |
| 292 | // TestInterceptMatrixBlock verifies all 17 points: a block ruling stops the |
| 293 | // operation and the reason is credential-redacted. |
| 294 | func TestInterceptMatrixBlock(t *testing.T) { |
| 295 | for _, tc := range pointCases() { |
| 296 | t.Run(string(tc.point), func(t *testing.T) { |
| 297 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 298 | return protocol.InterceptResult{Decision: protocol.DecisionBlock, Reason: "denied, token " + testSecret}, nil |
| 299 | }} |
| 300 | warns := &warnRecorder{} |
| 301 | d := buildDispatcher( |
| 302 | map[extension.InterceptorPoint][]extension.Contribution{tc.point: {interceptor("p1", tc.point, 0)}}, |
| 303 | nil, map[string]*fakeClient{"p1": fake}, nil, warns) |
| 304 | payload := tc.sample() |
| 305 | result, err := d.Intercept(context.Background(), tc.point, payload) |
| 306 | if err != nil { |
| 307 | t.Fatalf("Intercept: %v", err) |
| 308 | } |
| 309 | if !result.Blocked { |
| 310 | t.Fatalf("result = %+v, want blocked", result) |
| 311 | } |
| 312 | if strings.Contains(result.BlockReason, testSecret) { |
| 313 | t.Fatalf("BlockReason %q leaks the credential", result.BlockReason) |
| 314 | } |
| 315 | if !strings.Contains(result.BlockReason, "denied, token") { |
| 316 | t.Fatalf("BlockReason %q lost the human-readable reason", result.BlockReason) |
| 317 | } |
| 318 | if result.BlockReason == "denied, token "+testSecret { |
| 319 | t.Fatalf("BlockReason was not redacted at all") |
| 320 | } |
| 321 | }) |
| 322 | } |
| 323 | } |
| 324 | |
| 325 | // TestInterceptMatrixReplace verifies all 17 points: a replace ruling |
| 326 | // substitutes the payload and the caller observes the new value. |
| 327 | func TestInterceptMatrixReplace(t *testing.T) { |
| 328 | for _, tc := range pointCases() { |
| 329 | t.Run(string(tc.point), func(t *testing.T) { |
| 330 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 331 | return protocol.InterceptResult{Decision: protocol.DecisionReplace, Replacement: json.RawMessage(tc.replaceJSON)}, nil |
| 332 | }} |
| 333 | warns := &warnRecorder{} |
| 334 | d := buildDispatcher( |
| 335 | map[extension.InterceptorPoint][]extension.Contribution{tc.point: {interceptor("p1", tc.point, 0)}}, |
| 336 | nil, map[string]*fakeClient{"p1": fake}, nil, warns) |
| 337 | payload := tc.sample() |
| 338 | result, err := d.Intercept(context.Background(), tc.point, payload) |
| 339 | if err != nil { |
| 340 | t.Fatalf("Intercept: %v", err) |
| 341 | } |
| 342 | tc.checkReplaced(t, payload) |
| 343 | if !reflect.DeepEqual(result.Applied, []string{"p1"}) { |
| 344 | t.Fatalf("Applied = %v, want [p1]", result.Applied) |
| 345 | } |
| 346 | if warns.count() != 0 { |
| 347 | t.Fatalf("warns = %v, want none", warns.msgs) |
| 348 | } |
| 349 | }) |
| 350 | } |
| 351 | } |
| 352 | |
| 353 | // TestInterceptMatrixInvalidReplace verifies all 17 points: a replacement |
| 354 | // with unknown fields or one that fails Validate is a protocol violation — |
| 355 | // optional extensions are warned about once and skipped (payload unchanged), |
| 356 | // required extensions fail the operation. |
| 357 | func TestInterceptMatrixInvalidReplace(t *testing.T) { |
| 358 | badPayloads := map[string]string{ |
| 359 | "unknown field": `{"bogusField":1}`, |
| 360 | "failed validate": "", // filled per point from violateJSON |
| 361 | } |
| 362 | for _, tc := range pointCases() { |
| 363 | for name, bad := range badPayloads { |
| 364 | if name == "failed validate" { |
| 365 | bad = tc.violateJSON |
| 366 | } |
| 367 | t.Run(string(tc.point)+"/"+name+"_optional", func(t *testing.T) { |
| 368 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 369 | return protocol.InterceptResult{Decision: protocol.DecisionReplace, Replacement: json.RawMessage(bad)}, nil |
| 370 | }} |
| 371 | warns := &warnRecorder{} |
| 372 | d := buildDispatcher( |
| 373 | map[extension.InterceptorPoint][]extension.Contribution{tc.point: {interceptor("p1", tc.point, 0)}}, |
| 374 | nil, map[string]*fakeClient{"p1": fake}, nil, warns) |
| 375 | payload := tc.sample() |
| 376 | result, err := d.Intercept(context.Background(), tc.point, payload) |
| 377 | if err != nil { |
| 378 | t.Fatalf("Intercept: optional violation must not fail, got %v", err) |
| 379 | } |
| 380 | if result.Blocked || len(result.Applied) != 0 { |
| 381 | t.Fatalf("result = %+v, want the ruling skipped", result) |
| 382 | } |
| 383 | if !reflect.DeepEqual(payload, tc.sample()) { |
| 384 | t.Fatalf("payload = %+v, want unchanged %+v", payload, tc.sample()) |
| 385 | } |
| 386 | if warns.count() != 1 || !warns.contains("p1") { |
| 387 | t.Fatalf("warns = %v, want one warning naming p1", warns.msgs) |
| 388 | } |
| 389 | }) |
| 390 | t.Run(string(tc.point)+"/"+name+"_required", func(t *testing.T) { |
| 391 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 392 | return protocol.InterceptResult{Decision: protocol.DecisionReplace, Replacement: json.RawMessage(bad)}, nil |
| 393 | }} |
| 394 | warns := &warnRecorder{} |
| 395 | d := buildDispatcher( |
| 396 | map[extension.InterceptorPoint][]extension.Contribution{tc.point: {interceptor("p1", tc.point, 0)}}, |
| 397 | nil, map[string]*fakeClient{"p1": fake}, map[string]bool{"p1": true}, warns) |
| 398 | payload := tc.sample() |
| 399 | _, err := d.Intercept(context.Background(), tc.point, payload) |
| 400 | var violation *ViolationError |
| 401 | if !errors.As(err, &violation) { |
| 402 | t.Fatalf("err = %v (%T), want *ViolationError", err, err) |
| 403 | } |
| 404 | if violation.Plugin != "p1" || violation.Point != tc.point { |
| 405 | t.Fatalf("violation = %+v, want p1 at %s", violation, tc.point) |
| 406 | } |
| 407 | }) |
| 408 | } |
| 409 | } |
| 410 | } |
| 411 | |
| 412 | // TestInterceptMatrixAllowDenyRejected verifies the 16 non-permission points: |
| 413 | // allow/deny rulings there are a protocol violation. |
| 414 | func TestInterceptMatrixAllowDenyRejected(t *testing.T) { |
| 415 | for _, tc := range pointCases() { |
| 416 | if tc.point == extension.PointPermissionDecision { |
| 417 | continue |
| 418 | } |
| 419 | for _, decision := range []protocol.InterceptDecision{protocol.DecisionAllow, protocol.DecisionDeny} { |
| 420 | t.Run(string(tc.point)+"/"+string(decision)+"_optional", func(t *testing.T) { |
| 421 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 422 | return protocol.InterceptResult{Decision: decision}, nil |
| 423 | }} |
| 424 | warns := &warnRecorder{} |
| 425 | d := buildDispatcher( |
| 426 | map[extension.InterceptorPoint][]extension.Contribution{tc.point: {interceptor("p1", tc.point, 0)}}, |
| 427 | nil, map[string]*fakeClient{"p1": fake}, nil, warns) |
| 428 | payload := tc.sample() |
| 429 | result, err := d.Intercept(context.Background(), tc.point, payload) |
| 430 | if err != nil { |
| 431 | t.Fatalf("Intercept: optional violation must not fail, got %v", err) |
| 432 | } |
| 433 | if result.Permission != nil { |
| 434 | t.Fatalf("Permission = %v, want nil outside permission.decision", *result.Permission) |
| 435 | } |
| 436 | if warns.count() != 1 || !warns.contains("only legal") { |
| 437 | t.Fatalf("warns = %v, want one warning about the illegal decision", warns.msgs) |
| 438 | } |
| 439 | }) |
| 440 | t.Run(string(tc.point)+"/"+string(decision)+"_required", func(t *testing.T) { |
| 441 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 442 | return protocol.InterceptResult{Decision: decision}, nil |
| 443 | }} |
| 444 | warns := &warnRecorder{} |
| 445 | d := buildDispatcher( |
| 446 | map[extension.InterceptorPoint][]extension.Contribution{tc.point: {interceptor("p1", tc.point, 0)}}, |
| 447 | nil, map[string]*fakeClient{"p1": fake}, map[string]bool{"p1": true}, warns) |
| 448 | payload := tc.sample() |
| 449 | _, err := d.Intercept(context.Background(), tc.point, payload) |
| 450 | var violation *ViolationError |
| 451 | if !errors.As(err, &violation) { |
| 452 | t.Fatalf("err = %v (%T), want *ViolationError", err, err) |
| 453 | } |
| 454 | if !strings.Contains(violation.Detail, "only legal") { |
| 455 | t.Fatalf("violation detail = %q, want the legality explanation", violation.Detail) |
| 456 | } |
| 457 | }) |
| 458 | } |
| 459 | } |
| 460 | } |
| 461 | |
| 462 | // TestInterceptChainOrder verifies three extensions observe replaced payloads |
| 463 | // in exact chain order (priority ascending dominates plugin ID). |
| 464 | func TestInterceptChainOrder(t *testing.T) { |
| 465 | point := extension.PointInputReceive |
| 466 | // Deliberately unordered, with priority order opposite to plugin-ID order. |
| 467 | contribs := extension.SortInterceptors([]extension.Contribution{ |
| 468 | interceptor("zeta", point, 5), |
| 469 | interceptor("alpha", point, -3), |
| 470 | interceptor("mid", point, 0), |
| 471 | }) |
| 472 | appendSelf := func(pluginID string) func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 473 | return func(_ protocol.InterceptEvent, raw json.RawMessage) (protocol.InterceptResult, error) { |
| 474 | var payload InputPayload |
| 475 | if err := json.Unmarshal(raw, &payload); err != nil { |
| 476 | return protocol.InterceptResult{}, err |
| 477 | } |
| 478 | replacement, _ := json.Marshal(InputPayload{Text: payload.Text + ">" + pluginID}) |
| 479 | return protocol.InterceptResult{Decision: protocol.DecisionReplace, Replacement: replacement}, nil |
| 480 | } |
| 481 | } |
| 482 | fakes := map[string]*fakeClient{ |
| 483 | "alpha": {interceptFn: appendSelf("alpha")}, |
| 484 | "mid": {interceptFn: appendSelf("mid")}, |
| 485 | "zeta": {interceptFn: appendSelf("zeta")}, |
| 486 | } |
| 487 | warns := &warnRecorder{} |
| 488 | d := buildDispatcher(map[extension.InterceptorPoint][]extension.Contribution{point: contribs}, nil, fakes, nil, warns) |
| 489 | |
| 490 | payload := &InputPayload{Text: "start"} |
| 491 | result, err := d.Intercept(context.Background(), point, payload) |
| 492 | if err != nil { |
| 493 | t.Fatalf("Intercept: %v", err) |
| 494 | } |
| 495 | if want := "start>alpha>mid>zeta"; payload.Text != want { |
| 496 | t.Fatalf("Text = %q, want %q", payload.Text, want) |
| 497 | } |
| 498 | if want := []string{"alpha", "mid", "zeta"}; !reflect.DeepEqual(result.Applied, want) { |
| 499 | t.Fatalf("Applied = %v, want %v", result.Applied, want) |
| 500 | } |
| 501 | // Each extension observed exactly the value its predecessor produced. |
| 502 | wantSeen := map[string]string{"alpha": "start", "mid": "start>alpha", "zeta": "start>alpha>mid"} |
| 503 | for pluginID, want := range wantSeen { |
| 504 | observed := fakes[pluginID].observedPayloads() |
| 505 | if len(observed) != 1 { |
| 506 | t.Fatalf("%s observed %d payloads, want 1", pluginID, len(observed)) |
| 507 | } |
| 508 | var seen InputPayload |
| 509 | if err := json.Unmarshal(observed[0], &seen); err != nil { |
| 510 | t.Fatalf("%s observed payload: %v", pluginID, err) |
| 511 | } |
| 512 | if seen.Text != want { |
| 513 | t.Fatalf("%s observed %q, want %q", pluginID, seen.Text, want) |
| 514 | } |
| 515 | } |
| 516 | } |
| 517 | |
| 518 | func TestPermissionAllowOverridesHostDeny(t *testing.T) { |
| 519 | point := extension.PointPermissionDecision |
| 520 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 521 | return protocol.InterceptResult{Decision: protocol.DecisionAllow}, nil |
| 522 | }} |
| 523 | warns := &warnRecorder{} |
| 524 | d := buildDispatcher( |
| 525 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("ext-sec", point, 0)}}, |
| 526 | nil, map[string]*fakeClient{"ext-sec": fake}, nil, warns) |
| 527 | payload := &PermissionPayload{Name: "bash", Arguments: `{"cmd":"rm -rf x"}`, HostDecision: "deny"} |
| 528 | result, err := d.Intercept(context.Background(), point, payload) |
| 529 | if err != nil { |
| 530 | t.Fatalf("Intercept: %v", err) |
| 531 | } |
| 532 | if result.Permission == nil || !*result.Permission { |
| 533 | t.Fatalf("Permission = %v, want allow", result.Permission) |
| 534 | } |
| 535 | if len(result.Audit) != 1 || !strings.Contains(result.Audit[0], "ext-sec") || !strings.Contains(result.Audit[0], "host deny") { |
| 536 | t.Fatalf("Audit = %v, want one override note naming ext-sec", result.Audit) |
| 537 | } |
| 538 | } |
| 539 | |
| 540 | func TestPermissionDeny(t *testing.T) { |
| 541 | point := extension.PointPermissionDecision |
| 542 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 543 | return protocol.InterceptResult{Decision: protocol.DecisionDeny}, nil |
| 544 | }} |
| 545 | warns := &warnRecorder{} |
| 546 | d := buildDispatcher( |
| 547 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("ext-sec", point, 0)}}, |
| 548 | nil, map[string]*fakeClient{"ext-sec": fake}, nil, warns) |
| 549 | payload := &PermissionPayload{Name: "bash", HostDecision: "allow"} |
| 550 | result, err := d.Intercept(context.Background(), point, payload) |
| 551 | if err != nil { |
| 552 | t.Fatalf("Intercept: %v", err) |
| 553 | } |
| 554 | if result.Permission == nil || *result.Permission { |
| 555 | t.Fatalf("Permission = %v, want deny", result.Permission) |
| 556 | } |
| 557 | if len(result.Audit) != 0 { |
| 558 | t.Fatalf("Audit = %v, want none for a deny", result.Audit) |
| 559 | } |
| 560 | } |
| 561 | |
| 562 | func TestPermissionContinueLeavesHostDecision(t *testing.T) { |
| 563 | point := extension.PointPermissionDecision |
| 564 | fake := &fakeClient{} |
| 565 | warns := &warnRecorder{} |
| 566 | d := buildDispatcher( |
| 567 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("ext-sec", point, 0)}}, |
| 568 | nil, map[string]*fakeClient{"ext-sec": fake}, nil, warns) |
| 569 | payload := &PermissionPayload{Name: "bash", HostDecision: "deny"} |
| 570 | result, err := d.Intercept(context.Background(), point, payload) |
| 571 | if err != nil { |
| 572 | t.Fatalf("Intercept: %v", err) |
| 573 | } |
| 574 | if result.Permission != nil { |
| 575 | t.Fatalf("Permission = %v, want nil (host decision stands)", *result.Permission) |
| 576 | } |
| 577 | } |
| 578 | |
| 579 | // TestPermissionFirstRulingTerminal verifies the first allow/deny ends the |
| 580 | // extension phase: later interceptors are never called. |
| 581 | func TestPermissionFirstRulingTerminal(t *testing.T) { |
| 582 | point := extension.PointPermissionDecision |
| 583 | first := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 584 | return protocol.InterceptResult{Decision: protocol.DecisionAllow}, nil |
| 585 | }} |
| 586 | second := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 587 | return protocol.InterceptResult{Decision: protocol.DecisionDeny}, nil |
| 588 | }} |
| 589 | warns := &warnRecorder{} |
| 590 | d := buildDispatcher( |
| 591 | map[extension.InterceptorPoint][]extension.Contribution{point: { |
| 592 | interceptor("aaa-first", point, 0), interceptor("zzz-second", point, 1), |
| 593 | }}, |
| 594 | nil, map[string]*fakeClient{"aaa-first": first, "zzz-second": second}, nil, warns) |
| 595 | payload := &PermissionPayload{Name: "bash", HostDecision: "deny"} |
| 596 | result, err := d.Intercept(context.Background(), point, payload) |
| 597 | if err != nil { |
| 598 | t.Fatalf("Intercept: %v", err) |
| 599 | } |
| 600 | if result.Permission == nil || !*result.Permission { |
| 601 | t.Fatalf("Permission = %v, want the first ruling (allow)", result.Permission) |
| 602 | } |
| 603 | if second.interceptCount() != 0 { |
| 604 | t.Fatalf("second interceptor called %d times after a terminal ruling", second.interceptCount()) |
| 605 | } |
| 606 | } |
| 607 | |
| 608 | // TestPermissionBlock verifies block remains legal at permission.decision and |
| 609 | // reports a redacted reason. |
| 610 | func TestPermissionBlock(t *testing.T) { |
| 611 | point := extension.PointPermissionDecision |
| 612 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 613 | return protocol.InterceptResult{Decision: protocol.DecisionBlock, Reason: "suspicious, token " + testSecret}, nil |
| 614 | }} |
| 615 | warns := &warnRecorder{} |
| 616 | d := buildDispatcher( |
| 617 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("ext-sec", point, 0)}}, |
| 618 | nil, map[string]*fakeClient{"ext-sec": fake}, nil, warns) |
| 619 | payload := &PermissionPayload{Name: "bash", HostDecision: "allow"} |
| 620 | result, err := d.Intercept(context.Background(), point, payload) |
| 621 | if err != nil { |
| 622 | t.Fatalf("Intercept: %v", err) |
| 623 | } |
| 624 | if !result.Blocked || result.Permission != nil { |
| 625 | t.Fatalf("result = %+v, want blocked with no permission ruling", result) |
| 626 | } |
| 627 | if strings.Contains(result.BlockReason, testSecret) { |
| 628 | t.Fatalf("BlockReason %q leaks the credential", result.BlockReason) |
| 629 | } |
| 630 | } |
| 631 | |
| 632 | func TestStrategyOwnerReplacesSystemPrompt(t *testing.T) { |
| 633 | fake := &fakeClient{interceptFn: func(event protocol.InterceptEvent, _ json.RawMessage) (protocol.InterceptResult, error) { |
| 634 | if event != protocol.EventSystemPromptBuild { |
| 635 | t.Errorf("strategy event = %q, want %q", event, protocol.EventSystemPromptBuild) |
| 636 | } |
| 637 | return protocol.InterceptResult{Decision: protocol.DecisionReplace, Replacement: json.RawMessage(`{"prompt":"owned","workspaceRoot":"/ws"}`)}, nil |
| 638 | }} |
| 639 | warns := &warnRecorder{} |
| 640 | d := buildDispatcher(nil, |
| 641 | map[extension.Slot]extension.ContributionSource{extension.SlotSystemPrompt: owner("prompt-owner")}, |
| 642 | map[string]*fakeClient{"prompt-owner": fake}, nil, warns) |
| 643 | payload := &SystemPromptPayload{Prompt: "host default", WorkspaceRoot: "/ws"} |
| 644 | if err := d.RunStrategy(context.Background(), extension.SlotSystemPrompt, extension.PointSystemPromptBuild, payload); err != nil { |
| 645 | t.Fatalf("RunStrategy: %v", err) |
| 646 | } |
| 647 | if payload.Prompt != "owned" { |
| 648 | t.Fatalf("Prompt = %q, want the owner's replacement", payload.Prompt) |
| 649 | } |
| 650 | } |
| 651 | |
| 652 | // TestStrategyOwnerTimeoutIsFatal verifies a strategy owner's timeout always |
| 653 | // fails the operation (slot owners are required-class even without |
| 654 | // required:true). |
| 655 | func TestStrategyOwnerTimeoutIsFatal(t *testing.T) { |
| 656 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 657 | return protocol.InterceptResult{}, timeoutError("prompt-owner", extension.PointSystemPromptBuild) |
| 658 | }} |
| 659 | warns := &warnRecorder{} |
| 660 | d := buildDispatcher(nil, |
| 661 | map[extension.Slot]extension.ContributionSource{extension.SlotSystemPrompt: owner("prompt-owner")}, |
| 662 | map[string]*fakeClient{"prompt-owner": fake}, nil, warns) |
| 663 | payload := &SystemPromptPayload{Prompt: "host default", WorkspaceRoot: "/ws"} |
| 664 | err := d.RunStrategy(context.Background(), extension.SlotSystemPrompt, extension.PointSystemPromptBuild, payload) |
| 665 | var failure *FailureError |
| 666 | if !errors.As(err, &failure) { |
| 667 | t.Fatalf("err = %v (%T), want *FailureError", err, err) |
| 668 | } |
| 669 | var protocolErr *protocol.ProtocolError |
| 670 | if !errors.As(err, &protocolErr) || protocolErr.Reason != protocol.ErrInterceptTimeout { |
| 671 | t.Fatalf("err = %v, want the wrapped intercept_timeout protocol error", err) |
| 672 | } |
| 673 | if payload.Prompt != "host default" { |
| 674 | t.Fatalf("Prompt = %q, want the host default untouched on failure", payload.Prompt) |
| 675 | } |
| 676 | } |
| 677 | |
| 678 | // TestStrategyNoOwnerKeepsHostDefault verifies an unowned slot is a no-op. |
| 679 | func TestStrategyNoOwnerKeepsHostDefault(t *testing.T) { |
| 680 | warns := &warnRecorder{} |
| 681 | d := buildDispatcher(nil, nil, nil, nil, warns) |
| 682 | if _, ok := d.Strategy(extension.SlotSystemPrompt); ok { |
| 683 | t.Fatalf("Strategy reported an owner for an unowned slot") |
| 684 | } |
| 685 | payload := &SystemPromptPayload{Prompt: "host default", WorkspaceRoot: "/ws"} |
| 686 | if err := d.RunStrategy(context.Background(), extension.SlotSystemPrompt, extension.PointSystemPromptBuild, payload); err != nil { |
| 687 | t.Fatalf("RunStrategy: %v", err) |
| 688 | } |
| 689 | if payload.Prompt != "host default" { |
| 690 | t.Fatalf("Prompt = %q, want the host default", payload.Prompt) |
| 691 | } |
| 692 | } |
| 693 | |
| 694 | // TestStrategyNonOwnerCannotClaim verifies chain membership at a point does |
| 695 | // not make an extension the strategy owner: only the Replacements owner gets |
| 696 | // the strategy call. |
| 697 | func TestStrategyNonOwnerCannotClaim(t *testing.T) { |
| 698 | point := extension.PointSystemPromptBuild |
| 699 | observer := &fakeClient{} |
| 700 | owned := &fakeClient{} |
| 701 | warns := &warnRecorder{} |
| 702 | d := buildDispatcher( |
| 703 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("observer", point, 0)}}, |
| 704 | map[extension.Slot]extension.ContributionSource{extension.SlotSystemPrompt: owner("prompt-owner")}, |
| 705 | map[string]*fakeClient{"observer": observer, "prompt-owner": owned}, nil, warns) |
| 706 | client, ok := d.Strategy(extension.SlotSystemPrompt) |
| 707 | if !ok { |
| 708 | t.Fatalf("Strategy reported no owner") |
| 709 | } |
| 710 | if client != owned { |
| 711 | t.Fatalf("Strategy returned the wrong client: the chain observer must not claim the slot") |
| 712 | } |
| 713 | payload := &SystemPromptPayload{Prompt: "host default", WorkspaceRoot: "/ws"} |
| 714 | if err := d.RunStrategy(context.Background(), extension.SlotSystemPrompt, point, payload); err != nil { |
| 715 | t.Fatalf("RunStrategy: %v", err) |
| 716 | } |
| 717 | if observer.interceptCount() != 0 { |
| 718 | t.Fatalf("non-owner received %d strategy calls", observer.interceptCount()) |
| 719 | } |
| 720 | if owned.interceptCount() != 1 { |
| 721 | t.Fatalf("owner received %d strategy calls, want 1", owned.interceptCount()) |
| 722 | } |
| 723 | } |
| 724 | |
| 725 | // TestStrategyRulingPolicy verifies strategy owners may only continue or |
| 726 | // replace; block is fatal with a redacted reason, allow/deny and invalid |
| 727 | // replacements are fatal contract violations. |
| 728 | func TestStrategyRulingPolicy(t *testing.T) { |
| 729 | point := extension.PointSystemPromptBuild |
| 730 | newDispatcher := func(answer protocol.InterceptResult) (*Dispatcher, *SystemPromptPayload) { |
| 731 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 732 | return answer, nil |
| 733 | }} |
| 734 | warns := &warnRecorder{} |
| 735 | d := buildDispatcher(nil, |
| 736 | map[extension.Slot]extension.ContributionSource{extension.SlotSystemPrompt: owner("prompt-owner")}, |
| 737 | map[string]*fakeClient{"prompt-owner": fake}, nil, warns) |
| 738 | return d, &SystemPromptPayload{Prompt: "host default", WorkspaceRoot: "/ws"} |
| 739 | } |
| 740 | |
| 741 | t.Run("continue_keeps_default", func(t *testing.T) { |
| 742 | d, payload := newDispatcher(protocol.InterceptResult{Decision: protocol.DecisionContinue}) |
| 743 | if err := d.RunStrategy(context.Background(), extension.SlotSystemPrompt, point, payload); err != nil { |
| 744 | t.Fatalf("RunStrategy: %v", err) |
| 745 | } |
| 746 | if payload.Prompt != "host default" { |
| 747 | t.Fatalf("Prompt = %q, want the host default", payload.Prompt) |
| 748 | } |
| 749 | }) |
| 750 | t.Run("block_is_fatal_and_redacted", func(t *testing.T) { |
| 751 | d, payload := newDispatcher(protocol.InterceptResult{Decision: protocol.DecisionBlock, Reason: "no, token " + testSecret}) |
| 752 | err := d.RunStrategy(context.Background(), extension.SlotSystemPrompt, point, payload) |
| 753 | var blocked *BlockError |
| 754 | if !errors.As(err, &blocked) { |
| 755 | t.Fatalf("err = %v (%T), want *BlockError", err, err) |
| 756 | } |
| 757 | if strings.Contains(err.Error(), testSecret) { |
| 758 | t.Fatalf("block error %q leaks the credential", err) |
| 759 | } |
| 760 | }) |
| 761 | t.Run("allow_is_violation", func(t *testing.T) { |
| 762 | d, payload := newDispatcher(protocol.InterceptResult{Decision: protocol.DecisionAllow}) |
| 763 | err := d.RunStrategy(context.Background(), extension.SlotSystemPrompt, point, payload) |
| 764 | var violation *ViolationError |
| 765 | if !errors.As(err, &violation) { |
| 766 | t.Fatalf("err = %v (%T), want *ViolationError", err, err) |
| 767 | } |
| 768 | }) |
| 769 | t.Run("invalid_replace_is_violation", func(t *testing.T) { |
| 770 | d, payload := newDispatcher(protocol.InterceptResult{Decision: protocol.DecisionReplace, Replacement: json.RawMessage(`{"bogus":1}`)}) |
| 771 | err := d.RunStrategy(context.Background(), extension.SlotSystemPrompt, point, payload) |
| 772 | var violation *ViolationError |
| 773 | if !errors.As(err, &violation) { |
| 774 | t.Fatalf("err = %v (%T), want *ViolationError", err, err) |
| 775 | } |
| 776 | if payload.Prompt != "host default" { |
| 777 | t.Fatalf("Prompt = %q, want the host default untouched", payload.Prompt) |
| 778 | } |
| 779 | }) |
| 780 | } |
| 781 | |
| 782 | // TestOptionalTimeoutWarnsOnce verifies an optional extension's timeout is |
| 783 | // warned about exactly once per process and skipped. |
| 784 | func TestOptionalTimeoutWarnsOnce(t *testing.T) { |
| 785 | point := extension.PointToolBefore |
| 786 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 787 | return protocol.InterceptResult{}, timeoutError("opt", point) |
| 788 | }} |
| 789 | warns := &warnRecorder{} |
| 790 | d := buildDispatcher( |
| 791 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("opt", point, 0)}}, |
| 792 | nil, map[string]*fakeClient{"opt": fake}, nil, warns) |
| 793 | for i := range 2 { |
| 794 | payload := &ToolBeforePayload{Name: "bash", Arguments: `{"cmd":"ls"}`} |
| 795 | result, err := d.Intercept(context.Background(), point, payload) |
| 796 | if err != nil { |
| 797 | t.Fatalf("call %d: optional timeout must not fail, got %v", i, err) |
| 798 | } |
| 799 | if result.Blocked || len(result.Applied) != 0 { |
| 800 | t.Fatalf("call %d: result = %+v, want the extension skipped", i, result) |
| 801 | } |
| 802 | if payload.Name != "bash" { |
| 803 | t.Fatalf("call %d: payload changed to %+v", i, payload) |
| 804 | } |
| 805 | } |
| 806 | if warns.count() != 1 { |
| 807 | t.Fatalf("warns = %v, want exactly one warning across two timeouts", warns.msgs) |
| 808 | } |
| 809 | if !warns.contains("opt") || !warns.contains("skipping") { |
| 810 | t.Fatalf("warn %v must name the plugin and the skip", warns.msgs) |
| 811 | } |
| 812 | } |
| 813 | |
| 814 | // TestRequiredTimeoutFails verifies a required extension's timeout fails the |
| 815 | // operation and preserves the frozen protocol error for errors.As. |
| 816 | func TestRequiredTimeoutFails(t *testing.T) { |
| 817 | point := extension.PointToolBefore |
| 818 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 819 | return protocol.InterceptResult{}, timeoutError("req", point) |
| 820 | }} |
| 821 | warns := &warnRecorder{} |
| 822 | d := buildDispatcher( |
| 823 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("req", point, 0)}}, |
| 824 | nil, map[string]*fakeClient{"req": fake}, map[string]bool{"req": true}, warns) |
| 825 | payload := &ToolBeforePayload{Name: "bash"} |
| 826 | _, err := d.Intercept(context.Background(), point, payload) |
| 827 | var failure *FailureError |
| 828 | if !errors.As(err, &failure) { |
| 829 | t.Fatalf("err = %v (%T), want *FailureError", err, err) |
| 830 | } |
| 831 | var protocolErr *protocol.ProtocolError |
| 832 | if !errors.As(err, &protocolErr) || protocolErr.Reason != protocol.ErrInterceptTimeout { |
| 833 | t.Fatalf("err = %v, want the wrapped intercept_timeout protocol error", err) |
| 834 | } |
| 835 | } |
| 836 | |
| 837 | // TestSlotOwnerTimeoutFails verifies slot ownership alone (no required:true) |
| 838 | // upgrades an extension to required-class error policy. |
| 839 | func TestSlotOwnerTimeoutFails(t *testing.T) { |
| 840 | point := extension.PointInputReceive |
| 841 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 842 | return protocol.InterceptResult{}, timeoutError("ctx-owner", point) |
| 843 | }} |
| 844 | warns := &warnRecorder{} |
| 845 | d := buildDispatcher( |
| 846 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("ctx-owner", point, 0)}}, |
| 847 | map[extension.Slot]extension.ContributionSource{extension.SlotContext: owner("ctx-owner")}, |
| 848 | map[string]*fakeClient{"ctx-owner": fake}, nil, warns) |
| 849 | payload := &InputPayload{Text: "hi"} |
| 850 | if _, err := d.Intercept(context.Background(), point, payload); err == nil { |
| 851 | t.Fatalf("slot owner's timeout must fail the operation") |
| 852 | } |
| 853 | if warns.count() != 0 { |
| 854 | t.Fatalf("warns = %v, want none for a required-class failure", warns.msgs) |
| 855 | } |
| 856 | } |
| 857 | |
| 858 | // TestMissingClientPolicy verifies a chain member with no live sidecar client |
| 859 | // follows the same optional/required policy. |
| 860 | func TestMissingClientPolicy(t *testing.T) { |
| 861 | point := extension.PointInputReceive |
| 862 | t.Run("optional", func(t *testing.T) { |
| 863 | warns := &warnRecorder{} |
| 864 | d := buildDispatcher( |
| 865 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("gone", point, 0)}}, |
| 866 | nil, nil, nil, warns) |
| 867 | payload := &InputPayload{Text: "hi"} |
| 868 | if _, err := d.Intercept(context.Background(), point, payload); err != nil { |
| 869 | t.Fatalf("optional missing client must not fail, got %v", err) |
| 870 | } |
| 871 | if warns.count() != 1 { |
| 872 | t.Fatalf("warns = %v, want one warning", warns.msgs) |
| 873 | } |
| 874 | }) |
| 875 | t.Run("required", func(t *testing.T) { |
| 876 | warns := &warnRecorder{} |
| 877 | d := buildDispatcher( |
| 878 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("gone", point, 0)}}, |
| 879 | nil, nil, map[string]bool{"gone": true}, warns) |
| 880 | payload := &InputPayload{Text: "hi"} |
| 881 | var failure *FailureError |
| 882 | if _, err := d.Intercept(context.Background(), point, payload); !errors.As(err, &failure) { |
| 883 | t.Fatalf("err = %v, want *FailureError", err) |
| 884 | } |
| 885 | }) |
| 886 | } |
| 887 | |
| 888 | // TestSessionPhaseMustMatchPoint verifies a session replacement whose phase |
| 889 | // disagrees with the dispatched point is a contract violation. |
| 890 | func TestSessionPhaseMustMatchPoint(t *testing.T) { |
| 891 | point := extension.PointSessionStart |
| 892 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 893 | return protocol.InterceptResult{Decision: protocol.DecisionReplace, Replacement: json.RawMessage(`{"sessionPath":"/x","phase":"end"}`)}, nil |
| 894 | }} |
| 895 | t.Run("optional", func(t *testing.T) { |
| 896 | warns := &warnRecorder{} |
| 897 | d := buildDispatcher( |
| 898 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("p1", point, 0)}}, |
| 899 | nil, map[string]*fakeClient{"p1": fake}, nil, warns) |
| 900 | payload := &SessionPayload{SessionPath: "/tmp/s.json", Phase: PhaseStart} |
| 901 | if _, err := d.Intercept(context.Background(), point, payload); err != nil { |
| 902 | t.Fatalf("optional violation must not fail, got %v", err) |
| 903 | } |
| 904 | if payload.SessionPath != "/tmp/s.json" { |
| 905 | t.Fatalf("payload = %+v, want unchanged", payload) |
| 906 | } |
| 907 | if !warns.contains("does not match") { |
| 908 | t.Fatalf("warns = %v, want the phase-mismatch explanation", warns.msgs) |
| 909 | } |
| 910 | }) |
| 911 | t.Run("required", func(t *testing.T) { |
| 912 | warns := &warnRecorder{} |
| 913 | d := buildDispatcher( |
| 914 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("p1", point, 0)}}, |
| 915 | nil, map[string]*fakeClient{"p1": fake}, map[string]bool{"p1": true}, warns) |
| 916 | payload := &SessionPayload{SessionPath: "/tmp/s.json", Phase: PhaseStart} |
| 917 | var violation *ViolationError |
| 918 | if _, err := d.Intercept(context.Background(), point, payload); !errors.As(err, &violation) { |
| 919 | t.Fatalf("err = %v, want *ViolationError", err) |
| 920 | } |
| 921 | }) |
| 922 | } |
| 923 | |
| 924 | // TestEventNotifiesChainAndSlotObservers verifies fire-and-forget delivery to |
| 925 | // chain members and slot observers (deduplicated), best-effort on error. |
| 926 | func TestEventNotifiesChainAndSlotObservers(t *testing.T) { |
| 927 | point := extension.PointSystemPromptBuild |
| 928 | p1 := &fakeClient{} |
| 929 | p2 := &fakeClient{notifyFn: func(protocol.InterceptEvent, json.RawMessage) error { |
| 930 | return errors.New("notify blew up, token " + testSecret) |
| 931 | }} |
| 932 | p3 := &fakeClient{} |
| 933 | warns := &warnRecorder{} |
| 934 | d := buildDispatcher( |
| 935 | map[extension.InterceptorPoint][]extension.Contribution{point: { |
| 936 | interceptor("p1", point, 0), interceptor("p2", point, 1), interceptor("p3", point, 2), |
| 937 | }}, |
| 938 | // p3 is both a chain member and the slot owner: it must be notified once. |
| 939 | map[extension.Slot]extension.ContributionSource{extension.SlotSystemPrompt: owner("p3")}, |
| 940 | map[string]*fakeClient{"p1": p1, "p2": p2, "p3": p3}, nil, warns) |
| 941 | d.Event(point, &SystemPromptPayload{Prompt: "p", WorkspaceRoot: "/ws"}) |
| 942 | for pluginID, fake := range map[string]*fakeClient{"p1": p1, "p2": p2, "p3": p3} { |
| 943 | if fake.notifyCount() != 1 { |
| 944 | t.Fatalf("%s notifyCount = %d, want 1", pluginID, fake.notifyCount()) |
| 945 | } |
| 946 | } |
| 947 | if warns.count() != 1 || !warns.contains("p2") { |
| 948 | t.Fatalf("warns = %v, want one warning naming p2", warns.msgs) |
| 949 | } |
| 950 | if warns.contains(testSecret) { |
| 951 | t.Fatalf("warning leaks the credential: %v", warns.msgs) |
| 952 | } |
| 953 | } |
| 954 | |
| 955 | // TestEventMarshalFailureWarns verifies an unmarshalable payload degrades to |
| 956 | // a warning instead of a panic. |
| 957 | func TestEventMarshalFailureWarns(t *testing.T) { |
| 958 | warns := &warnRecorder{} |
| 959 | d := buildDispatcher(nil, nil, nil, nil, warns) |
| 960 | d.Event(extension.PointInputReceive, make(chan int)) |
| 961 | if warns.count() != 1 { |
| 962 | t.Fatalf("warns = %v, want one marshal-failure warning", warns.msgs) |
| 963 | } |
| 964 | } |
| 965 | |
| 966 | // TestConcurrentDispatch hammers one Dispatcher from 32 goroutines; run with |
| 967 | // -race to prove the read-only dispatch path and the warn-once dedup are |
| 968 | // safe. |
| 969 | func TestConcurrentDispatch(t *testing.T) { |
| 970 | inputPoint := extension.PointInputReceive |
| 971 | toolPoint := extension.PointToolBefore |
| 972 | replacer := &fakeClient{interceptFn: func(_ protocol.InterceptEvent, raw json.RawMessage) (protocol.InterceptResult, error) { |
| 973 | var payload InputPayload |
| 974 | if err := json.Unmarshal(raw, &payload); err != nil { |
| 975 | return protocol.InterceptResult{}, err |
| 976 | } |
| 977 | replacement, _ := json.Marshal(InputPayload{Text: payload.Text + ">p2"}) |
| 978 | return protocol.InterceptResult{Decision: protocol.DecisionReplace, Replacement: replacement}, nil |
| 979 | }} |
| 980 | failing := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 981 | return protocol.InterceptResult{}, timeoutError("p3", inputPoint) |
| 982 | }} |
| 983 | fakes := map[string]*fakeClient{"p1": {}, "p2": replacer, "p3": failing} |
| 984 | warns := &warnRecorder{} |
| 985 | d := buildDispatcher(map[extension.InterceptorPoint][]extension.Contribution{ |
| 986 | inputPoint: {interceptor("p1", inputPoint, 0), interceptor("p2", inputPoint, 1), interceptor("p3", inputPoint, 2)}, |
| 987 | toolPoint: {interceptor("p1", toolPoint, 0)}, |
| 988 | }, nil, fakes, nil, warns) |
| 989 | |
| 990 | var wg sync.WaitGroup |
| 991 | errs := make(chan error, 32) |
| 992 | for i := range 32 { |
| 993 | wg.Add(1) |
| 994 | go func(i int) { |
| 995 | defer wg.Done() |
| 996 | if i%2 == 0 { |
| 997 | payload := &InputPayload{Text: fmt.Sprintf("turn-%d", i)} |
| 998 | result, err := d.Intercept(context.Background(), inputPoint, payload) |
| 999 | if err != nil { |
| 1000 | errs <- err |
| 1001 | return |
| 1002 | } |
| 1003 | if want := fmt.Sprintf("turn-%d>p2", i); payload.Text != want { |
| 1004 | errs <- fmt.Errorf("payload = %q, want %q", payload.Text, want) |
| 1005 | } |
| 1006 | if !reflect.DeepEqual(result.Applied, []string{"p2"}) { |
| 1007 | errs <- fmt.Errorf("Applied = %v, want [p2]", result.Applied) |
| 1008 | } |
| 1009 | } else { |
| 1010 | payload := &ToolBeforePayload{Name: "bash", Arguments: `{}`} |
| 1011 | if _, err := d.Intercept(context.Background(), toolPoint, payload); err != nil { |
| 1012 | errs <- err |
| 1013 | } |
| 1014 | } |
| 1015 | d.Event(inputPoint, &InputPayload{Text: "observed"}) |
| 1016 | }(i) |
| 1017 | } |
| 1018 | wg.Wait() |
| 1019 | close(errs) |
| 1020 | for err := range errs { |
| 1021 | t.Fatal(err) |
| 1022 | } |
| 1023 | // p3 timed out from 16 goroutines but warns exactly once. |
| 1024 | if warns.count() != 1 { |
| 1025 | t.Fatalf("warns = %v, want one deduplicated warning", warns.msgs) |
| 1026 | } |
| 1027 | } |
| 1028 | |
| 1029 | // TestFrozenInputs verifies New deep-copies its inputs: mutating the caller's |
| 1030 | // chain, replacements, or required map afterwards cannot change dispatch |
| 1031 | // behavior, and per-turn payloads never touch the frozen chain. |
| 1032 | func TestFrozenInputs(t *testing.T) { |
| 1033 | point := extension.PointInputReceive |
| 1034 | real := &fakeClient{} |
| 1035 | evil := &fakeClient{} |
| 1036 | chain := map[extension.InterceptorPoint][]extension.Contribution{ |
| 1037 | point: {interceptor("real", point, 0)}, |
| 1038 | } |
| 1039 | replacements := map[extension.Slot]extension.ContributionSource{extension.SlotSystemPrompt: owner("real")} |
| 1040 | required := map[string]bool{"real": true} |
| 1041 | warns := &warnRecorder{} |
| 1042 | d := New(chain, replacements, func(pluginID string) Client { |
| 1043 | if pluginID == "evil" { |
| 1044 | return evil |
| 1045 | } |
| 1046 | return real |
| 1047 | }, required, Options{Warn: warns.warn}) |
| 1048 | |
| 1049 | // Mutate every input after construction. |
| 1050 | chain[point][0] = interceptor("evil", point, 0) |
| 1051 | chain[point] = append(chain[point], interceptor("evil", point, 1)) |
| 1052 | replacements[extension.SlotSystemPrompt] = owner("evil") |
| 1053 | delete(required, "real") |
| 1054 | |
| 1055 | payload := &InputPayload{Text: "hi"} |
| 1056 | if _, err := d.Intercept(context.Background(), point, payload); err != nil { |
| 1057 | t.Fatalf("Intercept: %v", err) |
| 1058 | } |
| 1059 | if real.interceptCount() != 1 || evil.interceptCount() != 0 { |
| 1060 | t.Fatalf("intercepts real=%d evil=%d, want 1 and 0", real.interceptCount(), evil.interceptCount()) |
| 1061 | } |
| 1062 | client, ok := d.Strategy(extension.SlotSystemPrompt) |
| 1063 | if !ok || client != real { |
| 1064 | t.Fatalf("Strategy owner changed after the replacements map was mutated") |
| 1065 | } |
| 1066 | |
| 1067 | // The required set is frozen too: "real" still fails rather than warns. |
| 1068 | failing := New( |
| 1069 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("real", point, 0)}}, |
| 1070 | nil, func(string) Client { |
| 1071 | return &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 1072 | return protocol.InterceptResult{}, timeoutError("real", point) |
| 1073 | }} |
| 1074 | }, map[string]bool{"real": true}, Options{Warn: warns.warn}) |
| 1075 | if _, err := failing.Intercept(context.Background(), point, &InputPayload{Text: "hi"}); err == nil { |
| 1076 | t.Fatalf("required-class failure must fail the operation") |
| 1077 | } |
| 1078 | } |
| 1079 | |
| 1080 | // TestPayloadTypeMismatch verifies a host programming error (wrong DTO for |
| 1081 | // the point) fails loudly instead of dispatching garbage. |
| 1082 | func TestPayloadTypeMismatch(t *testing.T) { |
| 1083 | warns := &warnRecorder{} |
| 1084 | d := buildDispatcher(nil, nil, nil, nil, warns) |
| 1085 | if _, err := d.Intercept(context.Background(), extension.PointInputReceive, &ToolBeforePayload{Name: "bash"}); err == nil { |
| 1086 | t.Fatalf("wrong payload type must fail") |
| 1087 | } |
| 1088 | if err := d.RunStrategy(context.Background(), extension.SlotSystemPrompt, extension.PointSystemPromptBuild, &InputPayload{}); err == nil { |
| 1089 | t.Fatalf("wrong strategy payload type must fail") |
| 1090 | } |
| 1091 | } |
| 1092 | |
| 1093 | // TestRedactionInWarnings verifies sidecar error text surfaced through |
| 1094 | // warnings is credential-redacted. |
| 1095 | func TestRedactionInWarnings(t *testing.T) { |
| 1096 | point := extension.PointToolBefore |
| 1097 | fake := &fakeClient{interceptFn: func(protocol.InterceptEvent, json.RawMessage) (protocol.InterceptResult, error) { |
| 1098 | return protocol.InterceptResult{}, errors.New("boom, token " + testSecret) |
| 1099 | }} |
| 1100 | warns := &warnRecorder{} |
| 1101 | d := buildDispatcher( |
| 1102 | map[extension.InterceptorPoint][]extension.Contribution{point: {interceptor("opt", point, 0)}}, |
| 1103 | nil, map[string]*fakeClient{"opt": fake}, nil, warns) |
| 1104 | if _, err := d.Intercept(context.Background(), point, &ToolBeforePayload{Name: "bash"}); err != nil { |
| 1105 | t.Fatalf("Intercept: %v", err) |
| 1106 | } |
| 1107 | if warns.contains(testSecret) { |
| 1108 | t.Fatalf("warning leaks the credential: %v", warns.msgs) |
| 1109 | } |
| 1110 | } |
| 1111 |