Skip to content

Commit 535235b

Browse files
m0verm0ver
committed
Checks for missing or invalid sessionId and returns a proper JSON-RPC error
1 parent 352a2f3 commit 535235b

File tree

1 file changed

+40
-13
lines changed

1 file changed

+40
-13
lines changed

‎src/main/java/org/tinystruct/mcp/MCPApplication.java

Lines changed: 40 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.tinystruct.http.SSEPushManager;
1111
import org.tinystruct.system.annotation.Action;
1212
import org.tinystruct.system.annotation.Argument;
13+
import org.tinystruct.http.Header;
1314

1415
import java.util.UUID;
1516
import java.util.logging.Level;
@@ -27,26 +28,26 @@
2728
public abstract class MCPApplication extends AbstractApplication {
2829
private static final Logger LOGGER = Logger.getLogger(MCPApplication.class.getName());
2930

30-
protected SSEHandler sseHandler;
3131
protected JsonRpcHandler jsonRpcHandler;
3232
protected AuthorizationHandler authHandler;
3333

3434
private boolean initialized = false;
3535
protected SessionState sessionState = SessionState.DISCONNECTED;
36-
private final String sessionId = UUID.randomUUID().toString();
36+
protected final Map<String, Object> sessionMap = new ConcurrentHashMap<>(); // sessionId -> user info or state
3737

3838
// Generic registries for tools, resources, prompts, and custom RPC handlers
3939
protected final Map<String, MCPTool> tools = new java.util.concurrent.ConcurrentHashMap<>();
4040
protected final Map<String, MCPDataResource> resources = new java.util.concurrent.ConcurrentHashMap<>();
4141
protected final Map<String, MCPPrompt> prompts = new java.util.concurrent.ConcurrentHashMap<>();
4242
protected final Map<String, RpcMethodHandler> rpcHandlers = new java.util.concurrent.ConcurrentHashMap<>();
4343

44-
@Override
44+
4545
/**
4646
* Initializes the MCP application, setting up authentication, SSE, JSON-RPC handler,
4747
* and registering core protocol handlers. Subclasses should call super.init() and
4848
* may register additional handlers for custom protocol extensions.
4949
*/
50+
@Override
5051
public void init() {
5152
this.setTemplateRequired(false);
5253
// Generate a random auth token if not configured
@@ -56,7 +57,6 @@ public void init() {
5657
LOGGER.info("Generated new MCP auth token: " + token);
5758
}
5859
this.authHandler = new AuthorizationHandler(getConfiguration().get(Config.AUTH_TOKEN));
59-
this.sseHandler = new SSEHandler();
6060
this.jsonRpcHandler = new JsonRpcHandler();
6161

6262
// Register core protocol handlers
@@ -92,7 +92,7 @@ private boolean isAllowedPreInitialization(String method) {
9292
}
9393

9494
/**
95-
* Main entry point for handling JSON-RPC requests.
95+
* Main entry point for handling JSON-RPC requests via Streamable HTTP POST.
9696
* Authenticates, parses, dispatches, and returns the response.
9797
*
9898
* @param request The HTTP request
@@ -105,11 +105,23 @@ public String handleRpcRequest(Request request, Response response) throws Applic
105105
try {
106106
// Validate authentication
107107
authHandler.validateAuthHeader(request);
108-
108+
// Session management: extract or assign sessionId
109+
String sessionId = (String) request.headers().get(Header.value0f("Mcp-Session-Id"));
110+
if (sessionId == null || sessionId.isEmpty()) {
111+
// If this is an initialize request, assign a new sessionId
112+
String jsonStr = request.body();
113+
if (jsonStr.contains("\"method\":\"initialize\"")) {
114+
sessionId = request.getSession().getId();
115+
response.addHeader("Mcp-Session-Id", sessionId);
116+
sessionMap.put(sessionId, System.currentTimeMillis()); // Store session state as needed
117+
} else {
118+
response.setStatus(ResponseStatus.UNAUTHORIZED);
119+
return jsonRpcHandler.createErrorResponse("Missing session ID", ErrorCodes.UNAUTHORIZED);
120+
}
121+
}
109122
// Parse the JSON-RPC request
110123
String jsonStr = request.body();
111124
LOGGER.fine("Received JSON: " + jsonStr);
112-
113125
// Add batch request support
114126
if (jsonStr.trim().startsWith("[")) {
115127
return jsonRpcHandler.handleBatchRequest(jsonStr, (rpcReq, rpcRes) -> {
@@ -121,15 +133,12 @@ public String handleRpcRequest(Request request, Response response) throws Applic
121133
}
122134
});
123135
}
124-
125136
if (!jsonRpcHandler.validateJsonRpcRequest(jsonStr)) {
126137
response.setStatus(ResponseStatus.BAD_REQUEST);
127138
return jsonRpcHandler.createErrorResponse("Invalid JSON-RPC request", ErrorCodes.INVALID_REQUEST);
128139
}
129-
130140
JsonRpcRequest rpcRequest = new JsonRpcRequest();
131141
rpcRequest.parse(jsonStr);
132-
133142
JsonRpcResponse jsonResponse = new JsonRpcResponse();
134143
// Restrict methods before READY
135144
String method = rpcRequest.getMethod();
@@ -143,6 +152,8 @@ public String handleRpcRequest(Request request, Response response) throws Applic
143152
jsonResponse.setError(new JsonRpcError(ErrorCodes.METHOD_NOT_FOUND, "Method not found: " + method));
144153
}
145154
}
155+
// For now, always return JSON. SSE streaming support to be added in GET endpoint.
156+
// response.addHeader("Content-Type", "application/json");
146157
return jsonResponse.toString();
147158
} catch (SecurityException e) {
148159
response.setStatus(ResponseStatus.UNAUTHORIZED);
@@ -235,7 +246,6 @@ protected void handleShutdown(JsonRpcRequest request, JsonRpcResponse response)
235246
return;
236247
}
237248

238-
sseHandler.closeAll();
239249
sessionState = SessionState.DISCONNECTED;
240250

241251
response.setId(request.getId());
@@ -250,11 +260,28 @@ protected void handleShutdown(JsonRpcRequest request, JsonRpcResponse response)
250260
* @param response The JSON-RPC response to populate
251261
*/
252262
protected void handleGetStatus(JsonRpcRequest request, JsonRpcResponse response) {
263+
String sessionId = null;
264+
if (request.getParams() != null && request.getParams().get("sessionId") != null) {
265+
sessionId = request.getParams().get("sessionId").toString();
266+
}
253267
Builder result = new Builder();
268+
if (sessionId == null || !sessionMap.containsKey(sessionId)) {
269+
response.setError(new JsonRpcError(ErrorCodes.UNAUTHORIZED, "Invalid or missing sessionId"));
270+
return;
271+
}
272+
Object sessionStartObj = sessionMap.get(sessionId);
273+
long sessionStart;
274+
if (sessionStartObj instanceof Long) {
275+
sessionStart = (Long) sessionStartObj;
276+
} else if (sessionStartObj instanceof Number) {
277+
sessionStart = ((Number) sessionStartObj).longValue();
278+
} else {
279+
response.setError(new JsonRpcError(ErrorCodes.INTERNAL_ERROR, "Session start time invalid"));
280+
return;
281+
}
254282
result.put("state", sessionState.toString());
255283
result.put("sessionId", sessionId);
256-
result.put("uptime", System.currentTimeMillis() - sseHandler.getFirstSessionCreatedAt());
257-
284+
result.put("uptime", System.currentTimeMillis() - sessionStart);
258285
response.setId(request.getId());
259286
response.setResult(result);
260287
}

0 commit comments

Comments
 (0)