Playa/src/Volcanic.AgenticCore/Features/Agent/AgentSession.cs

using Microsoft.Extensions.Logging;
using Volcanic.AgenticCore.Abstractions.Events;
using Volcanic.AgenticCore.Abstractions.Extensions;
using Volcanic.AgenticCore.Abstractions.Extensions.Events;
using Volcanic.AgenticCore.Abstractions.Messages;
using Volcanic.AgenticCore.Abstractions.Processes;
using Volcanic.AgenticCore.Abstractions.Providers;
using Volcanic.AgenticCore.Abstractions.Sessions;
using Volcanic.AgenticCore.Abstractions.Ui;
using Volcanic.AgenticCore.Features.Commands;
using Volcanic.AgenticCore.Features.Compaction;
using Volcanic.AgenticCore.Features.Extensions;
using Volcanic.AgenticCore.Features.Models;
using Volcanic.AgenticCore.Features.Sessions;
using Volcanic.AgenticCore.Features.SystemPrompt;
using Volcanic.AgenticCore.Features.Tools;
using Volcanic.AgenticCore.Hosting;

namespace Volcanic.AgenticCore.Features.Agent;

public sealed record AgentSessionState(
    string SessionId,
    string Cwd,
    Model? Model,
    ThinkingLevel ThinkingLevel,
    bool IsStreaming,
    bool IsCompacting,
    int MessageCount,
    IReadOnlyList<string> PendingSteering,
    IReadOnlyList<string> PendingFollowUp,
    string? SessionName,
    bool AutoCompaction,
    IReadOnlyList<ExtensionLoadError> ExtensionErrors);

/// <summary>
/// One conversation: the facade a host (UI) talks to. Owns the session tree, the extensions, the tool and model
/// registries, and runs the agent loop. Create it through <see cref="IAgentSessionFactory"/>.
/// </summary>
public sealed partial class AgentSession : IAsyncDisposable
{
    private readonly ISessionStore _store;
    private readonly ExtensionRunner _runner;
    private readonly ToolRegistry _tools;
    private readonly CommandRegistry _commands;
    private readonly ModelRegistry _models;
    private readonly Compactor _compactor;
    private readonly SystemPromptBuilder _promptBuilder;
    private readonly AgentSessionOptions _options;
    private readonly AgenticCoreOptions _coreOptions;
    private readonly RetrySettings _retry;
    private readonly ILogger _logger;
    private readonly ISecretRedactor _redactor;
    private readonly IUserInterface _ui;
    private readonly PendingMessageQueue _steering;
    private readonly PendingMessageQueue _followUp;
    private readonly List<CustomMessage> _nextTurnMessages = [];
    private readonly List<Action<AgentSessionEvent>> _subscribers = [];
    private readonly SemaphoreSlim _runGate = new(1, 1);
    private readonly Lock _stateLock = new();

    private SessionManager _session;
    private Model? _model;
    private ThinkingLevel _thinkingLevel;
    private bool _autoCompaction;
    private CancellationTokenSource? _abortCts;
    private Task? _runTask;
    private bool _isCompacting;
    private CompactOptions? _pendingCompaction;
    private string? _systemPromptForRun;
    private bool _disposed;

    internal AgentSession(
        SessionManager session, ISessionStore store, ExtensionRunner runner, ToolRegistry tools, CommandRegistry commands,
        ModelRegistry models, Compactor compactor, SystemPromptBuilder promptBuilder, AgentSessionOptions options,
        AgenticCoreOptions coreOptions, Model? model, ThinkingLevel thinkingLevel, ILogger logger, ISecretRedactor redactor)
    {
        _session = session;
        _store = store;
        _runner = runner;
        _tools = tools;
        _commands = commands;
        _models = models;
        _compactor = compactor;
        _promptBuilder = promptBuilder;
        _options = options;
        _coreOptions = coreOptions;
        _retry = options.Retry ?? coreOptions.Retry;
        _logger = logger;
        _redactor = redactor;
        _ui = options.Ui ?? NullUserInterface.Instance;
        _steering = new PendingMessageQueue(options.SteeringMode);
        _followUp = new PendingMessageQueue(options.FollowUpMode);
        _model = model;
        _thinkingLevel = thinkingLevel;
        _autoCompaction = options.AutoCompaction ?? coreOptions.Compaction.Enabled;
        _runner.OnError += error => Emit(new AgentSessionEvent.ExtensionError(error.ExtensionName, error.Event, error.Exception.Message));
    }

    // ---- State ----

