path: root/src/Agent.Core/Runtime/AgentRuntime.cs

using System.Text.Json;
using Agent.Core.Ollama;
using Agent.Core.Prompts;
using Agent.Core.Sessions;
using Agent.Core.Skills;
using Agent.Core.Tasks;
using Agent.Core.Tools;
using Agent.Planner;

namespace Agent.Core.Runtime;

public sealed class AgentRuntime(TaskRegistry registry, IOllamaClient ollama, AgentRuntimeSettings settings, InternalToolDispatcher? toolDispatcher = null, IAgentRuntimeObserver? observer = null, PlannerService? plannerService = null)
{
    private readonly InternalToolDispatcher _tools = toolDispatcher ?? new(
        new DenyByDefaultPolicy(InternalToolCatalog.All.Select(static item => item.Name).ToHashSet(StringComparer.Ordinal)),
        new FakeInternalToolBackend());
    private readonly IAgentRuntimeObserver _observer = observer ?? NullAgentRuntimeObserver.Instance;
    private readonly PlannerService _planner = plannerService ?? new(ollama, settings);

    /// <summary>
    /// Plan → Run lifecycle for a single MCP Task.
    /// 1. Planning: generate Initial TaskGraph via structured output.
    /// 2. Running: execute ready Tasks (dependencies satisfied) in critical-path / creation order.
    /// Sandbox Task Resource is applied at RunAsync time, not PlanAsync time.
    /// </summary>
    public async Task ExecuteAsync(Guid taskId, CancellationToken cancellationToken)
    {
        if (!registry.TryGet(taskId, null, out TaskSnapshot? task) || task is null) return;
        if (!registry.Sessions.TryGet(task.SessionId, out AgentSession? session) || session is null)
        {
            registry.Transition(taskId, AgentTaskState.Failed, result: "Session not found.");
            return;
        }

        try
        {
            // Vision capability check before any Ollama call.
            if (task.Images.Count > 0 && await ollama.SupportsVisionAsync(cancellationToken).ConfigureAwait(false) is false)
                throw new InvalidOperationException(settings.VisionUnsupportedMessage);

            // Phase 1: Planning — generate TaskGraph via structured output.
            TaskGraph graph;
            try
            {
                graph = await _planner.PlanAsync(taskId, task.SessionId, session, task.Request, task.Temperature, task.Images, cancellationToken).ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) { throw; }
            catch (Exception ex)
            {
                registry.Transition(taskId, AgentTaskState.Failed, result: $"Planning failed: {ex.Message}");
                return;
            }

            // Phase 2: Running — execute ready Tasks sequentially.
            registry.Transition(taskId, AgentTaskState.Running, "running");
            await RunAsync(taskId, task, session, graph, cancellationToken).ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            string? cleanupFailure = await EndTaskSafelyAsync(taskId, false, CancellationToken.None).ConfigureAwait(false);
            string result = cleanupFailure is null ? "Cancelled." : $"Cancelled. {cleanupFailure}";
            registry.Transition(taskId, AgentTaskState.Cancelled, result: result);
        }
        catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
        {
            string? cleanupFailure = await EndTaskSafelyAsync(taskId, false, CancellationToken.None).ConfigureAwait(false);
            string result = cleanupFailure is null ? ex.Message : $"{ex.Message} {cleanupFailure}";
            registry.Transition(taskId, AgentTaskState.Failed, result: result);
        }
    }

    /// <summary>
    /// Executes ready Tasks (dependencies satisfied) in critical-path then creation order.
    /// Each PlannerTask is executed as a single LLM chat-with-tools iteration loop.
    /// Sandbox Task Resource (BeginTask/EndTask) wraps the entire RunAsync phase.
    /// </summary>
    private async Task RunAsync(Guid taskId, TaskSnapshot task, AgentSession session, TaskGraph graph, CancellationToken cancellationToken)
    {
        // Apply Sandbox Task Resource at RunAsync time.
        await _tools.BeginTaskAsync(taskId, task.SessionId, cancellationToken).ConfigureAwait(false);

        var finalContents = new List<string>();
        int replanCount = 0;
        const int maxReplans = 3;
        try
        {
            // Execute tasks in ready order: dependencies satisfied → critical path → creation order.
            while (true)
            {
                PlannerTask? next = SelectNextReadyTask(graph);
                if (next is null) break;

                graph.SetState(next.Id, PlannerTaskState.Running);
                try
                {
                    string content = await ExecutePlannerTaskAsync(taskId, task, session, graph, next, cancellationToken).ConfigureAwait(false);
                    graph.SetState(next.Id, PlannerTaskState.Completed);
                    finalContents.Add(content);
                }
                catch (InvalidOperationException ex) when (ex.Message.Contains("denied by sandbox policy", StringComparison.OrdinalIgnoreCase) && replanCount < maxReplans)
                {
                    // Policy deny → replan possible (Design §8.4).
                    replanCount++;
                    graph.SetState(next.Id, PlannerTaskState.Failed);
                    graph = await _planner.ReplanAsync(taskId, task.SessionId, session, task.Request, task.Temperature, ex.Message, cancellationToken).ConfigureAwait(false);
                    // Transition back to Running for the replan iteration.
                    registry.Transition(taskId, AgentTaskState.Running, "replan");
                    continue;
                }
            }

            // All tasks completed — assemble final result.
            string finalContent = finalContents.Count > 0
                ? string.Join("\n\n", finalContents)
                : "No tasks were executed.";
            string? cleanupFailure = await EndTaskSafelyAsync(taskId, true, CancellationToken.None).ConfigureAwait(false);
            if (cleanupFailure is not null) throw new InvalidOperationException(cleanupFailure);
            session.Append(new ConversationMessage("user", task.Request));
            session.Append(new ConversationMessage("assistant", finalContent));
            registry.Transition(taskId, AgentTaskState.Completed, result: finalContent);
        }
        catch
        {
            // Re-throw without double-cleanup; ExecuteAsync's catch handles EndTaskSafelyAsync.
            throw;
        }
    }

    /// <summary>
    /// Selects the next ready Task: dependencies satisfied, not yet completed/failed.
    /// Priority: critical path first, then creation order.
    /// </summary>
    private static PlannerTask? SelectNextReadyTask(TaskGraph graph)
    {
        var completedIds = graph.Tasks
            .Where(static t => t.State is PlannerTaskState.Completed or PlannerTaskState.Failed)
            .Select(static t => t.Id)
            .ToHashSet();

        // Critical path task IDs in order.
        var criticalPath = new HashSet<Guid>(graph.CriticalPath.TaskIds);

        // Ready = not completed/failed/running and all dependencies completed.
        var ready = graph.Tasks
            .Where(t => t.State is not (PlannerTaskState.Completed or PlannerTaskState.Failed or PlannerTaskState.Running))
            .Where(t => t.Dependencies.All(depId => completedIds.Contains(depId)))
            .ToList();

        if (ready.Count == 0) return null;

        // Prefer critical path tasks, then creation order (Tasks list is in insertion order).
        return ready.FirstOrDefault(t => criticalPath.Contains(t.Id)) ?? ready[0];
    }

    /// <summary>
    /// Executes a single PlannerTask: LLM chat-with-tools iteration loop.
    /// Uses PlannerTask.RequiresExternalEffect (not keyword heuristic) for ToolUseGuard.
    /// </summary>
    private async Task<string> ExecutePlannerTaskAsync(Guid taskId, TaskSnapshot task, AgentSession session, TaskGraph graph, PlannerTask plannerTask, CancellationToken cancellationToken)
    {
        // Progressive Disclosure Stage 2: LLM Skill Selection.
        IReadOnlyList<AgentSkill> selectedSkills = await SelectSkillsAsync(task.Request, cancellationToken).ConfigureAwait(false);

        var messages = PromptComposer.Compose(new PromptContext(
            FixedPrompt.Load(),
            settings.UserSystemPrompt,
            settings.Locale,
            settings.Skills,
            selectedSkills,
            settings.SandboxCapabilitySummary,
            BuildTaskGraphSummary(graph, plannerTask),
            session.Summary ?? string.Empty,
            session.Conversation,
            plannerTask.Description,
            task.Images)).ToList();

        var guard = new ToolUseGuard(settings.SameFailureLimit);
        bool requiresExternalEffect = plannerTask.RequiresExternalEffect;
        bool externalEffectEvidence = false;
        var unresolvedToolFailures = new Dictionary<string, string>(StringComparer.Ordinal);
        string? finalContent = null;

        for (int iteration = 0; iteration < settings.MaxIterations; iteration++)
        {
            _observer.LlmRequest(taskId, iteration, messages);
            OllamaResult result = await ollama.ChatWithToolsAsync(messages, task.Temperature, InternalToolCatalog.All, cancellationToken).ConfigureAwait(false);
            _observer.LlmResponse(taskId, iteration, result);
            registry.AddUsage(taskId, result.InputTokens, result.OutputTokens);

            // Record usage in TaskGraph for CPM tracking.
            graph.AddUsage(plannerTask.Id, result.InputTokens, result.OutputTokens);

            // Context Summary: compact old conversation/results when context utilization exceeds threshold.
            // Usage = latest prompt_eval_count / num_ctx (NOT cumulative InputTokens).
            TryContextSummary(session, result.PromptEvalCount);

            IReadOnlyList<OllamaToolCall> calls = result.ToolCalls ?? [];
            if (calls.Count > 0)
            {
                messages.Add(new OllamaMessage("assistant", result.Content, ToolCalls: calls));
                string? failureFingerprint = null;
                foreach (OllamaToolCall call in calls)
                {
                    ToolExecutionResult toolResult;
                    try
                    {
                        toolResult = await _tools.DispatchAsync(taskId, call.Function.Name, call.Function.Arguments, cancellationToken).ConfigureAwait(false);
                    }
                    catch (UnauthorizedAccessException ex)
                    {
                        toolResult = new(false, ex.Message, false, $"{call.Function.Name}:{ex.GetType().Name}:{ex.Message}", RetryAllowed: false);
                    }
                    catch (Exception ex) when (ex is ToolValidationException or ArgumentException)
                    {
                        toolResult = new(false, ex.Message, false, $"{call.Function.Name}:{ex.GetType().Name}:{ex.Message}");
                    }

                    externalEffectEvidence |= toolResult.Success && toolResult.ExternalEffect;
                    if (toolResult.Success) unresolvedToolFailures.Remove(call.Function.Name);
                    else unresolvedToolFailures[call.Function.Name] = toolResult.FailureFingerprint ?? toolResult.Output;

                    // Record ExecutionResult in session for context summary.
                    session.AppendExecutionResult(new ExecutionResult(taskId, call.Function.Name, toolResult.Success, toolResult.Output));

                    _observer.ToolResult(taskId, call.Function.Name, toolResult);
                    if (!toolResult.Success && !toolResult.RetryAllowed)
                        throw new InvalidOperationException($"Tool '{call.Function.Name}' was denied by sandbox policy and will not be retried: {toolResult.Output}");
                    failureFingerprint = toolResult.FailureFingerprint ?? failureFingerprint;
                    messages.Add(new OllamaMessage("tool", toolResult.Output, ToolName: call.Function.Name));
                }
                ToolUseGuardResult callGuard = guard.Evaluate(requiresExternalEffect, true, externalEffectEvidence, failureFingerprint);
                if (callGuard.Decision == ToolUseGuardDecision.Fail) throw new InvalidOperationException(callGuard.Observation);
                continue;
            }

            if (unresolvedToolFailures.Count > 0)
            {
                messages.Add(new OllamaMessage("system", "Completion is not allowed while tool failures remain unresolved. Retry or correct: " +
                    string.Join("; ", unresolvedToolFailures.Select(static item => $"{item.Key}={item.Value}"))));
                continue;
            }

            ToolUseGuardResult completion = guard.Evaluate(requiresExternalEffect, false, externalEffectEvidence);
            if (completion.Decision == ToolUseGuardDecision.Retry)
            {
                messages.Add(new OllamaMessage("system", completion.Observation));
                continue;
            }
            if (completion.Decision == ToolUseGuardDecision.Fail) throw new InvalidOperationException(completion.Observation);
            finalContent = result.Content;
            break;
        }

        if (finalContent is null) throw new InvalidOperationException($"Agent exceeded the iteration limit of {settings.MaxIterations}.");
        return finalContent;
    }

    /// <summary>
    /// Builds a TaskGraph summary for the prompt context (Design §9.2 item 6).
    /// Shows the current task position within the overall graph.
    /// </summary>
    private static string BuildTaskGraphSummary(TaskGraph graph, PlannerTask currentTask)
    {
        var sb = new System.Text.StringBuilder();
        sb.AppendLine("Current TaskGraph:");
        for (int i = 0; i < graph.Tasks.Count; i++)
        {
            PlannerTask t = graph.Tasks[i];
            string marker = t.Id == currentTask.Id ? "→ " : "  ";
            sb.AppendLine($"{marker}[{t.State}] {t.Description}");
        }
        sb.AppendLine($"Critical path tokens: {graph.CriticalPath.CumulativeTokenCost:F0}");
        return sb.ToString();
    }

    /// <summary>
    /// Context Summary: compacts old Conversation/ExecutionResults when context utilization exceeds threshold.
    /// Usage = latest prompt_eval_count / num_ctx (Design §9.3). NOT cumulative InputTokens.
    /// Preserves: user constraints, incomplete tasks, completed task results, failure causes, security denials.
    /// </summary>
    private void TryContextSummary(AgentSession session, long promptEvalCount)
    {
        if (settings.NumCtx <= 0) return;
        double utilization = (double)promptEvalCount / settings.NumCtx;
        if (utilization < settings.ContextSummaryThreshold) return;

        // Compact old conversation entries, keeping only the most recent N.
        // The summary itself is generated by the LLM in a separate call (future enhancement);
        // for now, we compact the active context set to stay within num_ctx.
        session.CompactConversation(settings.ConversationKeepLastN);
        session.CompactExecutionResults(settings.ExecutionResultsKeepLastN);
    }

    private async Task<IReadOnlyList<AgentSkill>> SelectSkillsAsync(string request, CancellationToken cancellationToken)
    {
        if (settings.Skills.Count == 0) return [];

        string catalog = string.Join("\n", settings.Skills.Select(static s => $"- {s.Name}: {s.Description}"));
        var selectionMessages = new List<OllamaMessage>
        {
            new("system", $"Select skills relevant to the user request. Available skills:\n{catalog}\n\nReturn a JSON object with a 'selected_skills' array containing skill names. Select only relevant skills. If none are relevant, return an empty array."),
            new("user", request)
        };

        string schema = """{"type":"object","properties":{"selected_skills":{"type":"array","items":{"type":"string"}}},"required":["selected_skills"]}""";
        using JsonDocument formatDoc = JsonDocument.Parse(schema);

        OllamaResult result;
        try
        {
            result = await ollama.ChatStructuredAsync(selectionMessages, 0.0, formatDoc.RootElement.Clone(), cancellationToken).ConfigureAwait(false);
        }
        catch
        {
            return [];
        }

        return ParseSelectedSkills(result.Content);
    }

    private IReadOnlyList<AgentSkill> ParseSelectedSkills(string content)
    {
        try
        {
            using JsonDocument doc = JsonDocument.Parse(content);
            if (!doc.RootElement.TryGetProperty("selected_skills", out JsonElement skillsElement)) return [];
            var selectedNames = new HashSet<string>(StringComparer.Ordinal);
            foreach (JsonElement item in skillsElement.EnumerateArray())
            {
                string? name = item.GetString();
                if (!string.IsNullOrWhiteSpace(name)) selectedNames.Add(name);
            }
            var valid = new List<AgentSkill>();
            foreach (AgentSkill skill in settings.Skills)
            {
                if (selectedNames.Contains(skill.Name))
                {
                    valid.Add(skill);
                    if (valid.Count >= 10) break;
                }
            }
            return valid;
        }
        catch
        {
            return [];
        }
    }

    private async Task<string?> EndTaskSafelyAsync(Guid taskId, bool success, CancellationToken cancellationToken)
    {
        try
        {
            await _tools.EndTaskAsync(taskId, success, cancellationToken).ConfigureAwait(false);
            return null;
        }
        catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
        {
            return $"Task resource cleanup failed: {ex.Message}";
        }
    }
}