Skip to content

Commit 91eed0b

Browse files
yasmewadyasmewad
authored
Add Prompts support for MCP Server (#825)
This adds prompt support to the MCP server, letting LLMs get helpful context directly from Smithy models. What's New Core Features: • Define prompts in Smithy using @prompts trait • MCP server automatically exposes them via prompts/list and prompts/get • Template support with {{placeholder}} syntax for dynamic content. Key Components: • PromptLoader - discovers prompts from Smithy models • PromptProcessor - handles template substitution and validation Testing Added unit tests for various cases of PromptLoader, PromptProcessor and the MCP Server itself. Added example in ProxyMcpServerExample and ran it with Q CLI. Will add more testing examples in next revision after discussion has settled.
1 parent 5a1ebdc commit 91eed0b

File tree

13 files changed

+1011
-86
lines changed

13 files changed

+1011
-86
lines changed

‎examples/mcp-server/build.gradle.kts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ dependencies {
1010
val smithyJavaVersion: String by project
1111

1212
smithyBuild("software.amazon.smithy.java:plugins:$smithyJavaVersion")
13+
implementation("software.amazon.smithy.java:mcp-traits:$smithyJavaVersion")
1314
implementation("software.amazon.smithy.java:mcp-server:$smithyJavaVersion")
1415
implementation("software.amazon.smithy.java:server-proxy:$smithyJavaVersion")
1516
implementation("software.amazon.smithy.java:server-netty:$smithyJavaVersion")

‎examples/mcp-server/smithy-build.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
"service": "smithy.example.mcp#EmployeeService",
77
"namespace": "software.amazon.smithy.java.example.server.mcp",
88
"headerFile": "license.txt",
9-
"runtimeTraits": ["smithy.api#documentation", "smithy.api#examples" ]
9+
"runtimeTraits": ["smithy.api#documentation", "smithy.api#examples", "amazon.smithy.llm#prompts" ]
1010
}
1111
}
1212
}

‎examples/mcp-server/src/main/resources/software/amazon/smithy/java/example/server/mcp/main.smithy

Lines changed: 40 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,30 +3,58 @@ $version: "2"
33
namespace smithy.example.mcp
44

55
use aws.protocols#restJson1
6+
use amazon.smithy.llm#prompts
67

78
@restJson1
9+
@prompts({
10+
get_employee_info: {
11+
description: "Retrieve detailed information about an employee by their login ID"
12+
template: "Get employee details for login ID {{loginId}}. This will return the employee's name and manager information."
13+
arguments: GetEmployeeDetailsInput
14+
preferWhen: "User needs to look up employee information, find who someone reports to, or verify employee details"
15+
}
16+
get_coding_stats: {
17+
description: "Retrieve coding statistics and commit information for an employee"
18+
template: "Get coding statistics for employee {{login}}. This returns commit counts by programming language."
19+
arguments: GetCodingStatisticsInput
20+
preferWhen: "User wants to analyze developer productivity, review coding activity, or understand technology usage patterns"
21+
}
22+
employee_lookup: {
23+
description: "General employee lookup and information retrieval service"
24+
template: "This service provides employee information including personal details and coding statistics. Use get_employee_info for basic details or get_coding_stats for development metrics."
25+
preferWhen: "User needs any employee-related information or wants to understand available employee data"
26+
}
27+
})
828
service EmployeeService {
929
operations: [
1030
GetEmployeeDetails
1131
GetCodingStatistics
1232
]
1333
}
1434

35+
@prompts({
36+
get_employee_info_operation: {
37+
description: "Retrieve detailed information about an employee by their login ID on the operation."
38+
template: "Get employee details for login ID {{loginId}}. This will return the employee's name and manager information."
39+
arguments: GetEmployeeDetailsInput
40+
preferWhen: "User needs to look up employee information, find who someone reports to, or verify employee details"
41+
}
42+
})
1543
@documentation("Get employee information by login id")
1644
@http(method: "POST", uri: "/get-employee-details")
1745
operation GetEmployeeDetails {
18-
input := {
19-
@required
20-
loginId: LoginId
21-
}
22-
46+
input: GetEmployeeDetailsInput
2347
output: Employee
24-
2548
errors: [
2649
NoSuchUserException
2750
]
2851
}
2952