    public string SessionId => _session.SessionId;
    public string Cwd => _options.Cwd;
    /// <summary>The session's time zone: the session option, else the host default, else the machine's zone.</summary>
    public TimeZoneInfo TimeZone => _options.TimeZone ?? _coreOptions.DefaultTimeZone ?? TimeZoneInfo.Local;
    /// <summary>The host's chat channel for this session (e.g. for scheduled deliveries), or null.</summary>
    public IChatChannel? Channel => _options.Channel;
    public Model? Model => _model;
    public ThinkingLevel ThinkingLevel => _thinkingLevel;
    public bool IsStreaming { get { lock (_stateLock) return _runTask is { IsCompleted: false }; } }
    public bool IsCompacting => _isCompacting;
    public IReadOnlyList<Message> Messages => _session.GetMessages();
    public IReadOnlySessionManager SessionManager => _session;
    public IModelRegistry Models => _models;
    public IReadOnlyList<ToolInfo> Tools => _tools.All;
    public IReadOnlyList<string> ActiveTools => _tools.ActiveNames;
    public IReadOnlyList<CommandInfo> Commands => _commands.All;
    public IReadOnlyList<ExtensionLoadError> ExtensionErrors => _runner.LoadErrors;
    public IReadOnlyList<string> LoadedExtensions => _runner.LoadedExtensions;
    public string? SessionName => _session.Name;
    public bool AutoCompaction { get => _autoCompaction; set => _autoCompaction = value; }

    public AgentSessionState State => new(
        SessionId, Cwd, _model, _thinkingLevel, IsStreaming, _isCompacting, _session.GetMessages().Count,
        _steering.Texts(), _followUp.Texts(), _session.Name, _autoCompaction, _runner.LoadErrors);

    public string SystemPrompt => _systemPromptForRun ?? _promptBuilder.Render(_promptBuilder.DefaultOptions(_options.SystemPrompt ?? _coreOptions.DefaultSystemPrompt));

    // ---- Events ----

    /// <summary>Subscribe to session events. Handlers run synchronously on the loop; keep them fast and hand off to your own queue.</summary>
    public IDisposable Subscribe(Action<AgentSessionEvent> handler)
    {
        lock (_subscribers) _subscribers.Add(handler);
        return new Subscription(() => { lock (_subscribers) _subscribers.Remove(handler); });
    }

    private void Emit(AgentSessionEvent e)
    {
        Action<AgentSessionEvent>[] snapshot;
        lock (_subscribers) snapshot = _subscribers.ToArray();
        foreach (var handler in snapshot)
        {
            try { handler(e); }
            catch (Exception ex) { _logger.LogError(ex, "Session event subscriber failed for {Event}", e.GetType().Name); }
        }
    }

    // ---- Prompting ----

    /// <summary>
    /// Send user input. Slash commands run immediately (even while streaming). While a run is active,
    /// <paramref name="behavior"/> decides whether the text steers the current run or follows up after it;
    /// without it the call throws. When idle, starts a run and returns after the run settles.
    /// </summary>
    public async Task PromptAsync(string text, IReadOnlyList<ImageContent>? images = null, StreamingBehavior? behavior = null, CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
        if (CommandRegistry.Parse(text) is { } parsed && _commands.Get(parsed.Name) is { } command)
        {
            try
            {
                await command.Handler(parsed.Args, _runner.CreateCommandContext(), cancellationToken).ConfigureAwait(false);
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                _runner.Report(new ExtensionError("command:" + parsed.Name, "command", ex));
            }
            return;
        }

        var input = await _runner.EmitAsync(new InputEvent(text, images, InputSource.User, behavior), cancellationToken).ConfigureAwait(false);
        if (input.Handled) return;
        if (string.IsNullOrWhiteSpace(input.Text) && input.Images is not { Count: > 0 }) return;

        var message = UserMessage.From(input.Text, input.Images);
        if (IsStreaming)
        {
            switch (behavior)
            {
                case StreamingBehavior.Steer: Steer(message); return;
                case StreamingBehavior.FollowUp: FollowUp(message); return;
                default: throw new InvalidOperationException("The agent is streaming; pass StreamingBehavior.Steer or StreamingBehavior.FollowUp to queue the prompt.");
            }
        }

        await StartRunAsync(message, images, cancellationToken).ConfigureAwait(false);
    }

    /// <summary>Queue a steering message; delivered after the current turn's tools, before the next model call.</summary>
    public Task SteerAsync(string text, IReadOnlyList<ImageContent>? images = null)
    {
        Steer(UserMessage.From(text, images));
        return Task.CompletedTask;
    }

