diff --git a/go/sdk/interceptors/chain/chain.go b/go/sdk/interceptors/chain/chain.go index 20081e7..471ad90 100644 --- a/go/sdk/interceptors/chain/chain.go +++ b/go/sdk/interceptors/chain/chain.go @@ -25,6 +25,24 @@ type ChainEntry struct { Server *mcp.ClientSession } +// Directive carries per-invocation decisions from an ExecutionHandler +// back to the chain. nil fields mean "use descriptor default". +type Directive struct { + Mode *interceptors.Mode // nil = use descriptor's static mode +} + +// ExecutionHandler wraps each interceptor/invoke RPC call in the chain. +// It receives the interceptor entry, the invoke params that will be sent, +// and a next function that performs the actual RPC. The handler can modify +// params before calling next, inspect/modify results after, control the +// effective mode via Directive, or short-circuit by not calling next. +type ExecutionHandler func( + ctx context.Context, + entry ChainEntry, + params *interceptors.InvokeParams, + next func(ctx context.Context, params *interceptors.InvokeParams) (interceptors.InvokeResult, error), +) (interceptors.InvokeResult, *Directive, error) + // Chain is the SEP-compliant interceptor chain orchestrator. It holds // ChainEntry objects (interceptor descriptors + MCP server connections), // discovers interceptors via interceptors/list, and invokes them via @@ -33,6 +51,7 @@ type Chain struct { mu sync.Mutex entries []ChainEntry logger *slog.Logger + handler ExecutionHandler } // ChainOption configures a Chain. @@ -46,6 +65,17 @@ func WithChainLogger(l *slog.Logger) ChainOption { } } +// WithExecutionHandler sets the ExecutionHandler for the chain. +// When set, every interceptor/invoke call is routed through the handler, +// which can modify invoke params, inspect results, and override the +// effective mode. If not set, the chain uses the interceptor descriptor's +// static mode. +func WithExecutionHandler(h ExecutionHandler) ChainOption { + return func(c *Chain) { + c.handler = h + } +} + // NewChain creates a new Chain with optional configuration. func NewChain(opts ...ChainOption) *Chain { c := &Chain{} @@ -288,7 +318,7 @@ func (c *Chain) recordValidation(entry ChainEntry, result invokeOutcome, cr *Exe // Tally validation summary and check for abort in a single pass. if result.result.Validation != nil { - shouldAbort := entry.Interceptor.Mode != interceptors.ModeAudit && !result.result.Validation.Valid + shouldAbort := result.mode != interceptors.ModeAudit && !result.result.Validation.Valid aborted := false for _, msg := range result.result.Validation.Messages { switch msg.Severity { @@ -368,7 +398,7 @@ func (c *Chain) runMutators( cr.Results = append(cr.Results, result.result) // For audit-mode mutators, don't apply the mutated payload. - if m.Interceptor.Mode == interceptors.ModeAudit { + if result.mode == interceptors.ModeAudit { continue } @@ -384,14 +414,18 @@ func (c *Chain) runMutators( } } -// invokeOutcome wraps the result of a single interceptor/invoke call. +// invokeOutcome wraps the result of a single interceptor/invoke call +// along with the resolved effective mode. type invokeOutcome struct { result interceptors.InvokeResult + mode interceptors.Mode // effective mode (from handler directive or descriptor) err error } // callInvoke calls interceptor/invoke on the appropriate server for -// a single chain entry. +// a single chain entry. When an ExecutionHandler is set, the RPC is +// routed through the handler which may modify params and override +// the effective mode. func (c *Chain) callInvoke( ctx context.Context, params *ExecutionParams, @@ -404,18 +438,36 @@ func (c *Chain) callInvoke( Payload: params.Payload, Context: params.Context, } - - // Apply per-interceptor config if provided. if cfg, ok := params.Config[entry.Interceptor.Name]; ok { invokeParams.Config = cfg } - var result interceptors.InvokeResult - err := entry.Server.CallCustom(ctx, interceptors.MethodInvoke, invokeParams, &result) - if err != nil { - return invokeOutcome{err: err} + next := func(ctx context.Context, p *interceptors.InvokeParams) (interceptors.InvokeResult, error) { + var result interceptors.InvokeResult + err := entry.Server.CallCustom(ctx, interceptors.MethodInvoke, p, &result) + return result, err } - return invokeOutcome{result: result} + + if c.handler != nil { + result, directive, err := c.handler(ctx, entry, invokeParams, next) + if err != nil { + return invokeOutcome{mode: entry.Interceptor.Mode, err: err} + } + return invokeOutcome{result: result, mode: resolveMode(entry, directive)} + } + + result, err := next(ctx, invokeParams) + if err != nil { + return invokeOutcome{mode: entry.Interceptor.Mode, err: err} + } + return invokeOutcome{result: result, mode: entry.Interceptor.Mode} +} + +func resolveMode(entry ChainEntry, d *Directive) interceptors.Mode { + if d != nil && d.Mode != nil { + return *d.Mode + } + return entry.Interceptor.Mode } // timeoutResult sets the chain result to timeout status. diff --git a/go/sdk/interceptors/chain/chain_test.go b/go/sdk/interceptors/chain/chain_test.go index 25faf6e..f6c8a97 100644 --- a/go/sdk/interceptors/chain/chain_test.go +++ b/go/sdk/interceptors/chain/chain_test.go @@ -23,6 +23,10 @@ import ( // interceptors registered via custom methods, connects a chain via // in-memory transport, and returns the chain ready for testing. func setupChainWithInterceptors(t *testing.T, is ...interceptors.Interceptor) *chain.Chain { + return setupChainWithOpts(t, nil, is...) +} + +func setupChainWithOpts(t *testing.T, opts []chain.ChainOption, is ...interceptors.Interceptor) *chain.Chain { t.Helper() mcpServer := mcp.NewServer(&mcp.Implementation{ @@ -30,7 +34,6 @@ func setupChainWithInterceptors(t *testing.T, is ...interceptors.Interceptor) *c Version: "0.1.0", }, nil) - // Register interceptors/list and interceptor/invoke handlers. registerInterceptorMethods(mcpServer, is) serverTransport, clientTransport := mcp.NewInMemoryTransports() @@ -47,7 +50,8 @@ func setupChainWithInterceptors(t *testing.T, is ...interceptors.Interceptor) *c require.NoError(t, err) t.Cleanup(func() { cs.Close() }) - ch := chain.NewChain(chain.WithChainLogger(slog.Default())) + allOpts := append([]chain.ChainOption{chain.WithChainLogger(slog.Default())}, opts...) + ch := chain.NewChain(allOpts...) err = ch.AddMCPServer(context.Background(), cs) require.NoError(t, err) @@ -136,6 +140,172 @@ func registerInterceptorMethods(server *mcp.Server, is []interceptors.Intercepto ) } +func TestChain_ExecutionHandler(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + interceptor interceptors.Interceptor + directive *chain.Directive + phase interceptors.InterceptionPhase + wantStatus chain.ChainStatus + wantAborted bool + wantNoPayload bool + }{ + { + name: "audit-to-enforce validator override aborts chain", + interceptor: &interceptors.Validator{ + Metadata: interceptors.Metadata{ + Name: "v", + Hooks: []interceptors.Hook{{Events: []string{"test/event"}, Phase: interceptors.PhaseRequest}}, + Mode: interceptors.ModeAudit, + }, + Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) { + return &interceptors.ValidationResult{ + Valid: false, Severity: interceptors.SeverityError, + Messages: []interceptors.ValidationMessage{{Message: "blocked", Severity: interceptors.SeverityError}}, + }, nil + }, + }, + directive: &chain.Directive{Mode: modePtr(interceptors.ModeEnforce)}, + phase: interceptors.PhaseRequest, + wantStatus: chain.ChainValidationFailed, + wantAborted: true, + }, + { + name: "enforce-to-audit validator override does not abort", + interceptor: &interceptors.Validator{ + Metadata: interceptors.Metadata{ + Name: "v", + Hooks: []interceptors.Hook{{Events: []string{"test/event"}, Phase: interceptors.PhaseRequest}}, + Mode: interceptors.ModeEnforce, + }, + Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) { + return &interceptors.ValidationResult{ + Valid: false, Severity: interceptors.SeverityError, + Messages: []interceptors.ValidationMessage{{Message: "would block", Severity: interceptors.SeverityError}}, + }, nil + }, + }, + directive: &chain.Directive{Mode: modePtr(interceptors.ModeAudit)}, + phase: interceptors.PhaseRequest, + wantStatus: chain.ChainSuccess, + }, + { + name: "mutator audit override skips payload application", + interceptor: &interceptors.Mutator{ + Metadata: interceptors.Metadata{ + Name: "m", + Hooks: []interceptors.Hook{{Events: []string{"test/event"}, Phase: interceptors.PhaseResponse}}, + Mode: interceptors.ModeEnforce, + }, + Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.MutationResult, error) { + modified, _ := json.Marshal(map[string]any{"value": "mutated"}) + return &interceptors.MutationResult{Modified: true, Payload: modified}, nil + }, + }, + directive: &chain.Directive{Mode: modePtr(interceptors.ModeAudit)}, + phase: interceptors.PhaseResponse, + wantStatus: chain.ChainSuccess, + wantNoPayload: true, + }, + { + name: "nil directive uses descriptor mode", + interceptor: &interceptors.Validator{ + Metadata: interceptors.Metadata{ + Name: "v", + Hooks: []interceptors.Hook{{Events: []string{"test/event"}, Phase: interceptors.PhaseRequest}}, + Mode: interceptors.ModeEnforce, + }, + Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) { + return &interceptors.ValidationResult{ + Valid: false, Severity: interceptors.SeverityError, + Messages: []interceptors.ValidationMessage{{Message: "blocked", Severity: interceptors.SeverityError}}, + }, nil + }, + }, + directive: nil, + phase: interceptors.PhaseRequest, + wantStatus: chain.ChainValidationFailed, + wantAborted: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + handler := chain.ExecutionHandler( + func(ctx context.Context, entry chain.ChainEntry, params *interceptors.InvokeParams, + next func(ctx context.Context, params *interceptors.InvokeParams) (interceptors.InvokeResult, error), + ) (interceptors.InvokeResult, *chain.Directive, error) { + result, err := next(ctx, params) + return result, tt.directive, err + }, + ) + + ch := setupChainWithOpts(t, []chain.ChainOption{chain.WithExecutionHandler(handler)}, tt.interceptor) + payload, _ := json.Marshal(map[string]any{"value": "test"}) + cr, err := ch.Execute(context.Background(), &chain.ExecutionParams{ + Event: "test/event", + Phase: tt.phase, + Payload: payload, + }) + require.NoError(t, err) + assert.Equal(t, tt.wantStatus, cr.Status) + if tt.wantAborted { + assert.NotEmpty(t, cr.AbortedAt) + } else { + assert.Empty(t, cr.AbortedAt) + } + if tt.wantNoPayload { + assert.Nil(t, cr.FinalPayload) + } + }) + } +} + +func modePtr(m interceptors.Mode) *interceptors.Mode { return &m } + +func TestChain_ExecutionHandler_ShortCircuit(t *testing.T) { + t.Parallel() + invoked := false + v := &interceptors.Validator{ + Metadata: interceptors.Metadata{ + Name: "v", + Hooks: []interceptors.Hook{{Events: []string{"test/event"}, Phase: interceptors.PhaseRequest}}, + Mode: interceptors.ModeEnforce, + }, + Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) { + invoked = true + return &interceptors.ValidationResult{Valid: true}, nil + }, + } + + handler := chain.ExecutionHandler( + func(ctx context.Context, entry chain.ChainEntry, params *interceptors.InvokeParams, + next func(ctx context.Context, params *interceptors.InvokeParams) (interceptors.InvokeResult, error), + ) (interceptors.InvokeResult, *chain.Directive, error) { + return interceptors.InvokeResult{ + Interceptor: entry.Interceptor.Name, + Type: interceptors.TypeValidation, + Phase: interceptors.PhaseRequest, + Validation: &interceptors.ValidationResult{Valid: true}, + }, nil, nil + }, + ) + + ch := setupChainWithOpts(t, []chain.ChainOption{chain.WithExecutionHandler(handler)}, v) + payload, _ := json.Marshal(map[string]any{"value": "test"}) + cr, err := ch.Execute(context.Background(), &chain.ExecutionParams{ + Event: "test/event", + Phase: interceptors.PhaseRequest, + Payload: payload, + }) + require.NoError(t, err) + assert.Equal(t, chain.ChainSuccess, cr.Status) + assert.False(t, invoked, "handler short-circuited, server should not be invoked") +} + func TestChain_FailOpenRecordsExecutionResult(t *testing.T) { t.Parallel()