Playa/hosts/Volcanic.AgenticCore.Telegram/Chats/ChatRegistry.cs

using System.Collections.Concurrent;
using Microsoft.Extensions.Logging;
using Volcanic.AgenticCore.Abstractions.Events;
using Volcanic.AgenticCore.Abstractions.Extensions.Events;
using Volcanic.AgenticCore.Abstractions.Messages;
using Volcanic.AgenticCore.Abstractions.Ui;
using Volcanic.AgenticCore.Extensions.GoogleCalendar;
using Volcanic.AgenticCore.Extensions.Scheduler;
using Volcanic.AgenticCore.Features.Agent;
using Volcanic.AgenticCore.Features.Extensions;
using Volcanic.AgenticCore.Hosting;
using Volcanic.AgenticCore.Telegram.Inbound;
using Volcanic.AgenticCore.Telegram.Rendering;

namespace Volcanic.AgenticCore.Telegram.Chats;

/// <summary>Everything that belongs to one Telegram chat: its agent session, renderer and dialog surface.</summary>
public sealed class ChatContext(long chatId, AgentSession session, TelegramUi ui, ChatRenderer renderer, IDisposable subscription, ChatInbox inbox, ILogger logger,
    IChatChannel channel) : IAsyncDisposable
{
    public long ChatId => chatId;
    public AgentSession Session => session;
    public TelegramUi Ui => ui;
    public ChatRenderer Renderer => renderer;

    /// <summary>Delivers files and texts into this chat in order with rendered output (the scheduler uses it too).</summary>
    public IChatChannel Channel => channel;

    /// <summary>Groups this chat's incoming messages into prompts (see <see cref="ChatInbox"/>).</summary>
    public ChatInbox Inbox => inbox;

    /// <summary>The most recent background prompt; completes when that run has settled.</summary>
    public Task LastPrompt { get; private set; } = Task.CompletedTask;

    /// <summary>
    /// Send user input without waiting for the run (it can take minutes, and the polling loop must keep receiving so
    /// <c>/stop</c> gets through). While a run is active the text steers it.
    /// </summary>
    public Task PromptInBackground(string text, IReadOnlyList<ImageContent>? images = null) => LastPrompt = Task.Run(async () =>
    {
        try
        {
            await session.PromptAsync(text, images, StreamingBehavior.Steer).ConfigureAwait(false);
        }
        catch (Exception ex)
        {
            logger.LogError(ex, "Prompt in chat {ChatId} failed", chatId);
            Report($"The request failed: {ex.Message}");
        }
    });

    /// <summary>Show an error in the chat, in order with the rendered events.</summary>
    public void Report(string message) => renderer.Post(new AgentSessionEvent.Notification(message, NotifyKind.Error));

    public async ValueTask DisposeAsync()
    {
        await inbox.DisposeAsync().ConfigureAwait(false);
        subscription.Dispose();
        await session.DisposeAsync().ConfigureAwait(false);
        await renderer.DisposeAsync().ConfigureAwait(false);
    }
}