    /// <summary>Queue a follow-up; delivered when the agent has no more tool calls or steering messages.</summary>
    public Task FollowUpAsync(string text, IReadOnlyList<ImageContent>? images = null)
    {
        FollowUp(UserMessage.From(text, images));
        return Task.CompletedTask;
    }

    private void Steer(Message message)
    {
        _steering.Enqueue(message);
        Emit(new AgentSessionEvent.QueueUpdate(_steering.Texts(), _followUp.Texts()));
        if (!IsStreaming) _ = StartRunAsync(null, null, CancellationToken.None);
    }

    private void FollowUp(Message message)
    {
        _followUp.Enqueue(message);
        Emit(new AgentSessionEvent.QueueUpdate(_steering.Texts(), _followUp.Texts()));
        if (!IsStreaming) _ = StartRunAsync(null, null, CancellationToken.None);
    }

    /// <summary>Remove queued steering and follow-up messages; returns their text so a UI can restore it.</summary>
    public (IReadOnlyList<string> Steering, IReadOnlyList<string> FollowUp) ClearQueue()
    {
        var s = _steering.Clear().Select(m => m.Text).ToList();
        var f = _followUp.Clear().Select(m => m.Text).ToList();
        Emit(new AgentSessionEvent.QueueUpdate([], []));
        return (s, f);
    }

    /// <summary>Abort the current run (and compaction) and wait until the session is idle.</summary>
    public async Task AbortAsync()
    {
        CancellationTokenSource? cts;
        Task? run;
        lock (_stateLock) { cts = _abortCts; run = _runTask; }
        cts?.Cancel();
        if (run is not null)
        {
            try { await run.ConfigureAwait(false); } catch (Exception) { /* surfaced through events */ }
        }
    }

    public async Task WaitForIdleAsync(CancellationToken cancellationToken = default)
    {
        while (true)
        {
            Task? run;
            lock (_stateLock) run = _runTask;
            if (run is null || run.IsCompleted) return;
            try { await run.WaitAsync(cancellationToken).ConfigureAwait(false); } catch (Exception) when (!cancellationToken.IsCancellationRequested) { }
        }
    }

    // ---- Model and thinking ----

    /// <summary>Switch model. Returns false when the provider has no usable credentials.</summary>
    public async Task<bool> SetModelAsync(Model model, CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
        if (!_models.IsAvailable(model)) return false;
        var previous = _model;
        _model = model;
        await _session.AppendModelChangeAsync(model, cancellationToken).ConfigureAwait(false);
        var clamped = ThinkingLevels.Clamp(model, _thinkingLevel);
        if (clamped != _thinkingLevel) await SetThinkingLevelAsync(clamped, cancellationToken).ConfigureAwait(false);
        await _runner.EmitAsync(new ModelSelectEvent(model, previous, ModelSelectSource.Set), cancellationToken).ConfigureAwait(false);
        Emit(new AgentSessionEvent.ModelChanged(model));
        return true;
    }

    /// <summary>Resolve <c>provider/model-id</c> and switch. Returns false when unknown or unavailable.</summary>
    public Task<bool> SetModelAsync(string key, CancellationToken cancellationToken = default)
    {
        var model = _models.Find(key);
        return model is null ? Task.FromResult(false) : SetModelAsync(model, cancellationToken);
    }

    public async Task SetThinkingLevelAsync(ThinkingLevel level, CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
        var effective = _model is null ? level : ThinkingLevels.Clamp(_model, level);
        var previous = _thinkingLevel;
        if (effective == previous) return;
        _thinkingLevel = effective;
        await _session.AppendThinkingLevelChangeAsync(effective, cancellationToken).ConfigureAwait(false);
        await _runner.EmitAsync(new ThinkingLevelSelectEvent(effective, previous), cancellationToken).ConfigureAwait(false);
        Emit(new AgentSessionEvent.ThinkingLevelChanged(effective));
    }

    public IReadOnlyList<ThinkingLevel> AvailableThinkingLevels => _model is null ? [ThinkingLevel.Off] : ThinkingLevels.Supported(_model);

    public void SetActiveTools(IEnumerable<string> names) => _tools.SetActive(names);

    // ---- Session management ----

