feat(a2a): pool one agent session per caller context
This commit is contained in:
EugeneTes
155
src/Volcanic.AgenticCore.A2A.Server/Agent/A2ASessionPool.cs
Normal file
155
src/Volcanic.AgenticCore.A2A.Server/Agent/A2ASessionPool.cs
Normal file
@@ -0,0 +1,155 @@
|
||||
using System.Collections.Concurrent;
|
||||
using A2A;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using Microsoft.Extensions.Options;
|
||||
using Volcanic.AgenticCore.A2A.Server.Callers;
|
||||
using Volcanic.AgenticCore.A2A.Server.Tasks;
|
||||
using Volcanic.AgenticCore.Features.Agent;
|
||||
using Volcanic.AgenticCore.Features.Extensions;
|
||||
using Volcanic.AgenticCore.Hosting;
|
||||
|
||||
namespace Volcanic.AgenticCore.A2A.Server.Agent;
|
||||
|
||||
/// <summary>
|
||||
/// One <see cref="AgentSession"/> per caller context (scope <c><prefix>-a2a-<caller>-<contextId></c>) in the
|
||||
/// caller's workspace, so it loads that workspace's AGENTS.md/CLAUDE.md, PERSONA.md and MEMORY.md. One run at a time per
|
||||
/// context. Idle sessions are disposed after <see cref="A2AHostOptions.SessionIdleTimeout"/> and resumed from JSONL later.
|
||||
/// </summary>
|
||||
internal sealed class A2ASessionPool : IA2ALiveRuns, IAsyncDisposable
|
||||
{
|
||||
internal sealed class Entry(AgentSession session, A2AChatChannel channel, A2ASessionExtension extension)
|
||||
{
|
||||
public AgentSession Session => session;
|
||||
public A2AChatChannel Channel => channel;
|
||||
public A2ASessionExtension Extension => extension;
|
||||
public string? ActiveTaskId { get; set; }
|
||||
public DateTimeOffset LastUsed { get; set; }
|
||||
}
|
||||
|
||||
private readonly ConcurrentDictionary<string, Lazy<Task<Entry>>> _entries = new();
|
||||
private readonly ConcurrentDictionary<string, Entry> _live = new();
|
||||
private readonly IAgentSessionFactory _factory;
|
||||
private readonly A2AHostOptions _options;
|
||||
private readonly IA2AOwnerNotifier? _notifier;
|
||||
private readonly TimeProvider _time;
|
||||
private readonly ILogger<A2ASessionPool> _logger;
|
||||
private readonly ITimer _evictTimer;
|
||||
|
||||
public A2ASessionPool(IAgentSessionFactory factory, IOptions<A2AHostOptions> options, TimeProvider time,
|
||||
ILogger<A2ASessionPool> logger, IA2AOwnerNotifier? notifier = null)
|
||||
{
|
||||
_factory = factory;
|
||||
_options = options.Value;
|
||||
_notifier = notifier;
|
||||
_time = time;
|
||||
_logger = logger;
|
||||
_evictTimer = time.CreateTimer(_ => _ = EvictIdleAsync(), null, TimeSpan.FromMinutes(1), TimeSpan.FromMinutes(1));
|
||||
}
|
||||
|
||||
public static string ScopeFor(A2ACaller caller, string contextId) => $"{caller.ScopePrefix}-a2a-{caller.Name}-{contextId}";
|
||||
|
||||
public bool IsLive(string taskId) => _live.ContainsKey(taskId);
|
||||
|
||||
public bool IsBusy(string scope) =>
|
||||
_entries.TryGetValue(scope, out var lazy) && lazy.IsValueCreated && lazy.Value.IsCompletedSuccessfully &&
|
||||
lazy.Value.Result.ActiveTaskId is not null;
|
||||
|
||||
/// <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
|
||||
{
|
||||
entry = await lazy.Value.WaitAsync(cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
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);
|
||||
}
|
||||
|
||||
public async Task AbortAsync(string taskId)
|
||||
{
|
||||
if (_live.TryGetValue(taskId, out var entry)) await entry.Session.AbortAsync().ConfigureAwait(false);
|
||||
}
|
||||
|
||||
private void Release(Entry entry, string taskId)
|
||||
{
|
||||
_live.TryRemove(taskId, out _);
|
||||
lock (entry)
|
||||
{
|
||||
entry.ActiveTaskId = null;
|
||||
entry.LastUsed = _time.GetUtcNow();
|
||||
}
|
||||
}
|
||||
|
||||
private async Task<Entry> OpenAsync(A2ACaller caller, string scope)
|
||||
{
|
||||
Directory.CreateDirectory(caller.Workspace);
|
||||
var last = (await _factory.ListSessionsAsync(caller.Workspace, scope).ConfigureAwait(false)).FirstOrDefault();
|
||||
var channel = new A2AChatChannel(_options.MaxFileBytes);
|
||||
var extension = new A2ASessionExtension(caller, _notifier);
|
||||
var session = await _factory.CreateAsync(new AgentSessionOptions
|
||||
{
|
||||
Cwd = caller.Workspace,
|
||||
Scope = scope,
|
||||
SessionId = last?.SessionId,
|
||||
Model = _options.Model,
|
||||
Channel = channel,
|
||||
TimeZone = _options.TimeZoneFor?.Invoke(caller),
|
||||
ExtraExtensions = [ExtensionRegistration.FromInstance(extension)],
|
||||
}).ConfigureAwait(false);
|
||||
_logger.LogInformation("A2A: {Action} session {SessionId} for {Caller} (owner {Owner})",
|
||||
last is null ? "new" : "resumed", session.SessionId, caller.Name, caller.OwnerId);
|
||||
return new Entry(session, channel, extension) { LastUsed = _time.GetUtcNow() };
|
||||
}
|
||||
|
||||
internal async Task EvictIdleAsync()
|
||||
{
|
||||
var cutoff = _time.GetUtcNow() - _options.SessionIdleTimeout;
|
||||
foreach (var (scope, lazy) in _entries)
|
||||
{
|
||||
if (!lazy.IsValueCreated || !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;
|
||||
try { await entry.Session.DisposeAsync().ConfigureAwait(false); }
|
||||
catch (Exception ex) { _logger.LogDebug(ex, "A2A: disposing idle session {Scope} failed", scope); }
|
||||
}
|
||||
}
|
||||
|
||||
internal int OpenCount => _entries.Count;
|
||||
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
await _evictTimer.DisposeAsync().ConfigureAwait(false);
|
||||
foreach (var lazy in _entries.Values)
|
||||
{
|
||||
if (!lazy.IsValueCreated || !lazy.Value.IsCompletedSuccessfully) continue;
|
||||
try { await lazy.Value.Result.Session.DisposeAsync().ConfigureAwait(false); }
|
||||
catch (Exception ex) { _logger.LogDebug(ex, "A2A: disposing a session failed"); }
|
||||
}
|
||||
_entries.Clear();
|
||||
}
|
||||
|
||||
/// <summary>A context's session held for one task; dispose to release it.</summary>
|
||||
internal sealed class Lease(A2ASessionPool pool, Entry entry, string taskId) : IDisposable
|
||||
{
|
||||
public AgentSession Session => entry.Session;
|
||||
public A2AChatChannel Channel => entry.Channel;
|
||||
public A2ASessionExtension Extension => entry.Extension;
|
||||
public void Dispose() => pool.Release(entry, taskId);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user