Skip to content

Commit a293207

Browse files
KehrlannKehrlann
committed
wip
1 parent b4fef52 commit a293207

13 files changed

+317
-107
lines changed

‎mcp/src/main/java/io/modelcontextprotocol/client/McpClient.java

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,25 +4,26 @@
44

55
package io.modelcontextprotocol.client;
66

7-
import java.time.Duration;
8-
import java.util.ArrayList;
9-
import java.util.HashMap;
10-
import java.util.List;
11-
import java.util.Map;
12-
import java.util.function.Consumer;
13-
import java.util.function.Function;
14-
7+
import io.modelcontextprotocol.server.McpTransportContext;
158
import io.modelcontextprotocol.spec.McpClientTransport;
169
import io.modelcontextprotocol.spec.McpSchema;
17-
import io.modelcontextprotocol.spec.McpTransport;
1810
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
1911
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
2012
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
2113
import io.modelcontextprotocol.spec.McpSchema.ElicitRequest;
2214
import io.modelcontextprotocol.spec.McpSchema.ElicitResult;
2315
import io.modelcontextprotocol.spec.McpSchema.Implementation;
2416
import io.modelcontextprotocol.spec.McpSchema.Root;
17+
import io.modelcontextprotocol.spec.McpTransport;
2518
import io.modelcontextprotocol.util.Assert;
19+
import java.time.Duration;
20+
import java.util.ArrayList;
21+
import java.util.HashMap;
22+
import java.util.List;
23+
import java.util.Map;
24+
import java.util.function.Consumer;
25+
import java.util.function.Function;
26+
import java.util.function.Supplier;
2627
import reactor.core.publisher.Mono;
2728

2829
/**
@@ -183,6 +184,8 @@ class SyncSpec {
183184

184185
private Function<ElicitRequest, ElicitResult> elicitationHandler;
185186

187+
private Supplier<McpTransportContext> contextProvider = McpTransportContext.EMPTY::copy;
188+
186189
private SyncSpec(McpClientTransport transport) {
187190
Assert.notNull(transport, "Transport must not be null");
188191
this.transport = transport;
@@ -409,6 +412,11 @@ public SyncSpec progressConsumers(List<Consumer<McpSchema.ProgressNotification>>
409412
return this;
410413
}
411414

415+
public SyncSpec transportContextProvider(Supplier<McpTransportContext> contextProvider) {
416+
this.contextProvider = contextProvider;
417+
return this;
418+
}
419+
412420
/**
413421
* Create an instance of {@link McpSyncClient} with the provided configurations or
414422
* sensible defaults.
@@ -423,7 +431,8 @@ public McpSyncClient build() {
423431
McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
424432

425433
return new McpSyncClient(
426-
new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures));
434+
new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures),
435+
this.contextProvider);
427436
}
428437

429438
}

‎mcp/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java

Lines changed: 91 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,18 @@
44

55
package io.modelcontextprotocol.client;
66

7-
import java.time.Duration;
8-
9-
import org.slf4j.Logger;
10-
import org.slf4j.LoggerFactory;
11-
7+
import io.modelcontextprotocol.server.McpTransportContext;
128
import io.modelcontextprotocol.spec.McpSchema;
139
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
1410
import io.modelcontextprotocol.spec.McpSchema.GetPromptRequest;
1511
import io.modelcontextprotocol.spec.McpSchema.GetPromptResult;
1612
import io.modelcontextprotocol.spec.McpSchema.ListPromptsResult;
13+
import io.modelcontextprotocol.spec.McpTransport;
1714
import io.modelcontextprotocol.util.Assert;
15+
import java.time.Duration;
16+
import java.util.function.Supplier;
17+
import org.slf4j.Logger;
18+
import org.slf4j.LoggerFactory;
1819

1920
/**
2021
* A synchronous client implementation for the Model Context Protocol (MCP) that wraps an
@@ -63,14 +64,18 @@ public class McpSyncClient implements AutoCloseable {
6364

6465
private final McpAsyncClient delegate;
6566

67+
private final Supplier<McpTransportContext> contextProvider;
68+
6669
/**
6770
* Create a new McpSyncClient with the given delegate.
6871
* @param delegate the asynchronous kernel on top of which this synchronous client
6972
* provides a blocking API.
7073
*/
71-
McpSyncClient(McpAsyncClient delegate) {
74+
McpSyncClient(McpAsyncClient delegate, Supplier<McpTransportContext> contextProvider) {
7275
Assert.notNull(delegate, "The delegate can not be null");
76+
Assert.notNull(contextProvider, "The contextProvider can not be null");
7377
this.delegate = delegate;
78+
this.contextProvider = contextProvider;
7479
}
7580

7681
/**
@@ -177,36 +182,43 @@ public boolean closeGracefully() {
177182
public McpSchema.InitializeResult initialize() {
178183
// TODO: block takes no argument here as we assume the async client is
179184
// configured with a requestTimeout at all times
180-
return this.delegate.initialize().block();
185+
var context = this.contextProvider.get();
186+
return this.delegate.initialize().contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context)).block();
181187
}
182188

183189
/**
184190
* Send a roots/list_changed notification.
185191
*/
186192
public void rootsListChangedNotification() {
187-
this.delegate.rootsListChangedNotification().block();
193+
var context = this.contextProvider.get();
194+
this.delegate.rootsListChangedNotification()
195+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
196+
.block();
188197
}
189198

190199
/**
191200
* Add a roots dynamically.
192201
*/
193202
public void addRoot(McpSchema.Root root) {
194-
this.delegate.addRoot(root).block();
203+
var context = this.contextProvider.get();
204+
this.delegate.addRoot(root).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context)).block();
195205
}
196206

