<!-- Provide a brief summary of your changes --> ## Motivation and Context The SEP-2624 spec defines interceptor modes as "active" | "audit". This change aligns the Go SDK's interceptor `Mode` value with the SEP spec. ## How Has This Been Tested? - go test ./... from go/sdk — all packages pass (chain, extension, integrations/gomiddleware). - gofmt -l . — clean. - go vet ./... — clean. ## Breaking Changes Yes: - The exported constant interceptors.ModeEnforce is removed and replaced by interceptors.ModeActive. Callers referencing ModeEnforce must rename. - The serialized mode value changes from "enforce" to "active". Any config expecting "enforce" must be updated. ## Types of changes <!-- What types of changes does your code introduce? Put an `x` in all the boxes that apply: --> - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [x] Breaking change (fix or feature that would cause existing functionality to change) - [x] Documentation update ## Checklist <!-- Go over all the following points, and put an `x` in all the boxes that apply. --> - [x] I have read the [MCP Documentation](https://modelcontextprotocol.io) - [x] My code follows the repository's style guidelines - [x] New and existing tests pass locally - [ ] I have added appropriate error handling - [x] I have added or updated documentation as needed ## Additional context <!-- Add any other context, implementation notes, or design decisions --> Signed-off-by: monikaicampbell <mcampbell350@bloomberg.net>
687 lines
22 KiB
Go
687 lines
22 KiB
Go
package extension_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors"
|
|
)
|
|
|
|
func TestServer_ValidatorRejectsBlockedTool(t *testing.T) {
|
|
t.Parallel()
|
|
cs := setup(t, blockToolValidator("echo"))
|
|
|
|
_, err := callEcho(t, cs)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestServer_ValidatorAllowsTool(t *testing.T) {
|
|
t.Parallel()
|
|
cs := setup(t, allowAllValidator("allow-all"))
|
|
|
|
result, err := callEcho(t, cs)
|
|
require.NoError(t, err)
|
|
assert.True(t, strings.HasPrefix(resultText(t, result), "echo:"))
|
|
}
|
|
|
|
func TestServer_MutatorModifiesResponse(t *testing.T) {
|
|
t.Parallel()
|
|
cs := setup(t, prefixMutator("prefix", "[mutated]", interceptors.PhaseResponse, 0, interceptors.ModeActive))
|
|
|
|
result, err := callEcho(t, cs)
|
|
require.NoError(t, err)
|
|
assert.True(t, strings.HasPrefix(resultText(t, result), "[mutated] echo:"))
|
|
}
|
|
|
|
func TestServer_ChainedMutatorsWithAuditMode(t *testing.T) {
|
|
t.Parallel()
|
|
// 5 response mutators in priority order. Mutator 3 (priority 30) is
|
|
// audit-mode: its "[AUDIT]" prefix must NOT appear in the final output.
|
|
// Expected final text (applied innermost-first by ascending priority):
|
|
//
|
|
// [M5] [M4] [M2] [M1] echo: ...
|
|
cs := setup(t,
|
|
prefixMutator("m1", "[M1]", interceptors.PhaseResponse, 10, interceptors.ModeActive),
|
|
prefixMutator("m2", "[M2]", interceptors.PhaseResponse, 20, interceptors.ModeActive),
|
|
prefixMutator("m3-audit", "[AUDIT]", interceptors.PhaseResponse, 30, interceptors.ModeAudit),
|
|
prefixMutator("m4", "[M4]", interceptors.PhaseResponse, 40, interceptors.ModeActive),
|
|
prefixMutator("m5", "[M5]", interceptors.PhaseResponse, 50, interceptors.ModeActive),
|
|
)
|
|
|
|
result, err := callEcho(t, cs)
|
|
require.NoError(t, err)
|
|
|
|
text := resultText(t, result)
|
|
assert.True(t, strings.HasPrefix(text, "[M5] [M4] [M2] [M1] echo:"))
|
|
assert.NotContains(t, text, "[AUDIT]")
|
|
}
|
|
|
|
func TestServer_MutatorFailureAbortsChain(t *testing.T) {
|
|
t.Parallel()
|
|
// 3 response mutators. Mutator 2 (priority 20) returns an error
|
|
// and is fail-closed. The chain aborts; client receives an error.
|
|
cs := setup(t,
|
|
prefixMutator("m1", "[M1]", interceptors.PhaseResponse, 10, interceptors.ModeActive),
|
|
failMutator("m2-fail", interceptors.PhaseResponse, 20),
|
|
prefixMutator("m3", "[M3]", interceptors.PhaseResponse, 30, interceptors.ModeActive),
|
|
)
|
|
|
|
_, err := callEcho(t, cs)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestServer_ValidatorWarnDoesNotBlock(t *testing.T) {
|
|
t.Parallel()
|
|
// A validator returns Valid=false with only SeverityWarn messages.
|
|
// The chain should continue — warn-only failures don't abort.
|
|
warnValidator := &interceptors.Validator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "warn-only",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseRequest,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
|
return &interceptors.ValidationResult{
|
|
Valid: false,
|
|
Messages: []interceptors.ValidationMessage{
|
|
{Message: "something looks off", Severity: interceptors.SeverityWarn},
|
|
},
|
|
}, nil
|
|
},
|
|
}
|
|
|
|
cs := setup(t, warnValidator)
|
|
|
|
result, err := callEcho(t, cs)
|
|
require.NoError(t, err)
|
|
assert.True(t, strings.HasPrefix(resultText(t, result), "echo:"))
|
|
}
|
|
|
|
func TestServer_ValidationFailurePreventsM(t *testing.T) {
|
|
t.Parallel()
|
|
// A validator rejects the request. A mutator is also registered but
|
|
// should never run because the chain aborts before mutators.
|
|
mutatorRan := atomic.Bool{}
|
|
spy := &interceptors.Mutator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "spy-mutator",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseRequest,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.MutationResult, error) {
|
|
mutatorRan.Store(true)
|
|
return &interceptors.MutationResult{Modified: false}, nil
|
|
},
|
|
}
|
|
|
|
cs := setup(t, blockToolValidator("echo"), spy)
|
|
|
|
_, err := callEcho(t, cs)
|
|
assert.Error(t, err)
|
|
assert.False(t, mutatorRan.Load(), "mutator should not run after validation failure")
|
|
}
|
|
|
|
func TestServer_ValidatorTimeoutAbortsChain(t *testing.T) {
|
|
t.Parallel()
|
|
slowValidator := &interceptors.Validator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "slow-validator",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseRequest,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(ctx context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case <-time.After(100 * time.Millisecond):
|
|
return &interceptors.ValidationResult{Valid: true}, nil
|
|
}
|
|
},
|
|
}
|
|
|
|
cs := setup(t, slowValidator)
|
|
|
|
// Per-invocation timeout: context expires before the 100ms handler completes.
|
|
_, err := callEchoWithTimeout(t, cs, 5*time.Millisecond)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "context deadline exceeded")
|
|
}
|
|
|
|
func TestServer_MutatorTimeoutAbortsChain(t *testing.T) {
|
|
t.Parallel()
|
|
// 3 response mutators. Mutator 2 (priority 20) blocks longer than
|
|
// the invocation timeout. The chain aborts; client receives an error.
|
|
slowMutator := &interceptors.Mutator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "m2-slow",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseResponse,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
PriorityHint: interceptors.NewPriority(20),
|
|
},
|
|
Handler: func(ctx context.Context, _ *interceptors.Invocation) (*interceptors.MutationResult, error) {
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case <-time.After(100 * time.Millisecond):
|
|
return &interceptors.MutationResult{Modified: false}, nil
|
|
}
|
|
},
|
|
}
|
|
|
|
cs := setup(t,
|
|
prefixMutator("m1", "[M1]", interceptors.PhaseResponse, 10, interceptors.ModeActive),
|
|
slowMutator,
|
|
prefixMutator("m3", "[M3]", interceptors.PhaseResponse, 30, interceptors.ModeActive),
|
|
)
|
|
|
|
// Per-invocation timeout: context expires before the 100ms handler completes.
|
|
_, err := callEchoWithTimeout(t, cs, 5*time.Millisecond)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "context deadline exceeded")
|
|
}
|
|
|
|
func TestServer_MutatorFailOpenContinuesChain(t *testing.T) {
|
|
t.Parallel()
|
|
// 3 response mutators. Mutator 2 (priority 20) returns an error but
|
|
// has FailOpen set. The chain continues; M1 and M3 both apply.
|
|
failOpenMutator := &interceptors.Mutator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "m2-failopen",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseResponse,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
FailOpen: true,
|
|
PriorityHint: interceptors.NewPriority(20),
|
|
},
|
|
Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.MutationResult, error) {
|
|
return nil, fmt.Errorf("simulated failure")
|
|
},
|
|
}
|
|
|
|
cs := setup(t,
|
|
prefixMutator("m1", "[M1]", interceptors.PhaseResponse, 10, interceptors.ModeActive),
|
|
failOpenMutator,
|
|
prefixMutator("m3", "[M3]", interceptors.PhaseResponse, 30, interceptors.ModeActive),
|
|
)
|
|
|
|
result, err := callEcho(t, cs)
|
|
require.NoError(t, err)
|
|
|
|
text := resultText(t, result)
|
|
assert.Contains(t, text, "[M1]")
|
|
assert.Contains(t, text, "[M3]")
|
|
}
|
|
|
|
func TestServer_ResponseMutatorSeesMutatedParams(t *testing.T) {
|
|
t.Parallel()
|
|
// A request mutator injects a marker into the arguments. A response
|
|
// mutator reads the request payload and copies the marker into the
|
|
// response text.
|
|
//
|
|
// NOTE: In the SEP model, interceptors receive json.RawMessage payloads.
|
|
// Response-phase interceptors only see the response payload, not the
|
|
// request params. This test is simplified to verify request mutation
|
|
// flows through to the echo handler's output.
|
|
reqMutator := &interceptors.Mutator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "inject-marker",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseRequest,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(_ context.Context, inv *interceptors.Invocation) (*interceptors.MutationResult, error) {
|
|
raw, ok := inv.Payload.(json.RawMessage)
|
|
if !ok {
|
|
return nil, fmt.Errorf("unexpected payload type %T", inv.Payload)
|
|
}
|
|
var params map[string]any
|
|
if err := json.Unmarshal(raw, ¶ms); err != nil {
|
|
return nil, err
|
|
}
|
|
args, _ := params["arguments"].(map[string]any)
|
|
if args == nil {
|
|
args = map[string]any{}
|
|
}
|
|
args["marker"] = "injected-by-request-mutator"
|
|
params["arguments"] = args
|
|
data, err := json.Marshal(params)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &interceptors.MutationResult{Modified: true, Payload: json.RawMessage(data)}, nil
|
|
},
|
|
}
|
|
|
|
cs := setup(t, reqMutator)
|
|
|
|
result, err := callEcho(t, cs)
|
|
require.NoError(t, err)
|
|
assert.Contains(t, resultText(t, result), "injected-by-request-mutator")
|
|
}
|
|
|
|
func TestServer_CombinedValidatorsAndMutators(t *testing.T) {
|
|
t.Parallel()
|
|
// Exercises the full chain with 3 validators and 3 mutators on each
|
|
// side, verifying:
|
|
//
|
|
// Request: Validate (parallel) → Mutate (sequential by priority)
|
|
// Response: Mutate (sequential by priority) → Validate (parallel)
|
|
//
|
|
// Request mutators inject fields into arguments (reflected in echo output).
|
|
// Response mutators prepend tags to text (visible in final result).
|
|
// Atomic counters confirm every interceptor actually ran.
|
|
|
|
var (
|
|
reqValCount atomic.Int32
|
|
reqMutCount atomic.Int32
|
|
respMutCount atomic.Int32
|
|
respValCount atomic.Int32
|
|
)
|
|
|
|
// --- Request validators (3, parallel) ---
|
|
|
|
// req-v1: passes if tool name is "echo".
|
|
reqV1 := &interceptors.Validator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "req-v1-tool-check",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseRequest,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(_ context.Context, inv *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
|
reqValCount.Add(1)
|
|
raw, ok := inv.Payload.(json.RawMessage)
|
|
if !ok {
|
|
return nil, fmt.Errorf("unexpected type %T", inv.Payload)
|
|
}
|
|
var params struct {
|
|
Name string `json:"name"`
|
|
}
|
|
if err := json.Unmarshal(raw, ¶ms); err != nil {
|
|
return nil, err
|
|
}
|
|
if params.Name != "echo" {
|
|
return &interceptors.ValidationResult{
|
|
Valid: false,
|
|
Messages: []interceptors.ValidationMessage{
|
|
{Message: "only echo allowed", Severity: interceptors.SeverityError},
|
|
},
|
|
}, nil
|
|
}
|
|
return &interceptors.ValidationResult{Valid: true}, nil
|
|
},
|
|
}
|
|
|
|
// req-v2: passes if arguments contain "text".
|
|
reqV2 := &interceptors.Validator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "req-v2-args-check",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseRequest,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(_ context.Context, inv *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
|
reqValCount.Add(1)
|
|
raw, ok := inv.Payload.(json.RawMessage)
|
|
if !ok {
|
|
return nil, fmt.Errorf("unexpected type %T", inv.Payload)
|
|
}
|
|
var params struct {
|
|
Arguments json.RawMessage `json:"arguments"`
|
|
}
|
|
if err := json.Unmarshal(raw, ¶ms); err != nil {
|
|
return nil, err
|
|
}
|
|
var args map[string]any
|
|
if err := json.Unmarshal(params.Arguments, &args); err != nil {
|
|
return nil, err
|
|
}
|
|
if _, ok := args["text"]; !ok {
|
|
return &interceptors.ValidationResult{
|
|
Valid: false,
|
|
Messages: []interceptors.ValidationMessage{
|
|
{Message: "missing text argument", Severity: interceptors.SeverityError},
|
|
},
|
|
}, nil
|
|
}
|
|
return &interceptors.ValidationResult{Valid: true}, nil
|
|
},
|
|
}
|
|
|
|
// req-v3: always passes with a warning (warn doesn't block).
|
|
reqV3 := &interceptors.Validator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "req-v3-warn",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseRequest,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
|
reqValCount.Add(1)
|
|
return &interceptors.ValidationResult{
|
|
Valid: true,
|
|
Messages: []interceptors.ValidationMessage{
|
|
{Message: "request noted", Severity: interceptors.SeverityWarn},
|
|
},
|
|
}, nil
|
|
},
|
|
}
|
|
|
|
// --- Request mutators (3, sequential by priority 10 → 20 → 30) ---
|
|
// Each adds a field to the arguments.
|
|
|
|
requestArgMutator := func(name, key, value string, priority int) *interceptors.Mutator {
|
|
return &interceptors.Mutator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: name,
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseRequest,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
PriorityHint: interceptors.NewPriority(priority),
|
|
},
|
|
Handler: func(_ context.Context, inv *interceptors.Invocation) (*interceptors.MutationResult, error) {
|
|
reqMutCount.Add(1)
|
|
raw, ok := inv.Payload.(json.RawMessage)
|
|
if !ok {
|
|
return nil, fmt.Errorf("unexpected type %T", inv.Payload)
|
|
}
|
|
var params map[string]any
|
|
if err := json.Unmarshal(raw, ¶ms); err != nil {
|
|
return nil, err
|
|
}
|
|
argsRaw := params["arguments"]
|
|
var args map[string]any
|
|
if argsRaw != nil {
|
|
argsBytes, _ := json.Marshal(argsRaw)
|
|
json.Unmarshal(argsBytes, &args)
|
|
}
|
|
if args == nil {
|
|
args = map[string]any{}
|
|
}
|
|
args[key] = value
|
|
params["arguments"] = args
|
|
data, err := json.Marshal(params)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &interceptors.MutationResult{Modified: true, Payload: json.RawMessage(data)}, nil
|
|
},
|
|
}
|
|
}
|
|
|
|
reqM1 := requestArgMutator("req-m1", "step1", "done", 10)
|
|
reqM2 := requestArgMutator("req-m2", "step2", "done", 20)
|
|
reqM3 := requestArgMutator("req-m3", "step3", "done", 30)
|
|
|
|
// --- Response mutators (3, sequential by priority 10 → 20 → 30) ---
|
|
// Each prepends a tag. Applied innermost-first, so final text is:
|
|
// [RM3] [RM2] [RM1] echo: ...
|
|
|
|
respPrefixMutator := func(name, tag string, priority int) *interceptors.Mutator {
|
|
return &interceptors.Mutator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: name,
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseResponse,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
PriorityHint: interceptors.NewPriority(priority),
|
|
},
|
|
Handler: func(_ context.Context, inv *interceptors.Invocation) (*interceptors.MutationResult, error) {
|
|
respMutCount.Add(1)
|
|
raw, ok := inv.Payload.(json.RawMessage)
|
|
if !ok {
|
|
return nil, fmt.Errorf("unexpected type %T", inv.Payload)
|
|
}
|
|
var result mcp.CallToolResult
|
|
if err := json.Unmarshal(raw, &result); err != nil {
|
|
return nil, err
|
|
}
|
|
for _, c := range result.Content {
|
|
if tc, ok := c.(*mcp.TextContent); ok {
|
|
tc.Text = tag + " " + tc.Text
|
|
}
|
|
}
|
|
data, err := json.Marshal(&result)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &interceptors.MutationResult{Modified: true, Payload: json.RawMessage(data)}, nil
|
|
},
|
|
}
|
|
}
|
|
|
|
respM1 := respPrefixMutator("resp-m1", "[RM1]", 10)
|
|
respM2 := respPrefixMutator("resp-m2", "[RM2]", 20)
|
|
respM3 := respPrefixMutator("resp-m3", "[RM3]", 30)
|
|
|
|
// --- Response validators (3, parallel) ---
|
|
// Run after response mutators.
|
|
|
|
// resp-v1: passes if response has content.
|
|
respV1 := &interceptors.Validator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "resp-v1-has-content",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseResponse,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(_ context.Context, inv *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
|
respValCount.Add(1)
|
|
raw, ok := inv.Payload.(json.RawMessage)
|
|
if !ok {
|
|
return nil, fmt.Errorf("unexpected type %T", inv.Payload)
|
|
}
|
|
var result mcp.CallToolResult
|
|
if err := json.Unmarshal(raw, &result); err != nil {
|
|
return nil, err
|
|
}
|
|
if len(result.Content) == 0 {
|
|
return &interceptors.ValidationResult{
|
|
Valid: false,
|
|
Messages: []interceptors.ValidationMessage{
|
|
{Message: "empty response", Severity: interceptors.SeverityError},
|
|
},
|
|
}, nil
|
|
}
|
|
return &interceptors.ValidationResult{Valid: true}, nil
|
|
},
|
|
}
|
|
|
|
// resp-v2: rejects if response text doesn't contain "[RM1]".
|
|
// This proves the response-side ordering (Mutate → Validate): if
|
|
// this validator ran before mutators, it wouldn't find the tag.
|
|
respV2 := &interceptors.Validator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "resp-v2-sees-mutation",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseResponse,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(_ context.Context, inv *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
|
respValCount.Add(1)
|
|
raw, ok := inv.Payload.(json.RawMessage)
|
|
if !ok {
|
|
return nil, fmt.Errorf("unexpected type %T", inv.Payload)
|
|
}
|
|
var result mcp.CallToolResult
|
|
if err := json.Unmarshal(raw, &result); err != nil {
|
|
return nil, err
|
|
}
|
|
for _, c := range result.Content {
|
|
if tc, ok := c.(*mcp.TextContent); ok {
|
|
if strings.Contains(tc.Text, "[RM1]") {
|
|
return &interceptors.ValidationResult{Valid: true}, nil
|
|
}
|
|
}
|
|
}
|
|
return &interceptors.ValidationResult{
|
|
Valid: false,
|
|
Messages: []interceptors.ValidationMessage{
|
|
{Message: "response missing mutator tag [RM1]", Severity: interceptors.SeverityError},
|
|
},
|
|
}, nil
|
|
},
|
|
}
|
|
|
|
// resp-v3: always passes with an info message.
|
|
respV3 := &interceptors.Validator{
|
|
Metadata: interceptors.Metadata{
|
|
Name: "resp-v3-info",
|
|
Hooks: []interceptors.Hook{{
|
|
Events: []string{interceptors.EventToolsCall},
|
|
Phase: interceptors.PhaseResponse,
|
|
}},
|
|
Mode: interceptors.ModeActive,
|
|
},
|
|
Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
|
respValCount.Add(1)
|
|
return &interceptors.ValidationResult{
|
|
Valid: true,
|
|
Messages: []interceptors.ValidationMessage{
|
|
{Message: "response looks good", Severity: interceptors.SeverityInfo},
|
|
},
|
|
}, nil
|
|
},
|
|
}
|
|
|
|
cs := setup(t,
|
|
reqV1, reqV2, reqV3,
|
|
reqM1, reqM2, reqM3,
|
|
respM1, respM2, respM3,
|
|
respV1, respV2, respV3,
|
|
)
|
|
|
|
result, err := callEcho(t, cs)
|
|
require.NoError(t, err)
|
|
|
|
text := resultText(t, result)
|
|
|
|
// Response mutators applied in priority order (outermost = highest priority).
|
|
assert.True(t, strings.HasPrefix(text, "[RM3] [RM2] [RM1] echo:"),
|
|
"expected response mutator tags in priority order, got: %s", text)
|
|
|
|
// Request mutators injected fields into arguments (visible in echo output).
|
|
assert.Contains(t, text, "step1")
|
|
assert.Contains(t, text, "step2")
|
|
assert.Contains(t, text, "step3")
|
|
|
|
// Every interceptor ran exactly once.
|
|
assert.Equal(t, int32(3), reqValCount.Load(), "expected 3 request validators to run")
|
|
assert.Equal(t, int32(3), reqMutCount.Load(), "expected 3 request mutators to run")
|
|
assert.Equal(t, int32(3), respMutCount.Load(), "expected 3 response mutators to run")
|
|
assert.Equal(t, int32(3), respValCount.Load(), "expected 3 response validators to run")
|
|
}
|
|
|
|
// TestCustomMethodsFlowThroughMiddleware verifies that interceptors/list and
|
|
// interceptor/invoke are routed through AddReceivingMiddleware, just like
|
|
// standard MCP methods.
|
|
func TestCustomMethodsFlowThroughMiddleware(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var count atomic.Int64
|
|
counting := func(next mcp.MethodHandler) mcp.MethodHandler {
|
|
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
|
|
count.Add(1)
|
|
return next(ctx, method, req)
|
|
}
|
|
}
|
|
|
|
mcpServer := buildServer(t, allowAllValidator("v1"))
|
|
mcpServer.AddReceivingMiddleware(counting)
|
|
cs := connectHTTPClient(t, mcpServer)
|
|
count.Store(0) // reset after initialization handshake
|
|
|
|
_, err := cs.CallTool(context.Background(), &mcp.CallToolParams{
|
|
Name: "echo",
|
|
Arguments: map[string]any{"text": "hello"},
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Greater(t, count.Swap(0), int64(0), "middleware must fire for tools/call")
|
|
|
|
listResult, err := mcp.CallCustomMethod[*interceptors.ListParams, *interceptors.ListResult](
|
|
context.Background(), cs, interceptors.MethodList, nil,
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Len(t, listResult.Interceptors, 1)
|
|
assert.Greater(t, count.Swap(0), int64(0), "middleware must fire for interceptors/list")
|
|
|
|
payload, _ := json.Marshal(map[string]any{"name": "echo", "arguments": map[string]any{}})
|
|
_, err = mcp.CallCustomMethod[*interceptors.InvokeParams, *interceptors.InvokeResult](
|
|
context.Background(), cs, interceptors.MethodInvoke,
|
|
&interceptors.InvokeParams{
|
|
Name: "v1",
|
|
Event: interceptors.EventToolsCall,
|
|
Phase: interceptors.PhaseRequest,
|
|
Payload: payload,
|
|
},
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Greater(t, count.Swap(0), int64(0), "middleware must fire for interceptor/invoke")
|
|
}
|
|
|
|
// TestMiddlewareCanRejectCustomMethods verifies that middleware installed via
|
|
// AddReceivingMiddleware can block interceptors/list and interceptor/invoke,
|
|
// enabling auth guards over the interceptor protocol.
|
|
func TestMiddlewareCanRejectCustomMethods(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var rejected atomic.Bool
|
|
blockingMiddleware := func(next mcp.MethodHandler) mcp.MethodHandler {
|
|
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
|
|
if method != "initialize" {
|
|
rejected.Store(true)
|
|
return nil, fmt.Errorf("unauthorized")
|
|
}
|
|
return next(ctx, method, req)
|
|
}
|
|
}
|
|
|
|
mcpServer := buildServer(t, allowAllValidator("v1"))
|
|
mcpServer.AddReceivingMiddleware(blockingMiddleware)
|
|
cs := connectHTTPClient(t, mcpServer)
|
|
|
|
_, err := mcp.CallCustomMethod[*interceptors.ListParams, *interceptors.ListResult](
|
|
context.Background(), cs, interceptors.MethodList, nil,
|
|
)
|
|
assert.Error(t, err, "blockingMiddleware should reject interceptors/list")
|
|
assert.True(t, rejected.Load())
|
|
}
|