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, CancellationToken cancellationToken)
    {
        // Build a Planning prompt with skill catalog and sandbox capability summary.
        IReadOnlyList<AgentSkill> selectedSkills = await SelectSkillsAsync(request, 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)
        };

        // Use structured output to get a plan.
        OllamaResult result;
        try
        {
            result = await _ollama.ChatStructuredAsync(messages, 0.1, PlanSchema.Element, cancellationToken).ConfigureAwait(false);
        }
        catch
        {
            // Fallback: single-task plan with conservative estimates.
            return CreateFallbackGraph(request);
        }

        PlanResult? plan = ParsePlan(result.Content);
        if (plan is null || plan.Tasks.Count == 0)
            return CreateFallbackGraph(request);

        TaskGraph graph;
        try
        {
            graph = BuildTaskGraph(plan, request);
        }
        catch
        {
            return CreateFallbackGraph(request);
        }
        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, CancellationToken cancellationToken)
    {
        if (!session.TaskGraphs.TryGetValue(taskId, out TaskGraph? existing))
            return await PlanAsync(taskId, sessionId, session, request, temperature, null, 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)
        };

        OllamaResult result;
        try
        {
            result = await _ollama.ChatStructuredAsync(messages, 0.1, PlanSchema.Element, cancellationToken).ConfigureAwait(false);
        }
        catch
        {
            return existing;
        }

        PlanResult? plan = ParsePlan(result.Content);
        if (plan is null || plan.Tasks.Count == 0) return existing;

        // Build a new graph, preserving completed tasks from the existing graph.
        return RebuildTaskGraph(existing, plan, request);
    }

    private TaskGraph BuildTaskGraph(PlanResult plan, string originalRequest)
    {
        var graph = new TaskGraph(32);
        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)
                    dependencies.Add(idMap[depIdx]);
            }

            string description = string.IsNullOrWhiteSpace(pt.Description) ? originalRequest : pt.Description;
            try
            {
                PlannerTask task = graph.AddTask(description, budgetClass, estimate, dependencies, pt.RequiresExternalEffect);
                idMap.Add(task.Id);
            }
            catch (PlannerValidationException)
            {
                // If a task fails validation (e.g., bad estimate, cycle), fall back to single-task graph.
                return CreateFallbackGraph(originalRequest);
            }
        }

        // If all tasks were filtered out, fall back.
        if (graph.Count == 0) return CreateFallbackGraph(originalRequest);

        return graph;
    }

    private TaskGraph RebuildTaskGraph(TaskGraph existing, PlanResult plan, string originalRequest)
    {
        // Collect completed task descriptions to skip them in the new graph.
        var completedDescriptions = new HashSet<string>(
            existing.Tasks.Where(static t => t.State == PlannerTaskState.Completed).Select(static t => t.Description),
            StringComparer.Ordinal);

        var graph = new TaskGraph(32);
        var idMap = new List<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 (completedDescriptions.Contains(pt.Description))
            {
                // Find the completed task and add it as completed to the new graph.
                PlannerTask? completedTask = existing.Tasks.FirstOrDefault(t => t.Description == pt.Description);
                if (completedTask is not null)
                {
                    TokenBudgetClass budgetClass = ParseBudgetClass(pt.BudgetClass);
                    TokenEstimate estimate = CreateEstimate(pt, budgetClass);
                    var deps = new HashSet<Guid>();
                    foreach (int depIdx in pt.Dependencies)
                        if (depIdx >= 0 && depIdx < idMap.Count) deps.Add(idMap[depIdx]);
                    PlannerTask added = graph.AddTask(pt.Description, budgetClass, estimate, deps, pt.RequiresExternalEffect);
                    graph.SetState(added.Id, PlannerTaskState.Completed);
                    idMap.Add(added.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) dependencies.Add(idMap[depIdx]);

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

        return graph;
    }

    private TaskGraph CreateFallbackGraph(string request)
    {
        var graph = new TaskGraph(32);
        TokenEstimate estimate = new(512, 2048, 8192);
        graph.AddTask(request, TokenBudgetClass.K8, estimate, requiresExternalEffect: false);
        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 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, 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, 0.0, formatDoc.RootElement.Clone(), cancellationToken).ConfigureAwait(false);
            return ParseSelectedSkills(result.Content);
        }
        catch
        {
            return [];
        }
    }

    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.
        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.)
        
        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.
        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();
    }
}