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:
@@ -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);
|
||||
|
||||
|
||||
@@ -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()
|
||||
{
|
||||
|
||||
@@ -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()
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user