1 files changed, 52 insertions, 10 deletions
diff --git a/src/Agent.Core/Tools/TaskGraphToolHandler.cs b/src/Agent.Core/Tools/TaskGraphToolHandler.cs
index 03daa56..edc1cc5 100644
--- a/src/Agent.Core/Tools/TaskGraphToolHandler.cs
+++ b/src/Agent.Core/Tools/TaskGraphToolHandler.cs
@@ -1,30 +1,46 @@
 using System.Collections.Concurrent;
 using System.Text.Json;
+using Agent.Core.Sessions;
 using Agent.Planner;
 
 namespace Agent.Core.Tools;
 
 public interface ITaskUpdateHandler
 {
-    void BeginTask(Guid taskId);
+    void BeginTask(Guid taskId, Guid sessionId);
     ToolExecutionResult Execute(Guid taskId, JsonElement arguments);
     void EndTask(Guid taskId);
 }
 
-public sealed class TaskGraphToolHandler(int maxTasks = 32) : ITaskUpdateHandler
+public sealed class TaskGraphToolHandler(int maxTasks = 32, SessionStore? sessions = null) : ITaskUpdateHandler
 {
-    private readonly ConcurrentDictionary<Guid, TaskGraph> _graphs = new();
+    private readonly ConcurrentDictionary<Guid, TaskGraph> _fallbackGraphs = new();
 
-    public void BeginTask(Guid taskId)
+    public void BeginTask(Guid taskId, Guid sessionId)
     {
-        if (!_graphs.TryAdd(taskId, new TaskGraph(maxTasks)))
-            throw new InvalidOperationException($"TaskGraph for '{taskId}' already exists.");
+        TaskGraph graph = new(maxTasks);
+        if (sessions is not null && sessions.TryGet(sessionId, out AgentSession? session) && session is not null)
+        {
+            if (!session.TaskGraphs.TryAdd(taskId, graph))
+                throw new InvalidOperationException($"TaskGraph for '{taskId}' already exists in session '{sessionId}'.");
+        }
+        else
+        {
+            if (!_fallbackGraphs.TryAdd(taskId, graph))
+                throw new InvalidOperationException($"TaskGraph for '{taskId}' already exists.");
+        }
+    }
+
+    public TaskGraph? GetGraph(Guid taskId, Guid? sessionId = null)
+    {
+        if (sessionId is not null && sessions is not null && sessions.TryGet(sessionId.Value, out AgentSession? session) && session is not null)
+            return session.TaskGraphs.TryGetValue(taskId, out TaskGraph? g) ? g : null;
+        return _fallbackGraphs.TryGetValue(taskId, out TaskGraph? fallback) ? fallback : null;
     }
 
     public ToolExecutionResult Execute(Guid taskId, JsonElement arguments)
     {
-        if (!_graphs.TryGetValue(taskId, out TaskGraph? graph))
-            throw new InvalidOperationException("TaskGraph scope is not active.");
+        TaskGraph? graph = ResolveGraph(taskId) ?? throw new InvalidOperationException("TaskGraph scope is not active.");
 
         string operation = arguments.GetProperty("operation").GetString()!;
         switch (operation)
@@ -35,6 +51,7 @@ public sealed class TaskGraphToolHandler(int maxTasks = 32) : ITaskUpdateHandler
                     Budget(arguments),
                     Estimate(arguments),
                     Ids(arguments, "depends_on"),
+                    requiresExternalEffect: OptionalBool(arguments, "requires_external_effect"),
                     id: OptionalId(arguments, "task_id"));
                 break;
             case "add_dependency":
@@ -61,7 +78,8 @@ public sealed class TaskGraphToolHandler(int maxTasks = 32) : ITaskUpdateHandler
                 description = task.Description,
                 depends_on = task.Dependencies,
                 state = task.State.ToString().ToLowerInvariant(),
-                expected_tokens = task.ExpectedTokens
+                expected_tokens = task.ExpectedTokens,
+                requires_external_effect = task.RequiresExternalEffect
             }),
             critical_path = graph.CriticalPath.TaskIds,
             critical_path_tokens = graph.CriticalPath.CumulativeTokenCost,
@@ -70,7 +88,31 @@ public sealed class TaskGraphToolHandler(int maxTasks = 32) : ITaskUpdateHandler
         return new ToolExecutionResult(true, output, false);
     }
 
-    public void EndTask(Guid taskId) => _graphs.TryRemove(taskId, out _);
+    public void EndTask(Guid taskId)
+    {
+        // Remove from fallback dict; session-owned graphs persist for context summary.
+        _fallbackGraphs.TryRemove(taskId, out _);
+    }
+
+    private TaskGraph? ResolveGraph(Guid taskId)
+    {
+        // Search all sessions first, then fallback.
+        if (sessions is not null)
+        {
+            foreach (AgentSession? session in sessions.Sessions)
+            {
+                if (session is not null && session.TaskGraphs.TryGetValue(taskId, out TaskGraph? g))
+                    return g;
+            }
+        }
+        return _fallbackGraphs.TryGetValue(taskId, out TaskGraph? fallback) ? fallback : null;
+    }
+
+    private static bool OptionalBool(JsonElement arguments, string property)
+    {
+        if (!arguments.TryGetProperty(property, out JsonElement value)) return false;
+        return value.ValueKind == JsonValueKind.True;
+    }
 
     private static TokenBudgetClass Budget(JsonElement arguments) =>
         TokenBudgetClassExtensions.Parse(RequiredString(arguments, "budget_class"));