summaryrefslogtreecommitdiff
path: root/src/Agent.Mcp/AgentTools.cs
blob: 623c6a4c81a4f24dba8e42994ad7c8861bbb680d (plain)
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
using System.ComponentModel;
using Agent.Core.Localization;
using Agent.Core.Tasks;
using ModelContextProtocol.Server;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Extensions.Tasks;
using System.Text.Json;
using System.Text.Json.Nodes;

namespace Agent.Mcp;

[McpServerToolType]
public sealed class AgentTools(TaskRegistry registry, LocaleCatalog locale, VisionLimits visionLimits, PersistentMcpTaskStore taskStore)
{
    [McpServerTool(Name = "agent_run", ReadOnly = false, Destructive = true, OpenWorld = false, UseStructuredContent = false)]
    [Description("Queue a natural-language request for asynchronous execution.")]
    public async Task<string> AgentRun(
        RequestContext<CallToolRequestParams> context,
        [Description("Non-empty natural-language work request.")] string request,
        [Description("Existing session UUID, or omit to create a session.")] Guid? session_id = null,
        [Description("Task-scoped temperature from 0.0 through 2.0.")] double? temperature = null,
        [Description("Optional PNG, JPEG, or WebP images.")] IReadOnlyList<McpImageInput>? images = null,
        CancellationToken cancellationToken = default)
    {
        RequireOnly(context, "request", "session_id", "temperature", "images");
        ValidateImages(images);
        Guid? mcpTaskId = GetRelatedTaskId(context.Params?.Meta) ?? GetRelatedTaskId(context.JsonRpcRequest) ??
            taskStore.ResolveTaskId(cancellationToken);
        TaskSnapshot task = registry.Enqueue(request, session_id, temperature, images?.Select(static image => image.DataBase64).ToArray(), mcpTaskId);
        string message = locale.Format("mcp.task_queued", new Dictionary<string, object?>
        {
            ["task"] = task.TaskId,
            ["session"] = task.SessionId
        });
        string queued = string.Join('\n',
            $"task_id: {task.TaskId}",
            $"session_id: {task.SessionId}",
            "status: queued",
            $"message: {message}");
        if (mcpTaskId is null) return queued;

        using CancellationTokenRegistration registration = cancellationToken.Register(() => registry.Cancel(task.TaskId));
        TaskSnapshot terminal = await registry.WaitForTerminalAsync(task.TaskId, CancellationToken.None).ConfigureAwait(false);
        if (terminal.State == AgentTaskState.Failed) throw new InvalidOperationException(terminal.Result ?? "Agent task failed.");
        if (terminal.State == AgentTaskState.Cancelled) throw new OperationCanceledException("Agent task was cancelled.", cancellationToken);
        return FormatStatus(terminal);
    }

    private void ValidateImages(IReadOnlyList<McpImageInput>? images)
    {
        if (images is null) return;
        if (images.Count > visionLimits.MaxCount) throw new ArgumentException($"Image count exceeds configured limit {visionLimits.MaxCount}.");
        long total = 0;
        foreach (McpImageInput image in images)
        {
            if (image.MimeType is not ("image/png" or "image/jpeg" or "image/webp")) throw new ArgumentException($"Unsupported image MIME type '{image.MimeType}'.");
            byte[] bytes;
            try { bytes = Convert.FromBase64String(image.DataBase64); }
            catch (FormatException ex) { throw new ArgumentException("Image data_base64 is invalid.", ex); }
            if (bytes.LongLength > visionLimits.MaxBytesEach) throw new ArgumentException($"Image exceeds configured per-image limit {visionLimits.MaxBytesEach} bytes.");
            total = checked(total + bytes.LongLength);
            if (total > visionLimits.MaxTotalBytes) throw new ArgumentException($"Images exceed configured total limit {visionLimits.MaxTotalBytes} bytes.");
        }
    }

    [McpServerTool(Name = "agent_status", ReadOnly = true, Destructive = false, OpenWorld = false, UseStructuredContent = false)]
    [Description("Get the current state and result of a queued agent task.")]
    public string AgentStatus(
        RequestContext<CallToolRequestParams> context,
        [Description("Task UUID returned by agent_run.")] Guid task_id,
        [Description("Optional session UUID used to verify task ownership.")] Guid? session_id = null)
    {
        RequireOnly(context, "task_id", "session_id");
        if (!registry.TryGet(task_id, session_id, out TaskSnapshot? task) || task is null)
        {
            return string.Join('\n',
                $"task_id: {task_id}",
                "status: failed",
                $"message: {locale.Format("mcp.task_not_found", new Dictionary<string, object?> { ["task"] = task_id })}");
        }

        return FormatStatus(task);
    }

    private string FormatStatus(TaskSnapshot task)
    {
        string state = task.State.ToString().ToLowerInvariant();
        bool terminal = task.State is AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled;
        return string.Join('\n',
            $"task_id: {task.TaskId}",
            $"session_id: {task.SessionId}",
            $"status: {state}",
            $"completed_tasks: {task.CompletedTasks}/{task.TotalTasks}",
            $"current_task: {task.CurrentTask ?? string.Empty}",
            $"token_usage: input={task.InputTokens} output={task.OutputTokens} total={task.InputTokens + task.OutputTokens}",
            $"critical_path_tokens: {task.CriticalPathTokens.ToString(System.Globalization.CultureInfo.InvariantCulture)}",
            $"iterations: {task.Iterations}",
            $"replans: {task.Replans}",
            $"last_error_code: {task.LastErrorCode ?? string.Empty}",
            $"last_error: {task.LastError ?? string.Empty}",
            $"message: {locale.Get("mcp.status." + state)}",
            terminal ? $"result: {task.Result}" : "result:");
    }

    private static Guid? GetRelatedTaskId(object? metadata)
    {
        if (metadata is null) return null;
        JsonElement json = JsonSerializer.SerializeToElement(metadata);
        return Find(json);

        static Guid? Find(JsonElement json)
        {
            if (json.ValueKind == JsonValueKind.Object)
            {
                foreach (JsonProperty property in json.EnumerateObject())
                {
                    if (property.NameEquals(TasksProtocol.MetaRelatedTask)) return Parse(property.Value);
                    Guid? nested = Find(property.Value);
                    if (nested is not null) return nested;
                }
            }
            else if (json.ValueKind == JsonValueKind.Array)
            {
                foreach (JsonElement item in json.EnumerateArray())
                {
                    Guid? nested = Find(item);
                    if (nested is not null) return nested;
                }
            }
            return null;
        }

        static Guid? Parse(JsonElement related)
        {
        if (related.ValueKind == JsonValueKind.String && Guid.TryParse(related.GetString(), out Guid direct)) return direct;
        if (related.ValueKind == JsonValueKind.Object && related.TryGetProperty("taskId", out JsonElement taskId) &&
            Guid.TryParse(taskId.GetString(), out Guid nested)) return nested;
        return null;
        }
    }

    private static void RequireOnly(RequestContext<CallToolRequestParams> context, params string[] allowed)
    {
        if (context.Params?.Arguments is null) return;
        HashSet<string> allowedSet = allowed.ToHashSet(StringComparer.Ordinal);
        string? unknown = context.Params.Arguments.Keys.FirstOrDefault(key => !allowedSet.Contains(key));
        if (unknown is not null) throw new ArgumentException($"Unknown tool property '{unknown}'.");
    }
}