    public async Task SetSessionNameAsync(string? name, CancellationToken cancellationToken = default)
    {
        await _session.AppendSessionInfoAsync(name, cancellationToken).ConfigureAwait(false);
        await _runner.EmitAsync(new SessionInfoChangedEvent(name), cancellationToken).ConfigureAwait(false);
        Emit(new AgentSessionEvent.SessionInfoChanged(name));
    }

    /// <summary>Start a fresh session in the same working directory. Extensions may cancel.</summary>
    public Task<SwitchResult> NewSessionAsync(NewSessionOptions? options = null, CancellationToken cancellationToken = default) =>
        ReplaceSessionAsync(SessionStartReason.New, SessionShutdownReason.New, null, options?.ParentSessionId, cancellationToken);

    /// <summary>Resume another session from the store. Extensions may cancel.</summary>
    public Task<SwitchResult> SwitchSessionAsync(string sessionId, CancellationToken cancellationToken = default) =>
        ReplaceSessionAsync(SessionStartReason.Resume, SessionShutdownReason.Resume, sessionId, null, cancellationToken);

    public Task<IReadOnlyList<SessionSummary>> ListSessionsAsync(CancellationToken cancellationToken = default) =>
        _store.ListAsync(_session.Scope, cancellationToken);

    private async Task<SwitchResult> ReplaceSessionAsync(SessionStartReason startReason, SessionShutdownReason shutdownReason, string? targetSessionId, string? parentSessionId, CancellationToken cancellationToken)
    {
        ThrowIfDisposed();
        var before = await _runner.EmitAsync(new SessionBeforeSwitchEvent(startReason, targetSessionId), cancellationToken).ConfigureAwait(false);
        if (before.Cancelled) return new SwitchResult(true);

        SessionManager next;
        if (targetSessionId is not null)
        {
            next = await Sessions.SessionManager.OpenAsync(_store, _session.Scope, targetSessionId, cancellationToken).ConfigureAwait(false)
                   ?? throw new KeyNotFoundException($"Session '{targetSessionId}' not found.");
        }
        else
        {
            next = await Sessions.SessionManager.CreateAsync(_store, _session.Scope, Cwd, parentSessionId, cancellationToken: cancellationToken).ConfigureAwait(false);
        }

        await AbortAsync().ConfigureAwait(false);
        var previousId = _session.SessionId;
        await _runner.EmitAsync(new SessionShutdownEvent(shutdownReason, targetSessionId), cancellationToken).ConfigureAwait(false);
        _steering.Clear();
        _followUp.Clear();
        _nextTurnMessages.Clear();
        _session = next;
        RestoreBranchSettings();
        await _runner.EmitAsync(new SessionStartEvent(startReason, previousId), cancellationToken).ConfigureAwait(false);
        Emit(new AgentSessionEvent.SessionInfoChanged(_session.Name));
        return new SwitchResult(false);
    }

    /// <summary>Apply the model and thinking level recorded on the branch (used on resume and switch).</summary>
    internal void RestoreBranchSettings()
    {
        var (provider, modelId, level) = _session.GetBranchSettings();
        if (provider is not null && modelId is not null && _models.Find(provider, modelId) is { } model) _model = model;
        if (level is { } l) _thinkingLevel = _model is null ? l : ThinkingLevels.Clamp(_model, l);
    }

    // ---- Compaction ----

    /// <summary>Compact now when idle; when streaming, compaction runs between turns and this returns null.</summary>
    public async Task<CompactionResult?> CompactAsync(string? customInstructions = null, CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
        if (IsStreaming)
        {
            _pendingCompaction = new CompactOptions(customInstructions);
            return null;
        }
        return await CompactCoreAsync(CompactionReason.Manual, customInstructions, cancellationToken).ConfigureAwait(false);
    }

    public ContextUsage? GetContextUsage()
    {
        if (_model is null) return null;
        var tokens = Compactor.ContextTokens(_session.BuildContext());
        return new ContextUsage(tokens, _model.ContextWindow, tokens is null ? null : Math.Round(100.0 * tokens.Value / _model.ContextWindow, 1));
    }

    // ---- Lifecycle ----

    public async ValueTask DisposeAsync()
    {
        if (_disposed) return;
        _disposed = true;
        await AbortAsync().ConfigureAwait(false);
        try { await _runner.EmitAsync(new SessionShutdownEvent(SessionShutdownReason.Quit)).ConfigureAwait(false); }
        catch (Exception ex) { _logger.LogError(ex, "session_shutdown failed"); }
        _runner.Invalidate();
    }

    private void ThrowIfDisposed()
    {
        if (_disposed) throw new ObjectDisposedException(nameof(AgentSession));
    }
}