Rename mcpserver package to extension, remove *mcp.Server field (#11)
## Summary - Rename `interceptors/mcpserver/` → `interceptors/extension/`, replacing the `Server` wrapper with a lightweight `Extension` that can install itself on any `*mcp.Server` - Remove the stored `*mcp.Server` field — `Extension` is now a pure interceptor container with no server reference - Add `Install(server)` method to register JSON-RPC methods and init middleware on a given server - Update `LocalChain(ctx)` → `LocalChain(ctx, server)` to accept the server as a parameter - Remove `MCPServer()`, `Run()`, and `NewStreamableHTTPHandler()` wrapper methods — users interact with `mcp.Server` directly - Update all examples, tests, and documentation to use the new API ## Why `mcpserver.Server` pretended to *be* the server — it had `Run()`, `NewStreamableHTTPHandler()`, etc. — but it could never cover the full `*mcp.Server` API, so users constantly needed `MCPServer()` to unwrap it. `Extension` doesn't pretend to be the server. It's a pure interceptor container that can install itself on any server (or multiple servers) on the fly. No stored server reference, no wrapping, no delegation surface to maintain. ### Before ```go mcpServer := mcp.NewServer(...) srv := mcpserver.NewServer(mcpServer) srv.AddInterceptor(v) chain, err := srv.LocalChain(ctx) ``` ### After ```go mcpServer := mcp.NewServer(...) ext := extension.New() ext.AddInterceptor(v) ext.Install(mcpServer) chain, err := ext.LocalChain(ctx, mcpServer) ``` Signed-off-by: Kurt Degiorgio <kdegiorgio@bloomberg.net>
This commit is contained in:
Kurt Degiorgio
GitHub
parent
c9900f02ab
commit
074dc6cdbd
14 changed files with 172 additions and 173 deletions
|
|
@ -13,11 +13,11 @@ mcpServer := mcp.NewServer(&mcp.Implementation{
|
|||
Version: "0.1.0",
|
||||
}, nil)
|
||||
|
||||
// Wrap with interceptor support.
|
||||
srv := mcpserver.NewServer(mcpServer)
|
||||
// Create an extension and register interceptors.
|
||||
ext := extension.New()
|
||||
|
||||
// Register a validator that blocks dangerous tool calls.
|
||||
srv.AddInterceptor(&interceptors.Validator{
|
||||
ext.AddInterceptor(&interceptors.Validator{
|
||||
Metadata: interceptors.Metadata{
|
||||
Name: "block-dangerous",
|
||||
Hook: interceptors.Hook{
|
||||
|
|
@ -35,11 +35,12 @@ srv.AddInterceptor(&interceptors.Validator{
|
|||
},
|
||||
})
|
||||
|
||||
// Create a chain and install middleware for automatic execution.
|
||||
chain, err := srv.LocalChain(ctx)
|
||||
// Install on the server and create a chain for middleware.
|
||||
ext.Install(mcpServer)
|
||||
chain, err := ext.LocalChain(ctx, mcpServer)
|
||||
mcpServer.AddReceivingMiddleware(gomiddleware.Middleware(chain))
|
||||
|
||||
srv.Run(context.Background(), &mcp.StdioTransport{})
|
||||
mcpServer.Run(context.Background(), &mcp.StdioTransport{})
|
||||
```
|
||||
|
||||
See [`examples/`](examples/) for complete working examples.
|
||||
|
|
|
|||
|
|
@ -5,9 +5,10 @@
|
|||
The Go SDK follows the SEP execution model where interceptors are
|
||||
first-class MCP primitives:
|
||||
|
||||
1. **Interceptor Server** (`mcpserver.Server`) — wraps an `mcp.Server`,
|
||||
registers interceptors, and exposes them via `interceptors/list` and
|
||||
`interceptor/invoke` JSON-RPC methods.
|
||||
1. **Interceptor Extension** (`extension.Extension`) — registers
|
||||
interceptors and installs them on one or more `mcp.Server` instances,
|
||||
exposing them via `interceptors/list` and `interceptor/invoke` JSON-RPC
|
||||
methods.
|
||||
2. **Chain** (`chain.Chain`) — the SDK-level orchestrator in the
|
||||
`interceptors/chain` sub-package. It holds `ChainEntry` objects
|
||||
(interceptor descriptor + MCP client session) and invokes interceptors
|
||||
|
|
@ -31,8 +32,8 @@ Transport (SSE / stdio / HTTP)
|
|||
|
||||
## LocalChain: In-Memory Transport
|
||||
|
||||
`Server.LocalChain(ctx)` creates a chain connected to the server via
|
||||
in-memory transport. Under the hood it:
|
||||
`Extension.LocalChain(ctx, server)` creates a chain connected to the given
|
||||
server via in-memory transport. Under the hood it:
|
||||
|
||||
1. Creates an `InMemoryTransport` pair via `mcp.NewInMemoryTransports()`
|
||||
2. Connects the server side via `mcp.Server.Connect(ctx, serverTransport)`
|
||||
|
|
@ -206,11 +207,11 @@ Response-phase validators see the post-mutation payload.
|
|||
| `chain.go` | `Chain`, `ChainEntry`, `ExecutionParams`, `ExecutionResult`, `NewChain`, `AddMCPServer`, `Execute`; filtering, sorting, parallel validation, sequential mutation, trust-boundary ordering |
|
||||
| `result.go` | Chain result types: ChainStatus, AbortType, AbortInfo, ValidationSummary |
|
||||
|
||||
### `interceptors/mcpserver/` — MCP server integration
|
||||
### `interceptors/extension/` — MCP server integration
|
||||
|
||||
| File | Responsibility |
|
||||
|------|---------------|
|
||||
| `server.go` | `Server` wrapper, interceptor registry, `LocalChain`, capability declaration, `NewStreamableHTTPHandler` |
|
||||
| `server.go` | `Extension`, interceptor management, `Install`, `LocalChain`, capability declaration |
|
||||
| `rpc.go` | `handleList` and `handleInvoke` JSON-RPC method handlers |
|
||||
| `events.go` | Event name constants for standard MCP methods |
|
||||
|
||||
|
|
|
|||
|
|
@ -15,8 +15,8 @@ import (
|
|||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/extension"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/integrations/gomiddleware"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/mcpserver"
|
||||
)
|
||||
|
||||
func main() {
|
||||
|
|
@ -73,13 +73,14 @@ func main() {
|
|||
},
|
||||
}
|
||||
|
||||
// Create the interceptor server (registers interceptors as discoverable
|
||||
// Create the interceptor extension (registers interceptors as discoverable
|
||||
// resources via interceptors/list and interceptor/invoke).
|
||||
srv := mcpserver.NewServer(mcpServer)
|
||||
srv.AddInterceptor(m)
|
||||
ext := extension.New()
|
||||
ext.AddInterceptor(m)
|
||||
ext.Install(mcpServer)
|
||||
|
||||
// Create a chain connected via in-memory transport.
|
||||
chain, err := srv.LocalChain(context.Background())
|
||||
chain, err := ext.LocalChain(context.Background(), mcpServer)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
|
@ -101,7 +102,10 @@ func main() {
|
|||
),
|
||||
)
|
||||
|
||||
handler := mcpserver.NewStreamableHTTPHandler(srv, nil)
|
||||
handler := mcp.NewStreamableHTTPHandler(
|
||||
func(r *http.Request) *mcp.Server { return mcpServer },
|
||||
nil,
|
||||
)
|
||||
log.Println("listening on :8080")
|
||||
if err := http.ListenAndServe(":8080", handler); err != nil {
|
||||
log.Fatal(err)
|
||||
|
|
|
|||
|
|
@ -15,8 +15,8 @@ import (
|
|||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/extension"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/integrations/gomiddleware"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/mcpserver"
|
||||
)
|
||||
|
||||
func main() {
|
||||
|
|
@ -70,13 +70,14 @@ func main() {
|
|||
},
|
||||
}
|
||||
|
||||
// Create the interceptor server (registers interceptors as discoverable
|
||||
// Create the interceptor extension (registers interceptors as discoverable
|
||||
// resources via interceptors/list and interceptor/invoke).
|
||||
srv := mcpserver.NewServer(mcpServer)
|
||||
srv.AddInterceptor(v)
|
||||
ext := extension.New()
|
||||
ext.AddInterceptor(v)
|
||||
ext.Install(mcpServer)
|
||||
|
||||
// Create a chain connected via in-memory transport.
|
||||
chain, err := srv.LocalChain(context.Background())
|
||||
chain, err := ext.LocalChain(context.Background(), mcpServer)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
|
@ -98,7 +99,10 @@ func main() {
|
|||
),
|
||||
)
|
||||
|
||||
handler := mcpserver.NewStreamableHTTPHandler(srv, nil)
|
||||
handler := mcp.NewStreamableHTTPHandler(
|
||||
func(r *http.Request) *mcp.Server { return mcpServer },
|
||||
nil,
|
||||
)
|
||||
log.Println("listening on :8080")
|
||||
if err := http.ListenAndServe(":8080", handler); err != nil {
|
||||
log.Fatal(err)
|
||||
|
|
|
|||
|
|
@ -32,8 +32,8 @@
|
|||
// gomiddleware.Middleware(chain),
|
||||
// )
|
||||
//
|
||||
// For convenience, [mcpserver.Server.LocalChain] creates a chain
|
||||
// For convenience, [extension.Extension.LocalChain] creates a chain
|
||||
// pre-configured with a local in-process MCP server:
|
||||
//
|
||||
// chain, err := srv.LocalChain(ctx)
|
||||
// chain, err := ext.LocalChain(ctx, mcpServer)
|
||||
package chain
|
||||
|
|
|
|||
|
|
@ -10,20 +10,21 @@
|
|||
// This package is transport-agnostic — it has no dependency on the
|
||||
// MCP SDK's server or client types. The chain orchestrator lives in
|
||||
// the [interceptors/chain] sub-package, MCP server integration in
|
||||
// [interceptors/mcpserver], and middleware in
|
||||
// [interceptors/extension], and middleware in
|
||||
// [interceptors/integrations/gomiddleware].
|
||||
//
|
||||
// # Usage with MCP
|
||||
//
|
||||
// Create an interceptor [mcpserver.Server], register interceptors,
|
||||
// then use [mcpserver.Server.LocalChain] to get a chain for
|
||||
// middleware:
|
||||
// Create an [extension.Extension], register interceptors, install on
|
||||
// an [mcp.Server], then use [extension.Extension.LocalChain] to get a
|
||||
// chain for middleware:
|
||||
//
|
||||
// srv := mcpserver.NewServer(mcpServer)
|
||||
// srv.AddInterceptor(myValidator)
|
||||
// srv.AddInterceptor(myMutator)
|
||||
// ext := extension.New()
|
||||
// ext.AddInterceptor(myValidator)
|
||||
// ext.AddInterceptor(myMutator)
|
||||
// ext.Install(mcpServer)
|
||||
//
|
||||
// chain, err := srv.LocalChain(ctx)
|
||||
// chain, err := ext.LocalChain(ctx, mcpServer)
|
||||
// mcpServer.AddReceivingMiddleware(
|
||||
// gomiddleware.Middleware(chain),
|
||||
// )
|
||||
|
|
|
|||
55
go/sdk/interceptors/extension/doc.go
Normal file
55
go/sdk/interceptors/extension/doc.go
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
// Copyright 2025 The MCP Interceptors Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache-2.0
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// Package extension registers interceptors as first-class MCP
|
||||
// primitives. Adding an interceptor to the extension makes it
|
||||
// discoverable via "interceptors/list" and invocable via
|
||||
// "interceptor/invoke", so any MCP client — local or remote — can
|
||||
// discover and call interceptors using standard JSON-RPC.
|
||||
//
|
||||
// # Getting Started
|
||||
//
|
||||
// Create an [Extension], register interceptors, and install on an
|
||||
// [mcp.Server]:
|
||||
//
|
||||
// ext := extension.New(extension.WithLogger(logger))
|
||||
// ext.AddInterceptor(myValidator)
|
||||
// ext.AddInterceptor(myMutator)
|
||||
// ext.Install(mcpServer)
|
||||
//
|
||||
// At this point any connected MCP client can call interceptors/list
|
||||
// and interceptor/invoke.
|
||||
//
|
||||
// For convenience, [Extension.LocalChain] creates a [chain.Chain] wired
|
||||
// to the server over an in-memory transport, ready for use with the
|
||||
// gomiddleware package:
|
||||
//
|
||||
// chain, err := ext.LocalChain(ctx, mcpServer)
|
||||
// mcpServer.AddReceivingMiddleware(
|
||||
// gomiddleware.Middleware(chain,
|
||||
// gomiddleware.WithContextProvider(myProvider),
|
||||
// ),
|
||||
// )
|
||||
//
|
||||
// # JSON-RPC Methods
|
||||
//
|
||||
// The extension installs two JSON-RPC methods on the server:
|
||||
//
|
||||
// - "interceptors/list": Returns all registered interceptors,
|
||||
// optionally filtered by event name.
|
||||
// - "interceptor/invoke": Invokes a single interceptor by name
|
||||
// for a given event, phase, and payload.
|
||||
//
|
||||
// # Transports
|
||||
//
|
||||
// Users interact with the [mcp.Server] directly for transport
|
||||
// configuration. For stdio, use [mcp.Server.Run]. For HTTP, use
|
||||
// [mcp.NewStreamableHTTPHandler]:
|
||||
//
|
||||
// handler := mcp.NewStreamableHTTPHandler(
|
||||
// func(r *http.Request) *mcp.Server { return mcpServer },
|
||||
// nil,
|
||||
// )
|
||||
// http.ListenAndServe(":8080", handler)
|
||||
package extension
|
||||
113
go/sdk/interceptors/extension/rpc.go
Normal file
113
go/sdk/interceptors/extension/rpc.go
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
// Copyright 2025 The MCP Interceptors Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache-2.0
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package extension
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"slices"
|
||||
"time"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors"
|
||||
)
|
||||
|
||||
// handleList implements the "interceptors/list" JSON-RPC method.
|
||||
// It returns all registered interceptors, optionally filtered by event.
|
||||
func (e *Extension) handleList(_ context.Context, _ *mcp.ServerSession, raw json.RawMessage) (any, error) {
|
||||
var params interceptors.ListParams
|
||||
if len(raw) > 0 {
|
||||
if err := json.Unmarshal(raw, ¶ms); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
all := e.getInterceptors()
|
||||
infos := make([]interceptors.InterceptorInfo, 0, len(all))
|
||||
for _, ri := range all {
|
||||
if params.Event != "" && !slices.Contains(ri.GetMetadata().Hook.Events, params.Event) {
|
||||
continue
|
||||
}
|
||||
infos = append(infos, interceptors.InfoFromInterceptor(ri))
|
||||
}
|
||||
|
||||
return &interceptors.ListResult{Interceptors: infos}, nil
|
||||
}
|
||||
|
||||
// handleInvoke implements the "interceptor/invoke" JSON-RPC method.
|
||||
// It invokes a single interceptor by name and returns its result.
|
||||
func (e *Extension) handleInvoke(ctx context.Context, _ *mcp.ServerSession, raw json.RawMessage) (any, error) {
|
||||
var params interceptors.InvokeParams
|
||||
if err := json.Unmarshal(raw, ¶ms); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if params.Name == "" {
|
||||
return nil, &jsonrpc.Error{
|
||||
Code: jsonrpc.CodeInvalidParams,
|
||||
Message: "name is required",
|
||||
}
|
||||
}
|
||||
|
||||
// Look up the interceptor by name.
|
||||
i := e.findByName(params.Name)
|
||||
if i == nil {
|
||||
return nil, &jsonrpc.Error{
|
||||
Code: jsonrpc.CodeInvalidParams,
|
||||
Message: fmt.Sprintf("interceptor %q not found", params.Name),
|
||||
}
|
||||
}
|
||||
|
||||
// Apply timeout if specified.
|
||||
if params.TimeoutMs > 0 {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, time.Duration(params.TimeoutMs)*time.Millisecond)
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
inv := &interceptors.Invocation{
|
||||
Event: params.Event,
|
||||
Phase: params.Phase,
|
||||
Payload: params.Payload,
|
||||
Config: params.Config,
|
||||
Context: params.Context,
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
|
||||
result := &interceptors.InvokeResult{
|
||||
Interceptor: params.Name,
|
||||
Type: i.GetType(),
|
||||
Phase: params.Phase,
|
||||
}
|
||||
defer func() { result.DurationMs = time.Since(start).Milliseconds() }()
|
||||
|
||||
switch v := i.(type) {
|
||||
case *interceptors.Validator:
|
||||
vr, err := v.Handler(ctx, inv)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.Validation = vr
|
||||
|
||||
case *interceptors.Mutator:
|
||||
mr, err := v.Handler(ctx, inv)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.Mutation = mr
|
||||
result.Payload = mr.Payload
|
||||
|
||||
default:
|
||||
return nil, &jsonrpc.Error{
|
||||
Code: jsonrpc.CodeInternalError,
|
||||
Message: fmt.Sprintf("unsupported interceptor type %T", i),
|
||||
}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
264
go/sdk/interceptors/extension/rpc_test.go
Normal file
264
go/sdk/interceptors/extension/rpc_test.go
Normal file
|
|
@ -0,0 +1,264 @@
|
|||
package extension_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors"
|
||||
)
|
||||
|
||||
// --- interceptors/list tests ---
|
||||
|
||||
func TestListInterceptors(t *testing.T) {
|
||||
t.Parallel()
|
||||
cs := setupRPCServer(t,
|
||||
allowAllValidator("v1"),
|
||||
blockToolValidator("dangerous"),
|
||||
)
|
||||
|
||||
var result interceptors.ListResult
|
||||
err := cs.CallCustom(context.Background(), interceptors.MethodList, nil, &result)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Len(t, result.Interceptors, 2)
|
||||
names := make([]string, len(result.Interceptors))
|
||||
for i, info := range result.Interceptors {
|
||||
names[i] = info.Name
|
||||
}
|
||||
assert.Contains(t, names, "v1")
|
||||
assert.Contains(t, names, "block-dangerous")
|
||||
}
|
||||
|
||||
func TestListWithEventFilter(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
toolsValidator := allowAllValidator("tools-v")
|
||||
promptsValidator := &interceptors.Validator{
|
||||
Metadata: interceptors.Metadata{
|
||||
Name: "prompts-v",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventPromptsGet},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
||||
return &interceptors.ValidationResult{Valid: true}, nil
|
||||
},
|
||||
}
|
||||
|
||||
cs := setupRPCServer(t, toolsValidator, promptsValidator)
|
||||
|
||||
// Filter for tools/call — should only return tools-v
|
||||
var result interceptors.ListResult
|
||||
err := cs.CallCustom(context.Background(), interceptors.MethodList,
|
||||
&interceptors.ListParams{Event: interceptors.EventToolsCall}, &result)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, result.Interceptors, 1)
|
||||
assert.Equal(t, "tools-v", result.Interceptors[0].Name)
|
||||
|
||||
// Filter for prompts/get — should only return prompts-v
|
||||
var result2 interceptors.ListResult
|
||||
err = cs.CallCustom(context.Background(), interceptors.MethodList,
|
||||
&interceptors.ListParams{Event: interceptors.EventPromptsGet}, &result2)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, result2.Interceptors, 1)
|
||||
assert.Equal(t, "prompts-v", result2.Interceptors[0].Name)
|
||||
|
||||
// Filter for non-existent event — empty
|
||||
var result3 interceptors.ListResult
|
||||
err = cs.CallCustom(context.Background(), interceptors.MethodList,
|
||||
&interceptors.ListParams{Event: "nonexistent/method"}, &result3)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, result3.Interceptors)
|
||||
}
|
||||
|
||||
// --- interceptor/invoke tests ---
|
||||
|
||||
func TestInvokeValidator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cs := setupRPCServer(t, allowAllValidator("v1"))
|
||||
|
||||
payload, _ := json.Marshal(map[string]any{"name": "echo", "arguments": map[string]any{"text": "hello"}})
|
||||
var result interceptors.InvokeResult
|
||||
err := cs.CallCustom(context.Background(), interceptors.MethodInvoke,
|
||||
&interceptors.InvokeParams{
|
||||
Name: "v1",
|
||||
Event: interceptors.EventToolsCall,
|
||||
Phase: interceptors.PhaseRequest,
|
||||
Payload: payload,
|
||||
}, &result)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "v1", result.Interceptor)
|
||||
assert.Equal(t, interceptors.TypeValidation, result.Type)
|
||||
assert.Equal(t, interceptors.PhaseRequest, result.Phase)
|
||||
require.NotNil(t, result.Validation)
|
||||
assert.True(t, result.Validation.Valid)
|
||||
assert.Nil(t, result.Mutation)
|
||||
assert.Nil(t, result.Payload)
|
||||
}
|
||||
|
||||
func TestInvokeValidatorRejects(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rejectAll := &interceptors.Validator{
|
||||
Metadata: interceptors.Metadata{
|
||||
Name: "reject-all",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
||||
return &interceptors.ValidationResult{
|
||||
Valid: false,
|
||||
Messages: []interceptors.ValidationMessage{
|
||||
{Message: "rejected", Severity: interceptors.SeverityError},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
cs := setupRPCServer(t, rejectAll)
|
||||
|
||||
payload, _ := json.Marshal(map[string]any{"name": "echo"})
|
||||
var result interceptors.InvokeResult
|
||||
err := cs.CallCustom(context.Background(), interceptors.MethodInvoke,
|
||||
&interceptors.InvokeParams{
|
||||
Name: "reject-all",
|
||||
Event: interceptors.EventToolsCall,
|
||||
Phase: interceptors.PhaseRequest,
|
||||
Payload: payload,
|
||||
}, &result)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "reject-all", result.Interceptor)
|
||||
assert.Equal(t, interceptors.TypeValidation, result.Type)
|
||||
require.NotNil(t, result.Validation)
|
||||
assert.False(t, result.Validation.Valid)
|
||||
assert.Len(t, result.Validation.Messages, 1)
|
||||
assert.Equal(t, "rejected", result.Validation.Messages[0].Message)
|
||||
}
|
||||
|
||||
func TestInvokeMutator(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
mutator := &interceptors.Mutator{
|
||||
Metadata: interceptors.Metadata{
|
||||
Name: "add-field",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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 m map[string]any
|
||||
if err := json.Unmarshal(raw, &m); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m["injected"] = true
|
||||
data, err := json.Marshal(m)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &interceptors.MutationResult{Modified: true, Info: map[string]any{"added": "injected"}, Payload: json.RawMessage(data)}, nil
|
||||
},
|
||||
}
|
||||
|
||||
cs := setupRPCServer(t, mutator)
|
||||
|
||||
payload, _ := json.Marshal(map[string]any{"name": "echo"})
|
||||
var result interceptors.InvokeResult
|
||||
err := cs.CallCustom(context.Background(), interceptors.MethodInvoke,
|
||||
&interceptors.InvokeParams{
|
||||
Name: "add-field",
|
||||
Event: interceptors.EventToolsCall,
|
||||
Phase: interceptors.PhaseRequest,
|
||||
Payload: payload,
|
||||
}, &result)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "add-field", result.Interceptor)
|
||||
assert.Equal(t, interceptors.TypeMutation, result.Type)
|
||||
require.NotNil(t, result.Mutation)
|
||||
assert.True(t, result.Mutation.Modified)
|
||||
assert.Equal(t, "injected", result.Mutation.Info["added"])
|
||||
|
||||
// Verify mutated payload contains the injected field.
|
||||
require.NotNil(t, result.Payload)
|
||||
var mutated map[string]any
|
||||
require.NoError(t, json.Unmarshal(result.Payload, &mutated))
|
||||
assert.Equal(t, true, mutated["injected"])
|
||||
assert.Equal(t, "echo", mutated["name"])
|
||||
}
|
||||
|
||||
func TestInvokeUnknownName(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cs := setupRPCServer(t, allowAllValidator("v1"))
|
||||
|
||||
payload, _ := json.Marshal(map[string]any{})
|
||||
var result interceptors.InvokeResult
|
||||
err := cs.CallCustom(context.Background(), interceptors.MethodInvoke,
|
||||
&interceptors.InvokeParams{
|
||||
Name: "nonexistent",
|
||||
Event: interceptors.EventToolsCall,
|
||||
Phase: interceptors.PhaseRequest,
|
||||
Payload: payload,
|
||||
}, &result)
|
||||
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "nonexistent")
|
||||
}
|
||||
|
||||
func TestInvokeTimeout(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
slowValidator := &interceptors.Validator{
|
||||
Metadata: interceptors.Metadata{
|
||||
Name: "slow-v",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
Handler: func(ctx context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
||||
select {
|
||||
case <-time.After(5 * time.Second):
|
||||
return &interceptors.ValidationResult{Valid: true}, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
cs := setupRPCServer(t, slowValidator)
|
||||
|
||||
payload, _ := json.Marshal(map[string]any{})
|
||||
var result interceptors.InvokeResult
|
||||
err := cs.CallCustom(context.Background(), interceptors.MethodInvoke,
|
||||
&interceptors.InvokeParams{
|
||||
Name: "slow-v",
|
||||
Event: interceptors.EventToolsCall,
|
||||
Phase: interceptors.PhaseRequest,
|
||||
Payload: payload,
|
||||
TimeoutMs: 50,
|
||||
}, &result)
|
||||
|
||||
require.Error(t, err)
|
||||
}
|
||||
198
go/sdk/interceptors/extension/server.go
Normal file
198
go/sdk/interceptors/extension/server.go
Normal file
|
|
@ -0,0 +1,198 @@
|
|||
// Copyright 2025 The MCP Interceptors Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache-2.0
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package extension
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"sync"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/chain"
|
||||
)
|
||||
|
||||
// extensionID is the capability key used in Capabilities.Experimental
|
||||
const extensionID = "io.modelcontextprotocol/interceptors"
|
||||
|
||||
// Option configures an Extension.
|
||||
type Option func(*Extension)
|
||||
|
||||
// WithLogger sets the logger for interceptor chain execution.
|
||||
// If not set, slog.Default() is used.
|
||||
func WithLogger(l *slog.Logger) Option {
|
||||
return func(e *Extension) {
|
||||
e.logger = l
|
||||
}
|
||||
}
|
||||
|
||||
// Extension manages interceptors and can install them on one or more
|
||||
// *mcp.Server instances. Interceptors are registered as first-class
|
||||
// MCP resources discoverable via the "interceptors/list" JSON-RPC
|
||||
// method and invocable via "interceptor/invoke".
|
||||
//
|
||||
// Extension is a pure interceptor container — it does not hold a
|
||||
// reference to any particular server. Call [Extension.Install] to
|
||||
// register the JSON-RPC methods and capability middleware on a server.
|
||||
type Extension struct {
|
||||
mu sync.RWMutex
|
||||
interceptors []interceptors.Interceptor
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
// New creates a new interceptor Extension.
|
||||
func New(opts ...Option) *Extension {
|
||||
e := &Extension{}
|
||||
for _, opt := range opts {
|
||||
opt(e)
|
||||
}
|
||||
if e.logger == nil {
|
||||
e.logger = slog.Default()
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// Install registers the "interceptors/list" and "interceptor/invoke"
|
||||
// JSON-RPC methods and installs a receiving middleware to enrich the
|
||||
// initialize response with interceptor capabilities on the given
|
||||
// server. It can be called on multiple servers.
|
||||
func (e *Extension) Install(server *mcp.Server) {
|
||||
server.AddCustomMethod(interceptors.MethodList, e.handleList)
|
||||
server.AddCustomMethod(interceptors.MethodInvoke, e.handleInvoke)
|
||||
server.AddReceivingMiddleware(e.initMiddleware())
|
||||
}
|
||||
|
||||
// AddInterceptor registers an interceptor. It panics if the interceptor
|
||||
// has a nil handler. Except for the panic it is safe to call while the server is running,
|
||||
// returns the receiver for chaining.
|
||||
func (e *Extension) AddInterceptor(i interceptors.Interceptor) *Extension {
|
||||
switch v := i.(type) {
|
||||
case *interceptors.Validator:
|
||||
if v.Handler == nil {
|
||||
panic("interceptors: validator " + v.Name + " has nil handler")
|
||||
}
|
||||
case *interceptors.Mutator:
|
||||
if v.Handler == nil {
|
||||
panic("interceptors: mutator " + v.Name + " has nil handler")
|
||||
}
|
||||
}
|
||||
e.mu.Lock()
|
||||
e.interceptors = append(e.interceptors, i)
|
||||
e.mu.Unlock()
|
||||
return e
|
||||
}
|
||||
|
||||
// findByName returns the interceptor with the given name, or nil if not found.
|
||||
func (e *Extension) findByName(name string) interceptors.Interceptor {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
for _, i := range e.interceptors {
|
||||
if i.GetMetadata().Name == name {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// getInterceptors returns a snapshot of the current interceptor list.
|
||||
func (e *Extension) getInterceptors() []interceptors.Interceptor {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
cp := make([]interceptors.Interceptor, len(e.interceptors))
|
||||
copy(cp, e.interceptors)
|
||||
return cp
|
||||
}
|
||||
|
||||
// LocalChain creates an interceptor Chain connected to the given server
|
||||
// via in-memory transport. This is the primary way to get a Chain
|
||||
// for use with middleware.
|
||||
//
|
||||
// The returned chain discovers interceptors from the server via
|
||||
// interceptors/list and invokes them via interceptor/invoke, exactly
|
||||
// as a remote chain would — but over an in-memory transport.
|
||||
func (e *Extension) LocalChain(ctx context.Context, server *mcp.Server) (*chain.Chain, error) {
|
||||
serverTransport, clientTransport := mcp.NewInMemoryTransports()
|
||||
|
||||
// Connect the server side.
|
||||
ss, err := server.Connect(ctx, serverTransport, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Connect a client through the other end.
|
||||
client := mcp.NewClient(&mcp.Implementation{
|
||||
Name: "interceptor-chain-client",
|
||||
Version: "internal",
|
||||
}, nil)
|
||||
cs, err := client.Connect(ctx, clientTransport, nil)
|
||||
if err != nil {
|
||||
ss.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ch := chain.NewChain(chain.WithChainLogger(e.logger))
|
||||
if err := ch.AddMCPServer(ctx, cs); err != nil {
|
||||
cs.Close()
|
||||
ss.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
// initMiddleware returns a lightweight mcp.Middleware that enriches the
|
||||
// initialize response with interceptor capability metadata.
|
||||
func (e *Extension) initMiddleware() mcp.Middleware {
|
||||
return func(next mcp.MethodHandler) mcp.MethodHandler {
|
||||
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
|
||||
result, err := next(ctx, method, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if method == "initialize" {
|
||||
return e.enrichInitResult(result)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// enrichInitResult injects interceptor capability metadata into the
|
||||
// InitializeResult, following the same pattern as the variants extension.
|
||||
// The capability is declared under Capabilities.Experimental so it works
|
||||
// without upstream go-sdk changes.
|
||||
func (e *Extension) enrichInitResult(result mcp.Result) (mcp.Result, error) {
|
||||
initResult, ok := result.(*mcp.InitializeResult)
|
||||
if !ok {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
if initResult.Capabilities == nil {
|
||||
initResult.Capabilities = &mcp.ServerCapabilities{}
|
||||
}
|
||||
if initResult.Capabilities.Experimental == nil {
|
||||
initResult.Capabilities.Experimental = make(map[string]any)
|
||||
}
|
||||
|
||||
all := e.getInterceptors()
|
||||
supportedEvents := map[string]bool{}
|
||||
for _, ri := range all {
|
||||
for _, ev := range ri.GetMetadata().Hook.Events {
|
||||
supportedEvents[ev] = true
|
||||
}
|
||||
}
|
||||
|
||||
events := make([]string, 0, len(supportedEvents))
|
||||
for ev := range supportedEvents {
|
||||
events = append(events, ev)
|
||||
}
|
||||
|
||||
initResult.Capabilities.Experimental[extensionID] = map[string]any{
|
||||
"supportedEvents": events,
|
||||
}
|
||||
|
||||
return initResult, nil
|
||||
}
|
||||
612
go/sdk/interceptors/extension/server_integration_test.go
Normal file
612
go/sdk/interceptors/extension/server_integration_test.go
Normal file
|
|
@ -0,0 +1,612 @@
|
|||
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.ModeEnforce))
|
||||
|
||||
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.ModeEnforce),
|
||||
prefixMutator("m2", "[M2]", interceptors.PhaseResponse, 20, interceptors.ModeEnforce),
|
||||
prefixMutator("m3-audit", "[AUDIT]", interceptors.PhaseResponse, 30, interceptors.ModeAudit),
|
||||
prefixMutator("m4", "[M4]", interceptors.PhaseResponse, 40, interceptors.ModeEnforce),
|
||||
prefixMutator("m5", "[M5]", interceptors.PhaseResponse, 50, interceptors.ModeEnforce),
|
||||
)
|
||||
|
||||
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.ModeEnforce),
|
||||
failMutator("m2-fail", interceptors.PhaseResponse, 20),
|
||||
prefixMutator("m3", "[M3]", interceptors.PhaseResponse, 30, interceptors.ModeEnforce),
|
||||
)
|
||||
|
||||
_, 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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseResponse,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
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.ModeEnforce),
|
||||
slowMutator,
|
||||
prefixMutator("m3", "[M3]", interceptors.PhaseResponse, 30, interceptors.ModeEnforce),
|
||||
)
|
||||
|
||||
// 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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseResponse,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
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.ModeEnforce),
|
||||
failOpenMutator,
|
||||
prefixMutator("m3", "[M3]", interceptors.PhaseResponse, 30, interceptors.ModeEnforce),
|
||||
)
|
||||
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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,
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
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,
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseResponse,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseResponse,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseResponse,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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",
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseResponse,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
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")
|
||||
}
|
||||
291
go/sdk/interceptors/extension/testharness_test.go
Normal file
291
go/sdk/interceptors/extension/testharness_test.go
Normal file
|
|
@ -0,0 +1,291 @@
|
|||
package extension_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/extension"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/integrations/gomiddleware"
|
||||
)
|
||||
|
||||
// --- Server setup ---
|
||||
|
||||
// setup creates an HTTP test server with interceptor support and a connected
|
||||
// client session. The server has a single "echo" tool registered for
|
||||
// EventToolsCall. Cleanup is registered via t.Cleanup.
|
||||
func setup(t *testing.T, is ...interceptors.Interceptor) *mcp.ClientSession {
|
||||
t.Helper()
|
||||
return setupWithTools(t, defaultTools(), is...)
|
||||
}
|
||||
|
||||
// setupWithTools creates an HTTP test server with the given tools and
|
||||
// interceptors. Uses LocalChain to create a chain that invokes
|
||||
// interceptors via interceptor/invoke over in-memory transport.
|
||||
// Cleanup is registered via t.Cleanup.
|
||||
func setupWithTools(t *testing.T, tools []testTool, is ...interceptors.Interceptor) *mcp.ClientSession {
|
||||
t.Helper()
|
||||
|
||||
mcpServer := mcp.NewServer(&mcp.Implementation{
|
||||
Name: "test-server",
|
||||
Version: "0.1.0",
|
||||
}, nil)
|
||||
|
||||
for _, tool := range tools {
|
||||
mcpServer.AddTool(tool.tool, tool.handler)
|
||||
}
|
||||
|
||||
ext := extension.New()
|
||||
for _, i := range is {
|
||||
ext.AddInterceptor(i)
|
||||
}
|
||||
ext.Install(mcpServer)
|
||||
|
||||
// Create a chain via LocalChain (in-memory transport).
|
||||
chain, err := ext.LocalChain(context.Background(), mcpServer)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Install the middleware for automatic interceptor execution.
|
||||
mcpServer.AddReceivingMiddleware(gomiddleware.Middleware(chain))
|
||||
|
||||
handler := mcp.NewStreamableHTTPHandler(
|
||||
func(r *http.Request) *mcp.Server { return mcpServer },
|
||||
nil,
|
||||
)
|
||||
httpServer := httptest.NewServer(handler)
|
||||
t.Cleanup(httpServer.Close)
|
||||
|
||||
client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "0.1.0"}, nil)
|
||||
cs, err := client.Connect(context.Background(), &mcp.StreamableClientTransport{
|
||||
Endpoint: httpServer.URL,
|
||||
}, nil)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { cs.Close() })
|
||||
|
||||
return cs
|
||||
}
|
||||
|
||||
// --- Tool definitions ---
|
||||
|
||||
// testTool pairs a tool definition with its handler for use with setupWithTools.
|
||||
type testTool struct {
|
||||
tool *mcp.Tool
|
||||
handler mcp.ToolHandler
|
||||
}
|
||||
|
||||
// defaultTools returns the standard "echo" tool used by most tests.
|
||||
func defaultTools() []testTool {
|
||||
return []testTool{
|
||||
{
|
||||
tool: &mcp.Tool{
|
||||
Name: "echo",
|
||||
Description: "echoes input",
|
||||
InputSchema: map[string]any{"type": "object"},
|
||||
},
|
||||
handler: func(_ context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
return &mcp.CallToolResult{
|
||||
Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("echo: %s", req.Params.Arguments)}},
|
||||
}, nil
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// --- Call helpers ---
|
||||
|
||||
// callEcho calls the "echo" tool and returns the result.
|
||||
func callEcho(t *testing.T, cs *mcp.ClientSession) (*mcp.CallToolResult, error) {
|
||||
t.Helper()
|
||||
return cs.CallTool(context.Background(), &mcp.CallToolParams{
|
||||
Name: "echo",
|
||||
Arguments: map[string]any{"text": "hello"},
|
||||
})
|
||||
}
|
||||
|
||||
// callEchoWithTimeout calls the "echo" tool with a context timeout.
|
||||
func callEchoWithTimeout(t *testing.T, cs *mcp.ClientSession, timeout time.Duration) (*mcp.CallToolResult, error) {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||
defer cancel()
|
||||
return cs.CallTool(ctx, &mcp.CallToolParams{
|
||||
Name: "echo",
|
||||
Arguments: map[string]any{"text": "hello"},
|
||||
})
|
||||
}
|
||||
|
||||
// --- Assertion helpers ---
|
||||
|
||||
// resultText extracts the first TextContent string from a CallToolResult.
|
||||
// Fails the test if the result has no content or the first item isn't TextContent.
|
||||
func resultText(t *testing.T, result *mcp.CallToolResult) string {
|
||||
t.Helper()
|
||||
require.NotEmpty(t, result.Content)
|
||||
tc, ok := result.Content[0].(*mcp.TextContent)
|
||||
require.True(t, ok, "expected TextContent, got %T", result.Content[0])
|
||||
return tc.Text
|
||||
}
|
||||
|
||||
// --- RPC server setup ---
|
||||
|
||||
// setupRPCServer creates an interceptor extension with the given interceptors
|
||||
// and returns a connected client session. The client can call custom methods
|
||||
// via raw JSON-RPC.
|
||||
func setupRPCServer(t *testing.T, is ...interceptors.Interceptor) *mcp.ClientSession {
|
||||
t.Helper()
|
||||
|
||||
mcpServer := mcp.NewServer(&mcp.Implementation{
|
||||
Name: "test-server",
|
||||
Version: "0.1.0",
|
||||
}, nil)
|
||||
|
||||
// A minimal tool so the server has something to initialize with.
|
||||
mcpServer.AddTool(&mcp.Tool{
|
||||
Name: "echo",
|
||||
Description: "echoes input",
|
||||
InputSchema: map[string]any{"type": "object"},
|
||||
}, func(_ context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
return &mcp.CallToolResult{
|
||||
Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("echo: %s", req.Params.Arguments)}},
|
||||
}, nil
|
||||
})
|
||||
|
||||
ext := extension.New()
|
||||
for _, i := range is {
|
||||
ext.AddInterceptor(i)
|
||||
}
|
||||
ext.Install(mcpServer)
|
||||
|
||||
handler := mcp.NewStreamableHTTPHandler(
|
||||
func(r *http.Request) *mcp.Server { return mcpServer },
|
||||
nil,
|
||||
)
|
||||
httpServer := httptest.NewServer(handler)
|
||||
t.Cleanup(httpServer.Close)
|
||||
|
||||
client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "0.1.0"}, nil)
|
||||
cs, err := client.Connect(context.Background(), &mcp.StreamableClientTransport{
|
||||
Endpoint: httpServer.URL,
|
||||
}, nil)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { cs.Close() })
|
||||
|
||||
return cs
|
||||
}
|
||||
|
||||
// --- Interceptor builders ---
|
||||
// All handlers work with json.RawMessage payloads since that's what
|
||||
// interceptor/invoke delivers.
|
||||
|
||||
// prefixMutator creates a mutator that prepends tag to each TextContent
|
||||
// in a CallToolResult (response phase).
|
||||
func prefixMutator(name, tag string, phase interceptors.InterceptionPhase, priority int, mode interceptors.Mode) *interceptors.Mutator {
|
||||
return &interceptors.Mutator{
|
||||
Metadata: interceptors.Metadata{
|
||||
Name: name,
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: phase,
|
||||
},
|
||||
Mode: mode,
|
||||
PriorityHint: interceptors.NewPriority(priority),
|
||||
},
|
||||
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 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
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// failMutator creates a mutator that always returns an error.
|
||||
func failMutator(name string, phase interceptors.InterceptionPhase, priority int) *interceptors.Mutator {
|
||||
return &interceptors.Mutator{
|
||||
Metadata: interceptors.Metadata{
|
||||
Name: name,
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: phase,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
PriorityHint: interceptors.NewPriority(priority),
|
||||
},
|
||||
Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.MutationResult, error) {
|
||||
return nil, fmt.Errorf("simulated failure")
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// blockToolValidator creates a request validator that rejects calls to the named tool.
|
||||
func blockToolValidator(toolName string) *interceptors.Validator {
|
||||
return &interceptors.Validator{
|
||||
Metadata: interceptors.Metadata{
|
||||
Name: "block-" + toolName,
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
Handler: func(_ context.Context, inv *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
||||
raw, ok := inv.Payload.(json.RawMessage)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unexpected payload 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 == toolName {
|
||||
return &interceptors.ValidationResult{
|
||||
Valid: false,
|
||||
Messages: []interceptors.ValidationMessage{
|
||||
{Message: toolName + " is blocked", Severity: interceptors.SeverityError},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
return &interceptors.ValidationResult{Valid: true}, nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// allowAllValidator creates a request validator that allows all calls.
|
||||
func allowAllValidator(name string) *interceptors.Validator {
|
||||
return &interceptors.Validator{
|
||||
Metadata: interceptors.Metadata{
|
||||
Name: name,
|
||||
Hook: interceptors.Hook{
|
||||
Events: []string{interceptors.EventToolsCall},
|
||||
Phase: interceptors.PhaseRequest,
|
||||
},
|
||||
Mode: interceptors.ModeEnforce,
|
||||
},
|
||||
Handler: func(_ context.Context, _ *interceptors.Invocation) (*interceptors.ValidationResult, error) {
|
||||
return &interceptors.ValidationResult{Valid: true}, nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
@ -16,13 +16,14 @@
|
|||
//
|
||||
// # Usage
|
||||
//
|
||||
// // Create interceptor server (registers interceptors as resources)
|
||||
// srv := mcpserver.NewServer(mcpServer)
|
||||
// srv.AddInterceptor(myValidator)
|
||||
// srv.AddInterceptor(myMutator)
|
||||
// // Create interceptor extension (registers interceptors as resources)
|
||||
// ext := extension.New()
|
||||
// ext.AddInterceptor(myValidator)
|
||||
// ext.AddInterceptor(myMutator)
|
||||
// ext.Install(mcpServer)
|
||||
//
|
||||
// // Create chain via in-memory transport
|
||||
// chain, err := srv.LocalChain(ctx)
|
||||
// chain, err := ext.LocalChain(ctx, mcpServer)
|
||||
//
|
||||
// // Install middleware for automatic execution
|
||||
// mcpServer.AddReceivingMiddleware(
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
|
@ -13,8 +14,8 @@ import (
|
|||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/extension"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/integrations/gomiddleware"
|
||||
"github.com/modelcontextprotocol/ext-interceptors/go/sdk/interceptors/mcpserver"
|
||||
)
|
||||
|
||||
// setup creates a test server with middleware installed and returns a connected client session.
|
||||
|
|
@ -36,19 +37,23 @@ func setup(t *testing.T, is ...interceptors.Interceptor) *mcp.ClientSession {
|
|||
}, nil
|
||||
})
|
||||
|
||||
srv := mcpserver.NewServer(mcpServer)
|
||||
ext := extension.New()
|
||||
for _, i := range is {
|
||||
srv.AddInterceptor(i)
|
||||
ext.AddInterceptor(i)
|
||||
}
|
||||
ext.Install(mcpServer)
|
||||
|
||||
// Create chain via LocalChain (in-memory transport).
|
||||
chain, err := srv.LocalChain(context.Background())
|
||||
chain, err := ext.LocalChain(context.Background(), mcpServer)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Install the middleware.
|
||||
mcpServer.AddReceivingMiddleware(gomiddleware.Middleware(chain))
|
||||
|
||||
handler := mcpserver.NewStreamableHTTPHandler(srv, nil)
|
||||
handler := mcp.NewStreamableHTTPHandler(
|
||||
func(r *http.Request) *mcp.Server { return mcpServer },
|
||||
nil,
|
||||
)
|
||||
httpServer := httptest.NewServer(handler)
|
||||
t.Cleanup(httpServer.Close)
|
||||
|
||||
|
|
@ -260,10 +265,11 @@ func TestMiddlewareWithContextProvider(t *testing.T) {
|
|||
}, nil
|
||||
})
|
||||
|
||||
srv := mcpserver.NewServer(mcpServer)
|
||||
srv.AddInterceptor(principalCheck)
|
||||
ext := extension.New()
|
||||
ext.AddInterceptor(principalCheck)
|
||||
ext.Install(mcpServer)
|
||||
|
||||
chain, err := srv.LocalChain(context.Background())
|
||||
chain, err := ext.LocalChain(context.Background(), mcpServer)
|
||||
require.NoError(t, err)
|
||||
|
||||
mcpServer.AddReceivingMiddleware(gomiddleware.Middleware(chain,
|
||||
|
|
@ -277,7 +283,10 @@ func TestMiddlewareWithContextProvider(t *testing.T) {
|
|||
}),
|
||||
))
|
||||
|
||||
handler := mcpserver.NewStreamableHTTPHandler(srv, nil)
|
||||
handler := mcp.NewStreamableHTTPHandler(
|
||||
func(r *http.Request) *mcp.Server { return mcpServer },
|
||||
nil,
|
||||
)
|
||||
httpServer := httptest.NewServer(handler)
|
||||
t.Cleanup(httpServer.Close)
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue