Go SDK: Add ExecutionHandler

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)
+	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
+	}
+
+	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{err: err}
+		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 invokeOutcome{result: result}
+	return entry.Interceptor.Mode
 }
 
 // timeoutResult sets the chain result to timeout status.

go/sdk/interceptors/chain/chain_test.go

@@ -23,14 +23,17 @@ 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{
 		Name:    "chain-test-server",
 		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()