VolcanicAgenticCore/src/Volcanic.AgenticCore.A2A.Server/Agent/A2ARunMapper.cs

using System.Threading.Channels;
using A2A;
using Microsoft.Extensions.Logging;
using Volcanic.AgenticCore.Abstractions.Events;
using Volcanic.AgenticCore.Abstractions.Messages;
using Volcanic.AgenticCore.Abstractions.Processes;
using Volcanic.AgenticCore.Abstractions.Providers;

namespace Volcanic.AgenticCore.A2A.Server.Agent;

/// <summary>
/// Maps session events (raised synchronously on the loop) to task updates, written in order by a pump so the loop never
/// waits. Text deltas append to artifact <c>response</c>; tool starts become working-status messages. Tracks how the run
/// ended: the last assistant message decides (a retried failure followed by success counts as success), and its text is
/// the final answer. Writes use the run token, so nothing is written after a cancel; disposal waits a bounded time.
/// </summary>
internal sealed class A2ARunMapper : IAsyncDisposable
{
    public const string ResponseArtifactId = "response";

    /// <summary>How long disposal waits for queued writes before cancelling them (the event queue may be unread).</summary>
    internal static TimeSpan DrainTimeout { get; set; } = TimeSpan.FromSeconds(30);

    private readonly Channel<Func<ValueTask>> _work = Channel.CreateUnbounded<Func<ValueTask>>(new UnboundedChannelOptions { SingleReader = true });
    private readonly TaskUpdater _updater;
    private readonly ISecretRedactor _redactor;
    private readonly ILogger _logger;
    private readonly CancellationTokenSource _cts;
    private readonly CancellationToken _ct;
    private readonly Task _pump;
    private bool _textStarted;

    public A2ARunMapper(TaskUpdater updater, ISecretRedactor redactor, ILogger logger, CancellationToken cancellationToken)
    {
        _updater = updater;
        _redactor = redactor;
        _logger = logger;
        _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        _ct = _cts.Token;
        _pump = Task.Run(PumpAsync, CancellationToken.None);
    }

    public string? FinalText { get; private set; }
    public string? Error { get; private set; }
    public bool Aborted { get; private set; }

    public void OnEvent(AgentSessionEvent e)
    {
        switch (e)
        {
            case AgentSessionEvent.MessageUpdate { AssistantMessageEvent: AssistantMessageEvent.TextDelta delta }:
                // The run's first chunk replaces whatever response artifact the task already has.
                var append = _textStarted;
                _textStarted = true;
                Post(() => _updater.AddArtifactAsync([Part.FromText(delta.Delta)], ResponseArtifactId, "response", null,
                    lastChunk: false, append: append, _ct));
                break;
            case AgentSessionEvent.MessageEnd { Message: AssistantMessage message }:
                if (message.StopReason == StopReason.Aborted)
                {
                    Aborted = true;
                }
                else if (message.IsFailure)
                {
                    Error = _redactor.Redact(message.ErrorMessage ?? "The model call failed.");
                }
                else
                {
                    Error = null;
                    FinalText = string.IsNullOrWhiteSpace(message.Text) ? null : message.Text;
                }
                break;
            case AgentSessionEvent.ToolExecutionStart start:
                Post(() => _updater.StartWorkAsync(A2AText.Agent($"using {start.ToolName}"), _ct));
                break;
        }
    }

    private void Post(Func<ValueTask> work) => _work.Writer.TryWrite(work);

    private async Task PumpAsync()
    {
        await foreach (var work in _work.Reader.ReadAllAsync().ConfigureAwait(false))
        {
            if (_ct.IsCancellationRequested) continue; // drop the rest after a cancel
            try { await work().ConfigureAwait(false); }
            catch (OperationCanceledException) when (_ct.IsCancellationRequested) { }
            catch (Exception ex) { _logger.LogWarning(ex, "A2A: writing a task update for {TaskId} failed", _updater.TaskId); }
        }
    }

    /// <summary>Wait until every posted update is written, or <see cref="DrainTimeout"/> passes and the rest is dropped.</summary>
    public async ValueTask DisposeAsync()
    {
        _work.Writer.TryComplete();
        try
        {
            await _pump.WaitAsync(DrainTimeout).ConfigureAwait(false);
        }
        catch (TimeoutException)
        {
            _logger.LogWarning("A2A: task {TaskId} updates were not consumed in {Timeout}; dropping them", _updater.TaskId, DrainTimeout);
            await _cts.CancelAsync().ConfigureAwait(false);
            await _pump.ConfigureAwait(false);
        }
        _cts.Dispose();
    }
}

internal static class A2AText
{
    public static global::A2A.Message Agent(string text) => new()
    {
        Role = Role.Agent, MessageId = Guid.NewGuid().ToString("N"), Parts = [Part.FromText(text)],
    };
}