197207
/**
198208
* Remove a root dynamically.
199209
*/
200210
public void removeRoot(String rootUri) {
201-
this.delegate.removeRoot(rootUri).block();
211+
var context = this.contextProvider.get();
212+
this.delegate.removeRoot(rootUri).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context)).block();
202213
}
203214

204215
/**
205216
* Send a synchronous ping request.
206217
* @return
207218
*/
208219
public Object ping() {
209-
return this.delegate.ping().block();
220+
var context = this.contextProvider.get();
221+
return this.delegate.ping().contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context)).block();
210222
}
211223

212224
// --------------------------
@@ -224,7 +236,11 @@ public Object ping() {
224236
* Boolean indicating if the execution failed (true) or succeeded (false/absent)
225237
*/
226238
public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolRequest) {
227-
return this.delegate.callTool(callToolRequest).block();
239+
var context = this.contextProvider.get();
240+
return this.delegate.callTool(callToolRequest)
241+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
242+
.block();
243+
228244
}
229245

230246
/**
@@ -234,7 +250,8 @@ public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolReque
234250
* pagination if more tools are available
235251
*/
236252
public McpSchema.ListToolsResult listTools() {
237-
return this.delegate.listTools().block();
253+
var context = this.contextProvider.get();
254+
return this.delegate.listTools().contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context)).block();
238255
}
239256

240257
/**
@@ -245,7 +262,9 @@ public McpSchema.ListToolsResult listTools() {
245262
* pagination if more tools are available
246263
*/
247264
public McpSchema.ListToolsResult listTools(String cursor) {
248-
return this.delegate.listTools(cursor).block();
265+
var context = this.contextProvider.get();
266+
return this.delegate.listTools(cursor).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context)).block();
267+
249268
}
250269

251270
// --------------------------
@@ -257,7 +276,9 @@ public McpSchema.ListToolsResult listTools(String cursor) {
257276
* @return The list of all resources result
258277
*/
259278
public McpSchema.ListResourcesResult listResources() {
260-
return this.delegate.listResources().block();
279+
var context = this.contextProvider.get();
280+
return this.delegate.listResources().contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context)).block();
281+
261282
}
262283

