Skip to content

Commit c82469e

Browse files
Clay Giffordadwsingh
authored andcommitted
feat(mcp): Add McpServerInterceptor with scoped hooks
Add hook-based server interceptor for McpService, following the established ClientInterceptor pattern. Interceptors observe or modify requests at specific pipeline stages without exposing async dispatch internals. Hooks: readBefore/modifyBefore Execution and ToolCall, readAfter/modifyAfter Execution and ToolCall. Also makes McpServerProxy's rpc(), start(), and shutdown() protected to allow custom proxy transport implementations.
1 parent b53f721 commit c82469e

8 files changed

Lines changed: 1603 additions & 25 deletions

File tree

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.mcp.server;
7+
8+
import software.amazon.smithy.java.context.Context;
9+
import software.amazon.smithy.java.mcp.model.JsonRpcRequest;
10+
import software.amazon.smithy.utils.SmithyUnstableApi;
11+
12+
/**
13+
* Hook data available at the execution level. Passed to execution-scoped hooks in
14+
* {@link McpServerInterceptor}.
15+
*
16+
* <p>The {@link #context()} provides a per-request key-value store for passing state
17+
* between hooks. For example, a telemetry interceptor can stash a start timestamp in
18+
* {@code readBeforeExecution} and retrieve it in {@code readAfterExecution}.
19+
*/
20+
@SmithyUnstableApi
21+
public class McpExecutionHook {
22+
23+
private final JsonRpcRequest request;
24+
private final ProtocolVersion protocolVersion;
25+
private final Context context;
26+
27+
McpExecutionHook(JsonRpcRequest request, ProtocolVersion protocolVersion, Context context) {
28+
this.request = request;
29+
this.protocolVersion = protocolVersion;
30+
this.context = context;
31+
}
32+
33+
/**
34+
* The JSON-RPC request being handled.
35+
*/
36+
public JsonRpcRequest request() {
37+
return request;
38+
}
39+
40+
/**
41+
* Returns a new hook with the given request, or the same hook if unchanged.
42+
*/
43+
public McpExecutionHook withRequest(JsonRpcRequest request) {
44+
return this.request == request ? this : new McpExecutionHook(request, protocolVersion, context);
45+
}
46+
47+
/**
48+
* The MCP protocol version for this request.
49+
*/
50+
public ProtocolVersion protocolVersion() {
51+
return protocolVersion;
52+
}
53+
54+
/**
55+
* Per-request context for passing state between hooks.
56+
*/
57+
public Context context() {
58+
return context;
59+
}
60+
}

mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerBuilder.java

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ public final class McpServerBuilder {
2424
OutputStream os;
2525
Map<String, Service> services = new HashMap<>();
2626
List<McpServerProxy> proxyList = new ArrayList<>();
27+
McpServerInterceptor interceptor;
2728
String name;
2829
String version;
2930
ToolFilter toolFilter = (server, tool) -> true;
@@ -72,6 +73,10 @@ public Server build() {
7273
builder.version(version);
7374
}
7475

76+
if (interceptor != null) {
77+
builder.interceptor(interceptor);
78+
}
79+
7580
this.mcpService = builder.build();
7681
return new McpServer(this);
7782
}
@@ -101,6 +106,17 @@ public McpServerBuilder metricsObserver(McpMetricsObserver observer) {
101106
return this;
102107
}
103108

