fix(a2a-client): no leaked claim on a failed store, retry a flush that skipped a failing delivery, accurate StopAsync doc

This commit is contained in:
EugeneTes
2026-10-08 13:49:06 +00:00
parent e1e0f35bb9
commit d7f5e0b900
3 changed files with 117 additions and 14 deletions

View File

@@ -9,9 +9,9 @@ public interface IA2AResultDelivery
/// <summary>Run <paramref name="text"/> as a follow-up prompt in the target's conversation. False: not deliverable now.</summary>
/// <remarks>
/// May run a whole agent turn before returning, and that turn may call back into the <see cref="A2AWatchRegistry"/>
/// (cancel a watch, start one): the registry never holds a lock while delivering. A delivery in flight can't be
/// interrupted, so shutting the registry down waits for it. A false result keeps the text in <c>pending.json</c> for
/// the next flush; delivery is at least once.
/// (cancel a watch, start one): the registry never holds a lock while delivering. A delivery in flight isn't
/// interrupted, and shutting the registry down doesn't wait for it. A false result keeps the text in
/// <c>pending.json</c> for the next flush; delivery is at least once.
/// </remarks>
Task<bool> DeliverAsync(string target, string text, CancellationToken cancellationToken);

View File

@@ -45,6 +45,12 @@ public sealed class A2AWatchRegistry : IAsyncDisposable
/// <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.
/// </summary>
private readonly ConcurrentDictionary<string, byte> _flushRequested = new(StringComparer.Ordinal);
private volatile IA2AResultDelivery _delivery;
public A2AWatchRegistry(TimeProvider time, HttpMessageHandler handler, ILogger? logger = null)
@@ -166,25 +172,45 @@ public sealed class A2AWatchRegistry : IAsyncDisposable
/// <summary>Deliver results that were kept because the target had no live conversation.</summary>
public async Task FlushUndeliveredAsync(string target, string workspace, CancellationToken cancellationToken)
{
List<WatchRecord> claimed;
var claimed = new List<WatchRecord>();
using (await LockAsync(workspace, cancellationToken).ConfigureAwait(false))
{
claimed = PendingFile.Read(workspace)
.Where(r => r.Target == target && r.Undelivered is not null && _delivering.TryAdd(r.WatchId, 0))
.ToList();
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 (!await TryDeliverAsync(target, record.Undelivered!, cancellationToken).ConfigureAwait(false)) continue;
if (!await TryDeliverAsync(target, record.Undelivered!, cancellationToken).ConfigureAwait(false)) { failed = true; continue; }
using (await LockAsync(workspace).ConfigureAwait(false))
PendingFile.Remove(workspace, record.WatchId);
}
}
finally
{
foreach (var record in claimed) _delivering.TryRemove(record.WatchId, out _);
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 _);
}
}
@@ -383,19 +409,34 @@ public sealed class A2AWatchRegistry : IAsyncDisposable
using (await LockAsync(watch.Workspace).ConfigureAwait(false))
{
if (!_watches.TryRemove(new(record.WatchId, watch))) return; // stopped (cancel or shutdown) meanwhile
_delivering[record.WatchId] = 0;
// 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.
if (await TryDeliverAsync(record.Target, text, CancellationToken.None).ConfigureAwait(false))
delivered = await TryDeliverAsync(record.Target, text, CancellationToken.None).ConfigureAwait(false);
if (delivered)
using (await LockAsync(watch.Workspace).ConfigureAwait(false))
PendingFile.Remove(watch.Workspace, record.WatchId);
}
finally
{
_delivering.TryRemove(record.WatchId, out _);
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
@@ -433,8 +474,9 @@ public sealed class A2AWatchRegistry : IAsyncDisposable
private static TimeSpan Min(TimeSpan a, TimeSpan b) => a < b ? a : b;
/// <summary>
/// Stops every watch (the records stay in pending.json for the next restore). Waits for in-flight deliveries: a
/// delivery can't be interrupted.
/// 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()
{

View File

@@ -46,6 +46,28 @@ public sealed class A2AWatchTests : IDisposable
public IEnumerable<string> KnownTargets() => ["1"];
}
/// <summary>First delivery blocks until <see cref="Release"/> and fails (no live session yet); later ones succeed at once.</summary>
private sealed class FailFirstDelivery(string workspace) : IA2AResultDelivery
{
private readonly TaskCompletionSource _gate = new(TaskCreationOptions.RunContinuationsAsynchronously);
private int _entered;
public ConcurrentQueue<string> Delivered { get; } = new();
public int Entered => Volatile.Read(ref _entered);
public void Release() => _gate.TrySetResult();
public async Task<bool> DeliverAsync(string target, string text, CancellationToken cancellationToken)
{
if (Interlocked.Increment(ref _entered) == 1)
{
await _gate.Task;
return false;
}
Delivered.Enqueue(text);
return true;
}
public string? WorkspaceFor(string target) => workspace;
public IEnumerable<string> KnownTargets() => ["1"];
}
/// <summary>Delivery that runs <see cref="OnDeliver"/> first, like a whole agent turn calling back into the registry.</summary>
private sealed class ReentrantDelivery(string workspace) : IA2AResultDelivery
{
@@ -383,6 +405,45 @@ public sealed class A2AWatchTests : IDisposable
Assert.Empty(PendingFile.Read(_ws));
}
[Fact]
public async Task Flush_SkippingAClaimedRecord_RetriedWhenThatDeliveryFails()
{
PendingFile.Upsert(_ws, new WatchRecord("w1", "1", "remote", "t1", "c1", null, DateTimeOffset.UtcNow, WatchSources.Poll, Undelivered: "the result"));
var delivery = new FailFirstDelivery(_ws);
await using var registry = new A2AWatchRegistry(TimeProvider.System, new HttpClientHandler());
registry.UseDelivery(delivery);
var first = registry.FlushUndeliveredAsync("1", _ws, CancellationToken.None);
await WaitUntilAsync(() => delivery.Entered >= 1);
// A session start meanwhile: its flush finds the record claimed and skips it.
await registry.FlushUndeliveredAsync("1", _ws, CancellationToken.None);
Assert.Empty(delivery.Delivered);
delivery.Release(); // the first delivery fails: the skipped flush is run again
await first;
Assert.Equal("the result", Assert.Single(delivery.Delivered));
Assert.Empty(PendingFile.Read(_ws));
}
[Fact]
public async Task Finish_FailedDelivery_FlushesWhenAFlushSkippedIt()
{
await using var remote = await RemoteAgentServer.StartAsync();
remote.Faux.Enqueue(FauxTurn.Say("Stranded?"));
WriteAgents();
var delivery = new FailFirstDelivery(_ws);
await using var registry = new A2AWatchRegistry(TimeProvider.System, remote.Handler) { SourceOverride = WatchSources.Poll, PollDelays = [TimeSpan.FromMilliseconds(50)] };
registry.UseDelivery(delivery);
var settings = new A2AClientSettings();
using var agent = await registry.Connector(settings).ConnectAsync(AgentsFile.Find(_ws, "remote")!, CancellationToken.None);
await registry.StartAsync("1", _ws, agent, await StartRemoteTaskAsync(agent), settings, CancellationToken.None);
await WaitUntilAsync(() => delivery.Entered >= 1); // the watch is finishing and its delivery is in flight
await registry.FlushUndeliveredAsync("1", _ws, CancellationToken.None); // skipped: claimed
Assert.Empty(delivery.Delivered);
delivery.Release();
await WaitUntilAsync(() => !delivery.Delivered.IsEmpty);
Assert.Contains("Stranded?", Assert.Single(delivery.Delivered));
await WaitUntilAsync(() => PendingFile.Read(_ws).Count == 0);
}
[Fact]
public async Task StopWatching_RacingFinish_DoesNotResurrectTheRecord()
{