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

using System.Net.Http.Headers;
using A2A;
using A2AMessage = global::A2A.Message;

namespace Volcanic.AgenticCore.Extensions.A2A;

/// <summary>
/// A file the remote agent sent: saved (<see cref="Path"/>, workspace-relative) or not (<see cref="Problem"/>, e.g. over
/// the size limit or a failed download, with the <see cref="Url"/> when there was one).
/// </summary>
internal sealed record ReceivedFile(string Name, long? Bytes, string? MediaType, string? Path, string? Problem = null, string? Url = null);

/// <summary>
/// Saves the file parts of a remote answer under <c>inbox/a2a/&lt;taskId&gt;/</c>: inline bytes are written, URLs are
/// downloaded (streamed, capped by <see cref="A2AClientSettings.MaxReceivedFileBytes"/>). Parts come from the artifacts,
/// the status message and the agent's messages after our last one in the history; the same part found twice is saved
/// once. Saving the same task again (status after a watch) overwrites its files, so it is idempotent.
/// </summary>
internal static class A2AFiles
{
    public static readonly TimeSpan DownloadTimeout = TimeSpan.FromMinutes(5);

    public static Task<IReadOnlyList<ReceivedFile>> SaveAsync(string workspace, AgentTask task, RemoteAgent? agent,
        HttpMessageHandler handler, A2AClientSettings settings, CancellationToken ct) =>
        SaveAsync(workspace, task.Id, PartsOf(task), agent, handler, settings, ct);

    public static Task<IReadOnlyList<ReceivedFile>> SaveAsync(string workspace, A2AMessage message, RemoteAgent? agent,
        HttpMessageHandler handler, A2AClientSettings settings, CancellationToken ct) =>
        SaveAsync(workspace, message.TaskId ?? message.MessageId, (message.Parts ?? []).Select(p => (p, (string?)null)), agent, handler, settings, ct);

    /// <summary>File parts with a fallback name (the artifact's), in order: artifacts, status message, recent agent messages.</summary>
    internal static IEnumerable<(Part Part, string? Fallback)> PartsOf(AgentTask task)
    {
        foreach (var artifact in task.Artifacts ?? [])
            foreach (var part in artifact.Parts ?? [])
                yield return (part, artifact.Name ?? artifact.ArtifactId);
        foreach (var part in task.Status.Message?.Parts ?? [])
            yield return (part, null);
        var history = task.History ?? [];
        var lastOurs = history.FindLastIndex(m => m.Role == Role.User);
        foreach (var message in history.Skip(lastOurs + 1).Where(m => m.Role == Role.Agent))
            foreach (var part in message.Parts ?? [])
                yield return (part, null);
    }

    private static async Task<IReadOnlyList<ReceivedFile>> SaveAsync(string workspace, string? folderId,
        IEnumerable<(Part Part, string? Fallback)> parts, RemoteAgent? agent, HttpMessageHandler handler,
        A2AClientSettings settings, CancellationToken ct)
    {
        var files = new List<ReceivedFile>();
        var seen = new HashSet<string>();
        var names = new HashSet<string>(StringComparer.Ordinal);
        var folder = $"inbox/a2a/{Sanitize(folderId)}";
        foreach (var (part, fallback) in parts)
        {
            if (part.ContentCase == PartContentCase.None &&
                part.Metadata?.GetValueOrDefault(A2AWireCompat.UnrecognizedKey) is { ValueKind: System.Text.Json.JsonValueKind.String } keys)
            {
                if (seen.Add($"unknown:{keys.GetString()}"))
                    files.Add(new ReceivedFile("part", null, null, null, $"not understood (fields: {keys.GetString()})"));
                continue;
            }
            if (part.ContentCase is not (PartContentCase.Raw or PartContentCase.Url)) continue;
            var name = Sanitize(part.Filename ?? (part.Url is { } u ? UrlName(u) : null) ?? fallback);
            var mediaType = string.IsNullOrWhiteSpace(part.MediaType) ? null : part.MediaType;
            if (!seen.Add(part.Raw is { } raw ? $"raw:{name}:{raw.Length}" : $"url:{part.Url}")) continue;
            var relative = UniqueName(folder, name, names);
            var full = Path.Combine(workspace, relative);
            if (part.Raw is { } bytes)
            {
                if (bytes.LongLength > settings.MaxReceivedFileBytes)
                {
                    files.Add(new ReceivedFile(name, bytes.LongLength, mediaType, null, $"over {settings.MaxReceivedFileBytes} bytes"));
                    continue;
                }
                try
                {
                    Directory.CreateDirectory(Path.GetDirectoryName(full)!);
                    await File.WriteAllBytesAsync(full, bytes, ct).ConfigureAwait(false);
                    files.Add(new ReceivedFile(name, bytes.LongLength, mediaType, relative));
                }
                // A file that can't be written must not cost the answer (or loop a watch): report it instead.
                catch (Exception ex) when (ex is IOException or UnauthorizedAccessException)
                {
                    files.Add(new ReceivedFile(name, bytes.LongLength, mediaType, null, $"could not be written ({ex.GetType().Name})"));
                }
            }
            else
            {
                files.Add(await DownloadAsync(part.Url!, name, mediaType, relative, full, agent, handler, settings, ct).ConfigureAwait(false));
            }
        }
        return files;
    }

