Skip to content

Commit 812e5f7

Browse files
stephentoubstephentoub
committed
Address feedback
1 parent da912e9 commit 812e5f7

File tree

6 files changed

+97
-44
lines changed

6 files changed

+97
-44
lines changed

‎src/ModelContextProtocol/Hosting/SingleSessionMcpServerHostedService.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
1717
{
1818
try
1919
{
20-
await session.RunAsync(stoppingToken);
20+
await session.RunAsync(stoppingToken).ConfigureAwait(false);
2121
}
2222
finally
2323
{

‎src/ModelContextProtocol/Protocol/Transport/StreamServerTransport.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ public override async Task SendMessageAsync(IJsonRpcMessage message, Cancellatio
8282

8383
await JsonSerializer.SerializeAsync(_outputStream, message, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(IJsonRpcMessage)), cancellationToken).ConfigureAwait(false);
8484
await _outputStream.WriteAsync(s_newlineBytes, cancellationToken).ConfigureAwait(false);
85-
await _outputStream.FlushAsync(cancellationToken).ConfigureAwait(false);;
85+
await _outputStream.FlushAsync(cancellationToken).ConfigureAwait(false);
8686

8787
_logger.TransportSentMessage(_endpointName, id);
8888
}

‎src/ModelContextProtocol/Server/McpServer.cs

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -502,15 +502,24 @@ async ValueTask<TResult> InvokeScopedAsync(
502502
TParams? args,
503503
CancellationToken cancellationToken)
504504
{
505-
using var scope = Services?.GetService<IServiceScopeFactory>()?.CreateScope();
506-
507-
return await handler(
508-
new RequestContext<TParams>(this)
505+
var scope = Services?.GetService<IServiceScopeFactory>()?.CreateAsyncScope();
506+
try
507+
{
508+
return await handler(
509+
new RequestContext<TParams>(this)
510+
{
511+
Services = scope?.ServiceProvider ?? Services,
512+
Params = args
513+
},
514+
cancellationToken).ConfigureAwait(false);
515+
}
516+
finally
517+
{
518+
if (scope is not null)
509519
{
510-
Services = scope?.ServiceProvider ?? Services,
511-
Params = args
512-
},
513-
cancellationToken).ConfigureAwait(false);
520+
await scope.Value.DisposeAsync().ConfigureAwait(false);
521+
}
522+
}
514523
}
515524
}
516525

‎tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ protected override void ConfigureServices(ServiceCollection services, IMcpServer
107107
throw new Exception($"Unknown tool '{request.Params?.Name}'");
108108
}
109109
})
110-
.WithTools<EchoTool>(serializerOptions: JsonContext.Default.Options);
110+
.WithTools<EchoTool>(serializerOptions: BuilderToolsJsonContext.Default.Options);
111111

112112
services.AddSingleton(new ObjectWithId());
113113
}
@@ -429,8 +429,13 @@ public void Empty_Enumerables_Is_Allowed()
429429
[Fact]
430430
public void Register_Tools_From_Current_Assembly()
431431
{
432+
if (!JsonSerializer.IsReflectionEnabledByDefault)
433+
{
434+
return;
435+
}
436+
432437
ServiceCollection sc = new();
433-
sc.AddMcpServer().WithToolsFromAssembly(serializerOptions: JsonContext.Default.Options);
438+
sc.AddMcpServer().WithToolsFromAssembly();
434439
IServiceProvider services = sc.BuildServiceProvider();
435440

436441
Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "Echo");
@@ -446,7 +451,7 @@ public void WithTools_Parameters_Satisfiable_From_DI(bool parameterInServices)
446451
{
447452
sc.AddSingleton(new ComplexObject());
448453
}
449-
sc.AddMcpServer().WithTools([typeof(EchoTool)], JsonContext.Default.Options);
454+
sc.AddMcpServer().WithTools([typeof(EchoTool)], BuilderToolsJsonContext.Default.Options);
450455
IServiceProvider services = sc.BuildServiceProvider();
451456

452457
McpServerTool tool = services.GetServices<McpServerTool>().First(t => t.ProtocolTool.Name == "EchoComplex");
@@ -460,13 +465,19 @@ public void WithTools_Parameters_Satisfiable_From_DI(bool parameterInServices)
460465
}
461466
}
462467

