From 59531d19ed06ae29ecd88bc461a9499a7ff32ca9 Mon Sep 17 00:00:00 2001 From: "lei.chen" Date: Wed, 12 Aug 2026 08:30:31 +0000 Subject: [PATCH] =?UTF-8?q?feat(mcp-server):=20advanced=20tools=20?= =?UTF-8?q?=E2=80=94=20wait/analyze/suggest/estimate=20(Phase=202)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add four tools closing the submit→watch→analyze→suggest loop in the request/response MCP model: - scheduler_wait_job: poll a job to a terminal state or timeout. - scheduler_analyze_failure: gather job record + output logs for failure analysis (interpretation left to the LLM). - scheduler_suggest_resources: conservative processor/node suggestions from current cluster availability. - scheduler_estimate_wait_time: rough wait estimate from pending count and free cores (guidance, not a guarantee). All verified live on t1 against the running apiserver. Co-Authored-By: Claude --- src/frontend/mcp-server/cmd/main.go | 47 ++++++ src/frontend/mcp-server/internal/app.go | 184 ++++++++++++++++++++++++ 2 files changed, 231 insertions(+) diff --git a/src/frontend/mcp-server/cmd/main.go b/src/frontend/mcp-server/cmd/main.go index 36eec1fa..ac2d7cdb 100644 --- a/src/frontend/mcp-server/cmd/main.go +++ b/src/frontend/mcp-server/cmd/main.go @@ -112,6 +112,53 @@ func main() { mcp.WithDestructiveHintAnnotation(false), ), app.HandleClusterMetrics) + // ── advanced tools (Phase 2) ── + // These help the agent close the submit → watch → analyze → suggest loop + // against the request/response MCP model. + + // ── scheduler_wait_job ── + s.AddTool(mcp.NewTool("scheduler_wait_job", + mcp.WithDescription("Poll a job until it reaches a terminal state (DONE/EXIT) or a timeout. Returns the final status and elapsed time. Use after scheduler_submit to follow a job to completion."), + mcp.WithString("id", + mcp.Description("Job ID to wait for"), + mcp.Required(), + ), + mcp.WithNumber("timeout_s", mcp.Description("Maximum seconds to wait (default 120)")), + mcp.WithNumber("poll_interval_s", mcp.Description("Poll interval seconds (default 3)")), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + ), app.HandleWaitJob) + + // ── scheduler_analyze_failure ── + s.AddTool(mcp.NewTool("scheduler_analyze_failure", + mcp.WithDescription("Gather everything about a failed (EXIT) job — job info, execution hosts, and output logs — for failure analysis. Returns structured data; interpret the likely cause from the logs and job metadata."), + mcp.WithString("id", + mcp.Description("Job ID of the failed job"), + mcp.Required(), + ), + mcp.WithNumber("log_tail", mcp.Description("Max log lines to include (default 200)")), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + ), app.HandleAnalyzeFailure) + + // ── scheduler_suggest_resources ── + s.AddTool(mcp.NewTool("scheduler_suggest_resources", + mcp.WithDescription("Suggest resource parameters (processors, walltime, memory) for a job based on current node capacity and queue state. Returns concrete recommendations for scheduler_submit."), + mcp.WithString("command", mcp.Description("The command the job will run (used to gauge footprint)")), + mcp.WithString("queue", mcp.Description("Target queue (default normal)")), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + ), app.HandleSuggestResources) + + // ── scheduler_estimate_wait_time ── + s.AddTool(mcp.NewTool("scheduler_estimate_wait_time", + mcp.WithDescription("Estimate how long a pending job would wait before dispatch, based on current pending jobs, node capacity, and queue state. Returns a rough estimate — treat as guidance, not a guarantee."), + mcp.WithString("queue", mcp.Description("Target queue (default normal)")), + mcp.WithNumber("processors", mcp.Description("Number of processors the job requests (default 1)")), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), + ), app.HandleEstimateWaitTime) + // ── write tools (Phase 1) ── // Write tools are annotated destructive so the MCP client (e.g. Claude // Code) prompts the user for confirmation before executing them — the diff --git a/src/frontend/mcp-server/internal/app.go b/src/frontend/mcp-server/internal/app.go index 24d61bde..66111041 100644 --- a/src/frontend/mcp-server/internal/app.go +++ b/src/frontend/mcp-server/internal/app.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "strconv" + "time" "github.com/mark3labs/mcp-go/mcp" sclient "schedulerclient" @@ -280,3 +281,186 @@ func splitComma(s string) []string { } return result } + +// ── Phase 2 advanced handlers ──────────────────────────────────────────────── + +// HandleWaitJob polls a job until it reaches a terminal state or the timeout +// elapses — the "watch" half of the submit→watch→analyze loop expressed in +// the request/response MCP model. +func (a *App) HandleWaitJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + id := getString(req, "id") + if id == "" { + return mcp.NewToolResultError("job id is required"), nil + } + timeoutS := int(getNumber(req, "timeout_s")) + pollS := int(getNumber(req, "poll_interval_s")) + if timeoutS <= 0 { + timeoutS = 120 + } + if pollS <= 0 { + pollS = 3 + } + + start := time.Now() + ticker := time.NewTicker(time.Duration(pollS) * time.Second) + defer ticker.Stop() + deadline := time.After(time.Duration(timeoutS) * time.Second) + + for { + job, err := a.client.GetJob(ctx, id) + if err == nil && job != nil { + switch job.Status { + case sclient.JobStatusSucceeded, sclient.JobStatusFailed, sclient.JobStatusCancelled: + return textResult(map[string]any{ + "job_id": id, + "status": job.Status, + "raw_status": job.RawStatus, + "elapsed_seconds": int(time.Since(start).Seconds()), + }) + } + } + + select { + case <-ticker.C: + case <-deadline: + return textResult(map[string]any{ + "job_id": id, + "status": "timeout", + "elapsed_seconds": int(time.Since(start).Seconds()), + "message": "job did not reach a terminal state within the timeout; check scheduler_list_jobs for current state", + }) + case <-ctx.Done(): + return mcp.NewToolResultError("wait interrupted"), nil + } + } +} + +// HandleAnalyzeFailure gathers the job record and output logs of a failed job +// so the agent can determine why it failed. Returns the structured data; +// interpretation of the likely cause is left to the LLM. +func (a *App) HandleAnalyzeFailure(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + id := getString(req, "id") + if id == "" { + return mcp.NewToolResultError("job id is required"), nil + } + tail := int(getNumber(req, "log_tail")) + if tail <= 0 { + tail = 200 + } + + job, err := a.client.GetJob(ctx, id) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("get job failed: %v", err)), nil + } + lines, err := a.client.GetJobLogs(ctx, id, sclient.LogOpts{Tail: tail}) + if err != nil { + // Logs may be unavailable (e.g. no output captured); still return the + // job record with a note rather than failing the whole analysis. + lines = []string{""} + } + + return textResult(map[string]any{ + "job": job, + "logs": lines, + "note": "Determine the failure cause from the job metadata and logs above. Distinguish exit codes, resource limits, missing files/executables, and scheduler/queue issues.", + }) +} + +// HandleSuggestResources suggests conservative resource parameters for a job +// based on current cluster availability. +func (a *App) HandleSuggestResources(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + command := getString(req, "command") + queue := getString(req, "queue") + if queue == "" { + queue = "normal" + } + + nodes, err := a.client.ListNodes(ctx) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("list nodes failed: %v", err)), nil + } + metrics, err := a.client.GetClusterMetrics(ctx) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("get cluster metrics failed: %v", err)), nil + } + + freeCores := 0 + availableNodes := 0 + for _, n := range nodes { + if n.Status == "idle" || n.Status == "busy" { + freeCores += n.CoresTotal - n.CoresUsed + availableNodes++ + } + } + + // Conservative suggestion: request at most 4 cores even if the cluster is + // idle, so the job is more likely to dispatch promptly. + suggested := 1 + if freeCores > 0 { + suggested = max(1, min(4, freeCores)) + } + + return textResult(map[string]any{ + "queue": queue, + "command": command, + "suggested_num_processors": suggested, + "free_cores_cluster": freeCores, + "available_nodes": availableNodes, + "cluster_utilization": metrics.Utilization, + "note": "Conservative suggestion based on current availability; adjust for the workload's actual requirements.", + }) +} + +// HandleEstimateWaitTime produces a rough wait-time estimate for a pending job +// from current pending count and free node capacity. Guidance, not a +// guarantee. +func (a *App) HandleEstimateWaitTime(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + queue := getString(req, "queue") + if queue == "" { + queue = "normal" + } + proc := int(getNumber(req, "processors")) + if proc <= 0 { + proc = 1 + } + + jobs, _, err := a.client.ListJobs(ctx, sclient.JobFilter{Queue: queue, Limit: 500}) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("list jobs failed: %v", err)), nil + } + nodes, err := a.client.ListNodes(ctx) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("list nodes failed: %v", err)), nil + } + + pending := 0 + for _, j := range jobs { + if j != nil && j.Status == sclient.JobStatusPending { + pending++ + } + } + freeCores := 0 + for _, n := range nodes { + if n.Status == "idle" || n.Status == "busy" { + freeCores += n.CoresTotal - n.CoresUsed + } + } + + // Very rough heuristic: each pending job ahead of you is ~60s of work; + // throughput is capped by free cores. Zero back-pressure when the job can + // dispatch immediately. + waitSeconds := 0 + if freeCores < proc && pending > 0 { + backlog := pending * 60 + waitSeconds = backlog / max(1, freeCores) + } + + return textResult(map[string]any{ + "queue": queue, + "processors": proc, + "pending_jobs": pending, + "free_cores_cluster": freeCores, + "estimated_wait_seconds": waitSeconds, + "estimate_basis": "Rough heuristic from current pending count and free cores; not a guarantee.", + }) +}