VolcanicAgenticCore/extensions/Volcanic.AgenticCore.Extensions.A2A/RemoteAgents.cs

using A2A;

namespace Volcanic.AgenticCore.Extensions.A2A;

internal static class RemoteAgentEntryExtensions
{
    /// <summary>Texts shown to the model or logged (exception messages) must never echo the agent's key.</summary>
    public static string Scrub(this RemoteAgentEntry entry, string text) =>
        entry.Key is { Length: > 0 } key ? text.Replace(key, "[key]", StringComparison.Ordinal) : text;
}

/// <summary>A connected remote agent: its (extended) card and a client that sends the key.</summary>
internal sealed class RemoteAgent(RemoteAgentEntry entry, AgentCard card, IA2AClient client, HttpClient http) : IDisposable
{
    public RemoteAgentEntry Entry => entry;
    public AgentCard Card => card;
    public IA2AClient Client => client;

    /// <summary>Sends the agent's key: use only for the agent's own origin.</summary>
    public HttpClient Http => http;
    public bool Streaming => card.Capabilities?.Streaming == true;
    public bool SupportsPush => card.Capabilities?.PushNotifications == true;

    public void Dispose()
    {
        (client as IDisposable)?.Dispose();
        http.Dispose();
    }
}

/// <summary>Resolves an agent's card (public, then extended when we have a key) and builds a client over one shared handler.</summary>
internal sealed class RemoteAgentConnector(HttpMessageHandler handler, TimeSpan cardTimeout)
{
    /// <summary>An explicit <c>….json</c> URL is the card itself; anything else is a base URL with the well-known path.</summary>
    public static Uri CardUrl(string url)
    {
        var uri = new Uri(url);
        return uri.AbsolutePath.EndsWith(".json", StringComparison.OrdinalIgnoreCase)
            ? uri
            : new Uri(new Uri(url.TrimEnd('/') + "/"), ".well-known/agent-card.json");
    }

    public async Task<RemoteAgent> ConnectAsync(RemoteAgentEntry entry, CancellationToken cancellationToken)
    {
        var cardUrl = CardUrl(entry.Url);
        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeout.CancelAfter(cardTimeout);
        AgentCard card;
        using (var publicHttp = new HttpClient(handler, disposeHandler: false))
            card = await new A2ACardResolver(new Uri(cardUrl.GetLeftPart(UriPartial.Authority)), publicHttp, cardUrl.PathAndQuery)
                .GetAgentCardAsync(timeout.Token).ConfigureAwait(false);

        var http = new HttpClient(handler, disposeHandler: false) { Timeout = Timeout.InfiniteTimeSpan };
        try
        {
            if (entry.Key is { Length: > 0 } key)
            {
                if (UsesBearer(card)) http.DefaultRequestHeaders.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", key);
                else http.DefaultRequestHeaders.Add(ApiKeyHeader(card), key);
            }
            var client = A2AClientFactory.Create(card, http);
            if (entry.Key is { Length: > 0 } && card.Capabilities?.ExtendedAgentCard == true)
            {
                try
                {
                    var extended = await client.GetExtendedAgentCardAsync(new GetExtendedAgentCardRequest(), timeout.Token).ConfigureAwait(false);
                    (client as IDisposable)?.Dispose();
                    card = extended;
                    client = A2AClientFactory.Create(card, http);
                }
                // Refused, unreachable or too slow (the card timeout, not the caller's cancel): keep the public card.
                catch (Exception ex) when (ex is A2AException or HttpRequestException ||
                                           (ex is OperationCanceledException && !cancellationToken.IsCancellationRequested)) { }
            }
            return new RemoteAgent(entry, card, client, http);
        }
        catch
        {
            http.Dispose();
            throw;
        }
    }

    private static bool UsesBearer(AgentCard card) =>
        card.SecuritySchemes is { Count: > 0 } schemes &&
        schemes.Values.All(s => s.HttpAuthSecurityScheme is { Scheme: var scheme } && string.Equals(scheme, "bearer", StringComparison.OrdinalIgnoreCase));

    private static string ApiKeyHeader(AgentCard card) =>
        card.SecuritySchemes?.Values.Select(s => s.ApiKeySecurityScheme)
            .FirstOrDefault(k => k is not null && string.Equals(k.Location, "header", StringComparison.OrdinalIgnoreCase))?.Name
        ?? "X-API-Key";
}