fix(a2a): track live runs from the start, atomic eviction, token-based cancel and safe final writes

src/Volcanic.AgenticCore.A2A.Server/Agent/A2AAgentHandler.cs

@@ -3,67 +3,107 @@ using A2A;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using Volcanic.AgenticCore.A2A.Server.Callers;
+using Volcanic.AgenticCore.Abstractions.Events;
 using Volcanic.AgenticCore.Abstractions.Processes;
 
 namespace Volcanic.AgenticCore.A2A.Server.Agent;
 
 /// <summary>
 /// One task = one run of the caller context's session. Ends completed (final text as artifact <c>response</c>),
-/// input-required (the agent called <c>request_input</c>), failed (model error, hop limit, bad file) or canceled.
+/// input-required (the agent called <c>request_input</c>), failed (model error, hop limit, bad file, busy context, any
+/// unexpected exception) or canceled (the request token or a CancelTask). The task is registered as live in the pool
+/// before its first event and stays live until its final status is written, so restart recovery never touches it.
 /// </summary>
 internal sealed class A2AAgentHandler(
     A2ACaller caller, A2ASessionPool pool, IOptions<A2AHostOptions> options, ISecretRedactor redactor, HttpClient http,
     TimeProvider time, ILogger<A2AAgentHandler> logger) : IAgentHandler
 {
+    /// <summary>Bound on the final status write: after a client disconnect the bounded event queue may never be read.</summary>
+    internal static TimeSpan FinalWriteTimeout { get; set; } = TimeSpan.FromSeconds(30);
+
+    private enum Outcome { Completed, InputRequired, Failed, Canceled }
+
+    private sealed record Result(Outcome Outcome, string Label, string? Message = null, string? FinalText = null)
+    {
+        public static Result Canceled { get; } = new(Outcome.Canceled, "canceled");
+        public static Result Failed(string message, string label = "failed") => new(Outcome.Failed, label, message);
+    }
+
     public async Task ExecuteAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken)
     {
-        var updater = new TaskUpdater(eventQueue, context.TaskId, context.ContextId);
         var started = Stopwatch.GetTimestamp();
-        if (context.Task is null) await updater.SubmitAsync(cancellationToken).ConfigureAwait(false);
-        await updater.StartWorkAsync(null, cancellationToken).ConfigureAwait(false);
+        var updater = new TaskUpdater(eventQueue, context.TaskId, context.ContextId);
+        using var tracking = pool.Track(context.TaskId, out var run);
+        Result result;
+        using (var runCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, run.Cancel.Token))
+        {
+            var token = runCts.Token;
+            try
+            {
+                result = await RunAsync(context, updater, token).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException) when (token.IsCancellationRequested)
+            {
+                result = Result.Canceled;
+            }
+            catch (Exception ex)
+            {
+                logger.LogWarning(ex, "A2A: task {TaskId} of {Caller} failed", context.TaskId, caller.Name);
+                result = Result.Failed("The run failed: " + redactor.Redact(ex.Message), "failed (error)");
+            }
+            if (token.IsCancellationRequested) result = Result.Canceled;
+        }
+        var label = await WriteFinalAsync(context, updater, result).ConfigureAwait(false);
+        Log(context, label, started);
+    }
+
+    private async Task<Result> RunAsync(RequestContext context, TaskUpdater updater, CancellationToken token)
+    {
+        if (context.Task is null) await updater.SubmitAsync(token).ConfigureAwait(false);
+        await updater.StartWorkAsync(null, token).ConfigureAwait(false);
 
         var hops = Hops.Read(context.Metadata, context.Message.Metadata);
         if (hops > options.Value.MaxHops)
-        {
-            await updater.FailAsync(A2AText.Agent($"hop limit reached ({hops} > {options.Value.MaxHops})"), CancellationToken.None).ConfigureAwait(false);
-            Log(context, "failed (hop limit)", started);
-            return;
-        }
+            return Result.Failed($"hop limit reached ({hops} > {options.Value.MaxHops})", "failed (hop limit)");
 
         A2ASessionPool.Lease lease;
         try
         {
-            lease = await pool.AcquireAsync(caller, context.ContextId, context.TaskId, cancellationToken).ConfigureAwait(false);
+            lease = await pool.AcquireAsync(caller, context.ContextId, context.TaskId, token).ConfigureAwait(false);
         }
         catch (A2AException ex)
         {
-            await updater.FailAsync(A2AText.Agent(ex.Message), CancellationToken.None).ConfigureAwait(false);
-            Log(context, "failed (busy)", started);
-            return;
+            return Result.Failed(ex.Message, "failed (busy)");
         }
 
+        // The lease is released before the final status is written, so a continuation can take the context at once.
         using (lease)
         {
             string prompt;
             try
             {
-                prompt = await A2APrompt.BuildAsync(caller, context.Message, options.Value.MaxFileBytes, http, time, cancellationToken).ConfigureAwait(false);
+                prompt = await A2APrompt.BuildAsync(caller, context.Message, options.Value.MaxFileBytes, http, time, token).ConfigureAwait(false);
             }
             catch (A2APromptException ex)
             {
-                await updater.FailAsync(A2AText.Agent(ex.Message), CancellationToken.None).ConfigureAwait(false);
-                Log(context, "failed (file)", started);
-                return;
+                return Result.Failed(ex.Message, "failed (file)");
             }
+            token.ThrowIfCancellationRequested();
 
-            var mapper = new A2ARunMapper(updater, redactor, cancellationToken);
-            lease.Extension.BeginRun(hops);
-            lease.Channel.Bind(updater);
+            var session = lease.Session;
+            var mapper = new A2ARunMapper(updater, redactor, logger, token);
             try
             {
-                using var subscription = lease.Session.Subscribe(mapper.OnEvent);
-                using var abort = cancellationToken.Register(() => _ = lease.Session.AbortAsync());
-                await lease.Session.PromptAsync(prompt, cancellationToken: CancellationToken.None).ConfigureAwait(false);
+                lease.Extension.BeginRun(hops);
+                lease.Channel.Bind(updater);
+                using var subscription = session.Subscribe(e =>
+                {
+                    // An abort that came before the session's run existed had nothing to cancel: repeat it now.
+                    if (e is AgentSessionEvent.AgentStart && token.IsCancellationRequested) _ = session.AbortAsync();
+                    mapper.OnEvent(e);
+                });
+                using var abort = token.Register(() => _ = session.AbortAsync());
+                await session.PromptAsync(prompt, cancellationToken: CancellationToken.None).ConfigureAwait(false);
             }
             finally
             {
@@ -71,37 +111,48 @@ internal sealed class A2AAgentHandler(
                 await mapper.DisposeAsync().ConfigureAwait(false);
             }
 
-            string state;
-            if (cancellationToken.IsCancellationRequested || mapper.Aborted)
-            {
-                await updater.CancelAsync(CancellationToken.None).ConfigureAwait(false);
-                state = "canceled";
-            }
-            else if (mapper.Error is { } error)
-            {
-                await updater.FailAsync(A2AText.Agent(error), CancellationToken.None).ConfigureAwait(false);
-                state = "failed";
-            }
-            else
-            {
-                if (mapper.FinalText is { } text)
-                    await updater.AddArtifactAsync([Part.FromText(text)], A2ARunMapper.ResponseArtifactId, "response", null,
-                        lastChunk: true, append: false, CancellationToken.None).ConfigureAwait(false);
-                if (lease.Extension.PendingQuestion is { } question)
-                {
-                    await updater.RequireInputAsync(A2AText.Agent(question), CancellationToken.None).ConfigureAwait(false);
-                    state = "input-required";
-                }
-                else
-                {
-                    await updater.CompleteAsync(null, CancellationToken.None).ConfigureAwait(false);
-                    state = "completed";
-                }
-            }
-            Log(context, state, started);
+            if (token.IsCancellationRequested || mapper.Aborted) return Result.Canceled;
+            if (mapper.Error is { } error) return Result.Failed(error);
+            return lease.Extension.PendingQuestion is { } question
+                ? new Result(Outcome.InputRequired, "input-required", question, mapper.FinalText)
+                : new Result(Outcome.Completed, "completed", null, mapper.FinalText);
         }
     }
 
+    /// <summary>Write the final status (and answer artifact) with a bounded token; returns the label to log.</summary>
+    private async Task<string> WriteFinalAsync(RequestContext context, TaskUpdater updater, Result result)
+    {
+        using var cts = new CancellationTokenSource(FinalWriteTimeout);
+        try
+        {
+            switch (result.Outcome)
+            {
+                case Outcome.Canceled:
+                    await updater.CancelAsync(cts.Token).ConfigureAwait(false);
+                    break;
+                case Outcome.Failed:
+                    await updater.FailAsync(A2AText.Agent(result.Message!), cts.Token).ConfigureAwait(false);
+                    break;
+                default:
+                    if (result.FinalText is { } text)
+                        await updater.AddArtifactAsync([Part.FromText(text)], A2ARunMapper.ResponseArtifactId, "response", null,
+                            lastChunk: true, append: false, cts.Token).ConfigureAwait(false);
+                    if (result.Outcome == Outcome.InputRequired)
+                        await updater.RequireInputAsync(A2AText.Agent(result.Message!), cts.Token).ConfigureAwait(false);
+                    else
+                        await updater.CompleteAsync(null, cts.Token).ConfigureAwait(false);
+                    break;
+            }
+            return result.Label;
+        }
+        catch (Exception ex)
+        {
+            logger.LogWarning(ex, "A2A: writing the final status of task {TaskId} failed", context.TaskId);
+            return result.Label + " (final status not written)";
+        }
+    }
+
+    /// <summary>Cancel the run (its token, then the session) and write Canceled.</summary>
     public async Task CancelAsync(RequestContext context, AgentEventQueue eventQueue, CancellationToken cancellationToken)
     {
         await pool.AbortAsync(context.TaskId).ConfigureAwait(false);

src/Volcanic.AgenticCore.A2A.Server/Agent/A2APrompt.cs

@@ -93,5 +93,9 @@ internal static class A2APrompt
         {
             throw new A2APromptException($"Could not download {url}: {ex.Message}");
         }
+        catch (OperationCanceledException) when (!ct.IsCancellationRequested)
+        {
+            throw new A2APromptException($"Could not download {url}: the request timed out.");
+        }
     }
 }

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

