From 0e80b4de5e04848de43b85b619e1d36cd9358f81 Mon Sep 17 00:00:00 2001 From: sanskritipatole Date: Fri, 19 Dec 2025 16:37:21 -0800 Subject: [PATCH] Add MCP proxy server support for prompts --- .../smithy/java/mcp/server/McpService.java | 37 ++++++++++++++++--- 1 file changed, 32 insertions(+), 5 deletions(-) diff --git a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java index e97a7ebb0..345d25db2 100644 --- a/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java +++ b/mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpService.java @@ -220,8 +220,23 @@ private JsonRpcResponse handleInitialize(JsonRpcRequest req) { } private JsonRpcResponse handlePromptsList(JsonRpcRequest req) { + var allPrompts = new ArrayList<>(prompts.values().stream().map(Prompt::promptInfo).toList()); + + // Add prompts from proxy servers + for (McpServerProxy proxy : proxies.values()) { + try { + var response = proxy.rpc(req).join(); + if (response.getError() == null) { + var proxyPrompts = response.getResult().asShape(ListPromptsResult.builder()).getPrompts(); + allPrompts.addAll(proxyPrompts); + } + } catch (Exception e) { + LOG.error("Failed to fetch prompts from proxy: " + proxy.name(), e); + } + } + var result = ListPromptsResult.builder() - .prompts(prompts.values().stream().map(Prompt::promptInfo).toList()) + .prompts(allPrompts) .build(); return createSuccessResponse(req.getId(), result); } @@ -232,12 +247,24 @@ private JsonRpcResponse handlePromptsGet(JsonRpcRequest req) { var prompt = prompts.get(normalize(promptName)); - if (prompt == null) { - throw new RuntimeException("Prompt not found: " + promptName); + if (prompt != null) { + var result = promptProcessor.buildPromptResult(prompt, promptArguments); + return createSuccessResponse(req.getId(), result); } - var result = promptProcessor.buildPromptResult(prompt, promptArguments); - return createSuccessResponse(req.getId(), result); + // Try proxy servers + for (McpServerProxy proxy : proxies.values()) { + try { + var response = proxy.rpc(req).join(); + if (response.getError() == null) { + return createSuccessResponse(req.getId(), response.getResult()); + } + } catch (Exception e) { + LOG.error("Failed to get prompt from proxy: " + proxy.name(), e); + } + } + + throw new RuntimeException("Prompt not found: " + promptName); } private JsonRpcResponse handleToolsList(JsonRpcRequest req, ProtocolVersion protocolVersion) {