forge-tools-proxmox/internal/tools/task.go

package tools

import (
	"context"
	"fmt"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// task.go — задачи Proxmox (UPID): list (read), status (read, с опц. wait),
// log (read), stop (confirm + write). Мутации возвращают UPID, поэтому
// task_status(wait=true) — ключ к подтверждению async-операций, при этом
// poll строго ограничен timeout (bounded §9.6).

func registerTaskTools(s *mcp.Server) {
	s.AddTool(&mcp.Tool{
		Name:        "tasks_list",
		Description: "List recent cluster/node tasks. Read-only.",
		InputSchema: schema(map[string]any{
			"node":  strProps("Node name (empty = whole cluster)", false),
			"limit": intProps("Number of tasks (default: 20)", false),
			"host":  strProps("Cluster alias (default: primary)", false),
		}, nil),
	}, tasksListHandler)

	s.AddTool(&mcp.Tool{
		Name:        "task_status",
		Description: "Get status of a task by UPID. Set wait=true to poll until it finishes (bounded by task_poll_max_sec). Read-only.",
		InputSchema: schema(map[string]any{
			"node": strProps("Node name", true),
			"upid": strProps("Task UPID", true),
			"wait": boolProps("Wait for completion (default: false)", false),
			"host": strProps("Cluster alias (default: primary)", false),
		}, []string{"node", "upid"}),
	}, taskStatusHandler)

	s.AddTool(&mcp.Tool{
		Name:        "task_log",
		Description: "Get the log of a task by UPID. Read-only.",
		InputSchema: schema(map[string]any{
			"node":  strProps("Node name", true),
			"upid":  strProps("Task UPID", true),
			"limit": intProps("Number of lines (default: 50)", false),
			"host":  strProps("Cluster alias (default: primary)", false),
		}, []string{"node", "upid"}),
	}, taskLogHandler)

	registerPatternTool(s, &mcp.Tool{
		Name:        "task_stop",
		Description: "Stop/cancel a running task by UPID. Requires confirm + write permission.",
		InputSchema: schema(map[string]any{
			"node":    strProps("Node name", true),
			"upid":    strProps("Task UPID", true),
			"confirm": strProps("Set to \"true\" to confirm", true, "true"),
			"host":    strProps("Cluster alias (default: primary)", false),
		}, []string{"node", "upid", "confirm"}),
	}, nodePatterns, taskStopHandler)
}

func tasksListHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	t, err := tenantFor(ctx, req)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	args := requestArgs(req)
	host, err := resolveHost(t, args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	limit := getInt(args, "limit", 20)
	cctx, cancel := timeout(ctx, t)
	defer cancel()
	data, err := t.TasksList(cctx, host, getString(args, "node", ""), limit)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(pretty(data)), nil
}

func taskStatusHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	t, err := tenantFor(ctx, req)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	args := requestArgs(req)
	node, err := requireNode(args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	upid, err := requireUPID(args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	host, err := resolveHost(t, args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	if getBool(args, "wait", false) {
		out, err := t.TaskWait(ctx, host, node, upid)
		if err != nil {
			return errorResult(err.Error()), nil
		}
		return textResult(out), nil
	}
	cctx, cancel := timeout(ctx, t)
	defer cancel()
	data, err := t.TaskStatus(cctx, host, node, upid)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(pretty(data)), nil
}

func taskLogHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	t, err := tenantFor(ctx, req)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	args := requestArgs(req)
	node, err := requireNode(args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	upid, err := requireUPID(args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	host, err := resolveHost(t, args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	limit := getInt(args, "limit", 50)
	cctx, cancel := timeout(ctx, t)
	defer cancel()
	data, err := t.TaskLog(cctx, host, node, upid, limit)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(pretty(data)), nil
}

func taskStopHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	t, err := tenantFor(ctx, req)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	args := requestArgs(req)
	node, err := requireNode(args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	upid, err := requireUPID(args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	host, err := resolveHost(t, args)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	if err := confirm(args, "task_stop"); err != nil {
		return errorResult(err.Error()), nil
	}
	// task_stop — отмена выполняющейся задачи (не ресурс кластера).
	if err := gateNodeWrite(t, host, node); err != nil {
		return errorResult(err.Error()), nil
	}
	upidMsg2, err := t.TaskStop(ctx, host, node, upid)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(upidMsg("task_stop", fmt.Sprintf("%s %s", node, upid), upidMsg2)), nil
}