|
7 | 7 | import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientRequestCustomizer; |
8 | 8 | import io.modelcontextprotocol.client.transport.customizer.McpSyncHttpClientRequestCustomizer; |
9 | 9 | import io.modelcontextprotocol.common.McpTransportContext; |
| 10 | +import io.modelcontextprotocol.spec.HttpHeaders; |
10 | 11 | import io.modelcontextprotocol.spec.McpSchema; |
11 | 12 | import io.modelcontextprotocol.spec.McpTransportSessionClosedException; |
12 | 13 | import io.modelcontextprotocol.spec.ProtocolVersions; |
13 | 14 | import java.net.URI; |
14 | 15 | import java.net.URISyntaxException; |
| 16 | +import java.net.http.HttpRequest; |
15 | 17 | import java.util.Map; |
16 | 18 | import java.util.function.Consumer; |
17 | 19 | import java.util.function.Function; |
18 | 20 |
|
19 | 21 | import org.junit.jupiter.api.AfterAll; |
20 | 22 | import org.junit.jupiter.api.BeforeAll; |
21 | 23 | import org.junit.jupiter.api.Test; |
| 24 | +import org.mockito.ArgumentCaptor; |
22 | 25 | import org.testcontainers.containers.GenericContainer; |
23 | 26 | import org.testcontainers.containers.wait.strategy.Wait; |
24 | 27 | import reactor.core.publisher.Mono; |
25 | 28 | import reactor.test.StepVerifier; |
26 | 29 |
|
| 30 | +import static org.assertj.core.api.Assertions.assertThat; |
27 | 31 | import static org.mockito.ArgumentMatchers.any; |
28 | 32 | import static org.mockito.ArgumentMatchers.eq; |
29 | 33 | import static org.mockito.Mockito.atLeastOnce; |
@@ -166,4 +170,99 @@ void testCloseInitialized() { |
166 | 170 | .verify(); |
167 | 171 | } |
168 | 172 |
|
| 173 | + @Test |
| 174 | + void testMcpMethodHeaderOnRequest() throws URISyntaxException { |
| 175 | + var uri = new URI(host + "/mcp"); |
| 176 | + var mockRequestCustomizer = mock(McpSyncHttpClientRequestCustomizer.class); |
| 177 | + |
| 178 | + var transport = HttpClientStreamableHttpTransport.builder(host) |
| 179 | + .httpRequestCustomizer(mockRequestCustomizer) |
| 180 | + .build(); |
| 181 | + |
| 182 | + withTransport(transport, (t) -> { |
| 183 | + var initializeRequest = McpSchema.InitializeRequest |
| 184 | + .builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ClientCapabilities.builder().roots(true).build(), |
| 185 | + McpSchema.Implementation.builder("MCP Client", "0.3.1").build()) |
| 186 | + .build(); |
| 187 | + var testMessage = new McpSchema.JSONRPCRequest(McpSchema.METHOD_INITIALIZE, "test-id", initializeRequest); |
| 188 | + |
| 189 | + StepVerifier |
| 190 | + .create(t.sendMessage(testMessage).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))) |
| 191 | + .verifyComplete(); |
| 192 | + |
| 193 | + var requestCaptor = ArgumentCaptor.forClass(HttpRequest.Builder.class); |
| 194 | + verify(mockRequestCustomizer, atLeastOnce()).customize(requestCaptor.capture(), eq("POST"), eq(uri), any(), |
| 195 | + eq(context)); |
| 196 | + assertThat(requestCaptor.getValue().build().headers().firstValue(HttpHeaders.MCP_METHOD)) |
| 197 | + .hasValue(McpSchema.METHOD_INITIALIZE); |
| 198 | + }); |
| 199 | + } |
| 200 | + |
| 201 | + @Test |
| 202 | + void testMcpMethodHeaderOnNotification() throws URISyntaxException { |
| 203 | + var uri = new URI(host + "/mcp"); |
| 204 | + var mockRequestCustomizer = mock(McpSyncHttpClientRequestCustomizer.class); |
| 205 | + |
| 206 | + var transport = HttpClientStreamableHttpTransport.builder(host) |
| 207 | + .httpRequestCustomizer(mockRequestCustomizer) |
| 208 | + .build(); |
| 209 | + |
| 210 | + withTransport(transport, (t) -> { |
| 211 | + var initializeRequest = McpSchema.InitializeRequest |
| 212 | + .builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ClientCapabilities.builder().roots(true).build(), |
| 213 | + McpSchema.Implementation.builder("MCP Client", "0.3.1").build()) |
| 214 | + .build(); |
| 215 | + var initializeMessage = new McpSchema.JSONRPCRequest(McpSchema.METHOD_INITIALIZE, "test-id", |
| 216 | + initializeRequest); |
| 217 | + StepVerifier |
| 218 | + .create(t.sendMessage(initializeMessage).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))) |
| 219 | + .verifyComplete(); |
| 220 | + |
| 221 | + var testMessage = new McpSchema.JSONRPCNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED); |
| 222 | + |
| 223 | + StepVerifier |
| 224 | + .create(t.sendMessage(testMessage).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))) |
| 225 | + .verifyComplete(); |
| 226 | + |
| 227 | + var requestCaptor = ArgumentCaptor.forClass(HttpRequest.Builder.class); |
| 228 | + verify(mockRequestCustomizer, atLeastOnce()).customize(requestCaptor.capture(), eq("POST"), eq(uri), any(), |
| 229 | + eq(context)); |
| 230 | + assertThat(requestCaptor.getValue().build().headers().firstValue(HttpHeaders.MCP_METHOD)) |
| 231 | + .hasValue(McpSchema.METHOD_NOTIFICATION_INITIALIZED); |
| 232 | + }); |
| 233 | + } |
| 234 | + |
| 235 | + @Test |
| 236 | + void testNoMcpMethodHeaderOnResponse() throws URISyntaxException { |
| 237 | + var uri = new URI(host + "/mcp"); |
| 238 | + var mockRequestCustomizer = mock(McpSyncHttpClientRequestCustomizer.class); |
| 239 | + |
| 240 | + var transport = HttpClientStreamableHttpTransport.builder(host) |
| 241 | + .httpRequestCustomizer(mockRequestCustomizer) |
| 242 | + .build(); |
| 243 | + |
| 244 | + withTransport(transport, (t) -> { |
| 245 | + var initializeRequest = McpSchema.InitializeRequest |
| 246 | + .builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ClientCapabilities.builder().roots(true).build(), |
| 247 | + McpSchema.Implementation.builder("MCP Client", "0.3.1").build()) |
| 248 | + .build(); |
| 249 | + var initializeMessage = new McpSchema.JSONRPCRequest(McpSchema.METHOD_INITIALIZE, "test-id", |
| 250 | + initializeRequest); |
| 251 | + StepVerifier |
| 252 | + .create(t.sendMessage(initializeMessage).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))) |
| 253 | + .verifyComplete(); |
| 254 | + |
| 255 | + var testMessage = McpSchema.JSONRPCResponse.result("test-id-2", Map.of()); |
| 256 | + |
| 257 | + StepVerifier |
| 258 | + .create(t.sendMessage(testMessage).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))) |
| 259 | + .verifyComplete(); |
| 260 | + |
| 261 | + var requestCaptor = ArgumentCaptor.forClass(HttpRequest.Builder.class); |
| 262 | + verify(mockRequestCustomizer, atLeastOnce()).customize(requestCaptor.capture(), eq("POST"), eq(uri), any(), |
| 263 | + eq(context)); |
| 264 | + assertThat(requestCaptor.getValue().build().headers().firstValue(HttpHeaders.MCP_METHOD)).isEmpty(); |
| 265 | + }); |
| 266 | + } |
| 267 | + |
169 | 268 | } |
0 commit comments