@@ -1,5 +1,6 @@
 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;
@@ -10,23 +11,31 @@ 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).
+/// 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;
 
-    public A2ARunMapper(TaskUpdater updater, ISecretRedactor redactor, CancellationToken cancellationToken)
+    public A2ARunMapper(TaskUpdater updater, ISecretRedactor redactor, ILogger logger, CancellationToken cancellationToken)
     {
         _updater = updater;
         _redactor = redactor;
-        _ct = cancellationToken;
+        _logger = logger;
+        _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        _ct = _cts.Token;
         _pump = Task.Run(PumpAsync, CancellationToken.None);
     }
 
@@ -54,7 +63,7 @@ internal sealed class A2ARunMapper : IAsyncDisposable
                 else
                 {
                     Error = null;
-                    if (!string.IsNullOrWhiteSpace(message.Text)) FinalText = message.Text;
+                    FinalText = string.IsNullOrWhiteSpace(message.Text) ? null : message.Text;
                 }
                 break;
             case AgentSessionEvent.ToolExecutionStart start:
@@ -69,17 +78,28 @@ internal sealed class A2ARunMapper : IAsyncDisposable
     {
         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) { }
-            catch (InvalidOperationException) { } // queue already completed (task cancelled)
+            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.</summary>
+    /// <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();
-        await _pump.ConfigureAwait(false);
+        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();
     }
 }
 

