Playa/extensions/Volcanic.AgenticCore.Extensions.A2A/A2AWireCompat.cs

using System.IO.Pipelines;
using System.Net.Http.Headers;
using System.Text;
using System.Text.Json.Nodes;
using Microsoft.Extensions.Logging;

namespace Volcanic.AgenticCore.Extensions.A2A;

/// <summary>
/// Sits between the A2A SDK and the network and rewrites file parts the SDK (protocol 1.0) would drop silently: the
/// 0.3 shape <c>{"kind":"file","file":{"bytes"|"uri","name","mimeType"}}</c> and the same fields directly on the part
/// become <c>{"raw"|"url","filename","mediaType"}</c>. A part with none of <c>text</c>/<c>raw</c>/<c>url</c>/<c>data</c>
/// afterwards gets its field names in metadata <see cref="UnrecognizedKey"/>, so the result can say so instead of
/// losing it. Works on JSON responses and on SSE streams (line by line, each <c>data:</c> line on its own).
/// </summary>
internal sealed class A2AWireCompat(HttpMessageHandler inner, Func<ILogger> logger) : DelegatingHandler(inner)
{
    public const string UnrecognizedKey = "volcanic.unrecognizedPart";

    private static readonly string[] V1Keys = ["text", "raw", "url", "data"];

    protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
    {
        var response = await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
        var type = response.Content.Headers.ContentType?.MediaType;
        if (type is "application/json")
        {
            var body = await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false);
            var rewritten = Rewrite(body, logger());
            if (rewritten is not null) Replace(response, new StringContent(rewritten, Encoding.UTF8, "application/json"), disposeOld: true);
        }
        else if (type is "text/event-stream")
        {
            var source = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
            var content = new StreamContent(RewriteEvents(source, logger()));
            content.Headers.ContentType = new MediaTypeHeaderValue("text/event-stream");
            Replace(response, content, disposeOld: false); // its stream now belongs to the new content
        }
        return response;
    }

    private static void Replace(HttpResponseMessage response, HttpContent content, bool disposeOld)
    {
        var old = response.Content;
        foreach (var header in old.Headers.Where(h => h.Key is not ("Content-Type" or "Content-Length")))
            content.Headers.TryAddWithoutValidation(header.Key, header.Value);
        response.Content = content;
        if (disposeOld) old.Dispose();
    }

    /// <summary>The JSON with its parts rewritten, or null when nothing changed (or it isn't JSON).</summary>
    internal static string? Rewrite(string json, ILogger logger)
    {
        if (json.IndexOf("\"parts\"", StringComparison.Ordinal) < 0) return null;
        JsonNode? root;
        try { root = JsonNode.Parse(json); }
        catch (System.Text.Json.JsonException) { return null; }
        var (converted, unknown) = (0, new List<string>());
        Walk(root, ref converted, unknown);
        if (converted == 0 && unknown.Count == 0) return null;
        if (converted > 0) logger.LogInformation("A2A: rewrote {Count} file part(s) from the 0.3 shape", converted);
        foreach (var keys in unknown) logger.LogWarning("A2A: a part with unknown fields ({Keys}) from a remote agent", keys);
        return root!.ToJsonString();
    }

    private static void Walk(JsonNode? node, ref int converted, List<string> unknown)
    {
        switch (node)
        {
            case JsonObject obj:
                if (obj["parts"] is JsonArray parts)
                    foreach (var part in parts.OfType<JsonObject>())
                        Normalize(part, ref converted, unknown);
                foreach (var (_, child) in obj.ToList()) Walk(child, ref converted, unknown);
                break;
            case JsonArray array:
                foreach (var child in array.ToList()) Walk(child, ref converted, unknown);
                break;
        }
    }

    private static void Normalize(JsonObject part, ref int converted, List<string> unknown)
    {
        if (V1Keys.Any(k => Has(part, k))) return;
        var file = part["file"] as JsonObject ?? (Has(part, "bytes") || Has(part, "uri") ? part : null);
        if (file is not null && (Has(file, "bytes") || Has(file, "uri")))
        {
            if (Has(file, "bytes")) part["raw"] = file["bytes"]!.DeepClone();
            else part["url"] = file["uri"]!.DeepClone();
            if (!Has(part, "filename") && Has(file, "name")) part["filename"] = file["name"]!.DeepClone();
            if (!Has(part, "mediaType") && Has(file, "mimeType")) part["mediaType"] = file["mimeType"]!.DeepClone();
            foreach (var key in new[] { "file", "bytes", "uri", "name", "mimeType", "kind" }) part.Remove(key);
            converted++;
            return;
        }
        var keys = string.Join(", ", part.Select(p => p.Value is JsonObject inner ? $"{p.Key}{{{string.Join(", ", inner.Select(i => i.Key))}}}" : p.Key));
        unknown.Add(keys);
        var metadata = part["metadata"] as JsonObject ?? new JsonObject();
        metadata[UnrecognizedKey] = keys;
        part["metadata"] = metadata;
    }

    private static bool Has(JsonObject obj, string key) => obj[key] is { } value && value.GetValueKind() != System.Text.Json.JsonValueKind.Null;

    /// <summary>The event stream, each <c>data:</c> line rewritten on its own; the rest passes through.</summary>
    private static Stream RewriteEvents(Stream source, ILogger logger)
    {
        var pipe = new Pipe();
        _ = Task.Run(async () =>
        {
            Exception? error = null;
            try
            {
                using var reader = new StreamReader(source, Encoding.UTF8);
                while (await reader.ReadLineAsync().ConfigureAwait(false) is { } line)
                {
                    if (line.StartsWith("data:", StringComparison.Ordinal) && Rewrite(line[5..].TrimStart(), logger) is { } rewritten)
                        line = "data: " + rewritten;
                    var result = await pipe.Writer.WriteAsync(Encoding.UTF8.GetBytes(line + "\n")).ConfigureAwait(false);
                    if (result.IsCompleted) break; // the reader is gone
                }
            }
            catch (Exception ex) { error = ex; }
            finally
            {
                await source.DisposeAsync().ConfigureAwait(false);
                await pipe.Writer.CompleteAsync(error).ConfigureAwait(false);
            }
        });
        return new OwningStream(pipe.Reader.AsStream(), source);
    }

    /// <summary>Reads the rewritten events; disposing it also closes the source, which stops the pump.</summary>
    private sealed class OwningStream(Stream inner, Stream source) : Stream
    {
        public override bool CanRead => true;
        public override bool CanSeek => false;
        public override bool CanWrite => false;
        public override long Length => throw new NotSupportedException();
        public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
        public override void Flush() { }
        public override int Read(byte[] buffer, int offset, int count) => inner.Read(buffer, offset, count);
        public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default) => inner.ReadAsync(buffer, cancellationToken);
        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) => inner.ReadAsync(buffer, offset, count, cancellationToken);
        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
        public override void SetLength(long value) => throw new NotSupportedException();
        public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                source.Dispose();
                inner.Dispose();
            }
            base.Dispose(disposing);
        }
    }
}