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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
|
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, Guid sessionId);
ToolExecutionResult Execute(Guid taskId, JsonElement arguments);
void EndTask(Guid taskId);
}
public sealed class TaskGraphToolHandler(int maxTasks = 32, SessionStore? sessions = null) : ITaskUpdateHandler
{
private readonly ConcurrentDictionary<Guid, TaskGraph> _fallbackGraphs = new();
public void BeginTask(Guid taskId, Guid sessionId)
{
TaskGraph graph = new(maxTasks);
if (sessions is not null && sessions.TryGet(sessionId, out AgentSession? session) && session is not null)
{
// PlanAsync may have already created a TaskGraph for this taskId.
// If so, keep the planned graph; otherwise add a new empty one.
session.TaskGraphs.TryAdd(taskId, graph);
}
else
{
_fallbackGraphs.TryAdd(taskId, graph);
}
}
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)
{
TaskGraph? graph = ResolveGraph(taskId) ?? 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"),
requiresExternalEffect: OptionalBool(arguments, "requires_external_effect"),
id: OptionalId(arguments, "task_id"),
requiredEffectTools: Strings(arguments, "required_effect_tools"));
break;
case "add_dependency":
Guid taskIdValue = RequiredId(arguments, "task_id");
DemandStructuralMutation(graph, taskIdValue);
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_estimate":
Guid estimateTaskId = RequiredId(arguments, "task_id");
DemandStructuralMutation(graph, estimateTaskId);
graph.SetEstimate(estimateTaskId, 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,
requires_external_effect = task.RequiresExternalEffect,
required_effect_tools = task.RequiredEffectTools
}),
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)
{
// 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 string[] Strings(JsonElement arguments, string property)
{
if (!arguments.TryGetProperty(property, out JsonElement value) || value.ValueKind is JsonValueKind.Null or JsonValueKind.Undefined) return [];
if (value.ValueKind != JsonValueKind.Array) throw new ToolValidationException($"'{property}' must be an array.");
return value.EnumerateArray().Select(item => item.GetString() ?? throw new ToolValidationException($"'{property}' entries must be strings.")).ToArray();
}
private static void DemandStructuralMutation(TaskGraph graph, Guid taskId)
{
PlannerTaskState state = graph.Get(taskId).State;
if (state is not (PlannerTaskState.Queued or PlannerTaskState.Planning))
throw new PlannerValidationException($"TaskGraph structure or estimate cannot be changed while Task '{taskId}' is {state}.");
}
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];
}
}
|
