1 files changed, 76 insertions, 101 deletions

diff --git a/src/Agent.Core/Runtime/PlannerService.cs b/src/Agent.Core/Runtime/PlannerService.cs
index bf896f0..99569e4 100644
--- a/src/Agent.Core/Runtime/PlannerService.cs
+++ b/src/Agent.Core/Runtime/PlannerService.cs
@@ -23,10 +23,11 @@ public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings se
     /// 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)
+    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, cancellationToken).ConfigureAwait(false);
+        IReadOnlyList<AgentSkill> selectedSkills = await SelectSkillsAsync(request, temperature, observeAttempt, cancellationToken).ConfigureAwait(false);
 
         var messages = new List<OllamaMessage>
         {
@@ -39,39 +40,8 @@ public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings se
             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.
-            TaskGraph fallback = CreateFallbackGraph(request);
-            session.TaskGraphs[taskId] = fallback;
-            return fallback;
-        }
-
-        PlanResult? plan = ParsePlan(result.Content);
-        if (plan is null || plan.Tasks.Count == 0)
-        {
-            TaskGraph fallback = CreateFallbackGraph(request);
-            session.TaskGraphs[taskId] = fallback;
-            return fallback;
-        }
-
-        TaskGraph graph;
-        try
-        {
-            graph = BuildTaskGraph(plan, request);
-        }
-        catch
-        {
-            TaskGraph fallback = CreateFallbackGraph(request);
-            session.TaskGraphs[taskId] = fallback;
-            return fallback;
-        }
+        TaskGraph graph = await RequestValidGraphAsync(messages, temperature, "planning", observeAttempt,
+            plan => BuildTaskGraph(plan, request), cancellationToken).ConfigureAwait(false);
         session.TaskGraphs[taskId] = graph;
         return graph;
     }
@@ -80,10 +50,11 @@ public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings se
     /// 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)
+    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, cancellationToken).ConfigureAwait(false);
+            return await PlanAsync(taskId, sessionId, session, request, temperature, null, observeAttempt, cancellationToken).ConfigureAwait(false);
 
         var messages = new List<OllamaMessage>
         {
@@ -95,26 +66,15 @@ public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings se
             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);
+        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(32);
+        var graph = new TaskGraph(_settings.MaxTasks);
         var idMap = new List<Guid>();
 
         for (int i = 0; i < plan.Tasks.Count; i++)
@@ -127,81 +87,65 @@ public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings se
             var dependencies = new HashSet<Guid>();
             foreach (int depIdx in pt.Dependencies)
             {
-                if (depIdx >= 0 && depIdx < idMap.Count)
-                    dependencies.Add(idMap[depIdx]);
+                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;
-            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);
-            }
+            PlannerTask task = graph.AddTask(description, budgetClass, estimate, dependencies, pt.RequiresExternalEffect);
+            idMap.Add(task.Id);
         }
 
-        // If all tasks were filtered out, fall back.
-        if (graph.Count == 0) return CreateFallbackGraph(originalRequest);
+        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)
     {
-        // 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);
+        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);
+            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 (completedDescriptions.Contains(pt.Description))
+            if (completedByDescription.TryGetValue(pt.Description, out PlannerTask? completedTask) && reusedCompleted.Add(completedTask.Id))
             {
-                // 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;
-                }
+                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) dependencies.Add(idMap[depIdx]);
+            {
+                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);
             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);
+        if (graph.Tasks.All(static task => task.State == PlannerTaskState.Completed))
+            throw new PlannerValidationException("Replan did not produce an incomplete executable Task.");
         return graph;
     }
 
@@ -238,7 +182,8 @@ public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings se
         }
     }
 
-    private async Task<IReadOnlyList<AgentSkill>> SelectSkillsAsync(string request, CancellationToken cancellationToken)
+    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 [];
 
@@ -254,15 +199,45 @@ public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings se
 
         try
         {
-            OllamaResult result = await _ollama.ChatStructuredAsync(selectionMessages, 0.0, formatDoc.RootElement.Clone(), cancellationToken).ConfigureAwait(false);
+            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
+        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;
+        for (int attempt = 0; attempt < 2; attempt++)
+        {
+            OllamaResult? result = null;
+            try
+            {
+                bool? thinkOverride = attempt == 0 ? null : false;
+                result = await _ollama.ChatStructuredAsync(messages, 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, messages, 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, messages, result, ex);
+                firstFailure ??= ex;
+            }
+        }
+        throw new PlannerValidationException($"Structured {phase} failed after think=false retry: {firstFailure?.Message}");
+    }
+
     private IReadOnlyList<AgentSkill> ParseSelectedSkills(string content)
     {
         try
@@ -331,4 +306,4 @@ public sealed class PlannerService(IOllamaClient ollama, AgentRuntimeSettings se
         sb.AppendLine("Return a complete revised plan including completed tasks.");
         return sb.ToString();
     }
-}
\ No newline at end of file
+}