src/Volcanic.AgenticCore.A2A.Server/Agent/A2ASessionPool.cs

@@ -24,10 +24,32 @@ internal sealed class A2ASessionPool : IA2ALiveRuns, IAsyncDisposable
         public A2ASessionExtension Extension => extension;
         public string? ActiveTaskId { get; set; }
         public DateTimeOffset LastUsed { get; set; }
+
+        /// <summary>Removed from the pool by eviction (set under the entry lock); its session is being disposed.</summary>
+        public bool Evicted { get; set; }
     }
 
+    /// <summary>
+    /// A task whose handler is running, from its first event to its final status write. <see cref="Entry"/> is set
+    /// once the task holds its context's session; <see cref="Cancel"/> is the run-level cancellation signal.
+    /// </summary>
+    internal sealed class LiveRun
+    {
+        public Entry? Entry { get; internal set; }
+        public CancellationTokenSource Cancel { get; } = new();
+
+        /// <summary>When the handler finished (its tracking handle was disposed); null while it runs.</summary>
+        public DateTimeOffset? FinishedAt { get; internal set; }
+    }
+
+    /// <summary>
+    /// How long a finished task still counts as live. The SDK applies the final status event after the handler returns
+    /// (reading the task through the store first), so recovery must not see the task as orphaned in that window.
+    /// </summary>
+    internal static TimeSpan FinishedGrace { get; set; } = TimeSpan.FromMinutes(1);
+
     private readonly ConcurrentDictionary<string, Lazy<Task<Entry>>> _entries = new();