468+
463469
[Theory]
464470
[InlineData(ServiceLifetime.Singleton)]
465471
[InlineData(ServiceLifetime.Scoped)]
466472
[InlineData(ServiceLifetime.Transient)]
467473
[InlineData(null)]
468474
public void WithToolsFromAssembly_Parameters_Satisfiable_From_DI(ServiceLifetime? lifetime)
469475
{
476+
if (!JsonSerializer.IsReflectionEnabledByDefault)
477+
{
478+
return;
479+
}
480+
470481
ServiceCollection sc = new();
471482
switch (lifetime)
472483
{
@@ -483,7 +494,7 @@ public void WithToolsFromAssembly_Parameters_Satisfiable_From_DI(ServiceLifetime
483494
break;
484495
}
485496

486-
sc.AddMcpServer().WithToolsFromAssembly(serializerOptions: JsonContext.Default.Options);
497+
sc.AddMcpServer().WithToolsFromAssembly();
487498
IServiceProvider services = sc.BuildServiceProvider();
488499

489500
McpServerTool tool = services.GetServices<McpServerTool>().First(t => t.ProtocolTool.Name == "EchoComplex");
@@ -526,9 +537,9 @@ public void Register_Tools_From_Multiple_Sources()
526537
{
527538
ServiceCollection sc = new();
528539
sc.AddMcpServer()
529-
.WithTools<EchoTool>(serializerOptions: JsonContext.Default.Options)
530-
.WithTools<AnotherToolType>(serializerOptions: JsonContext.Default.Options)
531-
.WithTools([typeof(ToolTypeWithNoAttribute)], JsonContext.Default.Options);
540+
.WithTools<EchoTool>(serializerOptions: BuilderToolsJsonContext.Default.Options)
541+
.WithTools<AnotherToolType>(serializerOptions: BuilderToolsJsonContext.Default.Options)
542+
.WithTools([typeof(ToolTypeWithNoAttribute)], BuilderToolsJsonContext.Default.Options);
532543
IServiceProvider services = sc.BuildServiceProvider();
533544

534545
Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "double_echo");
@@ -759,4 +770,22 @@ public class ObjectWithId
759770
{
760771
public string Id { get; set; } = Guid.NewGuid().ToString("N");
761772
}
762-
}
773+
774+
public class ComplexObject
775+
{
776+
public string? Name { get; set; }
777+
public int Age { get; set; }
778+
}
779+
780+
[JsonSerializable(typeof(bool))]
781+
[JsonSerializable(typeof(int))]
782+
[JsonSerializable(typeof(long))]
783+
[JsonSerializable(typeof(double))]
784+
[JsonSerializable(typeof(string))]
785+
[JsonSerializable(typeof(DateTime))]
786+
[JsonSerializable(typeof(DateTimeOffset))]
787+
[JsonSerializable(typeof(ComplexObject))]
788+
[JsonSerializable(typeof(string[]))]
789+
[JsonSerializable(typeof(JsonElement))]
790+
partial class BuilderToolsJsonContext : JsonSerializerContext;
791+
}

‎tests/ModelContextProtocol.Tests/Configuration/McpServerScopedTests.cs

Lines changed: 41 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
using Microsoft.Extensions.AI;
2-
using Microsoft.Extensions.DependencyInjection;
1+
using Microsoft.Extensions.DependencyInjection;
32
using ModelContextProtocol.Client;
43
using ModelContextProtocol.Server;
4+
using ModelContextProtocol.Utils.Json;
55
using System.Text.Json;
6+
using System.Text.Json.Serialization;
67

78
namespace ModelContextProtocol.Tests.Configuration;
89

@@ -15,7 +16,7 @@ public McpServerScopedTests(ITestOutputHelper testOutputHelper)
1516

1617
protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
1718
{
18-
mcpServerBuilder.WithTools<EchoTool>(serializerOptions: JsonContext.Default.Options);
19+
mcpServerBuilder.WithTools<EchoTool>(serializerOptions: McpServerScopedTestsJsonContext.Default.Options);
1920
services.AddScoped(_ => new ComplexObject() { Name = "Scoped" });
2021
}
2122

@@ -24,11 +25,20 @@ public async Task InjectScopedServiceAsArgument()
2425
{
2526
IMcpClient client = await CreateMcpClientForServer();
2627

27-
var tools = await client.ListToolsAsync(JsonContext.Default.Options, TestContext.Current.CancellationToken);
28+
var tools = await client.ListToolsAsync(McpServerScopedTestsJsonContext.Default.Options, TestContext.Current.CancellationToken);
2829
var tool = tools.First(t => t.Name == nameof(EchoTool.EchoComplex));
29-
Assert.DoesNotContain("\"complex\"", JsonSerializer.Serialize(tool.JsonSchema, AIJsonUtilities.DefaultOptions));
30+
Assert.DoesNotContain("\"complex\"", JsonSerializer.Serialize(tool.JsonSchema, McpJsonUtilities.DefaultOptions));
3031

31-
Assert.Contains("\"Scoped\"", JsonSerializer.Serialize(await tool.InvokeAsync(cancellationToken: TestContext.Current.CancellationToken), AIJsonUtilities.DefaultOptions));
32+
int startingConstructed = ComplexObject.Constructed;
33+
int startingDisposed = ComplexObject.Disposed;
34+
35+
for (int i = 1; i <= 10; i++)
36+
{
37+
Assert.Contains("\"Scoped\"", JsonSerializer.Serialize(await tool.InvokeAsync(cancellationToken: TestContext.Current.CancellationToken), McpJsonUtilities.DefaultOptions));
38+
39+
Assert.Equal(startingConstructed + i, ComplexObject.Constructed);
40+
Assert.Equal(startingDisposed + i, ComplexObject.Disposed);
41+
}
3242
}
3343

3444
[McpServerToolType]
@@ -37,4 +47,29 @@ public sealed class EchoTool()
3747
[McpServerTool]
3848
public static string EchoComplex(ComplexObject complex) => complex.Name!;
3949
}
50+
51+
public class ComplexObject : IAsyncDisposable
52+
{
53+
public static int Constructed;
54+
public static int Disposed;
55+
56+
public ComplexObject()
57+
{
58+
Interlocked.Increment(ref Constructed);
59+
}
60+
61+
public ValueTask DisposeAsync()
62+
{
63+
Interlocked.Increment(ref Disposed);
64+
return default;
65+
}
66+
67+
public string? Name { get; set; }
68+
public int Age { get; set; }
69+
}
70+
71+
[JsonSerializable(typeof(string))]
72+
[JsonSerializable(typeof(ComplexObject))]
73+
[JsonSerializable(typeof(JsonElement))]
74+
partial class McpServerScopedTestsJsonContext : JsonSerializerContext;
4075
}

‎tests/ModelContextProtocol.Tests/JsonContext.cs

Lines changed: 0 additions & 20 deletions
This file was deleted.

0 commit comments

Comments
 (0)