MCPGateway/src/ManagedCode.MCPGateway/Internal/Catalog/McpGatewayRegistry.cs at main · managedcode/MCPGateway · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
using ManagedCode.MCPGateway.Abstractions;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.Options;
using ModelContextProtocol.Client;

namespace ManagedCode.MCPGateway;

internal sealed class McpGatewayRegistry(IOptions<McpGatewayOptions> options) : IMcpGatewayRegistry, IMcpGatewayCatalogSource, IAsyncDisposable
{
    private readonly McpGatewayRegistrationCollection _registrations = CreateRegistrations(options);
    private readonly McpGatewayOperationGate _operationGate = new();
    private int _version;

    public void AddTool(string sourceId, AITool tool, string? displayName = null)
        => Mutate(registrations => registrations.AddTool(sourceId, tool, displayName));

    public void AddTool(AITool tool, string sourceId = McpGatewayDefaults.DefaultSourceId, string? displayName = null)
        => Mutate(registrations => registrations.AddTool(tool, sourceId, displayName));

    public void AddTools(string sourceId, IEnumerable<AITool> tools, string? displayName = null)
        => Mutate(registrations => registrations.AddTools(sourceId, tools, displayName));

    public void AddTools(IEnumerable<AITool> tools, string sourceId = McpGatewayDefaults.DefaultSourceId, string? displayName = null)
        => Mutate(registrations => registrations.AddTools(tools, sourceId, displayName));

    public void AddHttpServer(
        string sourceId,
        Uri endpoint,
        IReadOnlyDictionary<string, string>? headers = null,
        string? displayName = null)
        => Mutate(registrations => registrations.AddHttpServer(sourceId, endpoint, headers, displayName));

    public void AddStdioServer(
        string sourceId,
        string command,
        IReadOnlyList<string>? arguments = null,
        string? workingDirectory = null,
        IReadOnlyDictionary<string, string?>? environmentVariables = null,
        string? displayName = null)
        => Mutate(registrations => registrations.AddStdioServer(sourceId, command, arguments, workingDirectory, environmentVariables, displayName));

    public void AddMcpClient(
        string sourceId,
        McpClient client,
        bool disposeClient = false,
        string? displayName = null)
        => Mutate(registrations => registrations.AddMcpClient(sourceId, client, disposeClient, displayName));

    public void AddMcpClientFactory(
        string sourceId,
        Func<CancellationToken, ValueTask<McpClient>> clientFactory,
        bool disposeClient = true,
        string? displayName = null)
        => Mutate(registrations => registrations.AddMcpClientFactory(sourceId, clientFactory, disposeClient, displayName));

    public McpGatewayCatalogSourceSnapshot CreateSnapshot()
    {
        _operationGate.Enter(this);
        try
        {
            _operationGate.ThrowIfDisposed(this);
            return new McpGatewayCatalogSourceSnapshot(
                Volatile.Read(ref _version),
                _registrations.Snapshot());
        }
        finally
        {
            _operationGate.Exit();
        }
    }

    public async ValueTask DisposeAsync()
    {
        if (!_operationGate.TryStartDispose(out var waitForDrain))
        {
            return;
        }

        Interlocked.Increment(ref _version);
        await waitForDrain;

        var registrations = _registrations.Drain();
        foreach (var registration in registrations)
        {
            await registration.DisposeAsync();
        }
    }

    private void Mutate(Action<McpGatewayRegistrationCollection> mutation)
    {
        _operationGate.Enter(this);
        try
        {
            _operationGate.ThrowIfDisposed(this);
            mutation(_registrations);
            Interlocked.Increment(ref _version);
        }
        finally
        {
            _operationGate.Exit();
        }
    }

    private static McpGatewayRegistrationCollection CreateRegistrations(IOptions<McpGatewayOptions> options)
    {
        ArgumentNullException.ThrowIfNull(options);
        return new McpGatewayRegistrationCollection(options.Value.SourceRegistrations);
    }
}