AI Management: MCP client sessions are cached for the process lifetime and never invalidated (expired sessions, configuration changes, tool list changes, shutdown) #10919

1 Answer(s)

• 

  1

  maliming created 7 days ago

  Support Team Fullstack Developer

  Hi,

  All four are bugs in the AI Management module. The MCP sessions are cached for the whole process and nothing refreshes or removes them.

  You can replace McpToolProvider in your Domain project for now. It creates a new session when the MCP server configuration changes. When a tool call gets 404, it starts a new session and retries the call once.

  using System;
  using System.Collections.Generic;
  using System.Linq;
  using System.Net;
  using System.Net.Http;
  using System.Security.Cryptography;
  using System.Text.Json;
  using System.Threading;
  using System.Threading.Tasks;
  using Microsoft.Extensions.AI;
  using Microsoft.Extensions.Logging;
  using Microsoft.Extensions.Options;
  using Volo.Abp.DependencyInjection;
  using Volo.AIManagement.Mcp;
  
  namespace YourProjectName;
  
  [Dependency(ReplaceServices = true)]
  [ExposeServices(typeof(IMcpToolProvider), typeof(McpToolProvider))]
  public class MyMcpToolProvider : McpToolProvider
  {
      public MyMcpToolProvider(
          ILogger<McpToolProvider> logger,
          IMcpServerConfigurationStore mcpServerConfigurationStore,
          IOptions<McpClientFactoryOptions> options,
          IServiceProvider serviceProvider,
          IMcpToolCatalogCache toolCatalogCache)
          : base(logger, mcpServerConfigurationStore, options, serviceProvider, toolCatalogCache)
      {
      }
  
      protected override async Task<IReadOnlyList<AIFunction>> GetToolsForConfigurationAsync(
          McpClientCreationConfiguration configuration,
          string workspaceName,
          CancellationToken cancellationToken)
      {
          IMcpClientSession? session = null;
          try
          {
              session = await GetOrCreateSessionAsync(configuration, cancellationToken);
              var tools = await session.GetToolsAsync(cancellationToken);
  
              await ToolCatalogCache.SetAsync(
                  configuration.McpServerId,
                  McpToolCatalogFingerprint.Compute(configuration),
                  McpToolCatalogMappingHelper.FromAIFunctions(tools),
                  cancellationToken: cancellationToken);
  
              return tools
                  .Select(tool => (AIFunction)new SessionRecoveringMcpTool(tool, this, configuration))
                  .ToList();
          }
          catch (Exception ex)
          {
              if (session != null)
              {
                  await RemoveSessionAsync(configuration, session);
              }
  
              Logger.LogWarning(ex, "Failed to get tools from MCP server '{ServerName}' for workspace '{WorkspaceName}'",
                  configuration.Name, workspaceName);
              return Array.Empty<AIFunction>();
          }
      }
  
      protected override async Task<IMcpClientSession> GetOrCreateSessionAsync(
          McpClientCreationConfiguration configuration,
          CancellationToken cancellationToken)
      {
          // The key contains the full configuration, so editing an MCP server creates a new session.
          var sessionKey = GetSessionKey(configuration);
  
          if (ActiveSessions.TryGetValue(sessionKey, out var existingSession) && existingSession.IsConnected)
          {
              return existingSession;
          }
  
          var gate = SessionCreateLocks.GetOrAdd(sessionKey, _ => new SemaphoreSlim(1, 1));
          await gate.WaitAsync(cancellationToken);
          try
          {
              if (ActiveSessions.TryGetValue(sessionKey, out existingSession) && existingSession.IsConnected)
              {
                  return existingSession;
              }
  
              if (existingSession != null)
              {
                  ActiveSessions.TryRemove(sessionKey, out _);
                  await DisposeSessionAsync(existingSession);
              }
  
              var session = await CreateSessionAsync(configuration, cancellationToken);
              ActiveSessions[sessionKey] = session;
              return session;
          }
          finally
          {
              gate.Release();
          }
      }
  
      public virtual async Task<object?> InvokeToolAsync(
          McpClientCreationConfiguration configuration,
          string toolName,
          AIFunctionArguments arguments,
          CancellationToken cancellationToken)
      {
          var session = await GetOrCreateSessionAsync(configuration, cancellationToken);
          try
          {
              return await InvokeToolAsync(session, toolName, arguments, cancellationToken);
          }
          catch (HttpRequestException ex) when (ex.StatusCode == HttpStatusCode.NotFound)
          {
              // The MCP server dropped the session (idle timeout or restart), start a new one and retry once.
              await RemoveSessionAsync(configuration, session);
              session = await GetOrCreateSessionAsync(configuration, cancellationToken);
              return await InvokeToolAsync(session, toolName, arguments, cancellationToken);
          }
      }
  
      protected virtual async Task<object?> InvokeToolAsync(
          IMcpClientSession session,
          string toolName,
          AIFunctionArguments arguments,
          CancellationToken cancellationToken)
      {
          var tools = await session.GetToolsAsync(cancellationToken);
          var tool = tools.FirstOrDefault(t => t.Name == toolName)
                     ?? throw new InvalidOperationException($"The MCP tool '{toolName}' is not available anymore.");
  
          return await tool.InvokeAsync(arguments, cancellationToken);
      }
  
      protected virtual async Task RemoveSessionAsync(McpClientCreationConfiguration configuration, IMcpClientSession session)
      {
          // Only remove the given session, another request may have replaced it already.
          if (ActiveSessions.TryRemove(new KeyValuePair<string, IMcpClientSession>(GetSessionKey(configuration), session)))
          {
              await DisposeSessionAsync(session);
          }
      }
  
      protected virtual string GetSessionKey(McpClientCreationConfiguration configuration)
      {
          var json = JsonSerializer.Serialize(new
          {
              Fingerprint = McpToolCatalogFingerprint.Compute(configuration),
              configuration.Name,
              configuration.ApiKey,
              configuration.CustomAuthHeaderValue
          });
  
          return $"{configuration.McpServerId}_{Convert.ToHexString(SHA256.HashData(System.Text.Encoding.UTF8.GetBytes(json)))}";
      }
  
      protected virtual async Task DisposeSessionAsync(IMcpClientSession session)
      {
          try
          {
              await session.DisposeAsync();
          }
          catch (Exception ex)
          {
              Logger.LogDebug(ex, "Failed to dispose MCP session '{SessionId}'", session.SessionId);
          }
      }
  }
  
  public class SessionRecoveringMcpTool : DelegatingAIFunction
  {
      protected MyMcpToolProvider Provider { get; }
      protected McpClientCreationConfiguration Configuration { get; }
  
      public SessionRecoveringMcpTool(
          AIFunction innerFunction,
          MyMcpToolProvider provider,
          McpClientCreationConfiguration configuration)
          : base(innerFunction)
      {
          Provider = provider;
          Configuration = configuration;
      }
  
      protected override async ValueTask<object?> InvokeCoreAsync(
          AIFunctionArguments arguments,
          CancellationToken cancellationToken)
      {
          return await Provider.InvokeToolAsync(Configuration, Name, arguments, cancellationToken);
      }
  }
  

  Then close the sessions on shutdown in your Domain module:

  using System.Threading.Tasks;
  using Volo.Abp;
  using Volo.AIManagement.Mcp;
  
  public override async Task OnApplicationShutdownAsync(ApplicationShutdownContext context)
  {
      await McpToolProvider.DisposeAllSessionsAsync();
  }
  

  This doesn't refresh the tool list of a running session. The updated tool list is used in the next chat request after a new session is created, for example after a 404 or a configuration change. Sessions of an old configuration are closed on shutdown.

  We will fix all four in the module, including handling notifications/tools/list_changed. Your ticket has been refunded.

  Thanks

  Markdown supported.

  Copy, paste, or drag & drop images and files (max 100 MB per file, 100 MB total per post)

  Save Cancel