From 57a526ec300f7caf9e587eea818722261e0fb21e Mon Sep 17 00:00:00 2001
From: Stephen Halter <halter73@gmail.com>
Date: Mon, 31 Mar 2025 14:35:28 -0700
Subject: [PATCH] Put MapMcp in Microsoft.AspNetCore.Builder namespace (#171)

* Still dispose McpSession when MessageProcessingTask throws non-OCE exceptions

* Put MapMcp in icrosoft.AspNetCore.Builder namespace
---
 samples/AspNetCoreSseServer/Program.cs        |  2 --
 .../McpEndpointRouteBuilderExtensions.cs      |  5 ++-
 src/ModelContextProtocol.AspNetCore/README.md |  1 -
 .../Shared/McpJsonRpcEndpoint.cs              | 31 +++++++++++--------
 .../Program.cs                                |  3 +-
 5 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/samples/AspNetCoreSseServer/Program.cs b/samples/AspNetCoreSseServer/Program.cs
index 774957e..9f94eb2 100644
--- a/samples/AspNetCoreSseServer/Program.cs
+++ b/samples/AspNetCoreSseServer/Program.cs
@@ -1,5 +1,3 @@
-using ModelContextProtocol.AspNetCore;
-
 var builder = WebApplication.CreateBuilder(args);
 builder.Services.AddMcpServer().WithToolsFromAssembly();
 var app = builder.Build();
diff --git a/src/ModelContextProtocol.AspNetCore/McpEndpointRouteBuilderExtensions.cs b/src/ModelContextProtocol.AspNetCore/McpEndpointRouteBuilderExtensions.cs
index 3dae6cd..6ddbc6c 100644
--- a/src/ModelContextProtocol.AspNetCore/McpEndpointRouteBuilderExtensions.cs
+++ b/src/ModelContextProtocol.AspNetCore/McpEndpointRouteBuilderExtensions.cs
@@ -1,5 +1,4 @@
-using Microsoft.AspNetCore.Builder;
-using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Routing;
 using Microsoft.AspNetCore.WebUtilities;
 using Microsoft.Extensions.DependencyInjection;
@@ -12,7 +11,7 @@ using ModelContextProtocol.Utils.Json;
 using System.Collections.Concurrent;
 using System.Security.Cryptography;
 
-namespace ModelContextProtocol.AspNetCore;
+namespace Microsoft.AspNetCore.Builder;
 
 /// <summary>
 /// Extension methods for <see cref="IEndpointRouteBuilder"/> to add MCP endpoints.
diff --git a/src/ModelContextProtocol.AspNetCore/README.md b/src/ModelContextProtocol.AspNetCore/README.md
index 76f7c2b..457321d 100644
--- a/src/ModelContextProtocol.AspNetCore/README.md
+++ b/src/ModelContextProtocol.AspNetCore/README.md
@@ -30,7 +30,6 @@ dotnet add package ModelContextProtocol.AspNetCore --prerelease
 
 ```csharp
 // Program.cs
-using ModelContextProtocol.AspNetCore;
 using ModelContextProtocol.Server;
 using System.ComponentModel;
 
diff --git a/src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs b/src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs
index ac0337f..3f7417a 100644
--- a/src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs
+++ b/src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs
@@ -94,25 +94,30 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
     {
         _logger.CleaningUpEndpoint(EndpointName);
 
-        if (_sessionCts is not null)
+        try
         {
-            await _sessionCts.CancelAsync().ConfigureAwait(false);
-        }
-
-        if (MessageProcessingTask is not null)
-        {
-            try
+            if (_sessionCts is not null)
             {
-                await MessageProcessingTask.ConfigureAwait(false);
+                await _sessionCts.CancelAsync().ConfigureAwait(false);
             }
-            catch (OperationCanceledException)
+
+            if (MessageProcessingTask is not null)
             {
-                // Ignore cancellation
+                try
+                {
+                    await MessageProcessingTask.ConfigureAwait(false);
+                }
+                catch (OperationCanceledException)
+                {
+                    // Ignore cancellation
+                }
             }
         }
-
-        _session?.Dispose();
-        _sessionCts?.Dispose();
+        finally
+        {
+            _session?.Dispose();
+            _sessionCts?.Dispose();
+        }
 
         _logger.EndpointCleanedUp(EndpointName);
     }
diff --git a/tests/ModelContextProtocol.TestSseServer/Program.cs b/tests/ModelContextProtocol.TestSseServer/Program.cs
index 4bbc5bf..23afc1d 100644
--- a/tests/ModelContextProtocol.TestSseServer/Program.cs
+++ b/tests/ModelContextProtocol.TestSseServer/Program.cs
@@ -1,5 +1,4 @@
-using ModelContextProtocol.AspNetCore;
-using ModelContextProtocol.Protocol.Types;
+using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
 using Serilog;
 using System.Text;