Skip to content

Commit 8736e35

Browse files
yasmewadyasmewad
committed
Add validation for prompt name uniqueness across the service.
- **New `PromptUniquenessValidator`**: Detects duplicate prompt names within services - **Updated examples**: Updated the examples to use the new naming convention. - **Only Service and Operations support**: Since we are not able to parse the prompts at resource level in the MCP Server, updated the trait selector to allow only service or operation level definition. Prevents runtime conflicts when services define prompts with the same name at different levels (e.g., service vs operation). This supports the established naming convention of `serviceName__prompt-name-as-defined-in-trait` which provides namespace disambiguation and aligns with MCP specification requirements.
1 parent 8c344f8 commit 8736e35

File tree

11 files changed

+335
-15
lines changed

11 files changed

+335
-15
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use smithy.ai#prompts
1919
arguments: GetCodingStatisticsInput
2020
preferWhen: "User wants to analyze developer productivity, review coding activity, or understand technology usage patterns"
2121
}
22-
employee_lookup: {
22+
Test_employee: {
2323
description: "General employee lookup and information retrieval service"
2424
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."
2525
preferWhen: "User needs any employee-related information or wants to understand available employee data"

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,9 @@ public static Map<String, Prompt> loadPrompts(Collection<Service> services) {
4040
Map<String, Prompt> prompts = new LinkedHashMap<>();
4141

4242
for (var service : services) {
43+
var serviceName = service.schema().id().getName();
4344
Map<String, PromptTemplateDefinition> promptDefinitions = new HashMap<>();
45+
4446
var servicePromptTrait = service.schema().getTrait(PROMPTS_TRAIT_KEY);
4547
if (servicePromptTrait != null) {
4648
promptDefinitions.putAll(servicePromptTrait.getValues());
@@ -50,13 +52,16 @@ public static Map<String, Prompt> loadPrompts(Collection<Service> services) {
5052
if (operationPromptsTrait != null) {
5153
promptDefinitions.putAll(operationPromptsTrait.getValues());
5254
}
53-
5455
});
56+
5557
for (Map.Entry<String, PromptTemplateDefinition> entry : promptDefinitions.entrySet()) {
56-
var promptName = entry.getKey().toLowerCase();
58+
var originalPromptName = entry.getKey();
5759
var promptTemplateDefinition = entry.getValue();
5860
var templateString = promptTemplateDefinition.getTemplate();
5961

62+
// Generate the new prompt name with service prefix
63+
var promptName = serviceName + "__" + originalPromptName;
64+
6065
var finalTemplateString = promptTemplateDefinition.getPreferWhen().isPresent()
6166
? templateString + TOOL_PREFERENCE_PREFIX
6267
+ promptTemplateDefinition.getPreferWhen().get()

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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -388,11 +388,11 @@ void testPromptsList() {
388388

389389
// Check the prompt (service and operation have same name, so only one is returned)
390390
var servicePrompt = prompts.stream()
391-
.filter(p -> p.asStringMap().get("name").asString().equals("search_users"))
391+
.filter(p -> p.asStringMap().get("name").asString().equals("TestService__search_users"))
392392
.findFirst()
393393
.orElseThrow();
394394
var servicePromptMap = servicePrompt.asStringMap();
395-
assertEquals("search_users", servicePromptMap.get("name").asString());
395+
assertEquals("TestService__search_users", servicePromptMap.get("name").asString());
396396
assertEquals("Test Template", servicePromptMap.get("description").asString());
397397
assertTrue(servicePromptMap.get("arguments").asList().isEmpty());
398398
}
@@ -415,7 +415,7 @@ void testPromptsGetWithValidPrompt() {
415415
write("prompts/get",
416416
Document.of(Map.of(
417417
"name",
418-
Document.of("search_users"))));
418+
Document.of("TestService__search_users"))));
419419
var response = read();
420420
var result = response.getResult().asStringMap();
421421

@@ -479,7 +479,7 @@ void testPromptsGetWithTemplateArguments() {
479479
write("prompts/get",
480480
Document.of(Map.of(
481481
"name",
482-
Document.of("search_with_args"),
482+
Document.of("TestServiceWithArgs__search_with_args"),
483483
"arguments",
484484
Document.of(Map.of(
485485
"query",
@@ -520,7 +520,7 @@ void testPromptsGetWithMissingRequiredArguments() {
520520
write("prompts/get",
521521
Document.of(Map.of(
522522
"name",
523-
Document.of("search_with_args"))));
523+
Document.of("TestServiceWithArgs__search_with_args"))));
524524
var response = read();
525525
var result = response.getResult().asStringMap();
526526

@@ -557,7 +557,7 @@ void testApplyTemplateArgumentsEdgeCases() {
557557
write("prompts/get",
558558
Document.of(Map.of(
559559
"name",
560-
Document.of("empty_template"))));
560+
Document.of("TestServiceEdgeCases__empty_template"))));
561561
var response = read();
562562
var result = response.getResult().asStringMap();
563563
var messages = result.get("messages").asList();
@@ -569,7 +569,7 @@ void testApplyTemplateArgumentsEdgeCases() {
569569
write("prompts/get",
570570
Document.of(Map.of(
571571
"name",
572-
Document.of("no_placeholders"))));
572+
Document.of("TestServiceEdgeCases__no_placeholders"))));
573573
response = read();
574574
result = response.getResult().asStringMap();
575575
messages = result.get("messages").asList();
@@ -581,7 +581,7 @@ void testApplyTemplateArgumentsEdgeCases() {
581581
write("prompts/get",
582582
Document.of(Map.of(
583583
"name",
584-
Document.of("duplicate_placeholders"),
584+
Document.of("TestServiceEdgeCases__duplicate_placeholders"),
585585
"arguments",
586586
Document.of(Map.of(
587587
"name",
@@ -597,7 +597,7 @@ void testApplyTemplateArgumentsEdgeCases() {
597597
write("prompts/get",
598598
Document.of(Map.of(
599599
"name",
600-
Document.of("missing_arg_template"))));
600+
Document.of("TestServiceEdgeCases__missing_arg_template"))));
601601
response = read();
602602
result = response.getResult().asStringMap();
603603
messages = result.get("messages").asList();

‎smithy-ai-traits/model/prompts.smithy

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,18 @@ $version: "2"
22

33
namespace smithy.ai
44

5-
// Prompt template trait - applied at operation level to provide guidance to LLMs
6-
@trait(selector: ":is(service, resource, operation)")
5+
@trait(selector: ":is(service, operation)")
76
map prompts {
87
/// Name of the prompt template
9-
key: String
8+
key: PromptName
109

1110
/// Definition of the prompt template
1211
value: PromptTemplateDefinition
1312
}
1413

14+
@pattern("^[a-zA-Z0-9]+(?:_[a-zA-Z0-9]+)*$")
15+
string PromptName
16+
1517
/// Defines the structure of the prompt
1618
@private
1719
structure PromptTemplateDefinition {

‎smithy-ai-traits/src/main/java/software/amazon/smithy/ai/PromptUniquenessValidator.java

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
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.ai;
7+
8+
import java.util.ArrayList;
9+
import java.util.HashMap;
10+
import java.util.List;
11+
import java.util.Map;
12+
import software.amazon.smithy.model.Model;
13+
import software.amazon.smithy.model.shapes.OperationShape;
14+
import software.amazon.smithy.model.shapes.ServiceShape;
15+
import software.amazon.smithy.model.shapes.Shape;
16+
import software.amazon.smithy.model.validation.AbstractValidator;
17+
import software.amazon.smithy.model.validation.ValidationEvent;
18+
19+
/**
20+
* Validates that prompt names are unique within each service when compared case-insensitively.
21+
* This validator checks both service-level and operation-level prompts to ensure no conflicts.
22+
*
23+
* For example, "search_users" and "Search_Users" would be considered conflicting prompt names
24+
* because they are identical when compared case-insensitively.
25+
*/
26+
public class PromptUniquenessValidator extends AbstractValidator {
27+
28+
@Override
29+
public List<ValidationEvent> validate(Model model) {
30+
List<ValidationEvent> events = new ArrayList<>();
31+
32+
for (ServiceShape service : model.getServiceShapes()) {
33+
validatePrompts(service, model, events);
34+
}
35+
36+
return events;
37+
}
38+
39+
private void validatePrompts(ServiceShape service, Model model, List<ValidationEvent> events) {
40+
// Track prompt names to detect case-insensitive conflicts
41+
Map<String, Shape> seenPromptNames = new HashMap<>();
42+
43+
service.getTrait(PromptsTrait.class).ifPresent(promptsTrait -> {
44+
for (String promptName : promptsTrait.getValues().keySet()) {
45+
checkPromptUniqueness(promptName, service, seenPromptNames, events);
46+
}
47+
});
48+
49+
for (OperationShape operation : service.getOperations()
50+
.stream()
51+
.map(shapeId -> model.expectShape(shapeId, OperationShape.class))
52+
.toList()) {
53+
54+
operation.getTrait(PromptsTrait.class).ifPresent(promptsTrait -> {
55+
for (String promptName : promptsTrait.getValues().keySet()) {
56+
checkPromptUniqueness(promptName, operation, seenPromptNames, events);
57+
}
58+
});
59+
}
60+
}
61+
62+
private void checkPromptUniqueness(
63+
String promptName,
64+
Shape shape,
65+
Map<String, Shape> seenPromptNames,
66+
List<ValidationEvent> events
67+
) {
68+
String normalizedName = promptName.toLowerCase();
69+
70+
if (seenPromptNames.containsKey(normalizedName)) {
71+
Shape conflictingShape = seenPromptNames.get(normalizedName);
72+
events.add(error(shape,
73+
String.format(
74+
"Duplicate prompt name detected: '%s' conflicts with an existing prompt " +
75+
"defined on %s when compared case-insensitively. Prompt names must be unique " +
76+
"within a service regardless of case.",
77+
promptName,
78+
conflictingShape.getId())));
79+
} else {
80+
seenPromptNames.put(normalizedName, shape);
81+
}
82+
}
83+
}

‎smithy-ai-traits/src/main/resources/META-INF/services/software.amazon.smithy.model.validation.Validator

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
software.amazon.smithy.ai.PromptUniquenessValidator

‎smithy-ai-traits/src/test/java/software/amazon/smithy/ai/PromptUniquenessValidatorTest.java

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
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.ai;
7+
8+
import static org.junit.jupiter.api.Assertions.assertEquals;
9+
import static org.junit.jupiter.api.Assertions.assertTrue;
10+
11+
import java.util.List;
12+
import java.util.Objects;
13+
import java.util.stream.Stream;
14+
import org.junit.jupiter.api.DisplayName;
15+
import org.junit.jupiter.params.ParameterizedTest;
16+
import org.junit.jupiter.params.provider.Arguments;
17+
import org.junit.jupiter.params.provider.MethodSource;
18+
import software.amazon.smithy.model.Model;
19+
import software.amazon.smithy.model.validation.ValidatedResult;
20+
import software.amazon.smithy.model.validation.ValidationEvent;
21+
22+
public class PromptUniquenessValidatorTest {
23+
24+
/**
25+
* Test case record for prompt uniqueness validation scenarios.
26+
*
27+
* @param testName descriptive name for the test scenario
28+
* @param resourcePath path to the Smithy model file
29+
* @param expectedErrorCount number of validation errors expected
30+
* @param expectedErrorContent specific content to verify in error messages (null if no errors expected)
31+
*/
32+
public record ValidationTestCase(
33+
String testName,
34+
String resourcePath,
35+
int expectedErrorCount,
36+
String expectedErrorContent) {
37+
@Override
38+
public String toString() {
39+
return testName;
40+
}
41+
}
42+
43+
@ParameterizedTest(name = "{0}")
44+
@MethodSource("provideValidationTestCases")
45+
@DisplayName("Prompt uniqueness validation scenarios")
46+
public void testPromptUniquenessValidation(ValidationTestCase testCase) {
47+
var result = assembleModelFromResource(testCase.resourcePath());
48+
List<ValidationEvent> duplicateEvents = filterDuplicatePromptEvents(result.getValidationEvents());
49+
50+
assertEquals(testCase.expectedErrorCount(),
51+
duplicateEvents.size(),
52+
String.format("Test '%s': Expected %d validation errors",
53+
testCase.testName(),
54+
testCase.expectedErrorCount()));
55+
56+
if (testCase.expectedErrorCount() > 0) {
57+
ValidationEvent event = duplicateEvents.get(0);
58+
assertTrue(event.getMessage().contains("Duplicate prompt name detected"),
59+
String.format("Test '%s': Should contain duplicate prompt error message", testCase.testName()));
60+
61+
if (testCase.expectedErrorContent() != null) {
62+
assertTrue(event.getMessage().contains(testCase.expectedErrorContent()),
63+
String.format("Test '%s': Should contain expected error content: %s",
64+
testCase.testName(),
65+
testCase.expectedErrorContent()));
66+
}
67+
68+
assertTrue(event.getMessage().contains("case-insensitively"),
69+
String.format("Test '%s': Should mention case-insensitive comparison", testCase.testName()));
70+
}
71+
}
72+
73+
/**
74+
* Provides test cases for prompt uniqueness validation scenarios.
75+
* Each test case includes: test name, resource path, expected error count, and expected error content.
76+
*/
77+
private static Stream<Arguments> provideValidationTestCases() {
78+
return Stream.of(
79+
new ValidationTestCase(
80+
"No duplicate prompts",
81+
"/unique-prompts.smithy",
82+
0,
83+
null),
84+
new ValidationTestCase(
85+
"Duplicate prompts within service",
86+
"/service-level-duplicates.smithy",
87+
1,
88+
"Search_Users"),
89+
new ValidationTestCase(
90+
"Duplicate prompts between service and operation",
91+
"/duplicate-prompts.smithy",
92+
1,
93+
"Search_Users"),
94+
new ValidationTestCase(
95+
"Same prompt names across different services (should be allowed)",
96+
"/cross-service-unique.smithy",
97+
0,
98+
null))
99+
.map(Arguments::of);
100+
}
101+
102+
/**
103+
* Helper method to assemble a Smithy model from a test resource file.
104+
* Reduces code duplication across test methods.
105+
*/
106+
private ValidatedResult<Model> assembleModelFromResource(String resourcePath) {
107+
return Model.assembler()
108+
.addImport(Objects.requireNonNull(getClass().getResource(resourcePath)))
109+
.discoverModels(getClass().getClassLoader())
110+
.assemble();
111+
}
112+
113+
/**
114+
* Helper method to filter validation events for duplicate prompt name errors.
115+
* Centralizes the filtering logic to reduce code duplication and improve maintainability.
116+
*/
117+
private List<ValidationEvent> filterDuplicatePromptEvents(List<ValidationEvent> events) {
118+
return events.stream()
119+
.filter(event -> event.getMessage().contains("Duplicate prompt name detected"))
120+
.toList();
121+
}
122+
}

‎smithy-ai-traits/src/test/resources/cross-service-unique.smithy

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
$version: "2.0"
2+
3+
namespace com.example.test.cross
4+
5+
use smithy.ai#prompts
6+
7+
// Test model for validating that same prompt names are allowed across different services
8+
@prompts(
9+
"search_users": {
10+
template: "Search for users in ServiceOne"
11+
description: "ServiceOne search template"
12+
}
13+
)
14+
service ServiceOne {
15+
operations: []
16+
}
17+
18+
@prompts(
19+
"search_users": {
20+
template: "Search for users in ServiceTwo"
21+
description: "ServiceTwo search template"
22+
}
23+
)
24+
service ServiceTwo {
25+
operations: []
26+
}

0 commit comments

Comments
 (0)