Skip to content

Commit aaf86c3

Browse files
fix: await async component registration during startup
Replace fire-and-forget subscribe() with centralized block-and-error handling so async tool, prompt, and resource registration completes before the server starts. Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 28420d4 commit aaf86c3

4 files changed

Lines changed: 40 additions & 8 deletions

File tree

src/main/java/com/github/thought2code/mcp/annotated/server/component/AbstractDualModeComponentRegistrar.java

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
package com.github.thought2code.mcp.annotated.server.component;
22

33
import com.github.thought2code.mcp.annotated.McpApplicationContext;
4+
import com.github.thought2code.mcp.annotated.exception.McpServerComponentRegistrationException;
45
import io.modelcontextprotocol.server.McpAsyncServer;
56
import io.modelcontextprotocol.server.McpSyncServer;
67
import java.lang.annotation.Annotation;
78
import java.lang.reflect.Method;
89
import java.util.Set;
910
import org.slf4j.Logger;
1011
import org.slf4j.LoggerFactory;
12+
import reactor.core.publisher.Mono;
1113

1214
/**
1315
* 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
6062
*
6163
* @param server target async server
6264
* @param specification specification to register
65+
* @return registration completion signal
6366
*/
64-
protected abstract void addAsyncSpecification(McpAsyncServer server, A specification);
67+
protected abstract Mono<Void> addAsyncSpecification(McpAsyncServer server, A specification);
6568

6669
/**
6770
* @param specification registered sync specification
@@ -95,9 +98,38 @@ public final void register(McpAsyncServer server, McpApplicationContext context)
9598
final String annotationTypeName = annotationType().getSimpleName();
9699
log.debug("Registering async {} method: {}", annotationTypeName, method.toGenericString());
97100
A specification = createAsyncSpecification(method, context);
98-
addAsyncSpecification(server, specification);
99101
final String specificationName = asyncSpecificationName(specification);
102+
Mono<Void> registration = addAsyncSpecification(server, specification);
103+
awaitAsyncRegistration(registration, annotationTypeName, specificationName);
100104
log.debug("Async {} {} registered successfully", annotationTypeName, specificationName);
101105
}
102106
}
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+
}
103135
}

src/main/java/com/github/thought2code/mcp/annotated/server/component/McpServerPrompt.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -130,9 +130,9 @@ protected void addSyncSpecification(
130130
}
131131

132132
@Override
133-
protected void addAsyncSpecification(
133+
protected Mono<Void> addAsyncSpecification(
134134
McpAsyncServer server, McpServerFeatures.AsyncPromptSpecification specification) {
135-
server.addPrompt(specification).subscribe();
135+
return server.addPrompt(specification);
136136
}
137137

138138
@Override

src/main/java/com/github/thought2code/mcp/annotated/server/component/McpServerResource.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -115,9 +115,9 @@ protected void addSyncSpecification(
115115
}
116116

117117
@Override
118-
protected void addAsyncSpecification(
118+
protected Mono<Void> addAsyncSpecification(
119119
McpAsyncServer server, McpServerFeatures.AsyncResourceSpecification specification) {
120-
server.addResource(specification).subscribe();
120+
return server.addResource(specification);
121121
}
122122

123123
@Override

src/main/java/com/github/thought2code/mcp/annotated/server/component/McpServerTool.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -143,9 +143,9 @@ protected void addSyncSpecification(
143143
}
144144

145145
@Override
146-
protected void addAsyncSpecification(
146+
protected Mono<Void> addAsyncSpecification(
147147
McpAsyncServer server, McpServerFeatures.AsyncToolSpecification specification) {
148-
server.addTool(specification).subscribe();
148+
return server.addTool(specification);
149149
}
150150

151151
@Override

0 commit comments

Comments
 (0)