cb877f697a
- A2AWatchRegistry.Shared was initialised before DefaultPollDelays (static fields run in source order), so the bot's registry had null poll delays and every poll watch crashed on its first wait, then sat in the registry doing nothing until a restart. The delays now come first; a test reads them through Shared. - TaskNotFound from the remote ends the watch at once with a "(lost)" notice instead of retrying for a day. - A crashed watch loop restarts by polling after a pause; the third crash in a row ends it with a "(stopped)" notice.
703 lines
36 KiB
C#
703 lines
36 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
|
|
{
|
|
// Static fields initialise in source order: everything an instance reads (DefaultPollDelays) must come before Shared,
|
|
// or Shared gets nulls (it did: every poll watch of the bot crashed on its first wait).
|
|
public static readonly TimeSpan[] DefaultPollDelays =
|
|
[TimeSpan.FromSeconds(15), TimeSpan.FromSeconds(30), TimeSpan.FromMinutes(1), TimeSpan.FromMinutes(2), TimeSpan.FromMinutes(5)];
|
|
|
|
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;
|
|
}
|
|
}
|
|
|
|
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;
|
|
|
|
/// <summary>A watch whose loop crashed this many times in a row ends with a <c>(stopped)</c> notice.</summary>
|
|
private const int CrashesBeforeGivingUp = 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));
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Runs a watch until it finishes or is stopped. A crash of its loop (a bug, never a remote failure: those are retried
|
|
/// inside) restarts the loop after a pause, by polling; after <see cref="CrashesBeforeGivingUp"/> crashes in a row the
|
|
/// watch ends with a <c>(stopped)</c> notice, so it never sits in the registry doing nothing.
|
|
/// </summary>
|
|
private async Task RunAsync(Watch watch, A2AClientSettings settings)
|
|
{
|
|
var ct = watch.Stop.Token;
|
|
try
|
|
{
|
|
for (var crashes = 1; ; crashes++)
|
|
{
|
|
try
|
|
{
|
|
await WatchAsync(watch, settings).ConfigureAwait(false);
|
|
return;
|
|
}
|
|
catch (Exception ex) when (ex is not OutOfMemoryException && !ct.IsCancellationRequested)
|
|
{
|
|
var record = watch.Record;
|
|
Logger.LogWarning(ex, "A2A: watch {WatchId} of task {TaskId} of {Agent} crashed ({Count} of {Max})",
|
|
record.WatchId, record.TaskId, record.Agent, crashes, CrashesBeforeGivingUp);
|
|
if (crashes >= CrashesBeforeGivingUp)
|
|
{
|
|
await FinishAsync(watch, $"[A2A result from \"{record.Agent}\", task {record.TaskId} (stopped)]\n" +
|
|
$"Watching this task failed ({ex.GetType().Name}), so it is no longer watched. " +
|
|
$"Check it with action=status, agent={record.Agent}, taskId={record.TaskId}.", null).ConfigureAwait(false);
|
|
return;
|
|
}
|
|
if (record.Source != WatchSources.Poll) await SetSourceAsync(watch, WatchSources.Poll).ConfigureAwait(false);
|
|
await WaitAsync(watch, DelayAt(ReconnectDelays, crashes - 1), ct).ConfigureAwait(false);
|
|
}
|
|
}
|
|
}
|
|
catch (OperationCanceledException) when (ct.IsCancellationRequested)
|
|
{
|
|
}
|
|
catch (Exception ex) when (ex is not OutOfMemoryException)
|
|
{
|
|
// Even the (stopped) notice failed: out of the live watches; its record stays for the next restore.
|
|
Logger.LogWarning(ex, "A2A: watch {WatchId} crashed", watch.Record.WatchId);
|
|
_watches.TryRemove(new(watch.Record.WatchId, watch));
|
|
}
|
|
finally
|
|
{
|
|
watch.DisposeStop();
|
|
}
|
|
}
|
|
|
|
/// <summary>The watch loop: returns when the watch finished or was stopped, throws only on a bug.</summary>
|
|
private async Task WatchAsync(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 (A2AException ex) when (ex.ErrorCode == A2AErrorCode.TaskNotFound)
|
|
{
|
|
// Definitive: the remote forgot the task (typically it restarted), so no check will ever find it.
|
|
Logger.LogInformation("A2A: task {TaskId} of {Agent} is unknown to the agent; the watch ends", record.TaskId, record.Agent);
|
|
await FinishAsync(watch, $"[A2A result from \"{record.Agent}\", task {record.TaskId} (lost)]\n" +
|
|
$"\"{record.Agent}\" no longer knows this task (it was probably restarted), so its result can't be fetched " +
|
|
"and the task is no longer watched. Send the request again with action=send if it is still needed.", agent).ConfigureAwait(false);
|
|
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)
|
|
{
|
|
}
|
|
finally
|
|
{
|
|
agent?.Dispose();
|
|
}
|
|
}
|
|
|
|
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();
|
|
}
|
|
}
|
|
}
|