|
1 | 1 | package com.github.thought2code.mcp.annotated.server.component; |
2 | 2 |
|
3 | 3 | import com.github.thought2code.mcp.annotated.McpApplicationContext; |
| 4 | +import com.github.thought2code.mcp.annotated.exception.McpServerComponentRegistrationException; |
4 | 5 | import io.modelcontextprotocol.server.McpAsyncServer; |
5 | 6 | import io.modelcontextprotocol.server.McpSyncServer; |
6 | 7 | import java.lang.annotation.Annotation; |
7 | 8 | import java.lang.reflect.Method; |
8 | 9 | import java.util.Set; |
9 | 10 | import org.slf4j.Logger; |
10 | 11 | import org.slf4j.LoggerFactory; |
| 12 | +import reactor.core.publisher.Mono; |
11 | 13 |
|
12 | 14 | /** |
13 | 15 | * Template base class for component registrars that support both sync and async server modes. |
@@ -60,8 +62,9 @@ public abstract class AbstractDualModeComponentRegistrar<S, A> implements McpCom |
60 | 62 | * |
61 | 63 | * @param server target async server |
62 | 64 | * @param specification specification to register |
| 65 | + * @return registration completion signal |
63 | 66 | */ |
64 | | - protected abstract void addAsyncSpecification(McpAsyncServer server, A specification); |
| 67 | + protected abstract Mono<Void> addAsyncSpecification(McpAsyncServer server, A specification); |
65 | 68 |
|
66 | 69 | /** |
67 | 70 | * @param specification registered sync specification |
@@ -95,9 +98,38 @@ public final void register(McpAsyncServer server, McpApplicationContext context) |
95 | 98 | final String annotationTypeName = annotationType().getSimpleName(); |
96 | 99 | log.debug("Registering async {} method: {}", annotationTypeName, method.toGenericString()); |
97 | 100 | A specification = createAsyncSpecification(method, context); |
98 | | - addAsyncSpecification(server, specification); |
99 | 101 | final String specificationName = asyncSpecificationName(specification); |
| 102 | + Mono<Void> registration = addAsyncSpecification(server, specification); |
| 103 | + awaitAsyncRegistration(registration, annotationTypeName, specificationName); |
100 | 104 | log.debug("Async {} {} registered successfully", annotationTypeName, specificationName); |
101 | 105 | } |
102 | 106 | } |
| 107 | + |
| 108 | + /** |
| 109 | + * Waits for an asynchronous component registration to complete during server startup. |
| 110 | + * |
| 111 | + * <p>Async MCP server APIs return a {@link Mono} that completes when registration finishes. This |
| 112 | + * method blocks on that signal so registration behaves like the synchronous path: each component |
| 113 | + * is fully registered before the next one is processed and before the server starts handling |
| 114 | + * requests. |
| 115 | + * |
| 116 | + * <p>If registration fails, the error is logged and rethrown as a {@link |
| 117 | + * McpServerComponentRegistrationException} with component context. |
| 118 | + * |
| 119 | + * @param registration registration completion signal from {@link #addAsyncSpecification} |
| 120 | + * @param componentType annotation type name used in error messages (for example, {@code McpTool}) |
| 121 | + * @param specificationName registered component name used in error messages |
| 122 | + * @throws McpServerComponentRegistrationException if registration fails |
| 123 | + */ |
| 124 | + private void awaitAsyncRegistration( |
| 125 | + Mono<Void> registration, String componentType, String specificationName) { |
| 126 | + try { |
| 127 | + registration.block(); |
| 128 | + } catch (RuntimeException e) { |
| 129 | + final String message = |
| 130 | + String.format("Failed to register async %s %s", componentType, specificationName); |
| 131 | + log.error(message, e); |
| 132 | + throw new McpServerComponentRegistrationException(message, e); |
| 133 | + } |
| 134 | + } |
103 | 135 | } |
0 commit comments