modelcontextprotocol/experimental-ext-interceptors
Watch
1
0
Fork
You've already forked experimental-ext-interceptors
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Go SDK: Add ExecutionHandler (#22)

Browse source 
This change adds an ExecutionHandler hook to the Go chain orchestrator,
allowing direct callers to wrap each interceptor/invoke RPC call. The
handler can inspect and modify invoke params before execution, inspect
results after execution, override the effective mode per-invocation, or
short-circuit execution entirely. The change is fully
backwards-compatible and requires no SEP/protocol modifications.

---------

Signed-off-by: Ukjae Jeong <jeongukjae@gmail.com>
This commit is contained in:
Ukjae Jeong 2026-06-17 14:07:21 +01:00 • committed by GitHub
commit 77a8464b33
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 234 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

74
go/sdk/interceptors/chain/chain.go
View file

@ -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.

174
go/sdk/interceptors/chain/chain_test.go
View file

@ -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()