MyAgent/MyAgent/Core/SessionManager.cs

using Microsoft.Extensions.Options;

namespace MyAgent.Core;

public sealed class SessionManager(IOptions<ShirobotOptions> options, TimeProvider timeProvider)
{
    private readonly ShirobotOptions _options = options.Value;
    private readonly Dictionary<string, PendingThread> _pending = new();
    private readonly SemaphoreSlim _lock = new(1, 1);

    public async Task<MessageBatch?> DebounceAsync(IncomingMessage message, CancellationToken ct = default)
    {
        var threadId = BuildThreadId(message);
        CancellationTokenSource supersede;
        long revision;

        await _lock.WaitAsync(ct);
        try
        {
            if (!_pending.TryGetValue(threadId, out var pending))
            {
                pending = new PendingThread();
                _pending[threadId] = pending;
            }

            var now = message.CreatedAt;
            pending.Messages.RemoveAll(m => (now - m.CreatedAt).TotalSeconds > _options.Fusion.SameUserMergeWindowSec);
            pending.Messages.Add(message);
            var max = Math.Max(1, _options.Fusion.MaxMergedMessages);
            if (pending.Messages.Count > max)
                pending.Messages.RemoveRange(0, pending.Messages.Count - max);

            pending.DelayCts?.Cancel();
            supersede = new CancellationTokenSource();
            pending.DelayCts = supersede;
            revision = ++pending.Revision;
        }
        finally
        {
            _lock.Release();
        }

        try
        {
            using var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, supersede.Token);
            await Task.Delay(TimeSpan.FromMilliseconds(Math.Max(0, _options.Fusion.DebounceMs)), timeProvider, linked.Token);
        }
        catch (OperationCanceledException) when (!ct.IsCancellationRequested)
        {
            supersede.Dispose();
            return null;
        }
        catch
        {
            await _lock.WaitAsync(CancellationToken.None);
            try
            {
                if (_pending.TryGetValue(threadId, out var pending) && pending.Revision == revision && ReferenceEquals(pending.DelayCts, supersede))
                    _pending.Remove(threadId);
            }
            finally
            {
                _lock.Release();
            }
            supersede.Dispose();
            throw;
        }

        await _lock.WaitAsync(CancellationToken.None);
        try
        {
            if (ct.IsCancellationRequested)
            {
                if (_pending.TryGetValue(threadId, out var cancelledPending) && cancelledPending.Revision == revision && ReferenceEquals(cancelledPending.DelayCts, supersede))
                    _pending.Remove(threadId);
                ct.ThrowIfCancellationRequested();
            }
            if (!_pending.TryGetValue(threadId, out var pending) || pending.Revision != revision || !ReferenceEquals(pending.DelayCts, supersede))
                return null;

            var stableMessages = pending.Messages.ToList();
            _pending.Remove(threadId);
            return CreateBatch(stableMessages, threadId);
        }
        finally
        {
            _lock.Release();
            supersede.Dispose();
        }
    }

    public Task<MessageBatch> MergeManyAsync(IReadOnlyList<IncomingMessage> messages, CancellationToken ct = default)
    {
        if (messages.Count == 0) throw new ArgumentException("messages is empty", nameof(messages));
        ct.ThrowIfCancellationRequested();
        var stableMessages = messages.ToList();
        return Task.FromResult(CreateBatch(stableMessages, BuildThreadId(stableMessages[^1])));
    }

    public async Task ClearGroupAsync(string groupId, CancellationToken ct = default)
    {
        if (string.IsNullOrWhiteSpace(groupId)) return;
        await _lock.WaitAsync(ct);
        try
        {
            foreach (var key in _pending.Keys.Where(k => k.StartsWith($"group:{groupId}:", StringComparison.OrdinalIgnoreCase)).ToList())
            {
                if (_pending.Remove(key, out var pending)) pending.DelayCts?.Cancel();
            }
        }
        finally
        {
            _lock.Release();
        }
    }

    public async Task RemoveMessageAsync(string messageId, CancellationToken ct = default)
    {
        if (string.IsNullOrWhiteSpace(messageId)) return;
        await _lock.WaitAsync(ct);
        try
        {
            var emptyThreads = new List<string>();
            foreach (var pair in _pending)
            {
                pair.Value.Messages.RemoveAll(m => string.Equals(m.MessageId, messageId, StringComparison.OrdinalIgnoreCase));
                if (pair.Value.Messages.Count == 0) emptyThreads.Add(pair.Key);
            }
            foreach (var threadId in emptyThreads)
            {
                if (_pending.Remove(threadId, out var pending)) pending.DelayCts?.Cancel();
            }
        }
        finally
        {
            _lock.Release();
        }
    }

    public static string BuildThreadId(IncomingMessage message)
    {
        if (!string.IsNullOrWhiteSpace(message.ReplyToMessageId))
            return $"group:{message.GroupId}:user:{message.UserId}:reply:{message.ReplyToMessageId}";
        return $"group:{message.GroupId}:user:{message.UserId}";
    }

    private static string RawJoin(IReadOnlyList<IncomingMessage> messages)
    {
        return string.Join("\n", messages.Select((m, i) => messages.Count == 1 ? m.Text.Trim() : $"{i + 1}. {m.Text.Trim()}"));
    }

    private static MessageBatch CreateBatch(IReadOnlyList<IncomingMessage> messages, string threadId)
    {
        var lastMessage = messages[^1];
        return new MessageBatch
        {
            GroupId = lastMessage.GroupId,
            UserId = lastMessage.UserId,
            ThreadId = threadId,
            Messages = messages.ToArray(),
            MergedText = RawJoin(messages)
        };
    }

    private sealed class PendingThread
    {
        public List<IncomingMessage> Messages { get; } = [];
        public long Revision { get; set; }
        public CancellationTokenSource? DelayCts { get; set; }
    }
}