263284
/**
@@ -266,7 +287,11 @@ public McpSchema.ListResourcesResult listResources() {
266287
* @return The list of resources result
267288
*/
268289
public McpSchema.ListResourcesResult listResources(String cursor) {
269-
return this.delegate.listResources(cursor).block();
290+
var context = this.contextProvider.get();
291+
return this.delegate.listResources(cursor)
292+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
293+
.block();
294+
270295
}
271296

272297
/**
@@ -275,7 +300,11 @@ public McpSchema.ListResourcesResult listResources(String cursor) {
275300
* @return the resource content.
276301
*/
277302
public McpSchema.ReadResourceResult readResource(McpSchema.Resource resource) {
278-
return this.delegate.readResource(resource).block();
303+
var context = this.contextProvider.get();
304+
return this.delegate.readResource(resource)
305+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
306+
.block();
307+
279308
}
280309

281310
/**
@@ -284,15 +313,23 @@ public McpSchema.ReadResourceResult readResource(McpSchema.Resource resource) {
284313
* @return the resource content.
285314
*/
286315
public McpSchema.ReadResourceResult readResource(McpSchema.ReadResourceRequest readResourceRequest) {
287-
return this.delegate.readResource(readResourceRequest).block();
316+
var context = this.contextProvider.get();
317+
return this.delegate.readResource(readResourceRequest)
318+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
319+
.block();
320+
288321
}
289322

290323
/**
291324
* Retrieves the list of all resource templates provided by the server.
292325
* @return The list of all resource templates result.
293326
*/
294327
public McpSchema.ListResourceTemplatesResult listResourceTemplates() {
295-
return this.delegate.listResourceTemplates().block();
328+
var context = this.contextProvider.get();
329+
return this.delegate.listResourceTemplates()
330+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
331+
.block();
332+
296333
}
297334

298335
/**
@@ -304,7 +341,11 @@ public McpSchema.ListResourceTemplatesResult listResourceTemplates() {
304341
* @return The list of resource templates result.
305342
*/
306343
public McpSchema.ListResourceTemplatesResult listResourceTemplates(String cursor) {
307-
return this.delegate.listResourceTemplates(cursor).block();
344+
var context = this.contextProvider.get();
345+
return this.delegate.listResourceTemplates(cursor)
346+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
347+
.block();
348+
308349
}
309350

310351
/**
@@ -317,7 +358,11 @@ public McpSchema.ListResourceTemplatesResult listResourceTemplates(String cursor
317358
* subscribe to.
318359
*/
319360
public void subscribeResource(McpSchema.SubscribeRequest subscribeRequest) {
320-
this.delegate.subscribeResource(subscribeRequest).block();
361+
var context = this.contextProvider.get();
362+
this.delegate.subscribeResource(subscribeRequest)
363+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
364+
.block();
365+
321366
}
322367

323368
/**
@@ -326,7 +371,11 @@ public void subscribeResource(McpSchema.SubscribeRequest subscribeRequest) {
326371
* to unsubscribe from.
327372
*/
328373
public void unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) {
329-
this.delegate.unsubscribeResource(unsubscribeRequest).block();
374+
var context = this.contextProvider.get();
375+
this.delegate.unsubscribeResource(unsubscribeRequest)
376+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
377+
.block();
378+
330379
}
331380

332381
// --------------------------
@@ -338,7 +387,8 @@ public void unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest)
338387
* @return The list of all prompts result.
339388
*/
340389
public ListPromptsResult listPrompts() {
341-
return this.delegate.listPrompts().block();
390+
var context = this.contextProvider.get();
391+
return this.delegate.listPrompts().contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context)).block();
342392
}
343393