-    private readonly ConcurrentDictionary<string, Entry> _live = new();
+    private readonly ConcurrentDictionary<string, LiveRun> _live = new();
     private readonly IAgentSessionFactory _factory;
     private readonly A2AHostOptions _options;
     private readonly IA2AOwnerNotifier? _notifier;
@@ -48,52 +70,91 @@ internal sealed class A2ASessionPool : IA2ALiveRuns, IAsyncDisposable
 
     public static string ScopeFor(A2ACaller caller, string contextId) => $"{caller.ScopePrefix}-a2a-{caller.Name}-{contextId}";
 
-    public bool IsLive(string taskId) => _live.ContainsKey(taskId);
+    public bool IsLive(string taskId)
+    {
+        if (!_live.TryGetValue(taskId, out var run)) return false;
+        if (!IsExpired(run)) return true;
+        _live.TryRemove(new KeyValuePair<string, LiveRun>(taskId, run));
+        return false;
+    }
+
+    private bool IsExpired(LiveRun run) => run.FinishedAt is { } at && _time.GetUtcNow() - at >= FinishedGrace;
 
     public bool IsBusy(string scope) =>
         _entries.TryGetValue(scope, out var lazy) && lazy.IsValueCreated && lazy.Value.IsCompletedSuccessfully &&
         lazy.Value.Result.ActiveTaskId is not null;
 
