ClaudeDo/src/ClaudeDo.Worker/External/AttachmentMcpTools.cs

using System.ComponentModel;
using System.Text;
using ClaudeDo.Data;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Hub;
using ModelContextProtocol.Server;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.External;

public sealed record AttachmentDto(string FileName, long ByteSize, DateTime CreatedAt);
public sealed record RemoveAttachmentResult(bool Removed, string TaskId, string FileName);

[McpServerToolType]
public sealed class AttachmentMcpTools
{
    private readonly TaskRepository _tasks;
    private readonly TaskAttachmentRepository _attachments;
    private readonly AttachmentStore _store;
    private readonly HubBroadcaster _broadcaster;

    public AttachmentMcpTools(
        TaskRepository tasks,
        TaskAttachmentRepository attachments,
        AttachmentStore store,
        HubBroadcaster broadcaster)
    {
        _tasks = tasks;
        _attachments = attachments;
        _store = store;
        _broadcaster = broadcaster;
    }

    [McpServerTool, Description(
        "Attach a read-only reference file to a task so the agent receives it at run time — use to prepare " +
        "context (plans, scripts, specs) for a task that will run later. Exactly one of textContent/" +
        "base64Content is required. Re-attaching the same fileName overwrites the previous version." +
        McpToolDocs.NotWhileRunning)]
    public async Task<AttachmentDto> AddTaskAttachment(
        string taskId,
        [Description("Name to store the attachment under; reusing an existing name overwrites it.")] string fileName,
        [Description("Plain-text content (plans, markdown, scripts). Provide this or base64Content, not both.")] string? textContent = null,
        [Description("Base64-encoded content for binary files (images, archives). Provide this or textContent, not both.")] string? base64Content = null,
        CancellationToken ct = default)
    {
        var task = await _tasks.GetByIdAsync(taskId, ct)
            ?? throw new InvalidOperationException($"Task {taskId} not found.");
        if (task.Status == TaskStatus.Running)
            throw new InvalidOperationException("Cannot add an attachment to a running task. Cancel it first.");

        if (textContent is null == base64Content is null)
            throw new InvalidOperationException(
                "Exactly one of textContent or base64Content must be provided, not both and not neither.");

        byte[] bytes;
        if (textContent is not null)
        {
            bytes = Encoding.UTF8.GetBytes(textContent);
        }
        else
        {
            try { bytes = Convert.FromBase64String(base64Content!); }
            catch (FormatException ex)
            {
                throw new InvalidOperationException("base64Content is not valid Base64.", ex);
            }
        }

        using var ms = new MemoryStream(bytes);
        var byteSize = await _store.SaveAsync(taskId, fileName, ms, ct);

        var existing = await _attachments.GetAsync(taskId, fileName, ct);
        if (existing is not null)
        {
            existing.ByteSize = byteSize;
            await _attachments.UpdateAsync(existing, ct);
        }
        else
        {
            await _attachments.AddAsync(new TaskAttachmentEntity
            {
                Id = Guid.NewGuid().ToString(),
                TaskId = taskId,
                FileName = fileName,
                ByteSize = byteSize,
                CreatedAt = DateTime.UtcNow,
            }, ct);
        }

        await _broadcaster.TaskUpdated(taskId);
        return new AttachmentDto(fileName, byteSize, existing?.CreatedAt ?? DateTime.UtcNow);
    }

    [McpServerTool, Description(
        "List all attachments on a task — use to check what reference files are already attached before adding more.")]
    public async Task<IReadOnlyList<AttachmentDto>> ListTaskAttachments(
        string taskId, CancellationToken ct = default)
    {
        var rows = await _attachments.ListByTaskIdAsync(taskId, ct);
        return rows.Select(r => new AttachmentDto(r.FileName, r.ByteSize, r.CreatedAt)).ToList();
    }

    [McpServerTool, Description(
        "Remove a single attachment from a task, deleting both the file on disk and its database record." +
        McpToolDocs.NotWhileRunning)]
    public async Task<RemoveAttachmentResult> RemoveTaskAttachment(
        string taskId, string fileName, CancellationToken ct = default)
    {
        var task = await _tasks.GetByIdAsync(taskId, ct)
            ?? throw new InvalidOperationException($"Task {taskId} not found.");
        if (task.Status == TaskStatus.Running)
            throw new InvalidOperationException("Cannot remove an attachment from a running task. Cancel it first.");

        _store.DeleteFile(taskId, fileName);
        await _attachments.DeleteAsync(taskId, fileName, ct);
        await _broadcaster.TaskUpdated(taskId);
        return new RemoveAttachmentResult(true, taskId, fileName);
    }
}