diff options
Diffstat (limited to 'src/Agent.Core/Tools')
| -rw-r--r-- | src/Agent.Core/Tools/InternalToolCatalog.cs | 2 | ||||
| -rw-r--r-- | src/Agent.Core/Tools/InternalToolDispatcher.cs | 2 | ||||
| -rw-r--r-- | src/Agent.Core/Tools/TaskGraphToolHandler.cs | 62 | ||||
| -rw-r--r-- | src/Agent.Core/Tools/ToolPolicy.cs | 2 |
4 files changed, 55 insertions, 13 deletions
diff --git a/src/Agent.Core/Tools/InternalToolCatalog.cs b/src/Agent.Core/Tools/InternalToolCatalog.cs index 14cba45..a92a829 100644 --- a/src/Agent.Core/Tools/InternalToolCatalog.cs +++ b/src/Agent.Core/Tools/InternalToolCatalog.cs @@ -19,7 +19,7 @@ public static class InternalToolCatalog Define("file_read", "Read an approved file.", """{"type":"object","additionalProperties":false,"required":["path","mode"],"properties":{"path":{"type":"string"},"mode":{"type":"string","enum":["whole","range","metadata"]},"start":{"type":"integer","minimum":0},"length":{"type":"integer","minimum":0}}}"""), Define("file_edit", "Edit an approved file transactionally.", """{"type":"object","additionalProperties":false,"required":["path","operation"],"properties":{"path":{"type":"string"},"operation":{"type":"string","enum":["overwrite","patch","atomic_replace","delete","rename"]},"content":{"type":["string","null"]},"patch":{"type":["string","null"]},"expected_sha256":{"type":["string","null"]},"new_path":{"type":["string","null"]}}}"""), Define("process_run", "Run one approved executable with separated arguments.", """{"type":"object","additionalProperties":false,"required":["executable","args","cwd"],"properties":{"executable":{"type":"string"},"args":{"type":"array","items":{"type":"string"}},"cwd":{"type":"string"},"timeout_seconds":{"type":"integer","minimum":1},"stdin":{"type":["string","null"]}}}"""), - Define("task_update", "Update the deterministic task graph.", """{"type":"object","additionalProperties":false,"required":["operation"],"properties":{"operation":{"type":"string","enum":["add","add_dependency","set_state","set_estimate"]},"task_id":{"type":["string","null"]},"description":{"type":["string","null"]},"depends_on":{"type":"array","items":{"type":"string"}},"budget_class":{"type":["string","null"],"enum":["1K","2K","4K","8K","16K","32K","64K","128K",null]},"optimistic_tokens":{"type":"integer","minimum":0},"most_likely_tokens":{"type":"integer","minimum":0},"pessimistic_tokens":{"type":"integer","minimum":0},"state":{"type":["string","null"],"enum":["queued","planning","running","completed","failed","cancelled",null]}}}"""), + Define("task_update", "Update the deterministic task graph.", """{"type":"object","additionalProperties":false,"required":["operation"],"properties":{"operation":{"type":"string","enum":["add","add_dependency","set_state","set_estimate"]},"task_id":{"type":["string","null"]},"description":{"type":["string","null"]},"depends_on":{"type":"array","items":{"type":"string"}},"budget_class":{"type":["string","null"],"enum":["1K","2K","4K","8K","16K","32K","64K","128K",null]},"optimistic_tokens":{"type":"integer","minimum":0},"most_likely_tokens":{"type":"integer","minimum":0},"pessimistic_tokens":{"type":"integer","minimum":0},"state":{"type":["string","null"],"enum":["queued","planning","running","completed","failed","cancelled",null]},"requires_external_effect":{"type":["boolean","null"]}}}"""), Define("dir", "Create, remove, or list directories under an approved root. Only empty directories can be removed.", """{"type":"object","additionalProperties":false,"required":["path","operation"],"properties":{"path":{"type":"string"},"operation":{"type":"string","enum":["mkdir","rmdir","list"]}}}""") ]; return definitions.ToDictionary(static definition => definition.Name, StringComparer.Ordinal); diff --git a/src/Agent.Core/Tools/InternalToolDispatcher.cs b/src/Agent.Core/Tools/InternalToolDispatcher.cs index c888705..0ae078c 100644 --- a/src/Agent.Core/Tools/InternalToolDispatcher.cs +++ b/src/Agent.Core/Tools/InternalToolDispatcher.cs @@ -4,7 +4,7 @@ namespace Agent.Core.Tools; public sealed class InternalToolDispatcher(IToolPolicyEngine policy, IInternalToolBackend backend) { - public Task BeginTaskAsync(Guid taskId, CancellationToken cancellationToken) => backend.BeginTaskAsync(taskId, cancellationToken); + public Task BeginTaskAsync(Guid taskId, Guid sessionId, CancellationToken cancellationToken) => backend.BeginTaskAsync(taskId, sessionId, cancellationToken); public Task EndTaskAsync(Guid taskId, bool success, CancellationToken cancellationToken) => backend.EndTaskAsync(taskId, success, cancellationToken); public async Task<ToolExecutionResult> DispatchAsync(Guid taskId, string toolName, JsonElement arguments, CancellationToken cancellationToken) 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")); diff --git a/src/Agent.Core/Tools/ToolPolicy.cs b/src/Agent.Core/Tools/ToolPolicy.cs index d478a2c..cefca8d 100644 --- a/src/Agent.Core/Tools/ToolPolicy.cs +++ b/src/Agent.Core/Tools/ToolPolicy.cs @@ -15,7 +15,7 @@ public interface IToolPolicyEngine public interface IInternalToolBackend { - Task BeginTaskAsync(Guid taskId, CancellationToken cancellationToken) => Task.CompletedTask; + Task BeginTaskAsync(Guid taskId, Guid sessionId, CancellationToken cancellationToken) => Task.CompletedTask; Task<ToolExecutionResult> ExecuteAsync(Guid taskId, string toolName, JsonElement arguments, CancellationToken cancellationToken); Task EndTaskAsync(Guid taskId, bool success, CancellationToken cancellationToken) => Task.CompletedTask; } |