+    /// <summary>
+    /// Mark the task live (so restart recovery leaves it alone) until <see cref="FinishedGrace"/> after the returned handle
+    /// is disposed. Call before the task's first event is queued and dispose after its final status is written.
+    /// </summary>
+    public IDisposable Track(string taskId, out LiveRun run)
+    {
+        foreach (var (id, old) in _live)
+            if (IsExpired(old)) _live.TryRemove(new KeyValuePair<string, LiveRun>(id, old));
+        var created = new LiveRun();
+        // A continuation can start while the previous run of the same task is still in its grace period: the newest run
+        // replaces it.
+        _live[taskId] = created;
+        run = created;
+        return new Tracking(this, created);
+    }
+
     /// <summary>Take the context's session for one task. Throws UnsupportedOperation while another task runs in it.</summary>
     public async Task<Lease> AcquireAsync(A2ACaller caller, string contextId, string taskId, CancellationToken cancellationToken)
     {
         var scope = ScopeFor(caller, contextId);
-        var lazy = _entries.GetOrAdd(scope, s => new Lazy<Task<Entry>>(() => OpenAsync(caller, s)));
-        Entry entry;
-        try
+        while (true)
         {
-            entry = await lazy.Value.WaitAsync(cancellationToken).ConfigureAwait(false);
+            var lazy = _entries.GetOrAdd(scope, s => new Lazy<Task<Entry>>(() => OpenAsync(caller, s)));
+            Entry entry;
+            try
+            {
+                entry = await lazy.Value.WaitAsync(cancellationToken).ConfigureAwait(false);
+            }
+            catch (Exception) when (lazy.Value.IsFaulted || lazy.Value.IsCanceled)
+            {
+                _entries.TryRemove(new KeyValuePair<string, Lazy<Task<Entry>>>(scope, lazy));
+                throw;
+            }
+            lock (entry)
+            {
+                if (entry.Evicted) continue; // disposed under us: open it again
+                if (entry.ActiveTaskId is not null)
+                    throw new A2AException("Context busy: retry when the current task ends.", A2AErrorCode.UnsupportedOperation);
+                entry.ActiveTaskId = taskId;
+            }
+            if (_live.TryGetValue(taskId, out var run)) run.Entry = entry;
+            return new Lease(this, entry, taskId);
         }
-        catch (Exception) when (lazy.Value.IsFaulted)
-        {
-            _entries.TryRemove(new KeyValuePair<string, Lazy<Task<Entry>>>(scope, lazy));
-            throw;
-        }
-        lock (entry)
-        {
-            if (entry.ActiveTaskId is not null)
-                throw new A2AException("Context busy: retry when the current task ends.", A2AErrorCode.UnsupportedOperation);
-            entry.ActiveTaskId = taskId;
-        }
-        _live[taskId] = entry;
-        return new Lease(this, entry, taskId);
     }
 
+    /// <summary>Cancel the task's run and wait until its session is idle.</summary>
     public async Task AbortAsync(string taskId)
     {
-        if (_live.TryGetValue(taskId, out var entry)) await entry.Session.AbortAsync().ConfigureAwait(false);
+        if (!_live.TryGetValue(taskId, out var run) || run.FinishedAt is not null) return;
+        try { await run.Cancel.CancelAsync().ConfigureAwait(false); }
+        catch (ObjectDisposedException) { return; }
+        if (run.Entry is { } entry && entry.ActiveTaskId == taskId)
+            await entry.Session.AbortAsync().ConfigureAwait(false);
     }
 
     private void Release(Entry entry, string taskId)
     {
-        _live.TryRemove(taskId, out _);
         lock (entry)
         {
-            entry.ActiveTaskId = null;
+            if (entry.ActiveTaskId == taskId) entry.ActiveTaskId = null;
             entry.LastUsed = _time.GetUtcNow();
         }
     }
 
+    /// <summary>The run ended: it stays live for <see cref="FinishedGrace"/>, then IsLive/Track drop it.</summary>
+    private void Untrack(LiveRun run)
+    {
+        run.FinishedAt = _time.GetUtcNow();
+        run.Cancel.Dispose();
+    }
+
     private async Task<Entry> OpenAsync(A2ACaller caller, string scope)
     {
         Directory.CreateDirectory(caller.Workspace);
@@ -120,11 +181,21 @@ internal sealed class A2ASessionPool : IA2ALiveRuns, IAsyncDisposable
         var cutoff = _time.GetUtcNow() - _options.SessionIdleTimeout;
         foreach (var (scope, lazy) in _entries)
         {
-            if (!lazy.IsValueCreated || !lazy.Value.IsCompletedSuccessfully) continue;
+            if (!lazy.IsValueCreated) continue;
+            if (lazy.Value.IsFaulted || lazy.Value.IsCanceled)
+            {
+                _entries.TryRemove(new KeyValuePair<string, Lazy<Task<Entry>>>(scope, lazy));
+                continue;
+            }
+            if (!lazy.Value.IsCompletedSuccessfully) continue;
             var entry = lazy.Value.Result;
             lock (entry)
-                if (entry.ActiveTaskId is not null || entry.LastUsed > cutoff) continue;
-            if (!_entries.TryRemove(new KeyValuePair<string, Lazy<Task<Entry>>>(scope, lazy))) continue;
+            {
+                // Check and remove atomically: AcquireAsync takes the same lock and retries when it sees Evicted.
+                if (entry.Evicted || entry.ActiveTaskId is not null || entry.LastUsed > cutoff) continue;
+                if (!_entries.TryRemove(new KeyValuePair<string, Lazy<Task<Entry>>>(scope, lazy))) continue;
+                entry.Evicted = true;
+            }
             try { await entry.Session.DisposeAsync().ConfigureAwait(false); }
             catch (Exception ex) { _logger.LogDebug(ex, "A2A: disposing idle session {Scope} failed", scope); }
         }
@@ -152,4 +223,14 @@ internal sealed class A2ASessionPool : IA2ALiveRuns, IAsyncDisposable
         public A2ASessionExtension Extension => entry.Extension;
         public void Dispose() => pool.Release(entry, taskId);
     }
