1 files changed, 123 insertions, 0 deletions
diff --git a/src/Agent.Core/Tools/TaskGraphToolHandler.cs b/src/Agent.Core/Tools/TaskGraphToolHandler.cs
new file mode 100644
index 0000000..03daa56
--- /dev/null
+++ b/src/Agent.Core/Tools/TaskGraphToolHandler.cs
@@ -0,0 +1,123 @@
+using System.Collections.Concurrent;
+using System.Text.Json;
+using Agent.Planner;
+
+namespace Agent.Core.Tools;
+
+public interface ITaskUpdateHandler
+{
+    void BeginTask(Guid taskId);
+    ToolExecutionResult Execute(Guid taskId, JsonElement arguments);
+    void EndTask(Guid taskId);
+}
+
+public sealed class TaskGraphToolHandler(int maxTasks = 32) : ITaskUpdateHandler
+{
+    private readonly ConcurrentDictionary<Guid, TaskGraph> _graphs = new();
+
+    public void BeginTask(Guid taskId)
+    {
+        if (!_graphs.TryAdd(taskId, new TaskGraph(maxTasks)))
+            throw new InvalidOperationException($"TaskGraph for '{taskId}' already exists.");
+    }
+
+    public ToolExecutionResult Execute(Guid taskId, JsonElement arguments)
+    {
+        if (!_graphs.TryGetValue(taskId, out TaskGraph? graph))
+            throw new InvalidOperationException("TaskGraph scope is not active.");
+
+        string operation = arguments.GetProperty("operation").GetString()!;
+        switch (operation)
+        {
+            case "add":
+                graph.AddTask(
+                    RequiredString(arguments, "description"),
+                    Budget(arguments),
+                    Estimate(arguments),
+                    Ids(arguments, "depends_on"),
+                    id: OptionalId(arguments, "task_id"));
+                break;
+            case "add_dependency":
+                Guid taskIdValue = RequiredId(arguments, "task_id");
+                Guid[] dependencies = Ids(arguments, "depends_on");
+                if (dependencies.Length == 0) throw new ToolValidationException("depends_on must contain at least one task id.");
+                foreach (Guid dependency in dependencies) graph.AddDependency(taskIdValue, dependency);
+                break;
+            case "set_state":
+                graph.SetState(RequiredId(arguments, "task_id"), ParseState(RequiredString(arguments, "state")));
+                break;
+            case "set_estimate":
+                graph.SetEstimate(RequiredId(arguments, "task_id"), Budget(arguments), Estimate(arguments));
+                break;
+            default:
+                throw new ToolValidationException($"Unsupported task_update operation '{operation}'.");
+        }
+
+        string output = JsonSerializer.Serialize(new
+        {
+            tasks = graph.Tasks.Select(static task => new
+            {
+                task_id = task.Id,
+                description = task.Description,
+                depends_on = task.Dependencies,
+                state = task.State.ToString().ToLowerInvariant(),
+                expected_tokens = task.ExpectedTokens
+            }),
+            critical_path = graph.CriticalPath.TaskIds,
+            critical_path_tokens = graph.CriticalPath.CumulativeTokenCost,
+            total_token_budget = graph.TotalTokenBudget
+        });
+        return new ToolExecutionResult(true, output, false);
+    }
+
+    public void EndTask(Guid taskId) => _graphs.TryRemove(taskId, out _);
+
+    private static TokenBudgetClass Budget(JsonElement arguments) =>
+        TokenBudgetClassExtensions.Parse(RequiredString(arguments, "budget_class"));
+
+    private static TokenEstimate Estimate(JsonElement arguments) => new(
+        RequiredInt64(arguments, "optimistic_tokens"),
+        RequiredInt64(arguments, "most_likely_tokens"),
+        RequiredInt64(arguments, "pessimistic_tokens"));
+
+    private static string RequiredString(JsonElement arguments, string property) =>
+        arguments.TryGetProperty(property, out JsonElement value) && value.ValueKind == JsonValueKind.String && !string.IsNullOrWhiteSpace(value.GetString())
+            ? value.GetString()!
+            : throw new ToolValidationException($"{property} is required for this operation.");
+
+    private static long RequiredInt64(JsonElement arguments, string property) =>
+        arguments.TryGetProperty(property, out JsonElement value) && value.ValueKind == JsonValueKind.Number
+            ? value.GetInt64()
+            : throw new ToolValidationException($"{property} is required for this operation.");
+
+    private static Guid RequiredId(JsonElement arguments, string property) =>
+        OptionalId(arguments, property) ?? throw new ToolValidationException($"{property} must be a UUID for this operation.");
+
+    private static Guid? OptionalId(JsonElement arguments, string property)
+    {
+        if (!arguments.TryGetProperty(property, out JsonElement value) || value.ValueKind == JsonValueKind.Null) return null;
+        return value.ValueKind == JsonValueKind.String && Guid.TryParse(value.GetString(), out Guid id)
+            ? id
+            : throw new ToolValidationException($"{property} must be a UUID.");
+    }
+
+    private static Guid[] Ids(JsonElement arguments, string property)
+    {
+        if (!arguments.TryGetProperty(property, out JsonElement value)) return [];
+        var ids = new List<Guid>();
+        foreach (JsonElement item in value.EnumerateArray())
+            ids.Add(Guid.TryParse(item.GetString(), out Guid id) ? id : throw new ToolValidationException($"{property} must contain only UUIDs."));
+        return [.. ids];
+    }
+
+    private static PlannerTaskState ParseState(string state) => state switch
+    {
+        "queued" => PlannerTaskState.Queued,
+        "planning" => PlannerTaskState.Planning,
+        "running" => PlannerTaskState.Running,
+        "completed" => PlannerTaskState.Completed,
+        "failed" => PlannerTaskState.Failed,
+        "cancelled" => PlannerTaskState.Cancelled,
+        _ => throw new ToolValidationException($"Unknown task state '{state}'.")
+    };
+}