79 lines
3.0 KiB
C#
79 lines
3.0 KiB
C#
|
using System.Text.Json;
|
||
|
|
using Microsoft.EntityFrameworkCore;
|
||
|
|
using Microsoft.Extensions.Logging;
|
||
|
|
using TeamUp.Modules.Integrations.Persistence;
|
||
|
|
using TeamUp.Modules.Integrations.Security;
|
||
|
|
using TeamUp.SharedKernel.Ai;
|
||
|
|
|
||
|
|
namespace TeamUp.Modules.Integrations.Mcp;
|
||
|
|
|
||
|
|
/// <summary>
|
||
|
|
/// Resolves an org's MCP server configs, decrypts their headers server-side, and talks to them via
|
||
|
|
/// <see cref="McpClient"/>. Discovery is best-effort: a server that fails to connect is logged and
|
||
|
|
/// skipped so it never fails the agent run. The decrypted headers never leave the server.
|
||
|
|
/// </summary>
|
||
|
|
internal sealed class McpGateway(
|
||
|
|
IntegrationsDbContext db,
|
||
|
|
McpClient client,
|
||
|
|
ISecretProtector protector,
|
||
|
|
ILogger<McpGateway> logger) : IMcpGateway
|
||
|
|
{
|
||
|
|
public async Task<IReadOnlyList<McpToolDescriptor>> ListToolsAsync(
|
||
|
|
Guid organizationId, IReadOnlyCollection<Guid> serverIds, CancellationToken cancellationToken = default)
|
||
|
|
{
|
||
|
|
if (serverIds.Count == 0)
|
||
|
|
{
|
||
|
|
return [];
|
||
|
|
}
|
||
|
|
|
||
|
|
var idSet = serverIds.ToHashSet();
|
||
|
|
var servers = await db.McpServers
|
||
|
|
.Where(s => s.OrganizationId == organizationId && s.Enabled && idSet.Contains(s.Id))
|
||
|
|
.ToListAsync(cancellationToken);
|
||
|
|
|
||
|
|
var tools = new List<McpToolDescriptor>();
|
||
|
|
foreach (var server in servers)
|
||
|
|
{
|
||
|
|
try
|
||
|
|
{
|
||
|
|
var discovered = await client.ListToolsAsync(server.Endpoint, DecryptHeaders(server.EncryptedHeaders), cancellationToken);
|
||
|
|
tools.AddRange(discovered.Select(t =>
|
||
|
|
new McpToolDescriptor(server.Id, server.Name, t.Name, t.Description, t.InputSchemaJson)));
|
||
|
|
}
|
||
|
|
catch (Exception ex)
|
||
|
|
{
|
||
|
|
logger.LogWarning(ex, "MCP server {Server} ({Endpoint}) unreachable; skipping its tools.", server.Name, server.Endpoint);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return tools;
|
||
|
|
}
|
||
|
|
|
||
|
|
public async Task<McpToolResult> CallToolAsync(
|
||
|
|
Guid organizationId, Guid serverId, string toolName, string argumentsJson, CancellationToken cancellationToken = default)
|
||
|
|
{
|
||
|
|
var server = await db.McpServers.FirstOrDefaultAsync(
|
||
|
|
s => s.Id == serverId && s.OrganizationId == organizationId && s.Enabled, cancellationToken);
|
||
|
|
if (server is null)
|
||
|
|
{
|
||
|
|
return new McpToolResult(false, null, "MCP server not found or disabled.");
|
||
|
|
}
|
||
|
|
|
||
|
|
try
|
||
|
|
{
|
||
|
|
var (success, content, error) = await client.CallToolAsync(
|
||
|
|
server.Endpoint, DecryptHeaders(server.EncryptedHeaders), toolName, argumentsJson, cancellationToken);
|
||
|
|
return new McpToolResult(success, content, error);
|
||
|
|
}
|
||
|
|
catch (Exception ex)
|
||
|
|
{
|
||
|
|
return new McpToolResult(false, null, ex.Message);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
private Dictionary<string, string>? DecryptHeaders(string? encrypted) =>
|
||
|
|
string.IsNullOrEmpty(encrypted)
|
||
|
|
? null
|
||
|
|
: JsonSerializer.Deserialize<Dictionary<string, string>>(protector.Unprotect(encrypted));
|
||
|
|
}
|