1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
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}'.")
};
}
|
