Phase 0: scheduleragent + mcp-server read-only agent framework
Complete the natural-language scheduler agent's Phase 0 read-only pipeline: Task #2 (API Key auth mechanism): - New store/apikey.go: APIKey store (hash, validate, expiry) - controllers/apikey.go: GenerateAPIKey finally persists + registers RBAC role - controllers/auth.go: GetUsernameFromToken recognizes xce_ak_ prefix - middleware/validatehandler.go: isValidToken handles API Key - config/rbac.yaml: mcp_readonly role (view-only permissions) - 6 new test files covering unit + integration full-chain Task #3 (schedulerclient Go library): - types.go: URM types (URMJob/URMQueue/URMNode/JobStatus/Metrics) - client.go: read-only Client interface (6 methods) - httpclient.go: resty-based HTTP implementation - statusmap.go + _test.go: TDD status mapping (33 tests) Task #4 (mcp-server binary): - cmd/main.go: entry, reads XCE_API_URL/XCE_API_KEY, registers 6 MCP tools - internal/app.go: tool handlers wiring schedulerclient to MCP stdio - Uses github.com/mark3labs/mcp-go@v0.57.0 for stdio transport All: go vet clean, go build clean, schedulerclient 33/33 tests pass. mcp-server binary ~11MB arm64, CGO_ENABLED=0. Architecture: Claude Code --stdio--> mcp-server --HTTP--> apiserver REST API --cgo--> C backend Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Lei Chen
parent
98066db68d
commit
e4af0cac73
27 changed files with 1760 additions and 7 deletions
86
src/frontend/mcp-server/cmd/main.go
Normal file
86
src/frontend/mcp-server/cmd/main.go
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"log"
|
||||
"os"
|
||||
|
||||
mcps "github.com/mark3labs/mcp-go/server"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"mcp-server/internal"
|
||||
)
|
||||
|
||||
func main() {
|
||||
baseURL := os.Getenv("XCE_API_URL")
|
||||
token := os.Getenv("XCE_API_KEY")
|
||||
|
||||
if baseURL == "" {
|
||||
baseURL = "https://localhost:8080"
|
||||
log.Printf("XCE_API_URL not set, defaulting to %s", baseURL)
|
||||
}
|
||||
if token == "" {
|
||||
log.Fatal("XCE_API_KEY environment variable is required")
|
||||
}
|
||||
|
||||
app, err := internal.NewApp(baseURL, token)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to initialize: %v", err)
|
||||
}
|
||||
|
||||
s := mcps.NewMCPServer(
|
||||
"scheduler-agent",
|
||||
"0.1.0",
|
||||
mcps.WithResourceCapabilities(true, true),
|
||||
mcps.WithLogging(),
|
||||
)
|
||||
|
||||
// ── scheduler_list_jobs ──
|
||||
s.AddTool(mcp.NewTool("scheduler_list_jobs",
|
||||
mcp.WithDescription("List jobs in the cluster. Supports filtering by queue, owner, status, and pagination."),
|
||||
mcp.WithString("queue", mcp.Description("Filter by queue name")),
|
||||
mcp.WithString("owner", mcp.Description("Filter by job owner username")),
|
||||
mcp.WithString("status", mcp.Description("Filter by status: pending|running|succeeded|failed|cancelled|suspended|unknown")),
|
||||
mcp.WithNumber("limit", mcp.Description("Max results to return")),
|
||||
mcp.WithString("page_token", mcp.Description("Pagination token from a prior response")),
|
||||
), app.HandleListJobs)
|
||||
|
||||
// ── scheduler_get_job ──
|
||||
s.AddTool(mcp.NewTool("scheduler_get_job",
|
||||
mcp.WithDescription("Get detailed information about a specific job by its ID."),
|
||||
mcp.WithString("id",
|
||||
mcp.Description("Job ID"),
|
||||
mcp.Required(),
|
||||
),
|
||||
), app.HandleGetJob)
|
||||
|
||||
// ── scheduler_list_queues ──
|
||||
s.AddTool(mcp.NewTool("scheduler_list_queues",
|
||||
mcp.WithDescription("List all queues/partitions and their current state (active, inactive, draining)."),
|
||||
), app.HandleListQueues)
|
||||
|
||||
// ── scheduler_list_nodes ──
|
||||
s.AddTool(mcp.NewTool("scheduler_list_nodes",
|
||||
mcp.WithDescription("List all compute nodes with their resource usage and status."),
|
||||
), app.HandleListNodes)
|
||||
|
||||
// ── scheduler_job_logs ──
|
||||
s.AddTool(mcp.NewTool("scheduler_job_logs",
|
||||
mcp.WithDescription("Retrieve the output logs of a job."),
|
||||
mcp.WithString("id",
|
||||
mcp.Description("Job ID"),
|
||||
mcp.Required(),
|
||||
),
|
||||
mcp.WithNumber("tail", mcp.Description("Number of lines from the end to return (0 = all)")),
|
||||
mcp.WithNumber("offset_bytes", mcp.Description("Byte offset to start reading from")),
|
||||
), app.HandleJobLogs)
|
||||
|
||||
// ── scheduler_cluster_metrics ──
|
||||
s.AddTool(mcp.NewTool("scheduler_cluster_metrics",
|
||||
mcp.WithDescription("Get a snapshot of overall cluster resource usage and utilization."),
|
||||
), app.HandleClusterMetrics)
|
||||
|
||||
log.Printf("Starting scheduler-agent MCP server (base-url: %s)", baseURL)
|
||||
if err := mcps.ServeStdio(s); err != nil {
|
||||
log.Fatalf("server error: %v", err)
|
||||
}
|
||||
}
|
||||
21
src/frontend/mcp-server/go.mod
Normal file
21
src/frontend/mcp-server/go.mod
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
module mcp-server
|
||||
|
||||
go 1.25.5
|
||||
|
||||
replace schedulerclient => ../schedulerclient
|
||||
|
||||
require (
|
||||
github.com/mark3labs/mcp-go v0.57.0
|
||||
schedulerclient v0.0.0-00010101000000-000000000000
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/google/jsonschema-go v0.4.2 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect
|
||||
github.com/spf13/cast v1.7.1 // indirect
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
|
||||
golang.org/x/net v0.33.0 // indirect
|
||||
golang.org/x/text v0.21.0 // indirect
|
||||
resty.dev/v3 v3.0.0-beta.2 // indirect
|
||||
)
|
||||
38
src/frontend/mcp-server/go.sum
Normal file
38
src/frontend/mcp-server/go.sum
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI=
|
||||
github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
|
||||
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
|
||||
github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8=
|
||||
github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
||||
github.com/mark3labs/mcp-go v0.57.0 h1:jzWKyCzdWnwnZt05cvcQQ+ngiUl2RnixXJa7Kj4qP1E=
|
||||
github.com/mark3labs/mcp-go v0.57.0/go.mod h1:+8WclSK1ZUweCP3hvktSji8n8ABG/95QaEkeVE/Uwas=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
||||
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
|
||||
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEVZGK7IN2kJkjTuQ=
|
||||
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU=
|
||||
github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
|
||||
github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
|
||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
|
||||
golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
|
||||
golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
|
||||
golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
|
||||
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
resty.dev/v3 v3.0.0-beta.2 h1:xu4mGAdbCLuc3kbk7eddWfWm4JfhwDtdapwss5nCjnQ=
|
||||
resty.dev/v3 v3.0.0-beta.2/go.mod h1:OgkqiPvTDtOuV4MGZuUDhwOpkY8enjOsjjMzeOHefy4=
|
||||
201
src/frontend/mcp-server/internal/app.go
Normal file
201
src/frontend/mcp-server/internal/app.go
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
// Package internal provides the scheduler-agent MCP tool handlers.
|
||||
package internal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
sclient "schedulerclient"
|
||||
)
|
||||
|
||||
// App wires the schedulerclient to the MCP tool handlers.
|
||||
type App struct {
|
||||
client sclient.Client
|
||||
}
|
||||
|
||||
// NewApp creates a new App with a schedulerclient connected to the given base URL.
|
||||
func NewApp(baseURL, token string) (*App, error) {
|
||||
hc := sclient.NewClient(baseURL, token)
|
||||
return &App{client: hc}, nil
|
||||
}
|
||||
|
||||
// ── helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
// textResult returns an MCP tool result containing the JSON representation of v.
|
||||
func textResult(v any) (*mcp.CallToolResult, error) {
|
||||
b, err := json.MarshalIndent(v, "", " ")
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("marshal error: %v", err)), nil
|
||||
}
|
||||
return mcp.NewToolResultText(string(b)), nil
|
||||
}
|
||||
|
||||
// textResultStr returns an MCP tool result wrapping a plain string.
|
||||
func textResultStr(s string) *mcp.CallToolResult {
|
||||
return mcp.NewToolResultText(s)
|
||||
}
|
||||
|
||||
// getString extracts a string parameter from an MCP tool request.
|
||||
func getString(req mcp.CallToolRequest, key string) string {
|
||||
args := req.GetArguments()
|
||||
raw, ok := args[key]
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
s, _ := raw.(string)
|
||||
return s
|
||||
}
|
||||
|
||||
// getNumber extracts a float64 parameter from an MCP tool request.
|
||||
func getNumber(req mcp.CallToolRequest, key string) float64 {
|
||||
args := req.GetArguments()
|
||||
raw, ok := args[key]
|
||||
if !ok {
|
||||
return 0
|
||||
}
|
||||
n, _ := raw.(float64)
|
||||
return n
|
||||
}
|
||||
|
||||
// ── tool handlers ───────────────────────────────────────────────────────────
|
||||
|
||||
func (a *App) HandleListJobs(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
filter := sclient.JobFilter{
|
||||
Queue: getString(req, "queue"),
|
||||
Owner: getString(req, "owner"),
|
||||
PageToken: getString(req, "page_token"),
|
||||
}
|
||||
if n := getNumber(req, "limit"); n > 0 {
|
||||
filter.Limit = int(n)
|
||||
}
|
||||
if s := getString(req, "status"); s != "" {
|
||||
filter.Statuses = parseStatuses(s)
|
||||
}
|
||||
|
||||
jobs, page, err := a.client.ListJobs(ctx, filter)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("list jobs failed: %v", err)), nil
|
||||
}
|
||||
|
||||
resp := map[string]any{
|
||||
"jobs": jobs,
|
||||
"page": page,
|
||||
}
|
||||
return textResult(resp)
|
||||
}
|
||||
|
||||
func (a *App) HandleGetJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
|
||||
job, err := a.client.GetJob(ctx, id)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("get job failed: %v", err)), nil
|
||||
}
|
||||
return textResult(job)
|
||||
}
|
||||
|
||||
func (a *App) HandleListQueues(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
queues, err := a.client.ListQueues(ctx)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("list queues failed: %v", err)), nil
|
||||
}
|
||||
return textResult(queues)
|
||||
}
|
||||
|
||||
func (a *App) HandleListNodes(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
nodes, err := a.client.ListNodes(ctx)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("list nodes failed: %v", err)), nil
|
||||
}
|
||||
return textResult(nodes)
|
||||
}
|
||||
|
||||
func (a *App) HandleJobLogs(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
opts := sclient.LogOpts{
|
||||
Tail: int(getNumber(req, "tail")),
|
||||
OffsetBytes: int64(getNumber(req, "offset_bytes")),
|
||||
}
|
||||
|
||||
lines, err := a.client.GetJobLogs(ctx, id, opts)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("get job logs failed: %v", err)), nil
|
||||
}
|
||||
return textResultStr(strconv.Itoa(len(lines)) + " lines\n" + joinLines(lines)), nil
|
||||
}
|
||||
|
||||
func (a *App) HandleClusterMetrics(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
metrics, err := a.client.GetClusterMetrics(ctx)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("get cluster metrics failed: %v", err)), nil
|
||||
}
|
||||
return textResult(metrics)
|
||||
}
|
||||
|
||||
// ── helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
// parseStatuses parses a comma-separated list of URM status strings.
|
||||
func parseStatuses(s string) []sclient.JobStatus {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
parts := splitComma(s)
|
||||
statuses := make([]sclient.JobStatus, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
statuses = append(statuses, sclient.JobStatus(p))
|
||||
}
|
||||
return statuses
|
||||
}
|
||||
|
||||
// joinLines joins log lines into a single string.
|
||||
func joinLines(lines []string) string {
|
||||
if len(lines) == 0 {
|
||||
return ""
|
||||
}
|
||||
if len(lines) == 1 {
|
||||
return lines[0]
|
||||
}
|
||||
// Estimate total length for capacity
|
||||
total := 0
|
||||
for _, l := range lines {
|
||||
total += len(l) + 1
|
||||
}
|
||||
b := make([]byte, 0, total)
|
||||
for i, l := range lines {
|
||||
if i > 0 {
|
||||
b = append(b, '\n')
|
||||
}
|
||||
b = append(b, []byte(l)...)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// splitComma splits a comma-separated string, trimming whitespace.
|
||||
func splitComma(s string) []string {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
var result []string
|
||||
start := 0
|
||||
for i := 0; i < len(s); i++ {
|
||||
if s[i] == ',' {
|
||||
if i > start {
|
||||
result = append(result, s[start:i])
|
||||
}
|
||||
start = i + 1
|
||||
}
|
||||
}
|
||||
if start < len(s) {
|
||||
result = append(result, s[start:])
|
||||
}
|
||||
return result
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue