path: root/src/Agent.Core/Runtime/PlannerService.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;

/// <summary>
/// Builds an initial TaskGraph from LLM structured output (Planning phase).
/// The Planner determines requires_external_effect per-task during Planning.
/// Pure analysis tasks get a TaskGraph with at least 1 Task (RequiresExternalEffect=false).
/// </summary>
public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings settings)
{
    private readonly IOllamaClient _ollama = ollama;
    private readonly AgentRuntimeSettings _settings = settings;

    /// <summary>
    /// Generates an initial TaskGraph for the given request via Ollama structured output.
    /// Returns a TaskGraph stored in the session, keyed by the MCP Task ID.
    /// </summary>
    public async Task<TaskGraph> PlanAsync(Guid taskId, Guid sessionId, AgentSession session, string request, double? temperature, IReadOnlyList<string>? images,
        Action<string, IReadOnlyList<OllamaMessage>, OllamaResult?, Exception?>? observeAttempt, CancellationToken cancellationToken)
    {
        // Build a Planning prompt with skill catalog and sandbox capability summary.
        IReadOnlyList<AgentSkill> selectedSkills = await SelectSkillsAsync(request, temperature, observeAttempt, cancellationToken).ConfigureAwait(false);

        var messages = new List<OllamaMessage>
        {
            new("system", FixedPrompt.Load()),
            new("system", _settings.UserSystemPrompt),
            new("system", $"Respond in locale: {_settings.Locale}."),
            new("system", ComposeSkillCatalog(_settings.Skills)),
            new("system", _settings.SandboxCapabilitySummary),
            new("system", BuildPlanningInstruction()),
            new("user", request, images)
        };

        TaskGraph graph = await RequestValidGraphAsync(messages, temperature, "planning", observeAttempt,
            plan => BuildTaskGraph(plan, request), cancellationToken).ConfigureAwait(false);
        session.TaskGraphs[taskId] = graph;
        return graph;
    }

    /// <summary>
    /// Replans after a policy deny or task failure by adding/replacing tasks in the existing TaskGraph.
    /// Completed tasks are immutable; only incomplete tasks are modified.
    /// </summary>
    public async Task<TaskGraph> ReplanAsync(Guid taskId, Guid sessionId, AgentSession session, string request, double? temperature, string replanReason,
        Action<string, IReadOnlyList<OllamaMessage>, OllamaResult?, Exception?>? observeAttempt, CancellationToken cancellationToken)
    {
        if (!session.TaskGraphs.TryGetValue(taskId, out TaskGraph? existing))
            return await PlanAsync(taskId, sessionId, session, request, temperature, null, observeAttempt, cancellationToken).ConfigureAwait(false);

        var messages = new List<OllamaMessage>
        {
            new("system", FixedPrompt.Load()),
            new("system", _settings.UserSystemPrompt),
            new("system", $"Respond in locale: {_settings.Locale}."),
            new("system", _settings.SandboxCapabilitySummary),
            new("system", BuildReplanInstruction(existing, replanReason)),
            new("user", request)
        };

        TaskGraph rebuilt = await RequestValidGraphAsync(messages, temperature, "replanning", observeAttempt,
            plan => RebuildTaskGraph(existing, plan, request), cancellationToken).ConfigureAwait(false);
        session.TaskGraphs[taskId] = rebuilt;
        return rebuilt;
    }

    private TaskGraph BuildTaskGraph(PlanResult plan, string originalRequest)
    {
        var graph = new TaskGraph(_settings.MaxTasks);
        var idMap = new List<Guid>();

        for (int i = 0; i < plan.Tasks.Count; i++)
        {
            PlannedTask pt = plan.Tasks[i];
            TokenBudgetClass budgetClass = ParseBudgetClass(pt.BudgetClass);
            TokenEstimate estimate = CreateEstimate(pt, budgetClass);

            // Map dependency indices to Guids.
            var dependencies = new HashSet<Guid>();
            foreach (int depIdx in pt.Dependencies)
            {
                if (depIdx < 0 || depIdx >= idMap.Count)
                    throw new PlannerValidationException($"Task {i} contains invalid dependency index {depIdx}.");
                dependencies.Add(idMap[depIdx]);
            }

            string description = string.IsNullOrWhiteSpace(pt.Description) ? originalRequest : pt.Description;
            PlannerTask task = graph.AddTask(description, budgetClass, estimate, dependencies, pt.RequiresExternalEffect,
                requiredTools: ValidateRequiredTools(pt));
            idMap.Add(task.Id);
        }

        if (graph.Count == 0) throw new PlannerValidationException("Structured plan did not contain an executable Task.");

        return graph;
    }

    private TaskGraph RebuildTaskGraph(TaskGraph existing, PlanResult plan, string originalRequest)
    {
        PlannerTask[] completed = existing.Tasks.Where(static t => t.State == PlannerTaskState.Completed).ToArray();
        var completedByDescription = completed.GroupBy(static t => t.Description, StringComparer.Ordinal)
            .ToDictionary(static group => group.Key, static group => group.First(), StringComparer.Ordinal);
        var graph = new TaskGraph(_settings.MaxTasks);
        foreach (PlannerTask item in completed)
        {
            Guid[] dependencies = item.Dependencies.Where(id => completed.Any(candidate => candidate.Id == id)).ToArray();
            PlannerTask copy = graph.AddTask(item.Description, item.BudgetClass, item.Estimate, dependencies, item.RequiresExternalEffect, item.Id, item.RequiredTools);
            graph.AddUsage(copy.Id, item.ActualInputTokens, item.ActualOutputTokens);
            graph.SetState(copy.Id, PlannerTaskState.Completed);
        }
        var idMap = new List<Guid>();
        var reusedCompleted = new HashSet<Guid>();

        for (int i = 0; i < plan.Tasks.Count; i++)
        {
            PlannedTask pt = plan.Tasks[i];

            // Skip tasks that are already completed in the existing graph.
            if (completedByDescription.TryGetValue(pt.Description, out PlannerTask? completedTask) && reusedCompleted.Add(completedTask.Id))
            {
                idMap.Add(completedTask.Id);
                continue;
            }

            TokenBudgetClass bc = ParseBudgetClass(pt.BudgetClass);
            TokenEstimate est = CreateEstimate(pt, bc);
            var dependencies = new HashSet<Guid>();
            foreach (int depIdx in pt.Dependencies)
            {
                if (depIdx < 0 || depIdx >= idMap.Count)
                    throw new PlannerValidationException($"Task {i} contains invalid dependency index {depIdx}.");
                dependencies.Add(idMap[depIdx]);
            }

            string description = string.IsNullOrWhiteSpace(pt.Description) ? originalRequest : pt.Description;
            PlannerTask task = graph.AddTask(description, bc, est, dependencies, pt.RequiresExternalEffect,
                requiredTools: ValidateRequiredTools(pt));
            idMap.Add(task.Id);
        }

        if (graph.Tasks.All(static task => task.State == PlannerTaskState.Completed))
            throw new PlannerValidationException("Replan did not produce an incomplete executable Task.");
        return graph;
    }

    private static TokenBudgetClass ParseBudgetClass(string value)
    {
        try { return TokenBudgetClassExtensions.Parse(value); }
        catch { return TokenBudgetClass.K4; }
    }

    private static TokenEstimate CreateEstimate(PlannedTask pt, TokenBudgetClass budgetClass)
    {
        long o = Math.Max(1, pt.OptimisticTokens);
        long m = Math.Max(o, pt.MostLikelyTokens);
        long p = Math.Max(m, pt.PessimisticTokens);
        long max = budgetClass.MaxTokens();
        if (p > max) p = max;
        if (m > p) m = p;
        if (o > m) o = m;
        return new TokenEstimate(o, m, p);
    }

    private static IReadOnlySet<string> ValidateRequiredTools(PlannedTask task)
    {
        HashSet<string> tools = task.RequiredTools.ToHashSet(StringComparer.Ordinal);
        string[] allowed = ["file_search", "file_read", "file_edit", "process_run", "task_update", "dir"];
        if (tools.Any(tool => !allowed.Contains(tool, StringComparer.Ordinal)))
            throw new PlannerValidationException("required_tools contains an unknown Internal Tool.");
        if (task.RequiresExternalEffect && tools.Count == 0)
            throw new PlannerValidationException("An external-effect Task must provide required_tools.");
        if (!task.RequiresExternalEffect && tools.Any(tool => tool is "file_edit" or "process_run"))
            throw new PlannerValidationException("A non-effect Task cannot require an effect-only Tool.");
        return tools;
    }

    private static PlanResult? ParsePlan(string content)
    {
        try
        {
            return JsonSerializer.Deserialize<PlanResult>(content, new JsonSerializerOptions
            {
                PropertyNameCaseInsensitive = true
            });
        }
        catch
        {
            return null;
        }
    }

    private async Task<IReadOnlyList<AgentSkill>> SelectSkillsAsync(string request, double? temperature,
        Action<string, IReadOnlyList<OllamaMessage>, OllamaResult?, Exception?>? observeAttempt, 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);

        try
        {
            OllamaResult result = await _ollama.ChatStructuredAsync(selectionMessages, temperature, formatDoc.RootElement.Clone(), cancellationToken).ConfigureAwait(false);
            observeAttempt?.Invoke("skill_selection", selectionMessages, result, null);
            return ParseSelectedSkills(result.Content);
        }
        catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
        {
            observeAttempt?.Invoke("skill_selection", selectionMessages, null, ex);
            return [];
        }
    }

    private async Task<TaskGraph> RequestValidGraphAsync(IReadOnlyList<OllamaMessage> messages, double? temperature, string phase,
        Action<string, IReadOnlyList<OllamaMessage>, OllamaResult?, Exception?>? observeAttempt, Func<PlanResult, TaskGraph> buildGraph, CancellationToken cancellationToken)
    {
        Exception? firstFailure = null;
        IReadOnlyList<OllamaMessage> attemptMessages = messages;
        for (int attempt = 0; attempt < 2; attempt++)
        {
            OllamaResult? result = null;
            try
            {
                bool? thinkOverride = attempt == 0 ? null : false;
                result = await _ollama.ChatStructuredAsync(attemptMessages, temperature, PlanSchema.Element, thinkOverride, cancellationToken).ConfigureAwait(false);
                PlanResult? plan = ParsePlan(result.Content);
                if (plan is null || plan.Tasks.Count == 0)
                    throw new PlannerValidationException("Ollama returned an empty or invalid structured plan.");
                TaskGraph graph = buildGraph(plan);
                observeAttempt?.Invoke(phase, attemptMessages, result, null);
                return graph;
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) { throw; }
            catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
            {
                observeAttempt?.Invoke(phase, attemptMessages, result, ex);
                firstFailure ??= ex;
                if (attempt == 0)
                {
                    attemptMessages = messages.Concat([
                        new OllamaMessage("system",
                            "The previous response was not a valid TaskGraph JSON object. Do not call or execute any Tool while planning. " +
                            "Return only one JSON object that matches the supplied schema, beginning with {\"tasks\":[ and containing no Markdown or commentary.")
                    ]).ToArray();
                }
            }
        }
        throw new PlannerValidationException($"Structured {phase} failed after think=false retry: {firstFailure?.Message}");
    }

    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 static string ComposeSkillCatalog(IReadOnlyList<AgentSkill> skills)
    {
        if (skills.Count == 0) return "No skills available.";
        var lines = new List<string> { "Available skills:" };
        foreach (AgentSkill skill in skills)
            lines.Add($"- {skill.Name}: {skill.Description}");
        return string.Join("\n", lines);
    }

    private static string BuildPlanningInstruction() =>
        """
        You are a Task Planner. Decompose the user request into a TaskGraph of 1-32 tasks.
        Planning is read-only: do not call or execute Tools. Return only one JSON object matching the supplied schema, with no Markdown or commentary.
        Each task must have:
        - description: a clear, self-contained sub-task description
        - dependencies: array of task indices (0-based) that must complete before this task
        - budget_class: one of 1K, 2K, 4K, 8K, 16K, 32K, 64K, 128K (token budget for this task)
        - optimistic_tokens, most_likely_tokens, pessimistic_tokens: PERT estimates (0 < O <= M <= P <= budget)
        - requires_external_effect: true if this task requires a verifiable external effect (file edit, process run, etc.)
        - required_tools: the exact Internal Tool names whose successful result is required to complete this task; use names from the supplied Tool catalog, or [] for a pure response task
        
        For pure analysis or informational tasks, requires_external_effect should be false.
        For tasks that modify files, run processes, or produce verifiable side effects, requires_external_effect should be true.
        Read-only operations (file_read, file_search, and dir list) do not create an external effect and must use requires_external_effect=false.
        A read/search/list task must name file_read, file_search, or dir in required_tools even though requires_external_effect is false. A pure response task uses an empty array. When requires_external_effect is true, required_tools must include the effect Tool that proves this task's own expected effect.
        Preserve every user-supplied absolute path, executable, cwd, argument array, operation name, and literal value exactly in the task that uses it; never abbreviate these values into a shell-like command.
        Keep each task limited to its own described operation. Do not put later dependent operations into an earlier task.
        A single task is valid for simple requests.
        """;

    private static string BuildReplanInstruction(TaskGraph existing, string reason)
    {
        var sb = new System.Text.StringBuilder();
        sb.AppendLine("You are a Task Planner performing replan. The previous plan encountered an issue.");
        sb.AppendLine($"Replan reason: {reason}");
        sb.AppendLine("Current task states:");
        foreach (PlannerTask t in existing.Tasks)
        {
            sb.AppendLine($"  - [{t.State}] {t.Description} (requires_external_effect={t.RequiresExternalEffect})");
        }
        sb.AppendLine();
        sb.AppendLine("Completed tasks must be preserved. Modify only incomplete/failed tasks.");
        sb.AppendLine("Return a complete revised plan including completed tasks.");
        return sb.ToString();
    }
}