f42b00cb1c
Hour-long remote tasks were never delivered: a subscribe watch held one SSE stream for up to a day, the remote sent nothing while its agent worked without text, and once the quiet connection died on the way (no FIN/RST) the watch waited on it with no error and no log line. - A subscription lasts at most 5 minutes (SubscriptionSlice), then GetTask and a new subscription: a pause, not a failure. A slice ended mid-read counts as ours even when it surfaces as an I/O error. - The shared handler connects with TCP keepalive (60 s, every 15 s, 4 probes) and HTTP/2 pings. - Log lines when a watch starts or is restored, when a slice ends on a running task, when a result is delivered or kept; the Telegram host's console log carries a UTC timestamp. - action=status on a settled task ends its watch, so the result isn't delivered a second time. - A "context busy" refusal of send says nothing was sent, names the watched task in that context and that its result will arrive by itself; refusals never touch a watch.
644 lines
33 KiB
C#
644 lines
33 KiB
C#
using System.Collections.Concurrent;
|
|
using System.Net;
|
|
using System.Net.Sockets;
|
|
using System.Security.Cryptography;
|
|
using System.Text;
|
|
using System.Text.Json;
|
|
using System.Threading.Channels;
|
|
using A2A;
|
|
using Microsoft.Extensions.Logging;
|
|
using Microsoft.Extensions.Logging.Abstractions;
|
|
|
|
namespace Volcanic.AgenticCore.Extensions.A2A;
|
|
|
|
/// <summary>
|
|
/// Process-wide (<see cref="Shared"/>): one HTTP handler for every remote call, and the watches of long remote tasks.
|
|
/// A watch follows its task by push (when the card supports it and the host gave us a push base URL), else by
|
|
/// subscribing (streaming agents), else by polling; it delivers the outcome through <see cref="IA2AResultDelivery"/> and
|
|
/// is persisted in the workspace's <c>pending.json</c> until delivered.
|
|
/// </summary>
|
|
/// <remarks>
|
|
/// Delivery is at least once. A finished watch first stores its text as <c>Undelivered</c> in <c>pending.json</c>, then
|
|
/// delivers outside any lock (a delivery may run a whole agent turn that calls back into this registry), then removes the
|
|
/// record; a crash in between delivers it again on the next flush. A per-workspace lock guards only the bookkeeping
|
|
/// (claiming a watch, writing <c>pending.json</c>), and records being delivered are claimed in <c>_delivering</c>, so
|
|
/// finishing, flushing and restoring deliver one result once, and a cancel during a delivery is never undone.
|
|
/// A push only means "check now": its body is never trusted.
|
|
/// </remarks>
|
|
public sealed class A2AWatchRegistry : IAsyncDisposable
|
|
{
|
|
public static A2AWatchRegistry Shared { get; } = new(TimeProvider.System, CreateSharedHandler());
|
|
|
|
/// <summary>
|
|
/// The shared handler: pooled connections renewed every 5 minutes, and TCP keepalive (plus HTTP/2 pings) so a
|
|
/// connection that died silently on the way (a proxy or NAT dropping an idle stream) fails instead of hanging.
|
|
/// </summary>
|
|
internal static SocketsHttpHandler CreateSharedHandler() => new()
|
|
{
|
|
PooledConnectionLifetime = TimeSpan.FromMinutes(5),
|
|
KeepAlivePingDelay = TimeSpan.FromSeconds(60),
|
|
KeepAlivePingTimeout = TimeSpan.FromSeconds(20),
|
|
ConnectCallback = (context, ct) => ConnectWithKeepAliveAsync(context.DnsEndPoint, ct),
|
|
};
|
|
|
|
/// <summary>A TCP connection probed after 60 s of silence, every 15 s, given up after 4 unanswered probes.</summary>
|
|
internal static async ValueTask<Stream> ConnectWithKeepAliveAsync(DnsEndPoint endPoint, CancellationToken cancellationToken)
|
|
{
|
|
var socket = new Socket(SocketType.Stream, ProtocolType.Tcp) { NoDelay = true };
|
|
try
|
|
{
|
|
socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.KeepAlive, true);
|
|
socket.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.TcpKeepAliveTime, 60);
|
|
socket.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.TcpKeepAliveInterval, 15);
|
|
socket.SetSocketOption(SocketOptionLevel.Tcp, SocketOptionName.TcpKeepAliveRetryCount, 4);
|
|
await socket.ConnectAsync(endPoint, cancellationToken).ConfigureAwait(false);
|
|
return new NetworkStream(socket, ownsSocket: true);
|
|
}
|
|
catch
|
|
{
|
|
socket.Dispose();
|
|
throw;
|
|
}
|
|
}
|
|
|
|
public static readonly TimeSpan[] DefaultPollDelays =
|
|
[TimeSpan.FromSeconds(15), TimeSpan.FromSeconds(30), TimeSpan.FromMinutes(1), TimeSpan.FromMinutes(2), TimeSpan.FromMinutes(5)];
|
|
|
|
private static readonly TimeSpan DeleteTimeout = TimeSpan.FromSeconds(10);
|
|
|
|
/// <summary>After this many failed checks in a row, a subscribe watch polls instead.</summary>
|
|
private const int SubscribeFailuresBeforePoll = 3;
|
|
|
|
private readonly HttpMessageHandler _handler;
|
|
|
|
/// <summary><see cref="_handler"/> behind <see cref="A2AWireCompat"/>, for every A2A call (not for downloads).</summary>
|
|
private readonly HttpMessageHandler _wire;
|
|
private readonly ConcurrentDictionary<string, Watch> _watches = new(StringComparer.Ordinal);
|
|
private readonly ConcurrentDictionary<string, SemaphoreSlim> _workspaceLocks = new(StringComparer.Ordinal);
|
|
|
|
/// <summary>Watch ids whose result is being delivered right now (claimed under the workspace lock).</summary>
|
|
private readonly ConcurrentDictionary<string, byte> _delivering = new(StringComparer.Ordinal);
|
|
|
|
/// <summary>
|
|
/// (workspace, target) pairs where a flush skipped a record because it was claimed: if that delivery fails, the flush
|
|
/// runs again (else the result would wait for the next session start). Set and cleared under the workspace lock.
|
|
/// Keyed by target as well as workspace because a flush delivers one target's records only, so the retry must flush
|
|
/// the target whose record was skipped.
|
|
/// </summary>
|
|
private readonly ConcurrentDictionary<string, byte> _flushRequested = new(StringComparer.Ordinal);
|
|
private volatile IA2AResultDelivery _delivery;
|
|
|
|
/// <summary>Set by <see cref="StopAsync"/>: later starts and re-arms still persist their record but run no loop.</summary>
|
|
private volatile bool _stopped;
|
|
|
|
public A2AWatchRegistry(TimeProvider time, HttpMessageHandler handler, ILogger? logger = null)
|
|
{
|
|
Time = time;
|
|
_handler = handler;
|
|
Logger = logger ?? NullLogger.Instance;
|
|
_wire = new A2AWireCompat(handler, () => Logger);
|
|
_delivery = Live;
|
|
}
|
|
|
|
public TimeProvider Time { get; }
|
|
public ILogger Logger { get; set; }
|
|
public LiveSessionDelivery Live { get; } = new();
|
|
|
|
internal string? SourceOverride { get; init; }
|
|
internal TimeSpan[] PollDelays { get; init; } = DefaultPollDelays;
|
|
internal TimeSpan[] ReconnectDelays { get; init; } = [TimeSpan.FromSeconds(1), TimeSpan.FromSeconds(5), TimeSpan.FromSeconds(15)];
|
|
internal TimeSpan SafetyPollInterval { get; init; } = TimeSpan.FromMinutes(10);
|
|
|
|
/// <summary>
|
|
/// Longest single subscription: the watch then reads the task with <c>GetTask</c> and subscribes again (a pause, not
|
|
/// a failure). Remotes send nothing while their agent works without text, and a stream that went quiet may have died
|
|
/// on the way without either side noticing, so no stream is trusted for longer than this.
|
|
/// </summary>
|
|
internal TimeSpan SubscriptionSlice { get; init; } = TimeSpan.FromMinutes(5);
|
|
|
|
/// <summary>Bound of one remote call (get, create/delete push config), never past the watch deadline.</summary>
|
|
internal TimeSpan CallTimeout { get; init; } = TimeSpan.FromSeconds(60);
|
|
|
|
/// <summary>How long after a session start its flush of kept results waits (see <see cref="A2AClientExtension"/>).</summary>
|
|
internal TimeSpan SessionStartFlushDelay { get; init; } = TimeSpan.FromSeconds(1);
|
|
|
|
public void UseDelivery(IA2AResultDelivery delivery) => _delivery = delivery;
|
|
|
|
/// <summary>The shared handler, also for downloading files the agents link to.</summary>
|
|
internal HttpMessageHandler Handler => _handler;
|
|
|
|
internal RemoteAgentConnector Connector(A2AClientSettings settings) =>
|
|
new(_wire, TimeSpan.FromSeconds(settings.CardTimeoutSeconds));
|
|
|
|
internal static TimeSpan DelayAt(TimeSpan[] delays, int attempt) => delays[Math.Clamp(attempt, 0, delays.Length - 1)];
|
|
|
|
public IReadOnlyList<WatchRecord> List(string target) =>
|
|
_watches.Values.Where(w => w.Record.Target == target).Select(w => w.Record).ToList();
|
|
|
|
/// <summary>Hand a still-running remote task to a watch. A task the target already watches keeps its watch.</summary>
|
|
internal async Task<WatchRecord> StartAsync(string target, string workspace, RemoteAgent agent, AgentTask task,
|
|
A2AClientSettings settings, CancellationToken cancellationToken)
|
|
{
|
|
var watchId = Guid.NewGuid().ToString("N");
|
|
var source = SourceOverride ?? (agent.Streaming ? WatchSources.Subscribe : WatchSources.Poll);
|
|
var watch = new Watch(new WatchRecord(watchId, target, agent.Entry.Name, task.Id, task.ContextId, null, Time.GetUtcNow(), source), workspace);
|
|
using (await LockAsync(workspace, cancellationToken).ConfigureAwait(false))
|
|
{
|
|
// A watch that is finishing (left the map, delivering) isn't seen here, so a start racing it may watch the
|
|
// task again: a rare duplicate result, within the registry's at-least-once delivery.
|
|
if (_watches.Values.FirstOrDefault(w => w.Record.Target == target && w.Record.TaskId == task.Id) is { } existing)
|
|
{
|
|
watch.DisposeStop();
|
|
return existing.Record;
|
|
}
|
|
// Reserved before the push config exists, so an early push finds its watch (it wakes it once the loop runs).
|
|
_watches[watchId] = watch;
|
|
}
|
|
try
|
|
{
|
|
if (SourceOverride is null && settings.PushBaseUrl is { Length: > 0 } pushBase && agent.SupportsPush)
|
|
{
|
|
var token = Convert.ToHexString(RandomNumberGenerator.GetBytes(32));
|
|
watch.Record = watch.Record with { PushToken = token };
|
|
try
|
|
{
|
|
using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
|
timeout.CancelAfter(CallTimeout);
|
|
await agent.Client.CreateTaskPushNotificationConfigAsync(new CreateTaskPushNotificationConfigRequest
|
|
{
|
|
TaskId = task.Id,
|
|
ConfigId = watchId,
|
|
Config = new PushNotificationConfig { Id = watchId, Url = $"{pushBase.TrimEnd('/')}/a2a/push/{watchId}", Token = token },
|
|
}, timeout.Token).ConfigureAwait(false);
|
|
watch.Record = watch.Record with { Source = WatchSources.Push };
|
|
}
|
|
catch (Exception ex) when (ex is not OutOfMemoryException && !cancellationToken.IsCancellationRequested)
|
|
{
|
|
Logger.LogInformation("A2A: {Agent} refused push for task {TaskId} ({Error}); watching by {Source}", agent.Entry.Name, task.Id, agent.Entry.Scrub(ex.Message), source);
|
|
watch.Record = watch.Record with { PushToken = null };
|
|
}
|
|
}
|
|
using (await LockAsync(workspace).ConfigureAwait(false))
|
|
PendingFile.Upsert(workspace, watch.Record);
|
|
}
|
|
catch
|
|
{
|
|
_watches.TryRemove(new(watchId, watch));
|
|
watch.DisposeStop();
|
|
throw;
|
|
}
|
|
if (_stopped)
|
|
{
|
|
// Shutting down: the record stays in pending.json for the next restore, no loop outlives the registry.
|
|
_watches.TryRemove(new(watchId, watch));
|
|
watch.DisposeStop();
|
|
return watch.Record;
|
|
}
|
|
Logger.LogInformation("A2A: watching task {TaskId} of {Agent} for {Target} by {Source} (watch {WatchId})",
|
|
task.Id, agent.Entry.Name, target, watch.Record.Source, watchId);
|
|
watch.Loop = Task.Run(() => RunAsync(watch, settings));
|
|
return watch.Record;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Push receiver: a known watch with the right token wakes up and checks its task (the body is never trusted).
|
|
/// False → 404.
|
|
/// </summary>
|
|
public Task<bool> OnPushAsync(string watchId, string? token, JsonElement body, CancellationToken cancellationToken)
|
|
{
|
|
if (!_watches.TryGetValue(watchId, out var watch) || watch.Record.PushToken is not { } expected || token is null ||
|
|
!CryptographicOperations.FixedTimeEquals(Encoding.UTF8.GetBytes(expected), Encoding.UTF8.GetBytes(token)))
|
|
return Task.FromResult(false);
|
|
watch.Wake.Writer.TryWrite(true);
|
|
return Task.FromResult(true);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Re-arm every persisted watch of the delivery's known targets, then deliver results kept undelivered (all watches
|
|
/// first: a delivery may run a whole chat turn). Returns the count re-armed.
|
|
/// </summary>
|
|
public async Task<int> RestoreAsync(A2AClientSettings settings, CancellationToken cancellationToken)
|
|
{
|
|
var count = 0;
|
|
var targets = new List<(string Target, string Workspace)>();
|
|
foreach (var target in _delivery.KnownTargets().Distinct(StringComparer.Ordinal))
|
|
if (_delivery.WorkspaceFor(target) is { } workspace) targets.Add((target, workspace));
|
|
foreach (var (target, workspace) in targets)
|
|
{
|
|
using (await LockAsync(workspace, cancellationToken).ConfigureAwait(false))
|
|
{
|
|
foreach (var stored in PendingFile.Read(workspace).Where(r => r.Target == target && r.Undelivered is null))
|
|
{
|
|
var record = stored;
|
|
// No push receiver now: nobody would wake a push watch, so follow the task another way.
|
|
if (record.Source == WatchSources.Push && string.IsNullOrEmpty(settings.PushBaseUrl))
|
|
record = record with { Source = WatchSources.Subscribe };
|
|
var watch = new Watch(record, workspace);
|
|
if (!_watches.TryAdd(record.WatchId, watch)) { watch.DisposeStop(); continue; }
|
|
if (!ReferenceEquals(record, stored)) PendingFile.Upsert(workspace, record);
|
|
// Checked after the watch is in the map: a StopAsync that set the flag later sees (and stops) it.
|
|
if (_stopped) { _watches.TryRemove(new(record.WatchId, watch)); watch.DisposeStop(); return count; }
|
|
Logger.LogInformation("A2A: restored the watch of task {TaskId} of {Agent} for {Target} by {Source} (watch {WatchId})",
|
|
record.TaskId, record.Agent, target, record.Source, record.WatchId);
|
|
watch.Loop = Task.Run(() => RunAsync(watch, settings));
|
|
count++;
|
|
}
|
|
}
|
|
}
|
|
foreach (var (target, workspace) in targets)
|
|
await FlushUndeliveredAsync(target, workspace, cancellationToken).ConfigureAwait(false);
|
|
return count;
|
|
}
|
|
|
|
/// <summary>Deliver results that were kept because the target had no live conversation. After <see cref="StopAsync"/> they stay kept.</summary>
|
|
public async Task FlushUndeliveredAsync(string target, string workspace, CancellationToken cancellationToken)
|
|
{
|
|
if (_stopped) return;
|
|
var claimed = new List<WatchRecord>();
|
|
using (await LockAsync(workspace, cancellationToken).ConfigureAwait(false))
|
|
{
|
|
foreach (var record in PendingFile.Read(workspace).Where(r => r.Target == target && r.Undelivered is not null))
|
|
{
|
|
if (_delivering.TryAdd(record.WatchId, 0)) claimed.Add(record);
|
|
else _flushRequested[FlushKey(workspace, target)] = 0; // being delivered elsewhere: retry here if that fails
|
|
}
|
|
}
|
|
var failed = false;
|
|
var retry = false;
|
|
try
|
|
{
|
|
foreach (var record in claimed)
|
|
{
|
|
if (_stopped) break; // shutting down: the rest stay in pending.json for the next start
|
|
if (!await TryDeliverAsync(target, record.Undelivered!, cancellationToken).ConfigureAwait(false)) { failed = true; continue; }
|
|
Logger.LogInformation("A2A: delivered the kept result of task {TaskId} of {Agent} to {Target}", record.TaskId, record.Agent, target);
|
|
using (await LockAsync(workspace).ConfigureAwait(false))
|
|
PendingFile.Remove(workspace, record.WatchId);
|
|
}
|
|
}
|
|
finally
|
|
{
|
|
retry = await ReleaseAsync(workspace, target, claimed.Select(r => r.WatchId), failed).ConfigureAwait(false);
|
|
}
|
|
if (retry) await FlushUndeliveredAsync(target, workspace, cancellationToken).ConfigureAwait(false);
|
|
}
|
|
|
|
private static string FlushKey(string workspace, string target) => Path.GetFullPath(workspace) + "\n" + target;
|
|
|
|
/// <summary>
|
|
/// Releases delivery claims under the workspace lock. True when a delivery failed and a flush skipped one of these
|
|
/// records meanwhile: the caller flushes again, so the result doesn't wait for the next session start.
|
|
/// </summary>
|
|
private async Task<bool> ReleaseAsync(string workspace, string target, IEnumerable<string> watchIds, bool failed)
|
|
{
|
|
using (await LockAsync(workspace).ConfigureAwait(false))
|
|
{
|
|
foreach (var id in watchIds) _delivering.TryRemove(id, out _);
|
|
return failed && _flushRequested.TryRemove(FlushKey(workspace, target), out _);
|
|
}
|
|
}
|
|
|
|
/// <summary>Stop watching a task (after <c>cancel</c>); the stopped watch's record, null when none was running.</summary>
|
|
public async Task<WatchRecord?> StopWatchingAsync(string target, string workspace, string taskId)
|
|
{
|
|
Watch? watch;
|
|
using (await LockAsync(workspace).ConfigureAwait(false))
|
|
{
|
|
watch = _watches.Values.FirstOrDefault(w => w.Record.Target == target && w.Record.TaskId == taskId);
|
|
if (watch is not null && !_watches.TryRemove(new(watch.Record.WatchId, watch))) watch = null;
|
|
foreach (var record in PendingFile.Read(workspace).Where(r => r.Target == target && r.TaskId == taskId))
|
|
PendingFile.Remove(workspace, record.WatchId);
|
|
}
|
|
if (watch is null) return null;
|
|
await watch.StopAsync().ConfigureAwait(false);
|
|
return watch.Record;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Resumes a watch stopped by <see cref="StopWatchingAsync"/> whose cancel didn't go through (same id, push token and
|
|
/// start). No-op when the task is watched again already.
|
|
/// </summary>
|
|
internal async Task RearmAsync(string workspace, WatchRecord record, A2AClientSettings settings)
|
|
{
|
|
using (await LockAsync(workspace).ConfigureAwait(false))
|
|
{
|
|
if (_watches.Values.Any(w => w.Record.Target == record.Target && w.Record.TaskId == record.TaskId)) return;
|
|
var watch = new Watch(record with { Undelivered = null }, workspace);
|
|
if (!_watches.TryAdd(record.WatchId, watch)) { watch.DisposeStop(); return; }
|
|
PendingFile.Upsert(workspace, watch.Record);
|
|
if (_stopped)
|
|
{
|
|
// Kept in pending.json for the next restore; no loop after StopAsync.
|
|
_watches.TryRemove(new(record.WatchId, watch));
|
|
watch.DisposeStop();
|
|
return;
|
|
}
|
|
watch.Loop = Task.Run(() => RunAsync(watch, settings));
|
|
}
|
|
}
|
|
|
|
private async Task RunAsync(Watch watch, A2AClientSettings settings)
|
|
{
|
|
var ct = watch.Stop.Token;
|
|
RemoteAgent? agent = null;
|
|
RemoteAgentEntry? lastEntry = null;
|
|
var failures = 0;
|
|
var polls = 0;
|
|
try
|
|
{
|
|
while (!ct.IsCancellationRequested)
|
|
{
|
|
var record = watch.Record;
|
|
var deadline = DeadlineOf(record.StartedAt, settings.WatchTimeout);
|
|
TimeSpan wait;
|
|
try
|
|
{
|
|
if (Time.GetUtcNow() >= deadline)
|
|
{
|
|
await FinishAsync(watch, $"[A2A result from \"{record.Agent}\", task {record.TaskId} (no result)]\n" +
|
|
$"\"{record.Agent}\" gave no result within {settings.WatchTimeout}. Check it with action=status, agent={record.Agent}, taskId={record.TaskId}.", agent).ConfigureAwait(false);
|
|
return;
|
|
}
|
|
var agents = AgentsFile.Read(watch.Workspace);
|
|
if (agents.Error is { } error) throw new IOException(error); // transient: the user may be editing it
|
|
var entry = agents.Agents.FirstOrDefault(a => a.Name == record.Agent);
|
|
if (entry is not null) lastEntry = entry;
|
|
if (entry is null)
|
|
{
|
|
await FinishAsync(watch, $"[A2A result from \"{record.Agent}\", task {record.TaskId} (stopped)]\n" +
|
|
$"\"{record.Agent}\" is no longer in agents.json, so its task is no longer watched.", null).ConfigureAwait(false);
|
|
return;
|
|
}
|
|
if (agent is null || agent.Entry != entry)
|
|
{
|
|
agent?.Dispose();
|
|
agent = null;
|
|
agent = await Connector(settings).ConnectAsync(entry, ct).ConfigureAwait(false);
|
|
}
|
|
if (record.Source == WatchSources.Subscribe && !agent.Streaming)
|
|
record = await SetSourceAsync(watch, WatchSources.Poll).ConfigureAwait(false);
|
|
|
|
AgentTask? task;
|
|
var pause = false;
|
|
if (record.Source == WatchSources.Subscribe)
|
|
{
|
|
(task, pause) = await SubscribeAsync(agent, record.TaskId, deadline, ct).ConfigureAwait(false);
|
|
if ((task is null || !A2AResults.IsSettled(task.Status.State)) && !pause)
|
|
throw new IOException("the task stream ended before the task settled");
|
|
if (pause && task is not null && !A2AResults.IsSettled(task.Status.State))
|
|
Logger.LogInformation("A2A: task {TaskId} of {Agent} is still {State}; subscribing again", record.TaskId, record.Agent,
|
|
A2AResults.StateName(task.Status.State));
|
|
}
|
|
else
|
|
{
|
|
task = await GetTaskAsync(agent, record.TaskId, deadline, ct).ConfigureAwait(false);
|
|
}
|
|
if (task is not null && A2AResults.IsSettled(task.Status.State))
|
|
{
|
|
var files = await A2AFiles.SaveAsync(watch.Workspace, task, agent, _handler, settings, ct).ConfigureAwait(false);
|
|
await FinishAsync(watch, A2AResults.DeliveryText(record.Agent, task, files), agent).ConfigureAwait(false);
|
|
return;
|
|
}
|
|
failures = 0;
|
|
wait = record.Source switch
|
|
{
|
|
WatchSources.Push => SafetyPollInterval,
|
|
WatchSources.Subscribe => TimeSpan.Zero, // our own cap or the deadline ended the stream
|
|
_ => DelayAt(PollDelays, polls++),
|
|
};
|
|
}
|
|
catch (OperationCanceledException) when (ct.IsCancellationRequested)
|
|
{
|
|
return;
|
|
}
|
|
catch (Exception ex) when (ex is not OutOfMemoryException)
|
|
{
|
|
failures++;
|
|
Logger.LogInformation("A2A: checking task {TaskId} of {Agent} failed ({Error})", record.TaskId, record.Agent,
|
|
lastEntry is null ? ex.Message : lastEntry.Scrub(ex.Message));
|
|
agent?.Dispose();
|
|
agent = null;
|
|
if (record.Source == WatchSources.Subscribe && failures >= SubscribeFailuresBeforePoll)
|
|
{
|
|
Logger.LogInformation("A2A: task {TaskId} of {Agent}: subscribing failed {Count} times; polling instead", record.TaskId, record.Agent, failures);
|
|
record = await SetSourceAsync(watch, WatchSources.Poll).ConfigureAwait(false);
|
|
}
|
|
wait = record.Source == WatchSources.Subscribe ? DelayAt(ReconnectDelays, failures - 1) : DelayAt(PollDelays, polls++);
|
|
}
|
|
await WaitAsync(watch, Min(wait, deadline - Time.GetUtcNow()), ct).ConfigureAwait(false);
|
|
}
|
|
}
|
|
catch (OperationCanceledException) when (ct.IsCancellationRequested)
|
|
{
|
|
}
|
|
catch (Exception ex) when (ex is not OutOfMemoryException)
|
|
{
|
|
Logger.LogWarning(ex, "A2A: watch {WatchId} crashed", watch.Record.WatchId);
|
|
}
|
|
finally
|
|
{
|
|
agent?.Dispose();
|
|
watch.DisposeStop();
|
|
}
|
|
}
|
|
|
|
private static DateTimeOffset DeadlineOf(DateTimeOffset start, TimeSpan timeout) =>
|
|
timeout >= DateTimeOffset.MaxValue - start ? DateTimeOffset.MaxValue : start + timeout;
|
|
|
|
private TimeSpan CallBound(DateTimeOffset deadline)
|
|
{
|
|
var remaining = deadline - Time.GetUtcNow();
|
|
return remaining < TimeSpan.FromSeconds(1) ? TimeSpan.FromSeconds(1) : Min(CallTimeout, remaining);
|
|
}
|
|
|
|
private async Task<AgentTask> GetTaskAsync(RemoteAgent agent, string taskId, DateTimeOffset deadline, CancellationToken ct)
|
|
{
|
|
using var bounded = CancellationTokenSource.CreateLinkedTokenSource(ct);
|
|
bounded.CancelAfter(CallBound(deadline));
|
|
return await agent.Client.GetTaskAsync(new GetTaskRequest { Id = taskId }, bounded.Token).ConfigureAwait(false);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Follows the task's stream until it settles or ends, then reads the task. <c>Paused</c>: the stream was ended by
|
|
/// us (the deadline or <see cref="SubscriptionSlice"/>), which is not a failure.
|
|
/// </summary>
|
|
private async Task<(AgentTask? Task, bool Paused)> SubscribeAsync(RemoteAgent agent, string taskId, DateTimeOffset deadline, CancellationToken ct)
|
|
{
|
|
AgentTask? current = null;
|
|
var paused = false;
|
|
using (var until = CancellationTokenSource.CreateLinkedTokenSource(ct))
|
|
{
|
|
var remaining = deadline - Time.GetUtcNow();
|
|
until.CancelAfter(remaining <= TimeSpan.Zero ? TimeSpan.Zero : Min(remaining, SubscriptionSlice));
|
|
try
|
|
{
|
|
await foreach (var ev in agent.Client.SubscribeToTaskAsync(new SubscribeToTaskRequest { Id = taskId }, until.Token).ConfigureAwait(false))
|
|
{
|
|
current = TaskProjection.Apply(current, ev);
|
|
if (current is not null && A2AResults.IsSettled(current.Status.State)) return (current, false);
|
|
}
|
|
}
|
|
catch (A2AException) // e.g. the task is already terminal: read it instead
|
|
{
|
|
}
|
|
// Ending a stream mid-read may surface as an I/O error rather than a cancel: ours either way.
|
|
catch (Exception ex) when (ex is not OutOfMemoryException && !ct.IsCancellationRequested && until.IsCancellationRequested)
|
|
{
|
|
paused = true;
|
|
}
|
|
}
|
|
return (await GetTaskAsync(agent, taskId, deadline, ct).ConfigureAwait(false), paused);
|
|
}
|
|
|
|
/// <summary>Waits <paramref name="wait"/>, or less when a push wakes the watch.</summary>
|
|
private async Task WaitAsync(Watch watch, TimeSpan wait, CancellationToken ct)
|
|
{
|
|
if (wait <= TimeSpan.Zero) return;
|
|
using var timeout = CancellationTokenSource.CreateLinkedTokenSource(ct);
|
|
var delay = Task.Delay(wait, Time, timeout.Token);
|
|
var woken = watch.Wake.Reader.WaitToReadAsync(timeout.Token).AsTask();
|
|
await Task.WhenAny(delay, woken).ConfigureAwait(false);
|
|
await timeout.CancelAsync().ConfigureAwait(false);
|
|
while (watch.Wake.Reader.TryRead(out _)) { }
|
|
ct.ThrowIfCancellationRequested();
|
|
}
|
|
|
|
private async Task<WatchRecord> SetSourceAsync(Watch watch, string source)
|
|
{
|
|
using (await LockAsync(watch.Workspace).ConfigureAwait(false))
|
|
{
|
|
watch.Record = watch.Record with { Source = source };
|
|
if (_watches.TryGetValue(watch.Record.WatchId, out var live) && ReferenceEquals(live, watch))
|
|
PendingFile.Upsert(watch.Workspace, watch.Record);
|
|
}
|
|
return watch.Record;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Claims the watch and stores its text as undelivered (under the lock), delivers outside the lock, removes the record
|
|
/// once delivered (under the lock), then drops the push config best-effort.
|
|
/// </summary>
|
|
private async Task FinishAsync(Watch watch, string text, RemoteAgent? agent)
|
|
{
|
|
var record = watch.Record;
|
|
using (await LockAsync(watch.Workspace).ConfigureAwait(false))
|
|
{
|
|
if (!_watches.TryRemove(new(record.WatchId, watch))) return; // stopped (cancel or shutdown) meanwhile
|
|
// Stored before the claim: a failed write leaves no claim behind (the plain record is re-armed on restore).
|
|
PendingFile.Upsert(watch.Workspace, record with { Undelivered = text });
|
|
_delivering[record.WatchId] = 0;
|
|
}
|
|
var delivered = false;
|
|
var retry = false;
|
|
try
|
|
{
|
|
// A failed delivery writes nothing: the text is stored, and a cancel meanwhile has removed it for good.
|
|
delivered = await TryDeliverAsync(record.Target, text, CancellationToken.None).ConfigureAwait(false);
|
|
Logger.LogInformation(delivered
|
|
? "A2A: delivered the result of task {TaskId} of {Agent} to {Target}"
|
|
: "A2A: kept the result of task {TaskId} of {Agent} for {Target} until it can be delivered",
|
|
record.TaskId, record.Agent, record.Target);
|
|
if (delivered)
|
|
using (await LockAsync(watch.Workspace).ConfigureAwait(false))
|
|
PendingFile.Remove(watch.Workspace, record.WatchId);
|
|
}
|
|
finally
|
|
{
|
|
retry = await ReleaseAsync(watch.Workspace, record.Target, [record.WatchId], !delivered).ConfigureAwait(false);
|
|
}
|
|
if (retry)
|
|
{
|
|
try
|
|
{
|
|
await FlushUndeliveredAsync(record.Target, watch.Workspace, CancellationToken.None).ConfigureAwait(false);
|
|
}
|
|
catch (Exception ex) when (ex is not OutOfMemoryException)
|
|
{
|
|
Logger.LogWarning(ex, "A2A: flushing results for {Target} failed", record.Target);
|
|
}
|
|
}
|
|
if (agent is null || record.PushToken is null) return;
|
|
try
|
|
{
|
|
using var timeout = new CancellationTokenSource(DeleteTimeout, Time);
|
|
await agent.Client.DeleteTaskPushNotificationConfigAsync(new DeleteTaskPushNotificationConfigRequest
|
|
{ TaskId = record.TaskId, Id = record.WatchId }, timeout.Token).ConfigureAwait(false);
|
|
}
|
|
catch (Exception ex) when (ex is not OutOfMemoryException)
|
|
{
|
|
Logger.LogInformation("A2A: removing the push config of task {TaskId} at {Agent} failed ({Error})", record.TaskId, record.Agent, agent.Entry.Scrub(ex.Message));
|
|
}
|
|
}
|
|
|
|
private async Task<bool> TryDeliverAsync(string target, string text, CancellationToken ct)
|
|
{
|
|
try
|
|
{
|
|
return await _delivery.DeliverAsync(target, text, ct).ConfigureAwait(false);
|
|
}
|
|
catch (Exception ex) when (ex is not OperationCanceledException and not OutOfMemoryException)
|
|
{
|
|
Logger.LogWarning(ex, "A2A: delivering a result to {Target} failed", target);
|
|
return false;
|
|
}
|
|
}
|
|
|
|
private async Task<IDisposable> LockAsync(string workspace, CancellationToken cancellationToken = default)
|
|
{
|
|
var gate = _workspaceLocks.GetOrAdd(Path.GetFullPath(workspace), _ => new SemaphoreSlim(1, 1));
|
|
await gate.WaitAsync(cancellationToken).ConfigureAwait(false);
|
|
return new Releaser(gate);
|
|
}
|
|
|
|
private static TimeSpan Min(TimeSpan a, TimeSpan b) => a < b ? a : b;
|
|
|
|
/// <summary>
|
|
/// Stops every watch still in the map (the records stay in pending.json for the next restore) and waits for their loops.
|
|
/// It does not wait for deliveries of watches that already left the map to deliver: such a delivery may still be in
|
|
/// flight; if it fails, its text stays <c>Undelivered</c> in pending.json and is flushed on the next start.
|
|
/// </summary>
|
|
public async Task StopAsync()
|
|
{
|
|
_stopped = true;
|
|
var watches = _watches.ToArray();
|
|
foreach (var pair in watches) _watches.TryRemove(pair);
|
|
await Task.WhenAll(watches.Select(p => p.Value.StopAsync())).ConfigureAwait(false);
|
|
}
|
|
|
|
public async ValueTask DisposeAsync() => await StopAsync().ConfigureAwait(false);
|
|
|
|
private sealed class Releaser(SemaphoreSlim gate) : IDisposable
|
|
{
|
|
private int _released;
|
|
public void Dispose()
|
|
{
|
|
if (Interlocked.Exchange(ref _released, 1) == 0) gate.Release();
|
|
}
|
|
}
|
|
|
|
private sealed class Watch(WatchRecord record, string workspace)
|
|
{
|
|
private volatile WatchRecord _record = record;
|
|
private int _disposed;
|
|
|
|
public WatchRecord Record { get => _record; set => _record = value; }
|
|
public string Workspace => workspace;
|
|
public CancellationTokenSource Stop { get; } = new();
|
|
public Channel<bool> Wake { get; } = Channel.CreateBounded<bool>(new BoundedChannelOptions(1) { FullMode = BoundedChannelFullMode.DropOldest });
|
|
public Task Loop { get; set; } = Task.CompletedTask;
|
|
|
|
public async Task StopAsync()
|
|
{
|
|
try { await Stop.CancelAsync().ConfigureAwait(false); }
|
|
catch (ObjectDisposedException) { } // the loop already ended
|
|
try { await Loop.ConfigureAwait(false); } catch (OperationCanceledException) { }
|
|
}
|
|
|
|
/// <summary>Called when the loop ends (its token's only user), or when the watch never started.</summary>
|
|
public void DisposeStop()
|
|
{
|
|
if (Interlocked.Exchange(ref _disposed, 1) == 0) Stop.Dispose();
|
|
}
|
|
}
|
|
}
|