109+
/**
110+
* Sets the server interceptor. Use {@link McpServerInterceptor#chain(List)} to compose
111+
* multiple interceptors into one.
112+
*
113+
* @see McpServerInterceptor for hook descriptions and the execution lifecycle
114+
*/
115+
public McpServerBuilder interceptor(McpServerInterceptor interceptor) {
116+
this.interceptor = Objects.requireNonNull(interceptor, "interceptor");
117+
return this;
118+
}
119+
104120
private void validate() {
105121
Objects.requireNonNull(is, "MCP server input stream is required");
106122
Objects.requireNonNull(os, "MCP server output stream is required");
Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.mcp.server;
7+
8+
import java.util.List;
9+
import software.amazon.smithy.java.mcp.model.JsonRpcRequest;
10+
import software.amazon.smithy.java.mcp.model.JsonRpcResponse;
11+
import software.amazon.smithy.utils.SmithyUnstableApi;
12+
13+
/**
14+
* Interceptor for MCP server request processing. Interceptors inject code into the
15+
* {@link McpService} request execution pipeline via hooks at specific stages.
16+
*
17+
* <p>Hooks are either "read" hooks (observe in-flight data) or "modify" hooks (transform
18+
* in-flight data). All hooks have default no-op implementations; override only the hooks
19+
* you need.
20+
*
21+
* <h2>Execution lifecycle</h2>
22+
*
23+
* <p>For every request:
24+
* <ol>
25+
* <li>{@link #readBeforeExecution} — observe the incoming request</li>
26+
* <li>{@link #modifyBeforeExecution} — optionally transform the request</li>
27+
* <li>For {@code tools/call} requests only:
28+
* <ol>
29+
* <li>{@link #readBeforeToolCall} — observe before tool dispatch</li>
30+
* <li>{@link #modifyBeforeToolCall} — optionally transform the request</li>
31+
* <li>Tool dispatch (local or proxy)</li>
32+
* <li>{@link #readAfterToolCall} — observe the tool result</li>
33+
* <li>{@link #modifyAfterToolCall} — optionally transform the response</li>
34+
* </ol>
35+
* </li>
36+
* <li>{@link #readAfterExecution} — observe the final result (ALWAYS fires)</li>
37+
* <li>{@link #modifyAfterExecution} — optionally transform the final response</li>
38+
* </ol>
39+
*
40+
* <h2>Error handling</h2>
41+
*
42+
* <p>Any hook may throw a {@link RuntimeException}. When a hook throws, remaining hooks
43+
* in that stage are skipped, and execution jumps to the after-execution hooks with the
44+
* error. The {@code readAfterExecution} and {@code modifyAfterExecution} hooks ALWAYS fire,
45+
* ensuring cleanup and telemetry logic runs regardless of errors.
46+
*
47+
* <h2>Async tool calls</h2>
48+
*
49+
* <p>For proxy tool calls, the after-tool-call and after-execution hooks fire on the
50+
* thread that receives the proxy response, not the original request thread. Hook
51+
* implementations must be thread-safe.
52+
*
53+
* <h2>Example: telemetry</h2>
54+
* <pre>{@code
55+
* public class TelemetryInterceptor implements McpServerInterceptor {
56+
* private static final Context.Key<Long> START = Context.key("start");
57+
*
58+
* @Override
59+
* public void readBeforeExecution(McpExecutionHook hook) {
60+
* hook.context().put(START, System.nanoTime());
61+
* }
62+
*
63+
* @Override
64+
* public void readAfterExecution(McpExecutionHook hook,
65+
* JsonRpcResponse response, RuntimeException error) {
66+
* long duration = System.nanoTime() - hook.context().get(START);
67+
* emitMetrics(hook.request().getMethod(), duration, error == null);
68+
* }
69+
* }
70+
* }</pre>
71+
*
72+
* <h2>Example: access control</h2>
73+
* <pre>{@code
74+
* public class AccessControlInterceptor implements McpServerInterceptor {
75+
* @Override
76+
* public void readBeforeToolCall(McpToolCallHook hook) {
77+
* if (isBlocked(hook.toolName(), hook.serverId())) {
78+
* throw new RuntimeException("Access denied: " + hook.toolName());
79+
* }
80+
* }
81+
* }
82+
* }</pre>
83+
*/
84+
@SmithyUnstableApi
85+
public interface McpServerInterceptor {
86+
87+
/**
88+
* An interceptor that does nothing.
89+
*/
90+
McpServerInterceptor NOOP = new McpServerInterceptor() {};
91+
92+
/**
93+
* Combines multiple interceptors into a single interceptor that invokes each one
94+
* in order. Hooks are called sequentially on each interceptor in list order.
95+
*
96+
* @param interceptors The interceptors to compose.
97+
* @return A single interceptor that delegates to all provided interceptors.
98+
*/
99+
static McpServerInterceptor chain(List<McpServerInterceptor> interceptors) {
100+
return switch (interceptors.size()) {
101+
case 0 -> NOOP;
102+
case 1 -> interceptors.get(0);
103+
default -> new McpServerInterceptorChain(List.copyOf(interceptors));
104+
};
105+
}
106+
107+
/**
108+
* Combines multiple interceptors into a single interceptor that invokes each one
109+
* in order. Convenience overload of {@link #chain(List)}.
110+
*
111+
* @param interceptors The interceptors to compose.
112+
* @return A single interceptor that delegates to all provided interceptors.
113+
*/
114+
static McpServerInterceptor chain(McpServerInterceptor... interceptors) {
115+
return chain(List.of(interceptors));
116+
}
117+
118+
// --- Execution-level hooks (fire for all requests) ---
119+
120+
/**
121+
* Called when a request is received, before any dispatch logic.
122+
*
123+
* @param hook Execution hook data containing the request, protocol version, and context.
124+
*/
125+
default void readBeforeExecution(McpExecutionHook hook) {}
126+
127+
/**
128+
* Called before dispatch. Can return a modified request.
129+
*
130+
* @param hook Execution hook data.
131+
* @return The request to dispatch, or {@code hook.request()} to pass through unmodified.
132+
*/
133+
default JsonRpcRequest modifyBeforeExecution(McpExecutionHook hook) {
134+
return hook.request();
135+
}
136+
137+
/**
138+
* Called when execution completes. ALWAYS fires, even if an earlier hook threw.
139+
*
140+
* @param hook Execution hook data.
141+
* @param response The response, or {@code null} for notifications and async proxy calls
142+
* still in flight.
143+
* @param error The error if one occurred, or {@code null} on success.
144+
*/
145+
default void readAfterExecution(McpExecutionHook hook, JsonRpcResponse response, RuntimeException error) {}
146+
147+
/**
148+
* Called when execution completes. Can modify the response or handle errors.
149+
* ALWAYS fires, even if an earlier hook threw.
150+
*
151+
* @param hook Execution hook data.
152+
* @param response The response, or {@code null} for notifications.
153+
* @param error The error if one occurred, or {@code null} on success.
154+
* @return The final response.
155+
* @throws RuntimeException to propagate or replace the error.
156+
*/
157+
default JsonRpcResponse modifyAfterExecution(
158+
McpExecutionHook hook,
159+
JsonRpcResponse response,
160+
RuntimeException error
161+
) {
162+
if (error != null) {
163+
throw error;
164+
}
165+
return response;
166+
}
167+
168+
// --- Tool-level hooks (fire only for tools/call) ---
169+
170+
/**
171+
* Called before a tool is invoked.
172+
*
173+
* @param hook Tool call hook data containing tool name, server ID, and proxy status.
174+
*/
175+
default void readBeforeToolCall(McpToolCallHook hook) {}
176+
177+
/**
178+
* Called before a tool is invoked. Can return a modified request.
179+
*
180+
* @param hook Tool call hook data.
181+
* @return The request to use for tool invocation, or {@code hook.request()} to pass
182+
* through unmodified.
183+
*/
184+
default JsonRpcRequest modifyBeforeToolCall(McpToolCallHook hook) {
185+
return hook.request();
186+
}
187+
188+
/**
189+
* Called after a tool completes. For proxy tools, this fires on the callback thread.
190+
*
191+
* @param hook Tool call hook data.
192+
* @param response The tool call response.
193+
* @param error The error if one occurred, or {@code null} on success.
194+
*/
195+
default void readAfterToolCall(McpToolCallHook hook, JsonRpcResponse response, RuntimeException error) {}
196+
197+
/**
198+
* Called after a tool completes. Can modify the response or handle errors.
199+
* For proxy tools, this fires on the callback thread.
200+
*
201+
* @param hook Tool call hook data.
202+
* @param response The tool call response.
203+
* @param error The error if one occurred, or {@code null} on success.
204+
* @return The response to return.
205+
* @throws RuntimeException to propagate or replace the error.
206+
*/
207+
default JsonRpcResponse modifyAfterToolCall(
208+
McpToolCallHook hook,
209+
JsonRpcResponse response,
210+
RuntimeException error
211+
) {
212+
if (error != null) {
213+
throw error;
214+
}
215+
return response;
216+
}
217+
}

0 commit comments

Comments
 (0)