/// <summary>
/// One <see cref="ChatContext"/> per chat, created on first use. Each chat works in <c>&lt;WorkspaceRoot&gt;/&lt;chatId&gt;</c>,
/// stores its sessions under scope <c>telegram-&lt;chatId&gt;</c> and resumes the most recent one after a restart.
/// </summary>
public sealed class ChatRegistry(
    IAgentSessionFactory factory,
    TelegramSender sender,
    TelegramOptions options,
    ILoggerFactory loggerFactory,
    TimeProvider? time = null,
    GoogleCalendarOptions? calendar = null,
    ChatSettingsStore? settings = null,
    ChatTimeZoneService? zones = null,
    JobScheduler? scheduler = null) : IAsyncDisposable
{
    private readonly ConcurrentDictionary<long, Lazy<Task<ChatContext>>> _chats = new();
    private readonly ILogger _logger = loggerFactory.CreateLogger<ChatRegistry>();
    private readonly TimeProvider _time = time ?? TimeProvider.System;

    public static string ScopeFor(long chatId) => $"telegram-{chatId}";

    public string WorkspaceFor(long chatId) => Path.Combine(options.WorkspaceRoot, chatId.ToString(System.Globalization.CultureInfo.InvariantCulture));

    /// <summary>The chat's context if it was already created, without creating it.</summary>
    public ChatContext? Find(long chatId) =>
        _chats.TryGetValue(chatId, out var lazy) && lazy.IsValueCreated && lazy.Value.IsCompletedSuccessfully ? lazy.Value.Result : null;

    public async Task<ChatContext> GetAsync(long chatId, CancellationToken cancellationToken = default)
    {
        var lazy = _chats.GetOrAdd(chatId, id => new Lazy<Task<ChatContext>>(() => CreateAsync(id)));
        try
        {
            return await lazy.Value.WaitAsync(cancellationToken).ConfigureAwait(false);
        }
        catch (Exception) when (lazy.Value.IsFaulted)
        {
            _chats.TryRemove(new KeyValuePair<long, Lazy<Task<ChatContext>>>(chatId, lazy)); // let the next message try again
            throw;
        }
    }

    private async Task<ChatContext> CreateAsync(long chatId)
    {
        var cwd = WorkspaceFor(chatId);
        Directory.CreateDirectory(cwd);
        var scope = ScopeFor(chatId);
        var last = (await factory.ListSessionsAsync(cwd, scope).ConfigureAwait(false)).FirstOrDefault();

        var logger = loggerFactory.CreateLogger<ChatContext>();
        var renderer = new ChatRenderer(chatId, sender, options, loggerFactory.CreateLogger<ChatRenderer>());
        var ui = new TelegramUi(chatId, sender, renderer);
        var channel = new TelegramChatChannel(chatId, sender, renderer);
        var extras = new List<ExtensionRegistration>();
        if (calendar is not null) extras.Add(GoogleCalendarSetup.Registration(calendar));
        if (zones is not null) extras.Add(ExtensionRegistration.FromInstance(new ChatSettingsExtension(chatId, zones)));
        if (scheduler is not null)
            extras.Add(ExtensionRegistration.FromInstance(new SchedulerExtension(scheduler, Scheduling.TelegramJobDelivery.TargetOf(chatId))));
        AgentSession session;
        try
        {
            session = await factory.CreateAsync(new AgentSessionOptions
            {
                Cwd = cwd,
                Scope = scope,
                SessionId = last?.SessionId,
                Model = options.Model,
                Ui = ui,
                Channel = channel,
                TimeZone = settings?.TimeZoneFor(chatId),
                ExtraExtensions = extras,
            }).ConfigureAwait(false);
        }
        catch
        {
            await renderer.DisposeAsync().ConfigureAwait(false);
            throw;
        }

        var subscription = session.Subscribe(renderer.Post);
        _logger.LogInformation("Chat {ChatId}: {Action} session {SessionId} in {Cwd}", chatId, last is null ? "new" : "resumed", session.SessionId, cwd);
        foreach (var error in session.ExtensionErrors) _logger.LogWarning("Chat {ChatId}: extension error {Source}: {Error}", chatId, error.Source, error.Error);
        ChatContext? context = null;
        var inbox = new ChatInbox(parts => FlushAsync(context!, parts), options.InputDebounce, _time,
            logger: loggerFactory.CreateLogger<ChatInbox>());
        context = new ChatContext(chatId, session, ui, renderer, subscription, inbox, logger, channel);
        return context; // no lease can be taken before this returns, so the flush always sees the context
    }

    /// <summary>
    /// The zone for a chat's timestamp headers: the chat's setting (<see cref="ChatSettingsStore.TimeZoneFor"/>), which also
    /// feeds the session's zone. Without a store (bare test setups) the session's zone, else the machine's. Read at flush
    /// time, so a <c>/timezone</c> change applies to the next batch. Never creates a session.
    /// </summary>
    public TimeZoneInfo TimeZoneFor(long chatId) => settings?.TimeZoneFor(chatId) ?? Find(chatId)?.Session.TimeZone ?? TimeZoneInfo.Local;

    /// <summary>A flushed batch becomes one prompt; while a run is active it steers it and each message gets a 👀 reaction.</summary>
    internal Task FlushAsync(ChatContext chat, IReadOnlyList<InboundPart> parts)
    {
        var (text, images) = InboundBatch.Render(parts, TimeZoneFor(chat.ChatId));
        // Tell the user their message joined the running request (best-effort, fire-and-forget). If the run ends between
        // this check and the steer, the message starts a new run and still shows 👀: harmless.
        if (chat.Session.IsStreaming)
            foreach (var part in parts) _ = sender.ReactAsync(chat.ChatId, part.MessageId, "👀");
        _ = chat.PromptInBackground(text, images.Count > 0 ? images : null);
        return Task.CompletedTask;
    }

    public async ValueTask DisposeAsync()
    {
        foreach (var lazy in _chats.Values)
        {
            if (!lazy.IsValueCreated) continue;
            try
            {
                await (await lazy.Value.ConfigureAwait(false)).DisposeAsync().ConfigureAwait(false);
            }
            catch (Exception ex)
            {
                _logger.LogDebug(ex, "Disposing a chat failed");
            }
        }
        _chats.Clear();
    }
}