344394
/**
@@ -347,19 +397,29 @@ public ListPromptsResult listPrompts() {
347397
* @return The list of prompts result.
348398
*/
349399
public ListPromptsResult listPrompts(String cursor) {
350-
return this.delegate.listPrompts(cursor).block();
400+
var context = this.contextProvider.get();
401+
return this.delegate.listPrompts(cursor).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context)).block();
402+
351403
}
352404

353405
public GetPromptResult getPrompt(GetPromptRequest getPromptRequest) {
354-
return this.delegate.getPrompt(getPromptRequest).block();
406+
var context = this.contextProvider.get();
407+
return this.delegate.getPrompt(getPromptRequest)
408+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
409+
.block();
410+
355411
}
356412

357413
/**
358414
* Client can set the minimum logging level it wants to receive from the server.
359415
* @param loggingLevel the min logging level
360416
*/
361417
public void setLoggingLevel(McpSchema.LoggingLevel loggingLevel) {
362-
this.delegate.setLoggingLevel(loggingLevel).block();
418+
var context = this.contextProvider.get();
419+
this.delegate.setLoggingLevel(loggingLevel)
420+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
421+
.block();
422+
363423
}
364424

365425
/**
@@ -369,7 +429,11 @@ public void setLoggingLevel(McpSchema.LoggingLevel loggingLevel) {
369429
* @return the completion result containing suggested values.
370430
*/
371431
public McpSchema.CompleteResult completeCompletion(McpSchema.CompleteRequest completeRequest) {
372-
return this.delegate.completeCompletion(completeRequest).block();
432+
var context = this.contextProvider.get();
433+
return this.delegate.completeCompletion(completeRequest)
434+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, context))
435+
.block();
436+
373437
}
374438

375439
}

‎mcp/src/main/java/io/modelcontextprotocol/client/transport/AsyncHttpRequestCustomizer.java

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44

55
package io.modelcontextprotocol.client.transport;
66

7+
import io.modelcontextprotocol.server.DefaultMcpTransportContext;
8+
import io.modelcontextprotocol.server.McpTransportContext;
79
import java.net.URI;
810
import java.net.http.HttpRequest;
911
import org.reactivestreams.Publisher;
@@ -21,8 +23,14 @@
2123
*/
2224
public interface AsyncHttpRequestCustomizer {
2325

26+
@Deprecated
27+
default Publisher<HttpRequest.Builder> customize(HttpRequest.Builder builder, String method, URI endpoint,
28+
@Nullable String body) {
29+
return customize(builder, method, endpoint, body, new DefaultMcpTransportContext());
30+
}
31+
2432
Publisher<HttpRequest.Builder> customize(HttpRequest.Builder builder, String method, URI endpoint,
25-
@Nullable String body);
33+
@Nullable String body, McpTransportContext context);
2634

2735
AsyncHttpRequestCustomizer NOOP = new Noop();
2836

@@ -33,8 +41,8 @@ Publisher<HttpRequest.Builder> customize(HttpRequest.Builder builder, String met
3341
* blocking implementation, consider using {@link Schedulers#boundedElastic()}.
3442
*/
3543
static AsyncHttpRequestCustomizer fromSync(SyncHttpRequestCustomizer customizer) {
36-
return (builder, method, uri, body) -> Mono.fromSupplier(() -> {
37-
customizer.customize(builder, method, uri, body);
44+
return (builder, method, uri, body, context) -> Mono.fromSupplier(() -> {
45+
customizer.customize(builder, method, uri, body, context);
3846
return builder;
3947
});
4048
}
@@ -43,7 +51,7 @@ class Noop implements AsyncHttpRequestCustomizer {
4351

4452
@Override
4553
public Publisher<HttpRequest.Builder> customize(HttpRequest.Builder builder, String method, URI endpoint,
46-
String body) {
54+
String body, McpTransportContext context) {
4755
return Mono.just(builder);
4856
}
4957

0 commit comments

Comments
 (0)