    /// <summary>
    /// Downloads a URL part to <paramref name="full"/> through a temp file (a failed download leaves nothing behind).
    /// The agent's key goes only to the agent's own origin; elsewhere only <c>https</c> is fetched, without the key.
    /// </summary>
    private static async Task<ReceivedFile> DownloadAsync(string url, string name, string? mediaType, string relative, string full,
        RemoteAgent? agent, HttpMessageHandler handler, A2AClientSettings settings, CancellationToken ct)
    {
        ReceivedFile Failed(string problem, long? size = null) => new(name, size, mediaType, null, problem, url);
        if (!Uri.TryCreate(url, UriKind.Absolute, out var uri) || uri.Scheme is not ("https" or "http"))
            return Failed("not an http(s) link, not downloaded");
        var sameOrigin = agent is not null && SameOrigin(uri, agent);
        if (!sameOrigin && uri.Scheme != "https") return Failed("not https, not downloaded");

        var temp = full + ".part";
        try
        {
            using var timeout = CancellationTokenSource.CreateLinkedTokenSource(ct);
            timeout.CancelAfter(DownloadTimeout);
            using var plain = sameOrigin ? null : new HttpClient(handler, disposeHandler: false) { Timeout = Timeout.InfiniteTimeSpan };
            var http = sameOrigin ? agent!.Http : plain!;
            using var response = await http.GetAsync(uri, HttpCompletionOption.ResponseHeadersRead, timeout.Token).ConfigureAwait(false);
            if (!response.IsSuccessStatusCode) return Failed($"download failed ({(int)response.StatusCode})");
            if (response.Content.Headers.ContentLength is { } length && length > settings.MaxReceivedFileBytes)
                return Failed($"over {settings.MaxReceivedFileBytes} bytes", length);
            mediaType ??= response.Content.Headers.ContentType is MediaTypeHeaderValue { MediaType: { } type } ? type : null;

            Directory.CreateDirectory(Path.GetDirectoryName(full)!);
            long total = 0;
            await using (var source = await response.Content.ReadAsStreamAsync(timeout.Token).ConfigureAwait(false))
            await using (var target = new FileStream(temp, FileMode.Create, FileAccess.Write, FileShare.None, 81920, FileOptions.Asynchronous))
            {
                var buffer = new byte[81920];
                int n;
                while ((n = await source.ReadAsync(buffer, timeout.Token).ConfigureAwait(false)) > 0)
                {
                    total += n;
                    if (total > settings.MaxReceivedFileBytes) break;
                    await target.WriteAsync(buffer.AsMemory(0, n), timeout.Token).ConfigureAwait(false);
                }
            }
            if (total > settings.MaxReceivedFileBytes)
            {
                File.Delete(temp);
                return Failed($"over {settings.MaxReceivedFileBytes} bytes");
            }
            File.Move(temp, full, overwrite: true);
            return new ReceivedFile(name, total, mediaType, relative, null, url);
        }
        catch (Exception ex) when (ex is HttpRequestException or IOException or UnauthorizedAccessException ||
                                   (ex is OperationCanceledException && !ct.IsCancellationRequested))
        {
            TryDelete(temp);
            var reason = ex is OperationCanceledException ? $"no complete download within {DownloadTimeout}"
                : agent is null ? ex.Message : agent.Entry.Scrub(ex.Message);
            return Failed($"download failed: {reason}");
        }
        catch
        {
            TryDelete(temp);
            throw;
        }
    }

    private static bool SameOrigin(Uri uri, RemoteAgent agent)
    {
        var origins = new List<string?> { agent.Entry.Url };
        origins.AddRange((agent.Card.SupportedInterfaces ?? []).Select(i => i.Url));
        return origins.Any(o => Uri.TryCreate(o, UriKind.Absolute, out var a) &&
                                Uri.Compare(a, uri, UriComponents.SchemeAndServer, UriFormat.Unescaped, StringComparison.OrdinalIgnoreCase) == 0);
    }

    private static void TryDelete(string path)
    {
        try { File.Delete(path); } catch (IOException) { } catch (UnauthorizedAccessException) { }
    }

    private static string UniqueName(string folder, string name, HashSet<string> taken)
    {
        var relative = $"{folder}/{name}";
        for (var n = 2; !taken.Add(relative); n++)
            relative = $"{folder}/{Path.GetFileNameWithoutExtension(name)}-{n}{Path.GetExtension(name)}";
        return relative;
    }

    private static string? UrlName(string url) =>
        Uri.TryCreate(url, UriKind.Absolute, out var uri) && Uri.UnescapeDataString(uri.Segments.LastOrDefault() ?? "").Trim('/') is { Length: > 0 } last
            ? last : null;

    internal static string Sanitize(string? name)
    {
        var file = Path.GetFileName(name ?? "");
        if (string.IsNullOrWhiteSpace(file) || file is "." or "..") return "file";
        var invalid = Path.GetInvalidFileNameChars();
        var clean = new string(file.Select(c => invalid.Contains(c) || char.IsWhiteSpace(c) ? '_' : c).Where(c => !char.IsControl(c)).ToArray());
        if (clean.Length == 0 || clean is "." or "..") return "file";
        return clean.Length > 80 ? clean[^80..] : clean;
    }
}