53+
structure GetEmployeeDetailsInput {
54+
@required
55+
loginId: LoginId
56+
}
57+
3058
structure Employee {
3159
@documentation("Name of the employee.")
3260
name: String
@@ -38,14 +66,15 @@ structure Employee {
3866
@documentation("Get coding statistics of an employee.")
3967
@http(method: "POST", uri: "/get-coding-statistics")
4068
operation GetCodingStatistics {
41-
input := {
42-
@required
43-
login: LoginId
44-
}
45-
69+
input: GetCodingStatisticsInput
4670
output: CodingStatistics
4771
}
4872

73+
structure GetCodingStatisticsInput {
74+
@required
75+
login: LoginId
76+
}
77+
4978
@documentation("Coding statistics of a user.")
5079
structure CodingStatistics {
5180
@documentation("Map of number of commits made per language. This can be empty or null.")

‎mcp/mcp-cli/src/main/java/software/amazon/smithy/java/mcp/cli/commands/StartServer.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,11 @@ public void execute(ExecutionContext context) throws IOException {
164164
}
165165

166166
this.mcpServer =
167-
(McpServer) McpServer.builder().stdio().addServices(services).name("smithy-mcp-server").build();
167+
(McpServer) McpServer.builder()
168+
.stdio()
169+
.addServices(services)
170+
.name("smithy-mcp-server")
171+
.build();
168172
mcpServer.start();
169173
awaitCompletion = mcpServer::awaitCompletion;
170174
shutdownMethod = mcpServer::shutdown;
@@ -185,6 +189,7 @@ public void execute(ExecutionContext context) throws IOException {
185189
private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
186190
Service service =
187191
McpBundles.getService(ConfigUtils.getMcpBundle(bundleConfig.getName()));
192+
188193
if (bundleConfig.hasAllowListedTools() || bundleConfig.hasBlockListedTools()) {
189194
var filter = OperationFilters.allowList(bundleConfig.getAllowListedTools())
190195
.and(OperationFilters.blockList(bundleConfig.getBlockListedTools()));

‎mcp/mcp-schemas/model/main.smithy

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,3 +182,68 @@ structure TextContent {
182182

183183
text: String
184184
}
185+
186+
structure ListPromptsResult {
187+
prompts: PromptInfoList
188+
}
189+
190+
list PromptInfoList {
191+
member: PromptInfo
192+
}
193+
194+
structure PromptArgument {
195+
@required
196+
name: String
197+
198+
description: String
199+
200+
required: Boolean
201+
}
202+
203+
list PromptArgumentList {
204+
member: PromptArgument
205+
}
206+
207+
structure PromptInfo {
208+
@required
209+
name: String
210+
211+
title: String
212+
213+
description: String
214+
215+
arguments: PromptArgumentList
216+
}
217+
218+
structure PromptMessageContent {
219+
@required
220+
type: PromptMessageContentType = "text"
221+
222+
text: String
223+
}
224+
225+
enum PromptMessageContentType {
226+
TEXT = "text"
227+
}
228+
229+
structure PromptMessage {
230+
@required
231+
role: String
232+
233+
@required
234+
content: PromptMessageContent
235+
}
236+
237+
enum PromptRole {
238+
USER = "user"
239+
ASSISTANT = "assistant"
240+
}
241+
242+
list PromptMessageList {
243+
member: PromptMessage
244+
}
245+
246+
structure GetPromptResult {
247+
description: String
248+
messages: PromptMessageList
249+
}

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,9 @@
4141
import software.amazon.smithy.java.mcp.model.JsonRpcErrorResponse;
4242
import software.amazon.smithy.java.mcp.model.JsonRpcRequest;
4343
import software.amazon.smithy.java.mcp.model.JsonRpcResponse;
44+
import software.amazon.smithy.java.mcp.model.ListPromptsResult;
4445
import software.amazon.smithy.java.mcp.model.ListToolsResult;
46+
import software.amazon.smithy.java.mcp.model.Prompts;
4547
import software.amazon.smithy.java.mcp.model.ServerInfo;
4648
import software.amazon.smithy.java.mcp.model.TextContent;
4749
import software.amazon.smithy.java.mcp.model.ToolInfo;
@@ -66,6 +68,8 @@ public final class McpServer implements Server {
6668
.build();
6769

6870
private final Map<String, Tool> tools;
71+
private final Map<String, Prompt> prompts;
72+
private final PromptProcessor promptProcessor;
6973
private final Thread listener;
7074
private final InputStream is;
7175
private final OutputStream os;
@@ -75,6 +79,8 @@ public final class McpServer implements Server {
7579

7680
McpServer(McpServerBuilder builder) {
7781
this.tools = createTools(builder.serviceList);
82+
this.prompts = PromptLoader.loadPrompts(builder.serviceList);
83+
this.promptProcessor = new PromptProcessor();
7884
this.is = builder.is;
7985
this.os = builder.os;
8086
this.name = builder.name;
@@ -113,12 +119,31 @@ private void handleRequest(JsonRpcRequest req) {
113119
InitializeResult.builder()
114120
.capabilities(Capabilities.builder()
115121
.tools(Tools.builder().listChanged(true).build())
122+
.prompts(Prompts.builder().listChanged(true).build())
116123
.build())
117124
.serverInfo(ServerInfo.builder()
118125
.name(name)
119126
.version("1.0.0")
120127
.build())
121128
.build());
129+
case "prompts/list" -> writeResponse(req.getId(),
130+
ListPromptsResult.builder()
131+
.prompts(prompts.values().stream().map(Prompt::promptInfo).toList())
132+
.build());
133+
case "prompts/get" -> {
134+
var promptName = req.getParams().getMember("name").asString();
135+
var promptArguments = req.getParams().getMember("arguments");
136+
137+
var prompt = prompts.get(promptName);
138+
139+
if (prompt == null) {
140+
internalError(req, new RuntimeException("Prompt not found: " + promptName));
141+
return;
142+
}
143+
144+
var result = promptProcessor.buildPromptResult(prompt, promptArguments);
145+
writeResponse(req.getId(), result);
146+
}
122147
case "tools/list" -> writeResponse(req.getId(),
123148
ListToolsResult.builder().tools(tools.values().stream().map(Tool::toolInfo).toList()).build());
124149
case "tools/call" -> {

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
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.mcp.model.PromptInfo;
9+
10+
public record Prompt(PromptInfo promptInfo, String promptTemplate) {}

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

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
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.ArrayList;
9+
import java.util.HashMap;
10+
import java.util.LinkedHashMap;
11+
import java.util.List;
12+
import java.util.Map;
13+
import software.amazon.smithy.ai.PromptTemplateDefinition;
14+
import software.amazon.smithy.ai.PromptsTrait;
15+
import software.amazon.smithy.java.core.schema.Schema;
16+
import software.amazon.smithy.java.core.schema.TraitKey;
17+
import software.amazon.smithy.java.mcp.model.PromptArgument;
18+
import software.amazon.smithy.java.mcp.model.PromptInfo;
19+
import software.amazon.smithy.java.server.Service;
20+
import software.amazon.smithy.utils.SmithyUnstableApi;
21+
import software.amazon.smithy.utils.StringUtils;
22+
23+
/**
24+
* Handles loading and parsing of prompts from Smithy models.
25+
*/
26+
@SmithyUnstableApi
27+
final class PromptLoader {
28+
29+
private static final TraitKey<PromptsTrait> PROMPTS_TRAIT_KEY = TraitKey.get(PromptsTrait.class);
30+
31+
public static final String TOOL_PREFERENCE_PREFIX = ".Tool preference: ";
32+
33+
/**
34+
* Loads prompts from the provided Smithy models.
35+
*
36+
* @return Map of prompt names to PromptInfo objects
37+
*/
38+
public static Map<String, Prompt> loadPrompts(List<Service> services) {
39+
Map<String, Prompt> prompts = new LinkedHashMap<>();
40+
41+
for (var service : services) {
42+
Map<String, PromptTemplateDefinition> promptDefinitions = new HashMap<>();
43+
var servicePromptTrait = service.schema().getTrait(PROMPTS_TRAIT_KEY);
44+
if (servicePromptTrait != null) {
45+
promptDefinitions.putAll(servicePromptTrait.getValues());
46+
}
47+
service.getAllOperations().forEach(operation -> {
48+
var operationPromptsTrait = operation.getApiOperation().schema().getTrait(PROMPTS_TRAIT_KEY);
49+
if (operationPromptsTrait != null) {
50+
promptDefinitions.putAll(operationPromptsTrait.getValues());
51+
}
52+
53+
});
54+
for (Map.Entry<String, PromptTemplateDefinition> entry : promptDefinitions.entrySet()) {
55+
var promptName = entry.getKey().toLowerCase();
56+
var promptTemplateDefinition = entry.getValue();
57+
var templateString = promptTemplateDefinition.getTemplate();
58+
59+
var finalTemplateString = promptTemplateDefinition.getPreferWhen().isPresent()
60+
? templateString + TOOL_PREFERENCE_PREFIX
61+
+ promptTemplateDefinition.getPreferWhen().get()
62+
: templateString;
63+
64+
var promptInfo = PromptInfo
65+
.builder()
66+
.name(promptName)
67+
.title(StringUtils.capitalize(promptName))
68+
.description(promptTemplateDefinition.getDescription())
69+
.arguments(promptTemplateDefinition.getArguments().isPresent()
70+
? convertArgumentShapeToPromptArgument(
71+
service.schemaIndex()
72+
.getSchema(promptTemplateDefinition.getArguments().get()))
73+
: List.of())
74+
.build();
75+
76+
prompts.put(
77+
promptName,
78+
new Prompt(promptInfo, finalTemplateString));
79+
}
80+
}
81+
return prompts;
82+
}
83+
84+
/**
85+
* Converts a Smithy structure shape to a list of PromptArgument objects.
86+
*
87+
* @param argument The ShapeId of the structure to convert
88+
* @return List of PromptArgument objects representing the structure members
89+
*/
90+
public static List<PromptArgument> convertArgumentShapeToPromptArgument(Schema argument) {
91+
List<PromptArgument> promptArguments = new ArrayList<>();
92+
93+
for (var member : argument.members()) {
94+
String memberName = member.memberName();
95+
96+
// Get description from documentation trait, use empty string if not present
97+
String description = "";
98+
var documentationTrait = member.getTrait(TraitKey.DOCUMENTATION_TRAIT);
99+
if (documentationTrait != null) {
100+
description = documentationTrait.getValue();
101+
}
102+
103+
// Check if member is required
104+
boolean isRequired = member.getTrait(TraitKey.REQUIRED_TRAIT) != null;
105+
106+
// Build the PromptArgument
107+
PromptArgument promptArgument = PromptArgument.builder()
108+
.name(memberName)
109+
.description(description)
110+
.required(isRequired)
111+
.build();
112+
113+
promptArguments.add(promptArgument);
114+
}
115+
116+
return promptArguments;
117+
}
118+
}

0 commit comments

Comments
 (0)