Skip to content

Commit ca2ccd2

Browse files
ejsmithejsmith
committed
Add MCP session migration cache (#2351)
1 parent 82837c6 commit ca2ccd2

3 files changed

Lines changed: 280 additions & 0 deletions

File tree

‎src/Exceptionless.Web/Mcp/McpSessionMigrationHandler.cs‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
using System.Security.Claims;
2+
using Exceptionless.Core.Authorization;
3+
using Exceptionless.Core.Extensions;
4+
using Foundatio.Caching;
5+
using Foundatio.Serializer;
6+
using Microsoft.Extensions.Options;
7+
using ModelContextProtocol.AspNetCore;
8+
using ModelContextProtocol.Protocol;
9+
10+
namespace Exceptionless.Web.Mcp;
11+
12+
public sealed class McpSessionMigrationHandler(
13+
ICacheClient cacheClient,
14+
ITextSerializer serializer,
15+
IOptions<HttpServerTransportOptions> transportOptions,
16+
TimeProvider timeProvider,
17+
ILogger<McpSessionMigrationHandler> logger) : ISessionMigrationHandler
18+
{
19+
private const string CacheKeyPrefix = "mcp:session:";
20+
private const string UserClientId = "user";
21+
private static readonly TimeSpan LifetimeBuffer = TimeSpan.FromMinutes(5);
22+
private static readonly TimeSpan DefaultIdleTimeout = TimeSpan.FromHours(2);
23+
24+
public async ValueTask OnSessionInitializedAsync(HttpContext context, string sessionId, InitializeRequestParams initializeParams, CancellationToken cancellationToken)
25+
{
26+
cancellationToken.ThrowIfCancellationRequested();
27+
28+
var sessionIdentity = GetSessionIdentity(context);
29+
if (sessionIdentity is null)
30+
{
31+
logger.LogWarning("Skipping MCP session migration persistence for unauthenticated session {SessionIdHash}", HashSessionId(sessionId));
32+
return;
33+
}
34+
35+
var state = new McpSessionMigrationState
36+
{
37+
UserId = sessionIdentity.UserId,
38+
ClientId = sessionIdentity.ClientId,
39+
Resource = sessionIdentity.Resource,
40+
InitializeParams = initializeParams,
41+
CreatedUtc = timeProvider.GetUtcNow().UtcDateTime
42+
};
43+
44+
string serializedState = serializer.SerializeToString(state);
45+
await cacheClient.SetAsync(GetCacheKey(sessionId), serializedState, GetCacheLifetime());
46+
}
47+
48+
public async ValueTask<InitializeRequestParams?> AllowSessionMigrationAsync(HttpContext context, string sessionId, CancellationToken cancellationToken)
49+
{
50+
cancellationToken.ThrowIfCancellationRequested();
51+
52+
var sessionIdentity = GetSessionIdentity(context);
53+
if (sessionIdentity is null)
54+
return null;
55+
56+
string cacheKey = GetCacheKey(sessionId);
57+
string? serializedState = await cacheClient.GetAsync<string?>(cacheKey, null);
58+
if (String.IsNullOrEmpty(serializedState))
59+
return null;
60+
61+
McpSessionMigrationState? state;
62+
try
63+
{
64+
state = serializer.Deserialize<McpSessionMigrationState>(serializedState);
65+
}
66+
catch (Exception ex)
67+
{
68+
logger.LogWarning(ex, "Failed to deserialize MCP session migration state for session {SessionIdHash}", HashSessionId(sessionId));
69+
await cacheClient.RemoveAsync(cacheKey);
70+
return null;
71+
}
72+
73+
if (state is null)
74+
{
75+
await cacheClient.RemoveAsync(cacheKey);
76+
return null;
77+
}
78+
79+
if (!Matches(state, sessionIdentity))
80+
{
81+
logger.LogWarning("Rejected MCP session migration for session {SessionIdHash} because the authenticated caller changed", HashSessionId(sessionId));
82+
return null;
83+
}
84+
85+
await cacheClient.SetExpirationAsync(cacheKey, GetCacheLifetime());
86+
return state.InitializeParams;
87+
}
88+
89+
private TimeSpan GetCacheLifetime()
90+
{
91+
TimeSpan idleTimeout = transportOptions.Value.IdleTimeout;
92+
if (idleTimeout <= TimeSpan.Zero)
93+
idleTimeout = DefaultIdleTimeout;
94+
95+
return idleTimeout + LifetimeBuffer;
96+
}
97+
98+
private static bool Matches(McpSessionMigrationState state, McpSessionIdentity sessionIdentity)
99+
{
100+
return String.Equals(state.UserId, sessionIdentity.UserId, StringComparison.Ordinal)
101+
&& String.Equals(state.ClientId, sessionIdentity.ClientId, StringComparison.Ordinal)
102+
&& String.Equals(state.Resource, sessionIdentity.Resource, StringComparison.Ordinal);
103+
}
104+
105+
private static McpSessionIdentity? GetSessionIdentity(HttpContext context)
106+
{
107+
var user = context.User;
108+
if (!user.IsAuthenticated() || !user.IsInRole(AuthorizationRoles.McpRead))
109+
return null;
110+
111+
string? userId = user.GetClaimValue(ClaimTypes.NameIdentifier);
112+
if (String.IsNullOrWhiteSpace(userId))
113+
return null;
114+
115+
string clientId = user.GetClaimValue(IdentityUtils.OAuthClientIdClaim) ?? UserClientId;
116+
string resource = user.GetClaimValue(IdentityUtils.OAuthResourceClaim) ?? context.Request.Path.ToString();
117+
return String.IsNullOrWhiteSpace(resource)
118+
? null
119+
: new McpSessionIdentity(userId, clientId, resource);
120+
}
121+
122+
private static string GetCacheKey(string sessionId)
123+
{
124+
return String.Concat(CacheKeyPrefix, sessionId.ToSHA1());
125+
}
126+
127+
private static string HashSessionId(string sessionId)
128+
{
129+
return sessionId.ToSHA1();
130+
}
131+
132+
private sealed record McpSessionIdentity(string UserId, string ClientId, string Resource);
133+
}
134+
135+
public sealed class McpSessionMigrationState
136+
{
137+
public required string UserId { get; init; }
138+
139+
public required string ClientId { get; init; }
140+
141+
public required string Resource { get; init; }
142+
143+
public required InitializeRequestParams InitializeParams { get; init; }
144+
145+
public DateTime CreatedUtc { get; init; }
146+
}

‎src/Exceptionless.Web/Startup.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
using Microsoft.AspNetCore.Mvc;
2525
using Microsoft.AspNetCore.Mvc.ModelBinding.Metadata;
2626
using Microsoft.Net.Http.Headers;
27+
using ModelContextProtocol.AspNetCore;
2728
using Scalar.AspNetCore;
2829
using Serilog;
2930
using Serilog.Events;
@@ -134,6 +135,7 @@ public void ConfigureServices(IServiceCollection services)
134135
var appOptions = AppOptions.ReadFromConfiguration(Configuration);
135136
Bootstrapper.RegisterServices(services, appOptions, Log.Logger.ToLoggerFactory());
136137
services.AddScoped<McpContextService>();
138+
services.AddSingleton<ISessionMigrationHandler, McpSessionMigrationHandler>();
137139
services.AddMcpServer()
138140
.WithHttpTransport(o => o.Stateless = false)
139141
.WithTools<ExceptionlessMcpTools>();

‎tests/Exceptionless.Tests/Mcp/McpSessionMigrationHandlerTests.cs‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
using System.Security.Claims;
2+
using Exceptionless.Core.Authorization;
3+
using Exceptionless.Core.Extensions;
4+
using Exceptionless.Web.Mcp;
5+
using Foundatio.Caching;
6+
using Foundatio.Serializer;
7+
using Microsoft.AspNetCore.Http;
8+
using Microsoft.Extensions.Logging;
9+
using Microsoft.Extensions.Options;
10+
using ModelContextProtocol.AspNetCore;
11+
using ModelContextProtocol.Protocol;
12+
using Xunit;
13+
14+
namespace Exceptionless.Tests.Mcp;
15+
16+
public sealed class McpSessionMigrationHandlerTests : TestWithServices
17+
{
18+
public McpSessionMigrationHandlerTests(ITestOutputHelper output) : base(output) { }
19+
20+
[Fact]
21+
public async Task AllowSessionMigrationAsync_SameCaller_ReturnsStoredInitializeParams()
22+
{
23+
var handler = CreateHandler();
24+
string sessionId = Guid.NewGuid().ToString("N");
25+
var context = CreateContext();
26+
var initializeParams = CreateInitializeParams();
27+
28+
await handler.OnSessionInitializedAsync(context, sessionId, initializeParams, CancellationToken.None);
29+
30+
var restored = await handler.AllowSessionMigrationAsync(context, sessionId, CancellationToken.None);
31+
32+
Assert.NotNull(restored);
33+
Assert.Equal(initializeParams.ProtocolVersion, restored.ProtocolVersion);
34+
Assert.Equal(initializeParams.ClientInfo.Name, restored.ClientInfo.Name);
35+
Assert.Equal(initializeParams.ClientInfo.Version, restored.ClientInfo.Version);
36+
}
37+
38+
[Fact]
39+
public async Task AllowSessionMigrationAsync_DifferentUser_ReturnsNull()
40+
{
41+
var handler = CreateHandler();
42+
string sessionId = Guid.NewGuid().ToString("N");
43+
var initializeContext = CreateContext(userId: "user-a");
44+
var migrationContext = CreateContext(userId: "user-b");
45+
46+
await handler.OnSessionInitializedAsync(initializeContext, sessionId, CreateInitializeParams(), CancellationToken.None);
47+
48+
var restored = await handler.AllowSessionMigrationAsync(migrationContext, sessionId, CancellationToken.None);
49+
50+
Assert.Null(restored);
51+
}
52+
53+
[Fact]
54+
public async Task AllowSessionMigrationAsync_DifferentOAuthClient_ReturnsNull()
55+
{
56+
var handler = CreateHandler();
57+
string sessionId = Guid.NewGuid().ToString("N");
58+
var initializeContext = CreateContext(clientId: "client-a");
59+
var migrationContext = CreateContext(clientId: "client-b");
60+
61+
await handler.OnSessionInitializedAsync(initializeContext, sessionId, CreateInitializeParams(), CancellationToken.None);
62+
63+
var restored = await handler.AllowSessionMigrationAsync(migrationContext, sessionId, CancellationToken.None);
64+
65+
Assert.Null(restored);
66+
}
67+
68+
[Fact]
69+
public async Task AllowSessionMigrationAsync_MissingMcpReadRole_ReturnsNull()
70+
{
71+
var handler = CreateHandler();
72+
string sessionId = Guid.NewGuid().ToString("N");
73+
var initializeContext = CreateContext();
74+
var migrationContext = CreateContext(includeMcpReadRole: false);
75+
76+
await handler.OnSessionInitializedAsync(initializeContext, sessionId, CreateInitializeParams(), CancellationToken.None);
77+
78+
var restored = await handler.AllowSessionMigrationAsync(migrationContext, sessionId, CancellationToken.None);
79+
80+
Assert.Null(restored);
81+
}
82+
83+
private McpSessionMigrationHandler CreateHandler()
84+
{
85+
return new McpSessionMigrationHandler(
86+
GetService<ICacheClient>(),
87+
GetService<ITextSerializer>(),
88+
Options.Create(new HttpServerTransportOptions { IdleTimeout = TimeSpan.FromMinutes(30) }),
89+
TimeProvider,
90+
GetService<ILogger<McpSessionMigrationHandler>>());
91+
}
92+
93+
private static DefaultHttpContext CreateContext(
94+
string userId = "user-1",
95+
string clientId = "test-client",
96+
string resource = "http://localhost/mcp",
97+
bool includeMcpReadRole = true)
98+
{
99+
var claims = new List<Claim>
100+
{
101+
new(ClaimTypes.NameIdentifier, userId),
102+
new(IdentityUtils.OAuthClientIdClaim, clientId),
103+
new(IdentityUtils.OAuthResourceClaim, resource)
104+
};
105+
106+
if (includeMcpReadRole)
107+
claims.Add(new Claim(ClaimTypes.Role, AuthorizationRoles.McpRead));
108+
109+
var context = new DefaultHttpContext
110+
{
111+
User = new ClaimsPrincipal(new ClaimsIdentity(claims, IdentityUtils.TokenAuthenticationType))
112+
};
113+
context.Request.Scheme = "http";
114+
context.Request.Host = new HostString("localhost");
115+
context.Request.Path = "/mcp";
116+
return context;
117+
}
118+
119+
private static InitializeRequestParams CreateInitializeParams()
120+
{
121+
return new InitializeRequestParams
122+
{
123+
ProtocolVersion = "2025-06-18",
124+
Capabilities = new ClientCapabilities(),
125+
ClientInfo = new Implementation
126+
{
127+
Name = "codex",
128+
Version = "1.0.0"
129+
}
130+
};
131+
}
132+
}

0 commit comments

Comments
 (0)