Skip to content

Commit 2997def

Browse files
adwsinghadwsingh
committed
Send back a notification after successfull initialize
1 parent 4b46bb3 commit 2997def

File tree

4 files changed

+88
-40
lines changed

4 files changed

+88
-40
lines changed

‎mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/HttpMcpProxy.java‎

Lines changed: 3 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import software.amazon.smithy.java.mcp.model.JsonRpcErrorResponse;
2626
import software.amazon.smithy.java.mcp.model.JsonRpcRequest;
2727
import software.amazon.smithy.java.mcp.model.JsonRpcResponse;
28-
import software.amazon.smithy.model.shapes.ShapeType;
2928
import software.amazon.smithy.utils.SmithyUnstableApi;
3029

3130
@SmithyUnstableApi
@@ -208,7 +207,7 @@ private JsonRpcResponse parseSseResponse(HttpResponse response, JsonRpcRequest r
208207
// This is a notification - convert Document to JsonRpcRequest and forward
209208
JsonRpcRequest notification = jsonDocument.asShape(JsonRpcRequest.builder());
210209
LOG.debug("Received notification from SSE stream: method={}", notification.getMethod());
211-
notifyRequest(notification);
210+
notify(notification);
212211
} else {
213212
// This is a response - convert Document to JsonRpcResponse
214213
finalResponse = jsonDocument.asShape(JsonRpcResponse.builder());
@@ -236,14 +235,14 @@ private JsonRpcResponse parseSseResponse(HttpResponse response, JsonRpcRequest r
236235
.build();
237236
LOG.debug("Received notification from remaining SSE buffer: method={}",
238237
notification.getMethod());
239-
notifyRequest(notification);
238+
notify(notification);
240239
} else {
241240
JsonRpcResponse message = JsonRpcResponse.builder()
242241
.deserialize(jsonDocument.createDeserializer())
243242
.build();
244243

245244
if (message.getId() == null) {
246-
notifyRequest(JsonRpcRequest.builder()
245+
notify(JsonRpcRequest.builder()
247246
.jsonrpc("2.0")
248247
.method("notifications/unknown")
249248
.build());
@@ -282,27 +281,6 @@ private JsonRpcResponse parseSseResponse(HttpResponse response, JsonRpcRequest r
282281
}
283282
}
284283

285-
/**
286-
* Determines if a Document represents a notification (has "method" but no "id")
287-
* rather than a response (has "id").
288-
*
289-
* - Responses have an "id" field at the top level
290-
* - Notifications have a "method" field but no "id" field at the top level
291-
*/
292-
private boolean isNotification(Document doc) {
293-
try {
294-
if (!doc.isType(ShapeType.STRUCTURE) && !doc.isType(ShapeType.MAP)) {
295-
return false;
296-
}
297-
298-
// If it has a "method" field but no "id", it's a notification
299-
return doc.getMember("id") == null && doc.getMember("method") != null;
300-
} catch (Exception e) {
301-
LOG.warn("Failed to determine if notification from Document", e);
302-
return false;
303-
}
304-
}
305-
306284
private JsonRpcResponse handleErrorResponse(HttpResponse response) {
307285
long contentLength = response.body().contentLength();
308286
String errorMessage = "HTTP " + response.statusCode();

‎mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerProxy.java‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import software.amazon.smithy.java.mcp.model.ListToolsResult;
2323
import software.amazon.smithy.java.mcp.model.PromptInfo;
2424
import software.amazon.smithy.java.mcp.model.ToolInfo;
25+
import software.amazon.smithy.model.shapes.ShapeType;
2526

2627
public abstract class McpServerProxy {
2728

@@ -81,6 +82,14 @@ public void initialize(
8182
if (result.getError() != null) {
8283
throw new RuntimeException("Error during initialization: " + result.getError().getMessage());
8384
}
85+
86+
// Send the initialized notification per MCP protocol spec
87+
JsonRpcRequest initializedNotification = JsonRpcRequest.builder()
88+
.method("notifications/initialized")
89+
.jsonrpc("2.0")
90+
.build();
91+
rpc(initializedNotification);
92+
8493
this.notificationConsumer.set(notificationConsumer);
8594
this.requestNotificationConsumer.set(requestNotificationConsumer);
8695
this.protocolVersion.set(protocolVersion);
@@ -127,7 +136,7 @@ protected void notify(JsonRpcResponse response) {
127136
* Forwards a notification request by converting it to a response format.
128137
* Notifications have a method field but no id.
129138
*/
130-
protected void notifyRequest(JsonRpcRequest notification) {
139+
protected void notify(JsonRpcRequest notification) {
131140
var rnc = requestNotificationConsumer.get();
132141
if (rnc != null) {
133142
LOG.debug("Forwarding notification to consumer: method={}", notification.getMethod());
@@ -138,5 +147,26 @@ protected void notifyRequest(JsonRpcRequest notification) {
138147
}
139148
}
140149

150+
/**
151+
* Determines if a Document represents a notification (has "method" but no "id")
152+
* rather than a response (has "id").
153+
*
154+
* - Responses have an "id" field at the top level
155+
* - Notifications have a "method" field but no "id" field at the top level
156+
*/
157+
protected static boolean isNotification(Document doc) {
158+
try {
159+
if (!doc.isType(ShapeType.STRUCTURE) && !doc.isType(ShapeType.MAP)) {
160+
return false;
161+
}
162+
163+
// If it has a "method" field but no "id", it's a notification
164+
return doc.getMember("id") == null && doc.getMember("method") != null;
165+
} catch (Exception e) {
166+
LOG.warn("Failed to determine if notification from Document", e);
167+
return false;
168+
}
169+
}
170+
141171
public abstract String name();
142172
}

‎mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/StdioProxy.java‎

Lines changed: 35 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,25 @@ public CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request) {
116116
return future;
117117
}
118118

119+
// Notifications don't have an ID and don't expect a response
120+
if (request.getId() == null) {
121+
try {
122+
writeLock.lock();
123+
String serializedRequest = JSON_CODEC.serializeToString(request);
124+
LOG.debug("Sending notification: {}", serializedRequest);
125+
writer.write(serializedRequest);
126+
writer.newLine();
127+
writer.flush();
128+
} catch (IOException e) {
129+
LOG.error("Error sending notification to MCP server", e);
130+
return CompletableFuture.failedFuture(
131+
new RuntimeException("Failed to send notification to MCP server: " + e.getMessage(), e));
132+
} finally {
133+
writeLock.unlock();
134+
}
135+
return CompletableFuture.completedFuture(null);
136+
}
137+
119138
String requestId = getStringRequestId(request.getId());
120139
CompletableFuture<JsonRpcResponse> responseFuture = new CompletableFuture<>();
121140
pendingRequests.put(requestId, responseFuture);
@@ -185,20 +204,23 @@ public synchronized void start() {
185204
}
186205

187206
LOG.debug("Received response: {}", responseLine);
188-
JsonRpcResponse response = JsonRpcResponse.builder()
189-
.deserialize(
190-
JSON_CODEC.createDeserializer(
191-
responseLine.getBytes(StandardCharsets.UTF_8)))
192-
.build();
193-
194-
String responseId = getStringRequestId(response.getId());
195-
LOG.debug("Processing response ID: {}", responseId);
196-
197-
CompletableFuture<JsonRpcResponse> future = pendingRequests.remove(responseId);
198-
if (future != null) {
199-
future.complete(response);
207+
var output =
208+
JSON_CODEC.createDeserializer(responseLine.getBytes(StandardCharsets.UTF_8))
209+
.readDocument();
210+
if (isNotification(output)) {
211+
notify(output.asShape(JsonRpcRequest.builder()));
200212
} else {
201-
notify(response);
213+
JsonRpcResponse response = output.asShape(JsonRpcResponse.builder());
214+
215+
String responseId = getStringRequestId(response.getId());
216+
LOG.debug("Processing response ID: {}", responseId);
217+
218+
CompletableFuture<JsonRpcResponse> future = pendingRequests.remove(responseId);
219+
if (future != null) {
220+
future.complete(response);
221+
} else {
222+
notify(response);
223+
}
202224
}
203225
} catch (IOException e) {
204226
if (running) {

‎mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1431,6 +1431,10 @@ void testToolsListChangedNotificationInvalidatesCache() {
14311431
.build();
14321432
service.handleRequest(initRequest, r -> {}, ProtocolVersion.defaultVersion());
14331433

1434+
// Verify notifications/initialized was sent during initialization
1435+
assertTrue(mockProxy.getSentNotifications().contains("notifications/initialized"),
1436+
"notifications/initialized should be sent during initialization");
1437+
14341438
// First tools/list - fetches from proxy
14351439
var toolsRequest = JsonRpcRequest.builder()
14361440
.method("tools/list")
@@ -1488,6 +1492,10 @@ void testOtherNotificationsDoNotInvalidateCache() {
14881492
.build();
14891493
service.handleRequest(initRequest, r -> {}, ProtocolVersion.defaultVersion());
14901494

1495+
// Verify notifications/initialized was sent during initialization
1496+
assertTrue(mockProxy.getSentNotifications().contains("notifications/initialized"),
1497+
"notifications/initialized should be sent during initialization");
1498+
14911499
// First tools/list
14921500
var toolsRequest = JsonRpcRequest.builder()
14931501
.method("tools/list")
@@ -1516,6 +1524,7 @@ void testOtherNotificationsDoNotInvalidateCache() {
15161524

15171525
private static class CacheTestProxy extends McpServerProxy {
15181526
private final AtomicInteger callCounter;
1527+
private final List<String> sentNotifications = new ArrayList<>();
15191528

15201529
CacheTestProxy(AtomicInteger callCounter) {
15211530
this.callCounter = callCounter;
@@ -1539,6 +1548,11 @@ public List<software.amazon.smithy.java.mcp.model.PromptInfo> listPrompts() {
15391548

15401549
@Override
15411550
CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request) {
1551+
// Notifications have no ID
1552+
if (request.getId() == null) {
1553+
sentNotifications.add(request.getMethod());
1554+
return CompletableFuture.completedFuture(null);
1555+
}
15421556
return CompletableFuture.completedFuture(
15431557
JsonRpcResponse.builder()
15441558
.id(request.getId())
@@ -1547,6 +1561,10 @@ CompletableFuture<JsonRpcResponse> rpc(JsonRpcRequest request) {
15471561
.build());
15481562
}
15491563

1564+
List<String> getSentNotifications() {
1565+
return sentNotifications;
1566+
}
1567+
15501568
@Override
15511569
void start() {}
15521570

@@ -1561,7 +1579,7 @@ public String name() {
15611579
}
15621580

15631581
void sendNotification(JsonRpcRequest notification) {
1564-
notifyRequest(notification);
1582+
notify(notification);
15651583
}
15661584
}
15671585
}

0 commit comments

Comments
 (0)