+
+    private sealed class Tracking(A2ASessionPool pool, LiveRun run) : IDisposable
+    {
+        private int _disposed;
+
+        public void Dispose()
+        {
+            if (Interlocked.Exchange(ref _disposed, 1) == 0) pool.Untrack(run);
+        }
+    }
 }

tests/Volcanic.AgenticCore.A2A.Server.Tests/SessionPoolTests.cs

@@ -0,0 +1,239 @@
+using A2A;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging.Abstractions;
+using Microsoft.Extensions.Options;
+using Microsoft.Extensions.Time.Testing;
+using Volcanic.AgenticCore.A2A.Server.Agent;
+using Volcanic.AgenticCore.A2A.Server.Callers;
+using Volcanic.AgenticCore.Features.Agent;
+using Volcanic.AgenticCore.Features.Extensions;
+using Volcanic.AgenticCore.Features.Processes;
+using Volcanic.AgenticCore.Hosting;
+using Volcanic.AgenticCore.Testing;
+
+namespace Volcanic.AgenticCore.A2A.Server.Tests;
+
+public sealed class SessionPoolTests : IAsyncDisposable
+{
+    private readonly string _root = Path.Combine(Path.GetTempPath(), "a2a-pool", Guid.NewGuid().ToString("N"));
+    private readonly FauxProvider _faux = new();
+    private readonly ServiceProvider _services;
+    private readonly FakeTimeProvider _time = new(DateTimeOffset.Parse("2026-10-07T10:00:00Z"));
+    private readonly A2AHostOptions _options = new() { SessionIdleTimeout = TimeSpan.FromMinutes(30) };
+    private readonly A2ASessionPool _pool;
+    private readonly A2ACaller _caller;
+
+    public SessionPoolTests()
+    {
+        Directory.CreateDirectory(_root);
+        _caller = new A2ACaller("alpha", "isolated", Path.Combine(_root, "ws"), "p", "1");
+        var services = new ServiceCollection();
+        services.AddAgenticCore(o =>
+        {
+            o.AgentDir = Path.Combine(_root, ".agent");
+            o.DiscoverExtensionAssemblies = false;
+            o.DefaultModel = FauxProvider.ModelKey;
+            o.Extensions.Add(ExtensionRegistration.FromInstance(_faux.Extension));
+        });
+        _services = services.BuildServiceProvider();
+        _pool = new A2ASessionPool(_services.GetRequiredService<IAgentSessionFactory>(), Options.Create(_options), _time,
+            NullLogger<A2ASessionPool>.Instance);
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await _pool.DisposeAsync();
+        await _services.DisposeAsync();
+        try { Directory.Delete(_root, true); } catch (Exception) { }
+    }
+
+    private A2AAgentHandler Handler() =>
+        new(_caller, _pool, Options.Create(_options), SecretRedactor.Empty, new HttpClient(), TimeProvider.System,
+            NullLogger<A2AAgentHandler>.Instance);
+
+    private static RequestContext Context(string text, string taskId = "t1", string contextId = "c1") => new()
+    {
+        Message = new Message { Role = Role.User, MessageId = Guid.NewGuid().ToString("N"), Parts = [Part.FromText(text)] },
+        TaskId = taskId,
+        ContextId = contextId,
+        StreamingResponse = false,
+    };
+
+    [Fact]
+    public void Pool_Track_MakesTaskLiveBeforeAcquire()
+    {
+        var tracking = _pool.Track("t1", out var run);
+        Assert.True(_pool.IsLive("t1"));
+        Assert.Null(run.Entry);
+        tracking.Dispose();
+        Assert.True(_pool.IsLive("t1")); // the SDK may still be applying the final event
+        _time.Advance(A2ASessionPool.FinishedGrace);
+        Assert.False(_pool.IsLive("t1"));
+    }
+
+    [Fact]
+    public void Pool_Track_FinishedRunIsReplacedByContinuation()
+    {
+        var first = _pool.Track("t1", out _);
+        first.Dispose();
+        using var second = _pool.Track("t1", out var run);
+        _time.Advance(A2ASessionPool.FinishedGrace);
+        Assert.True(_pool.IsLive("t1"));
+        Assert.Null(run.FinishedAt);
+    }
+
+    [Fact]
+    public async Task Pool_AbortAsync_FinishedRun_IsNoOp()
+    {
+        _pool.Track("t1", out var run).Dispose();
+        await _pool.AbortAsync("t1"); // the run's token source is disposed; must not throw
+        Assert.NotNull(run.FinishedAt);
+    }
+
+    [Fact]
+    public async Task Pool_Acquire_SetsTrackedEntry()
+    {
+        using var tracking = _pool.Track("t1", out var run);
+        using var lease = await _pool.AcquireAsync(_caller, "c1", "t1", CancellationToken.None);
+        Assert.Same(lease.Session, run.Entry!.Session);
+        Assert.True(_pool.IsBusy(A2ASessionPool.ScopeFor(_caller, "c1")));
+    }
+
+    [Fact]
+    public async Task Pool_SecondAcquire_SameContext_Throws()
+    {
+        using var lease = await _pool.AcquireAsync(_caller, "c1", "t1", CancellationToken.None);
+        var ex = await Assert.ThrowsAsync<A2AException>(() => _pool.AcquireAsync(_caller, "c1", "t2", CancellationToken.None));
+        Assert.Equal(A2AErrorCode.UnsupportedOperation, ex.ErrorCode);
+    }
+
+    [Fact]
+    public async Task Pool_AbortAsync_CancelsRunToken()
+    {
+        using var tracking = _pool.Track("t1", out var run);
+        await _pool.AbortAsync("t1");
+        Assert.True(run.Cancel.IsCancellationRequested);
+    }
+
+    [Fact]
+    public async Task Pool_EvictedEntry_IsReopenedOnAcquire()
+    {
+        AgentSession first;
+        using (var lease = await _pool.AcquireAsync(_caller, "c1", "t1", CancellationToken.None)) first = lease.Session;
+        _time.Advance(TimeSpan.FromMinutes(31));
+        await _pool.EvictIdleAsync();
+        Assert.Equal(0, _pool.OpenCount);
+
+        _faux.Enqueue(FauxTurn.Say("hello again"));
+        using var again = await _pool.AcquireAsync(_caller, "c1", "t2", CancellationToken.None);
+        Assert.NotSame(first, again.Session);
+        await again.Session.PromptAsync("hi"); // a disposed session would throw ObjectDisposedException
+    }
+
+    [Fact]
+    public async Task Pool_Eviction_SkipsBusyAndRecentSessions()
+    {
+        using var busy = await _pool.AcquireAsync(_caller, "busy", "t1", CancellationToken.None);
+        using (await _pool.AcquireAsync(_caller, "idle", "t2", CancellationToken.None)) { }
+        await _pool.EvictIdleAsync();
+        Assert.Equal(2, _pool.OpenCount);
+        _time.Advance(TimeSpan.FromMinutes(31));
+        await _pool.EvictIdleAsync();
+        Assert.Equal(1, _pool.OpenCount);
+        Assert.True(_pool.IsBusy(A2ASessionPool.ScopeFor(_caller, "busy")));
+    }
+
+    [Fact]
+    public async Task Handler_TaskIsLive_ForEveryQueuedEvent()
+    {
+        _faux.Enqueue(FauxTurn.Say("done here"));
+        var queue = new AgentEventQueue();
+        var exec = Task.Run(async () =>
+        {
+            try { await Handler().ExecuteAsync(Context("go"), queue, CancellationToken.None); }
+            finally { queue.Complete(); }
+        });
+        var liveness = new List<bool>();
+        var states = new List<TaskState>();
+        await foreach (var e in queue)
+        {
+            liveness.Add(_pool.IsLive("t1"));
+            if (e.StatusUpdate is { } s) states.Add(s.Status.State);
+            if (e.Task is { } t) states.Add(t.Status.State);
+        }
+        await exec;
+        Assert.NotEmpty(liveness);
+        Assert.All(liveness, Assert.True);
+        Assert.Equal(TaskState.Submitted, states[0]);
+        Assert.Equal(TaskState.Completed, states[^1]);
+        _time.Advance(A2ASessionPool.FinishedGrace);
+        Assert.False(_pool.IsLive("t1"));
+    }
+
+    [Fact]
+    public async Task Handler_HopLimit_FailsWhileTracked()
+    {
+        _options.MaxHops = -1;
+        var queue = new AgentEventQueue();
+        var exec = Task.Run(async () =>
+        {
+            try { await Handler().ExecuteAsync(Context("go"), queue, CancellationToken.None); }
+            finally { queue.Complete(); }
+        });
+        TaskState? last = null;
+        var allLive = true;
+        await foreach (var e in queue)
+        {
+            allLive &= _pool.IsLive("t1");
+            if (e.StatusUpdate is { } s) last = s.Status.State;
+        }
+        await exec;
+        Assert.True(allLive);
+        Assert.Equal(TaskState.Failed, last);
+    }
+
+    [Fact]
+    public async Task Handler_Cancel_EndsCanceled()
+    {
+        _faux.Enqueue(FauxTurn.Say("slow").WithDelay(TimeSpan.FromSeconds(30)));
+        var handler = Handler();
+        var queue = new AgentEventQueue();
+        var exec = Task.Run(async () =>
+        {
+            try { await handler.ExecuteAsync(Context("go"), queue, CancellationToken.None); }
+            finally { queue.Complete(); }
+        });
+        var states = new List<TaskState>();
+        var reader = Task.Run(async () =>
+        {
+            await foreach (var e in queue)
+                if (e.StatusUpdate is { } s) lock (states) states.Add(s.Status.State);
+        });
+        var scope = A2ASessionPool.ScopeFor(_caller, "c1");
+        while (!_pool.IsBusy(scope)) await Task.Delay(10);
+        await Task.Delay(100); // let the model call start
+        await handler.CancelAsync(Context("go"), new AgentEventQueue(), CancellationToken.None);
+        await reader.WaitAsync(TimeSpan.FromSeconds(10));
+        await exec.WaitAsync(TimeSpan.FromSeconds(10));
+        Assert.Equal(TaskState.Canceled, states[^1]);
+        Assert.DoesNotContain(TaskState.Completed, states);
+    }
+
+    [Fact]
+    public async Task Handler_RequestTokenCancelledBeforeStart_EndsCanceledWithoutHanging()
+    {
+        using var cts = new CancellationTokenSource();
+        await cts.CancelAsync();
+        var queue = new AgentEventQueue();
+        var exec = Task.Run(async () =>
+        {
+            try { await Handler().ExecuteAsync(Context("go"), queue, cts.Token); }
+            finally { queue.Complete(); }
+        });
+        var states = new List<TaskState>();
+        await foreach (var e in queue)
+            if (e.StatusUpdate is { } s) states.Add(s.Status.State);
+        await exec.WaitAsync(TimeSpan.FromSeconds(10));
+        Assert.Equal([TaskState.Canceled], states);
+    }
+}