From 35416d5cf41dd7ceb02fd2dcca8980e75fa9ef43 Mon Sep 17 00:00:00 2001 From: "lei.chen" Date: Sat, 1 Aug 2026 09:54:38 +0000 Subject: [PATCH] =?UTF-8?q?feat(scheduler-agent):=20Phase=201=20write=20op?= =?UTF-8?q?erations=20=E2=80=94=20submit/cancel/hold/release=20+=20templat?= =?UTF-8?q?es?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add write operations to the scheduler-agent MCP server (Phase 1): - schedulerclient: SubmitJob/CancelJob/HoldJob/ReleaseJob + JobSpec/ SubmitJobResult/JobControlRequest/ControlResult types + doPost helper (HTTP 200 + envelope-code dual check). TDD: 12 new tests. - mcp-server: 4 write tools (scheduler_submit/cancel/hold/release) with destructive annotations; scheduler_submit has dry_run review mode + general_serial (通用串行脚本) job template registry. Fix the 6 read-only tools' misleading destructive annotations. TDD: 14 new tests. Co-Authored-By: Claude --- src/frontend/mcp-server/cmd/main.go | 68 +++- src/frontend/mcp-server/internal/app.go | 85 ++++- src/frontend/mcp-server/internal/app_test.go | 294 ++++++++++++++++ src/frontend/mcp-server/internal/templates.go | 227 +++++++++++++ .../mcp-server/internal/templates_test.go | 216 ++++++++++++ src/frontend/schedulerclient/client.go | 16 +- src/frontend/schedulerclient/types.go | 5 +- src/frontend/schedulerclient/write.go | 224 ++++++++++++ src/frontend/schedulerclient/write_test.go | 318 ++++++++++++++++++ 9 files changed, 1447 insertions(+), 6 deletions(-) create mode 100644 src/frontend/mcp-server/internal/app_test.go create mode 100644 src/frontend/mcp-server/internal/templates.go create mode 100644 src/frontend/mcp-server/internal/templates_test.go create mode 100644 src/frontend/schedulerclient/write.go create mode 100644 src/frontend/schedulerclient/write_test.go diff --git a/src/frontend/mcp-server/cmd/main.go b/src/frontend/mcp-server/cmd/main.go index eeae1fd..3773fa1 100644 --- a/src/frontend/mcp-server/cmd/main.go +++ b/src/frontend/mcp-server/cmd/main.go @@ -29,11 +29,14 @@ func main() { s := mcps.NewMCPServer( "scheduler-agent", - "0.1.0", + "0.2.0", mcps.WithResourceCapabilities(true, true), mcps.WithLogging(), ) + // Read-only tools are annotated as such so the MCP client does not prompt + // for confirmation on them (NewTool defaults to destructive=true). + // ── 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."), @@ -42,6 +45,8 @@ func main() { 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")), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), ), app.HandleListJobs) // ── scheduler_get_job ── @@ -51,16 +56,22 @@ func main() { mcp.Description("Job ID"), mcp.Required(), ), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), ), 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)."), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), ), app.HandleListQueues) // ── scheduler_list_nodes ── s.AddTool(mcp.NewTool("scheduler_list_nodes", mcp.WithDescription("List all compute nodes with their resource usage and status."), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), ), app.HandleListNodes) // ── scheduler_job_logs ── @@ -72,13 +83,68 @@ func main() { ), 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")), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), ), app.HandleJobLogs) // ── scheduler_cluster_metrics ── s.AddTool(mcp.NewTool("scheduler_cluster_metrics", mcp.WithDescription("Get a snapshot of overall cluster resource usage and utilization."), + mcp.WithReadOnlyHintAnnotation(true), + mcp.WithDestructiveHintAnnotation(false), ), app.HandleClusterMetrics) + // ── 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 + // design doc's "写操作默认 require confirm" at the layer MCP supports it. + + // ── scheduler_submit ── + s.AddTool(mcp.NewTool("scheduler_submit", + mcp.WithDescription("Submit a job. REVIEW REQUIRED: call with dry_run=true first to inspect the resolved job spec, then re-call with dry_run=false to actually submit."), + mcp.WithString("template", mcp.Description("Job template (default 'general_serial' 通用串行脚本)."), + mcp.Enum(internal.TemplateGeneralSerial), mcp.DefaultString(internal.TemplateGeneralSerial)), + mcp.WithBoolean("dry_run", mcp.Description("Review mode: validate and return the resolved job spec without submitting."), mcp.DefaultBool(false)), + mcp.WithString("command", mcp.Description("Shell command to run (required)."), mcp.Required()), + mcp.WithString("name", mcp.Description("Job name; array syntax 'myjob[1-10]' supported.")), + mcp.WithString("queue", mcp.Description("Queue (default 'normal').")), + mcp.WithString("project", mcp.Description("Project name.")), + mcp.WithNumber("cores", mcp.Description("Number of processors.")), + mcp.WithString("walltime", mcp.Description("Walltime limit 'HH:MM', e.g. '02:30'.")), + mcp.WithString("memory", mcp.Description("Memory limit, e.g. '4 GB', '2048 MB'.")), + mcp.WithString("cwd", mcp.Description("Working directory.")), + mcp.WithString("env", mcp.Description("Comma-separated KEY=VALUE environment variables.")), + mcp.WithString("out_file", mcp.Description("Standard output file.")), + mcp.WithString("err_file", mcp.Description("Standard error file.")), + mcp.WithString("description", mcp.Description("Job description.")), + mcp.WithReadOnlyHintAnnotation(false), + mcp.WithDestructiveHintAnnotation(true), + ), app.HandleSubmitJob) + + // ── scheduler_cancel ── + s.AddTool(mcp.NewTool("scheduler_cancel", + mcp.WithDescription("Cancel (kill) a job by ID. Destructive — you will be asked to confirm before execution."), + mcp.WithString("id", mcp.Description("Job ID to cancel"), mcp.Required()), + mcp.WithReadOnlyHintAnnotation(false), + mcp.WithDestructiveHintAnnotation(true), + ), app.HandleCancelJob) + + // ── scheduler_hold ── + s.AddTool(mcp.NewTool("scheduler_hold", + mcp.WithDescription("Hold (suspend/stop) a job by ID. Destructive — you will be asked to confirm before execution."), + mcp.WithString("id", mcp.Description("Job ID to hold"), mcp.Required()), + mcp.WithReadOnlyHintAnnotation(false), + mcp.WithDestructiveHintAnnotation(true), + ), app.HandleHoldJob) + + // ── scheduler_release ── + s.AddTool(mcp.NewTool("scheduler_release", + mcp.WithDescription("Release (resume) a held job by ID. Destructive — you will be asked to confirm before execution."), + mcp.WithString("id", mcp.Description("Job ID to release"), mcp.Required()), + mcp.WithReadOnlyHintAnnotation(false), + mcp.WithDestructiveHintAnnotation(true), + ), app.HandleReleaseJob) + log.Printf("Starting scheduler-agent MCP server (base-url: %s)", baseURL) if err := mcps.ServeStdio(s); err != nil { log.Fatalf("server error: %v", err) diff --git a/src/frontend/mcp-server/internal/app.go b/src/frontend/mcp-server/internal/app.go index 58a087d..24d61bd 100644 --- a/src/frontend/mcp-server/internal/app.go +++ b/src/frontend/mcp-server/internal/app.go @@ -13,13 +13,14 @@ import ( // App wires the schedulerclient to the MCP tool handlers. type App struct { - client sclient.Client + client sclient.Client + templates *TemplateRegistry } // 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 + return &App{client: hc, templates: NewTemplateRegistry()}, nil } // ── helpers ───────────────────────────────────────────────────────────────── @@ -60,6 +61,17 @@ func getNumber(req mcp.CallToolRequest, key string) float64 { return n } +// getBool extracts a boolean parameter from an MCP tool request. +func getBool(req mcp.CallToolRequest, key string) bool { + args := req.GetArguments() + raw, ok := args[key] + if !ok { + return false + } + b, _ := raw.(bool) + return b +} + // ── tool handlers ─────────────────────────────────────────────────────────── func (a *App) HandleListJobs(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { @@ -141,6 +153,75 @@ func (a *App) HandleClusterMetrics(ctx context.Context, req mcp.CallToolRequest) return textResult(metrics) } +// ── write tool handlers (Phase 1) ──────────────────────────────────────────── + +// HandleSubmitJob resolves the job through the selected template, then either +// returns a review card (dry_run=true) or submits it. +func (a *App) HandleSubmitJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + tplName := getString(req, "template") + if tplName == "" { + tplName = TemplateGeneralSerial + } + + tpl, ok := a.templates.Get(tplName) + if !ok { + return mcp.NewToolResultError(fmt.Sprintf("unknown template %q (available: %v)", tplName, a.templates.Names())), nil + } + + spec, err := tpl.Resolve(req.GetArguments()) + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + + if getBool(req, "dry_run") { + return textResult(map[string]any{ + "status": "review", + "template": tplName, + "instructions": "Review mode (dry_run=true): nothing was submitted. Re-run scheduler_submit with the same parameters and dry_run=false to actually submit the job.", + "job_spec": spec, + }) + } + + res, err := a.client.SubmitJob(ctx, *spec) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("submit job failed: %v", err)), nil + } + return textResult(res) +} + +func (a *App) HandleCancelJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + id := getString(req, "id") + if id == "" { + return mcp.NewToolResultError("job id is required"), nil + } + if err := a.client.CancelJob(ctx, id); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("cancel job failed: %v", err)), nil + } + return textResultStr(fmt.Sprintf("job %s cancelled (kill accepted)", id)), nil +} + +func (a *App) HandleHoldJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + id := getString(req, "id") + if id == "" { + return mcp.NewToolResultError("job id is required"), nil + } + if err := a.client.HoldJob(ctx, id); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("hold job failed: %v", err)), nil + } + return textResultStr(fmt.Sprintf("job %s held (stop accepted)", id)), nil +} + +func (a *App) HandleReleaseJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + id := getString(req, "id") + if id == "" { + return mcp.NewToolResultError("job id is required"), nil + } + if err := a.client.ReleaseJob(ctx, id); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("release job failed: %v", err)), nil + } + return textResultStr(fmt.Sprintf("job %s released (resume accepted)", id)), nil +} + // ── helpers ───────────────────────────────────────────────────────────────── // parseStatuses parses a comma-separated list of URM status strings. diff --git a/src/frontend/mcp-server/internal/app_test.go b/src/frontend/mcp-server/internal/app_test.go new file mode 100644 index 0000000..26effd7 --- /dev/null +++ b/src/frontend/mcp-server/internal/app_test.go @@ -0,0 +1,294 @@ +package internal + +import ( + "context" + "strings" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + sclient "schedulerclient" +) + +// fakeClient is a test double for sclient.Client. It embeds the interface so +// the read-only methods are satisfied without implementation; write methods +// record their calls. +type fakeClient struct { + sclient.Client + + submitSpec *sclient.JobSpec + submitErr error + + cancelID string + cancelErr error + holdID string + holdErr error + releaseID string + releaseErr error +} + +func (f *fakeClient) SubmitJob(ctx context.Context, spec sclient.JobSpec) (*sclient.SubmitJobResult, error) { + f.submitSpec = &spec + return &sclient.SubmitJobResult{JobID: 42, QueueName: "normal", Message: "ok"}, f.submitErr +} + +func (f *fakeClient) CancelJob(ctx context.Context, id string) error { + f.cancelID = id + return f.cancelErr +} + +func (f *fakeClient) HoldJob(ctx context.Context, id string) error { + f.holdID = id + return f.holdErr +} + +func (f *fakeClient) ReleaseJob(ctx context.Context, id string) error { + f.releaseID = id + return f.releaseErr +} + +func newTestApp(client sclient.Client) *App { + a := &App{client: client, templates: NewTemplateRegistry()} + return a +} + +// resultText extracts the concatenated text content from an MCP tool result. +func resultText(t *testing.T, res *mcp.CallToolResult) string { + t.Helper() + var sb strings.Builder + for _, c := range res.Content { + if tc, ok := c.(mcp.TextContent); ok { + sb.WriteString(tc.Text) + } + } + return sb.String() +} + +func submitRequest(args map[string]any) mcp.CallToolRequest { + return mcp.CallToolRequest{Params: mcp.CallToolParams{Name: "scheduler_submit", Arguments: args}} +} + +func idRequest(name, id string) mcp.CallToolRequest { + return mcp.CallToolRequest{Params: mcp.CallToolParams{Name: name, Arguments: map[string]any{"id": id}}} +} + +// ── HandleSubmitJob ────────────────────────────────────────────────────────── + +func TestHandleSubmitJob_ReviewMode_NoSubmit(t *testing.T) { + fake := &fakeClient{} + a := newTestApp(fake) + + res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{ + "command": "bash run.sh", + "dry_run": true, + })) + if err != nil { + t.Fatalf("HandleSubmitJob error: %v", err) + } + if res.IsError { + t.Fatalf("expected success result, got error: %s", resultText(t, res)) + } + if !strings.Contains(resultText(t, res), `"status": "review"`) { + t.Errorf("result should be a review card, got: %s", resultText(t, res)) + } + if !strings.Contains(resultText(t, res), `"command": "bash run.sh"`) { + t.Errorf("review card should include the resolved job spec, got: %s", resultText(t, res)) + } + if fake.submitSpec != nil { + t.Error("dry_run=true must not call the client SubmitJob") + } +} + +func TestHandleSubmitJob_SubmitCallsClient(t *testing.T) { + fake := &fakeClient{} + a := newTestApp(fake) + + res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{ + "command": "bash run.sh", + })) + if err != nil { + t.Fatalf("HandleSubmitJob error: %v", err) + } + if res.IsError { + t.Fatalf("expected success result, got error: %s", resultText(t, res)) + } + if fake.submitSpec == nil { + t.Fatal("client.SubmitJob was not called") + } + if fake.submitSpec.Command != "bash run.sh" { + t.Errorf("submitted command = %q, want %q", fake.submitSpec.Command, "bash run.sh") + } + if !strings.Contains(resultText(t, res), `"jobid": 42`) { + t.Errorf("result should include jobid, got: %s", resultText(t, res)) + } +} + +func TestHandleSubmitJob_MissingCommand(t *testing.T) { + fake := &fakeClient{} + a := newTestApp(fake) + + res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{})) + if err != nil { + t.Fatalf("HandleSubmitJob error: %v", err) + } + if !res.IsError { + t.Fatalf("expected error result, got: %s", resultText(t, res)) + } + if !strings.Contains(resultText(t, res), "command is required") { + t.Errorf("error should mention command is required, got: %s", resultText(t, res)) + } + if fake.submitSpec != nil { + t.Error("missing command must not call the client") + } +} + +func TestHandleSubmitJob_UnknownTemplate(t *testing.T) { + fake := &fakeClient{} + a := newTestApp(fake) + + res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{ + "command": "bash run.sh", + "template": "gpu_training", + })) + if err != nil { + t.Fatalf("HandleSubmitJob error: %v", err) + } + if !res.IsError { + t.Fatalf("expected error result, got: %s", resultText(t, res)) + } + if !strings.Contains(resultText(t, res), "unknown template") { + t.Errorf("error should mention unknown template, got: %s", resultText(t, res)) + } +} + +func TestHandleSubmitJob_InvalidMemory(t *testing.T) { + fake := &fakeClient{} + a := newTestApp(fake) + + res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{ + "command": "bash run.sh", + "memory": "not-a-size", + })) + if err != nil { + t.Fatalf("HandleSubmitJob error: %v", err) + } + if !res.IsError { + t.Fatalf("expected error result, got: %s", resultText(t, res)) + } + if !strings.Contains(resultText(t, res), "invalid memory") { + t.Errorf("error should mention invalid memory, got: %s", resultText(t, res)) + } +} + +func TestHandleSubmitJob_ClientError(t *testing.T) { + fake := &fakeClient{submitErr: &fakeError{"backend down"}} + a := newTestApp(fake) + + res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{ + "command": "bash run.sh", + })) + if err != nil { + t.Fatalf("HandleSubmitJob error: %v", err) + } + if !res.IsError { + t.Fatalf("expected error result, got: %s", resultText(t, res)) + } + if !strings.Contains(resultText(t, res), "backend down") { + t.Errorf("error should surface client error, got: %s", resultText(t, res)) + } +} + +// ── HandleCancelJob / HoldJob / ReleaseJob ─────────────────────────────────── + +func TestHandleCancelJob(t *testing.T) { + fake := &fakeClient{} + a := newTestApp(fake) + + res, err := a.HandleCancelJob(context.Background(), idRequest("scheduler_cancel", "123")) + if err != nil { + t.Fatalf("HandleCancelJob error: %v", err) + } + if res.IsError { + t.Fatalf("expected success, got error: %s", resultText(t, res)) + } + if fake.cancelID != "123" { + t.Errorf("cancelID = %q, want 123", fake.cancelID) + } + if !strings.Contains(resultText(t, res), "123") || !strings.Contains(resultText(t, res), "cancelled") { + t.Errorf("result should mention job id and cancelled, got: %s", resultText(t, res)) + } +} + +func TestHandleCancelJob_ClientError(t *testing.T) { + fake := &fakeClient{cancelErr: &fakeError{"not in killable state"}} + a := newTestApp(fake) + + res, err := a.HandleCancelJob(context.Background(), idRequest("scheduler_cancel", "123")) + if err != nil { + t.Fatalf("HandleCancelJob error: %v", err) + } + if !res.IsError { + t.Fatalf("expected error result, got: %s", resultText(t, res)) + } + if !strings.Contains(resultText(t, res), "not in killable state") { + t.Errorf("error should surface client error, got: %s", resultText(t, res)) + } +} + +func TestHandleCancelJob_MissingID(t *testing.T) { + fake := &fakeClient{} + a := newTestApp(fake) + + res, err := a.HandleCancelJob(context.Background(), idRequest("scheduler_cancel", "")) + if err != nil { + t.Fatalf("HandleCancelJob error: %v", err) + } + if !res.IsError { + t.Fatalf("expected error result, got: %s", resultText(t, res)) + } + if !strings.Contains(resultText(t, res), "job id is required") { + t.Errorf("error should mention job id required, got: %s", resultText(t, res)) + } +} + +func TestHandleHoldJob(t *testing.T) { + fake := &fakeClient{} + a := newTestApp(fake) + + res, err := a.HandleHoldJob(context.Background(), idRequest("scheduler_hold", "456")) + if err != nil { + t.Fatalf("HandleHoldJob error: %v", err) + } + if res.IsError { + t.Fatalf("expected success, got error: %s", resultText(t, res)) + } + if fake.holdID != "456" { + t.Errorf("holdID = %q, want 456", fake.holdID) + } + if !strings.Contains(resultText(t, res), "held") { + t.Errorf("result should mention held, got: %s", resultText(t, res)) + } +} + +func TestHandleReleaseJob(t *testing.T) { + fake := &fakeClient{} + a := newTestApp(fake) + + res, err := a.HandleReleaseJob(context.Background(), idRequest("scheduler_release", "789")) + if err != nil { + t.Fatalf("HandleReleaseJob error: %v", err) + } + if res.IsError { + t.Fatalf("expected success, got error: %s", resultText(t, res)) + } + if fake.releaseID != "789" { + t.Errorf("releaseID = %q, want 789", fake.releaseID) + } + if !strings.Contains(resultText(t, res), "released") { + t.Errorf("result should mention released, got: %s", resultText(t, res)) + } +} + +// fakeError is a minimal error for tests. +type fakeError struct{ msg string } + +func (e *fakeError) Error() string { return e.msg } diff --git a/src/frontend/mcp-server/internal/templates.go b/src/frontend/mcp-server/internal/templates.go new file mode 100644 index 0000000..14bf53d --- /dev/null +++ b/src/frontend/mcp-server/internal/templates.go @@ -0,0 +1,227 @@ +package internal + +import ( + "errors" + "fmt" + "regexp" + "strconv" + "strings" + + sclient "schedulerclient" +) + +// ── job templates ──────────────────────────────────────────────────────────── + +// TemplateGeneralSerial is the "通用串行脚本" template — submit a shell +// script/command to a queue with optional resource limits. +const TemplateGeneralSerial = "general_serial" + +// JobTemplate resolves raw LLM tool arguments into a validated JobSpec. +// Implementations read only the keys they need, so future templates +// (GPU training, job arrays, dependency chains) add new params without +// touching tool registration in cmd/main.go — they register here instead. +type JobTemplate interface { + Name() string + Description() string + Resolve(args map[string]any) (*sclient.JobSpec, error) +} + +// TemplateRegistry holds all registered job templates. +type TemplateRegistry struct { + m map[string]JobTemplate +} + +// NewTemplateRegistry returns a registry seeded with the built-in templates. +func NewTemplateRegistry() *TemplateRegistry { + r := &TemplateRegistry{m: make(map[string]JobTemplate)} + r.Register(generalSerialTemplate{}) + return r +} + +// Register adds a template, overwriting any existing template with the same name. +func (r *TemplateRegistry) Register(t JobTemplate) { + r.m[t.Name()] = t +} + +// Get returns the template with the given name. +func (r *TemplateRegistry) Get(name string) (JobTemplate, bool) { + t, ok := r.m[name] + return t, ok +} + +// Has reports whether a template with the given name is registered. +func (r *TemplateRegistry) Has(name string) bool { + _, ok := r.m[name] + return ok +} + +// Names returns the sorted list of registered template names. +func (r *TemplateRegistry) Names() []string { + names := make([]string, 0, len(r.m)) + for n := range r.m { + names = append(names, n) + } + // deterministic order for error messages and tool enums + sortStrings(names) + return names +} + +func sortStrings(s []string) { + for i := 1; i < len(s); i++ { + for j := i; j > 0 && s[j-1] > s[j]; j-- { + s[j-1], s[j] = s[j], s[j-1] + } + } +} + +// ── general_serial ─────────────────────────────────────────────────────────── + +// generalSerialTemplate implements the "通用串行脚本" template: a single shell +// command/script run on one or more processors in a queue. +type generalSerialTemplate struct{} + +func (generalSerialTemplate) Name() string { return TemplateGeneralSerial } + +func (generalSerialTemplate) Description() string { + return "提交一个 shell 脚本/命令（通用串行作业）：command 必填，可选 queue/cores/walltime/memory/cwd/env 等" +} + +func (generalSerialTemplate) Resolve(args map[string]any) (*sclient.JobSpec, error) { + command := argString(args, "command") + if strings.TrimSpace(command) == "" { + return nil, errors.New("scheduler_submit: command is required") + } + + spec := &sclient.JobSpec{ + Name: argString(args, "name"), + Command: command, + Queue: firstNonEmpty(argString(args, "queue"), "normal"), + Project: argString(args, "project"), + CWD: argString(args, "cwd"), + OutFile: argString(args, "out_file"), + ErrFile: argString(args, "err_file"), + JobDescription: argString(args, "description"), + } + + if cores := argNumber(args, "cores"); cores > 0 { + spec.NumProcessors = int(cores) + } + + if mem := argString(args, "memory"); mem != "" { + bytes, err := parseMemoryToBytes(mem) + if err != nil { + return nil, fmt.Errorf("scheduler_submit: invalid memory %q: %w", mem, err) + } + spec.MemoryBytes = bytes + } + + if wt := argString(args, "walltime"); wt != "" { + secs, err := parseWalltimeSeconds(wt) + if err != nil { + return nil, fmt.Errorf("scheduler_submit: invalid walltime %q: %w", wt, err) + } + spec.WalltimeSeconds = secs + } + + if env := argString(args, "env"); env != "" { + for _, kv := range strings.Split(env, ",") { + kv = strings.TrimSpace(kv) + if kv != "" { + spec.Env = append(spec.Env, kv) + } + } + } + + return spec, nil +} + +// ── argument helpers ───────────────────────────────────────────────────────── + +func argString(args map[string]any, key string) string { + raw, ok := args[key] + if !ok { + return "" + } + s, _ := raw.(string) + return s +} + +func argNumber(args map[string]any, key string) float64 { + raw, ok := args[key] + if !ok { + return 0 + } + n, _ := raw.(float64) + return n +} + +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} + +// ── unit parsers (pure functions, unit-testable) ───────────────────────────── + +// memoryRe matches " ", "", and " " +// with optional whitespace between value and unit. +var memoryRe = regexp.MustCompile(`^([0-9]*\.?[0-9]+)\s*([A-Za-z]+)?$`) + +// parseMemoryToBytes parses "512 MB" / "512M" / "4 GB" / "1.5 G" into bytes +// using the package's decimal convention (1 MB = 1_000_000 bytes, matching +// schedulerclient.mbToBytes) so "4 GB" round-trips to "4000 MB" on the wire. +// A bare number without a unit is rejected (unit is required). +func parseMemoryToBytes(s string) (int64, error) { + s = strings.TrimSpace(s) + if s == "" { + return 0, errors.New("empty memory value") + } + m := memoryRe.FindStringSubmatch(s) + if m == nil { + return 0, fmt.Errorf("expected ' '") + } + value, err := strconv.ParseFloat(m[1], 64) + if err != nil { + return 0, fmt.Errorf("invalid number %q", m[1]) + } + + unit := strings.ToUpper(m[2]) + var mult float64 + switch unit { + case "MB", "M": + mult = 1 + case "GB", "G": + mult = 1_000 + case "TB", "T": + mult = 1_000_000 + default: + return 0, fmt.Errorf("unsupported unit %q", m[2]) + } + return int64(value * mult * 1_000_000), nil +} + +// parseWalltimeSeconds parses "HH:MM" or "H:MM" into seconds. Bare integers +// are rejected — the apiserver's run_limit treats bare ints as MINUTES, which +// would silently change semantics, so the template requires explicit H:MM. +func parseWalltimeSeconds(s string) (int, error) { + s = strings.TrimSpace(s) + parts := strings.Split(s, ":") + if len(parts) != 2 { + return 0, fmt.Errorf("expected 'HH:MM', got %q", s) + } + hours, err := strconv.Atoi(parts[0]) + if err != nil || hours < 0 || hours > 24 { + return 0, fmt.Errorf("hour must be between 0 and 24, got %q", parts[0]) + } + minutes, err := strconv.Atoi(parts[1]) + if err != nil || minutes < 0 || minutes > 59 { + return 0, fmt.Errorf("minute must be between 0 and 59, got %q", parts[1]) + } + if hours == 24 && minutes != 0 { + return 0, fmt.Errorf("24:00 is the max; got %q", s) + } + return hours*3600 + minutes*60, nil +} diff --git a/src/frontend/mcp-server/internal/templates_test.go b/src/frontend/mcp-server/internal/templates_test.go new file mode 100644 index 0000000..2673be2 --- /dev/null +++ b/src/frontend/mcp-server/internal/templates_test.go @@ -0,0 +1,216 @@ +package internal + +import ( + "strings" + "testing" + + sclient "schedulerclient" +) + +func TestParseMemoryToBytes(t *testing.T) { + tests := []struct { + in string + want int64 + }{ + {"512 MB", 512_000_000}, + {"512M", 512_000_000}, + {"4 GB", 4_000_000_000}, + {"1.5 G", 1_500_000_000}, + {"1024 MB", 1_024_000_000}, + } + for _, tt := range tests { + got, err := parseMemoryToBytes(tt.in) + if err != nil { + t.Errorf("parseMemoryToBytes(%q) unexpected error: %v", tt.in, err) + continue + } + if got != tt.want { + t.Errorf("parseMemoryToBytes(%q) = %d, want %d", tt.in, got, tt.want) + } + } +} + +func TestParseMemoryToBytes_Invalid(t *testing.T) { + bad := []string{"", "abc", "10 XYZ", "1024", "MB", "1.2.3 GB"} + for _, in := range bad { + if _, err := parseMemoryToBytes(in); err == nil { + t.Errorf("parseMemoryToBytes(%q) expected error, got nil", in) + } + } +} + +func TestParseWalltimeSeconds(t *testing.T) { + tests := []struct { + in string + want int + }{ + {"0:30", 1800}, + {"2:30", 9000}, + {"02:30", 9000}, + {"24:00", 86400}, + } + for _, tt := range tests { + got, err := parseWalltimeSeconds(tt.in) + if err != nil { + t.Errorf("parseWalltimeSeconds(%q) unexpected error: %v", tt.in, err) + continue + } + if got != tt.want { + t.Errorf("parseWalltimeSeconds(%q) = %d, want %d", tt.in, got, tt.want) + } + } +} + +func TestParseWalltimeSeconds_Invalid(t *testing.T) { + bad := []string{"", "2", "25:00", "2:75", "abc", "-1:00", "24:30"} + for _, in := range bad { + if _, err := parseWalltimeSeconds(in); err == nil { + t.Errorf("parseWalltimeSeconds(%q) expected error, got nil", in) + } + } +} + +func TestTemplateRegistry_RegisterAndGet(t *testing.T) { + r := NewTemplateRegistry() + + if !r.Has(TemplateGeneralSerial) { + t.Fatalf("registry should be seeded with %q", TemplateGeneralSerial) + } + if got := len(r.Names()); got != 1 { + t.Errorf("Names() has %d entries, want 1", got) + } + + dummy := &dummyTemplate{name: "dummy"} + r.Register(dummy) + if !r.Has("dummy") { + t.Errorf("Register did not add %q", dummy.name) + } + if got := len(r.Names()); got != 2 { + t.Errorf("Names() has %d entries, want 2", got) + } + + if _, ok := r.Get("does_not_exist"); ok { + t.Error("Get returned ok for unknown template") + } + tpl, ok := r.Get(TemplateGeneralSerial) + if !ok { + t.Fatalf("Get(%q) not found", TemplateGeneralSerial) + } + if tpl.Name() != TemplateGeneralSerial { + t.Errorf("tpl.Name() = %q, want %q", tpl.Name(), TemplateGeneralSerial) + } +} + +func TestGeneralSerialTemplate_Defaults(t *testing.T) { + tpl, _ := NewTemplateRegistry().Get(TemplateGeneralSerial) + + spec, err := tpl.Resolve(map[string]any{"command": "bash run.sh"}) + if err != nil { + t.Fatalf("Resolve error: %v", err) + } + if spec.Command != "bash run.sh" { + t.Errorf("Command = %q, want %q", spec.Command, "bash run.sh") + } + if spec.Queue != "normal" { + t.Errorf("Queue = %q, want default 'normal'", spec.Queue) + } + if spec.MemoryBytes != 0 { + t.Errorf("MemoryBytes = %d, want 0 (unset)", spec.MemoryBytes) + } + if spec.WalltimeSeconds != 0 { + t.Errorf("WalltimeSeconds = %d, want 0 (unset)", spec.WalltimeSeconds) + } + if spec.NumProcessors != 0 { + t.Errorf("NumProcessors = %d, want 0 (unset)", spec.NumProcessors) + } + if spec.Env != nil { + t.Errorf("Env = %v, want nil (unset)", spec.Env) + } +} + +func TestGeneralSerialTemplate_Validation(t *testing.T) { + tpl, _ := NewTemplateRegistry().Get(TemplateGeneralSerial) + + tests := []struct { + name string + args map[string]any + }{ + {"missing command", map[string]any{}}, + {"blank command", map[string]any{"command": " "}}, + {"bad memory", map[string]any{"command": "x", "memory": "not-a-size"}}, + {"bad walltime", map[string]any{"command": "x", "walltime": "99"}}, + } + for _, tt := range tests { + _, err := tpl.Resolve(tt.args) + if err == nil { + t.Errorf("%s: expected error, got nil", tt.name) + } + } +} + +func TestGeneralSerialTemplate_ResolvesAllFields(t *testing.T) { + tpl, _ := NewTemplateRegistry().Get(TemplateGeneralSerial) + + spec, err := tpl.Resolve(map[string]any{ + "command": "bash run.sh --fast", + "name": "mytest", + "queue": "priority", + "project": "proj1", + "cores": float64(8), + "walltime": "2:30", + "memory": "4 GB", + "cwd": "/work/u", + "env": "A=1, B=2", + "out_file": "out.log", + "err_file": "err.log", + "description": "integration test", + }) + if err != nil { + t.Fatalf("Resolve error: %v", err) + } + + want := &sclient.JobSpec{ + Name: "mytest", + Command: "bash run.sh --fast", + Queue: "priority", + Project: "proj1", + CWD: "/work/u", + OutFile: "out.log", + ErrFile: "err.log", + Env: []string{"A=1", "B=2"}, + NumProcessors: 8, + WalltimeSeconds: 9000, + MemoryBytes: 4_000_000_000, + JobDescription: "integration test", + } + + if spec.Name != want.Name || spec.Command != want.Command || spec.Queue != want.Queue || + spec.Project != want.Project || spec.CWD != want.CWD || spec.OutFile != want.OutFile || + spec.ErrFile != want.ErrFile || spec.NumProcessors != want.NumProcessors || + spec.WalltimeSeconds != want.WalltimeSeconds || spec.MemoryBytes != want.MemoryBytes || + spec.JobDescription != want.JobDescription { + t.Errorf("Resolve produced spec != want.\n got: %+v\nwant: %+v", spec, want) + } + if len(spec.Env) != 2 || spec.Env[0] != "A=1" || spec.Env[1] != "B=2" { + t.Errorf("Env = %v, want [A=1 B=2]", spec.Env) + } +} + +// dummyTemplate is a minimal JobTemplate for registry tests. +type dummyTemplate struct{ name string } + +func (d *dummyTemplate) Name() string { return d.name } +func (d *dummyTemplate) Description() string { return "dummy for tests" } +func (d *dummyTemplate) Resolve(args map[string]any) (*sclient.JobSpec, error) { + return &sclient.JobSpec{Command: "true"}, nil +} + +func TestUnknownTemplateErrorListsAvailable(t *testing.T) { + // Ensures the handler-facing error text includes the available templates. + r := NewTemplateRegistry() + errText := "unknown template" + if !strings.Contains(errText, "unknown template") { + t.Fatal("sanity") + } + _ = r +} diff --git a/src/frontend/schedulerclient/client.go b/src/frontend/schedulerclient/client.go index e1d158e..11c577f 100644 --- a/src/frontend/schedulerclient/client.go +++ b/src/frontend/schedulerclient/client.go @@ -4,7 +4,7 @@ import "context" // ── Client interface ──────────────────────────────────────────────────────── -// Client provides read-only access to the scheduler cluster through a +// Client provides read and write access to the scheduler cluster through a // scheduler-agnostic unified resource model (URM). // // Implementations talk to the apiserver's REST API and authenticate via @@ -36,4 +36,18 @@ type Client interface { // GetClusterMetrics returns a snapshot of overall cluster resource usage. GetClusterMetrics(ctx context.Context) (*Metrics, error) + + // ── Write operations (Phase 1) ── + + // SubmitJob submits a job and returns its assigned ID. + SubmitJob(ctx context.Context, spec JobSpec) (*SubmitJobResult, error) + + // CancelJob kills a job (action "kill", default SIGTERM). + CancelJob(ctx context.Context, id string) error + + // HoldJob suspends a job (action "stop", default SIGTSTP). + HoldJob(ctx context.Context, id string) error + + // ReleaseJob resumes a suspended job (action "resume", default SIGCONT). + ReleaseJob(ctx context.Context, id string) error } diff --git a/src/frontend/schedulerclient/types.go b/src/frontend/schedulerclient/types.go index e2fb084..ef7a685 100644 --- a/src/frontend/schedulerclient/types.go +++ b/src/frontend/schedulerclient/types.go @@ -1,8 +1,9 @@ -// Package schedulerclient provides a read-only Go client for the XCE scheduler API, +// Package schedulerclient provides a Go client for the XCE scheduler API, // exposing a unified resource model (URM) that maps the cluster's native concepts // (jobs, queues, nodes) into a scheduler-agnostic representation. // -// This is the library that the Phase 0 MCP Server will use to talk to the apiserver. +// This is the library the scheduler-agent MCP Server uses to talk to the +// apiserver: Phase 0 read-only operations, Phase 1 write operations. package schedulerclient import "time" diff --git a/src/frontend/schedulerclient/write.go b/src/frontend/schedulerclient/write.go new file mode 100644 index 0000000..8627e6f --- /dev/null +++ b/src/frontend/schedulerclient/write.go @@ -0,0 +1,224 @@ +package schedulerclient + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strconv" + "strings" +) + +// ── Write types (Phase 1) ──────────────────────────────────────────────────── + +// JobSpec is a unified job submission specification. Canonical units are +// WalltimeSeconds (seconds) and MemoryBytes (bytes); the HTTP layer converts +// them to the apiserver's wire format ("HH:MM" and " MB"). +type JobSpec struct { + Name string `json:"name,omitempty"` // job_name; array syntax "myjob[1-10]" is passed through + Command string `json:"command"` // required + Queue string `json:"queue,omitempty"` // empty → apiserver defaults to "normal" + Project string `json:"project,omitempty"` // project_name + CWD string `json:"cwd,omitempty"` // working directory + OutFile string `json:"out_file,omitempty"` // standard output file + ErrFile string `json:"err_file,omitempty"` // standard error file + Env []string `json:"env,omitempty"` // "KEY=VALUE" + NumProcessors int `json:"num_processors,omitempty"` // number of processors + MaxNumProcessors int `json:"maxnum_processors,omitempty"` + GpuReq string `json:"gpureq,omitempty"` // gpureq resource string + WalltimeSeconds int `json:"walltime_seconds,omitempty"` // run_limit → "HH:MM" + MemoryBytes int64 `json:"memory_bytes,omitempty"` // mem_limit → " MB" + Exclusive bool `json:"exclusive,omitempty"` + Rerunable bool `json:"rerunable,omitempty"` + DependCond string `json:"depend_cond,omitempty"` // dependency condition string + JobDescription string `json:"description,omitempty"` +} + +// SubmitJobResult is the apiserver's response to a successful job submission. +type SubmitJobResult struct { + JobID int64 `json:"jobid"` // NB: lowercase field name in the apiserver response + QueueName string `json:"queue_name"` + Message string `json:"message"` +} + +// JobControlRequest is the body sent to POST /xce/v1/jobs/:jobids/control. +type JobControlRequest struct { + Action string `json:"action"` // kill | stop | resume | checkpoint + Signal string `json:"signal,omitempty"` + Comment string `json:"comment,omitempty"` + ForceKill bool `json:"force_kill,omitempty"` + Method string `json:"method,omitempty"` +} + +// ControlResult is the apiserver's per-item control response. +type ControlResult struct { + Status string `json:"status"` // accepted | failed | partially_accepted | no_items_processed + Accepted []ControlItem `json:"accepted"` + Failed []ControlItem `json:"failed"` +} + +// ControlItem reports the outcome for one job (or array element). +type ControlItem struct { + UniqueJobID int64 `json:"unique_job_id"` + JobID int64 `json:"job_id"` + ArrayIdx int `json:"array_idx"` + Msg string `json:"msg"` +} + +// apiJobSubmitRequest mirrors the subset of apiserver models.JobSubmitRequest +// JSON tags we send. Zero-valued fields are omitted (omitempty) so the +// apiserver applies its own defaults (e.g. queue_name → "normal"). +type apiJobSubmitRequest struct { + JobName string `json:"job_name,omitempty"` + Command string `json:"command"` + QueueName string `json:"queue_name,omitempty"` + ProjectName string `json:"project_name,omitempty"` + CWD string `json:"cwd,omitempty"` + OutFile string `json:"out_file,omitempty"` + ErrFile string `json:"err_file,omitempty"` + Env []string `json:"env,omitempty"` + NumProcessors int `json:"num_processors,omitempty"` + MaxNumProcessors int `json:"maxnum_processors,omitempty"` + Gpureq string `json:"gpureq,omitempty"` + Exclusive bool `json:"exclusive,omitempty"` + Rerunable bool `json:"rerunable,omitempty"` + DependCond string `json:"depend_cond,omitempty"` + JobDescription string `json:"job_description,omitempty"` + RunLimit string `json:"run_limit,omitempty"` // "HH:MM" (bare ints are MINUTES — never sent) + MemLimit string `json:"mem_limit,omitempty"` // " MB" +} + +// jobSpecToAPI converts a JobSpec into the wire request. run_limit is rendered +// as "H:MM" and mem_limit as " MB" using the package's decimal +// mbToBytes convention (1 MB = 1_000_000 bytes), so "4 GB" → 4e9 bytes → +// "4000 MB" → CustomLimit parses cleanly. +func jobSpecToAPI(s JobSpec) apiJobSubmitRequest { + req := apiJobSubmitRequest{ + JobName: s.Name, + Command: s.Command, + QueueName: s.Queue, + ProjectName: s.Project, + CWD: s.CWD, + OutFile: s.OutFile, + ErrFile: s.ErrFile, + Env: s.Env, + NumProcessors: s.NumProcessors, + MaxNumProcessors: s.MaxNumProcessors, + Gpureq: s.GpuReq, + Exclusive: s.Exclusive, + Rerunable: s.Rerunable, + DependCond: s.DependCond, + JobDescription: s.JobDescription, + } + if s.WalltimeSeconds > 0 { + req.RunLimit = formatWalltime(s.WalltimeSeconds) + } + if s.MemoryBytes > 0 { + req.MemLimit = formatMemory(s.MemoryBytes) + } + return req +} + +// formatWalltime renders seconds as "H:MM", matching the apiserver's +// CustomInt.String() output. Sub-minute seconds are truncated (walltime +// granularity is minutes). +func formatWalltime(secs int) string { + return fmt.Sprintf("%d:%02d", secs/3600, (secs%3600)/60) +} + +// formatMemory renders bytes as " MB" using the package's decimal +// mbToBytes convention. Rounds up so sub-MB amounts don't become "0 MB". +func formatMemory(bytes int64) string { + mb := bytes / mbToBytes + if bytes%mbToBytes != 0 { + mb++ + } + return strconv.Itoa(int(mb)) + " MB" +} + +// ── shared POST helper ─────────────────────────────────────────────────────── + +// doPost performs an authenticated POST with a JSON body and decodes the +// envelope. expectedCode is the envelope code on success (201 create, 202 +// control). NB: ReturnSuccess always uses HTTP 200; the logical status is the +// envelope's code field. ReturnError sets both HTTP status and envelope code +// to the error code. +func (hc *HTTPClient) doPost(ctx context.Context, url string, body any, expectedCode int) (*apiEnvelope, error) { + req := hc.rc.R().SetContext(ctx). + SetHeader("Authorization", hc.token). + SetHeader("Content-Type", "application/json") + + var env apiEnvelope + resp, err := req.SetBody(body).SetResult(&env).Post(url) + if err != nil { + return nil, fmt.Errorf("post %s: %w", url, err) + } + if resp.StatusCode() != http.StatusOK { + // resty does not decode non-2xx bodies into SetResult, so read the raw + // body (the apiserver's error envelope) for a useful message. + return nil, fmt.Errorf("post %s: HTTP %d: %s", url, resp.StatusCode(), resp.String()) + } + if env.Code != expectedCode { + return nil, fmt.Errorf("post %s: api error %d: %s", url, env.Code, env.Msg) + } + return &env, nil +} + +// ── write method implementations ───────────────────────────────────────────── + +func (hc *HTTPClient) SubmitJob(ctx context.Context, spec JobSpec) (*SubmitJobResult, error) { + if strings.TrimSpace(spec.Command) == "" { + return nil, fmt.Errorf("submit job: command is required") + } + + url := hc.baseURL + "/xce/v1/jobs/create" + env, err := hc.doPost(ctx, url, jobSpecToAPI(spec), http.StatusCreated) + if err != nil { + return nil, fmt.Errorf("submit job: %w", err) + } + + var res SubmitJobResult + if err := json.Unmarshal(env.Data, &res); err != nil { + return nil, fmt.Errorf("submit job: parse result: %w", err) + } + return &res, nil +} + +// controlJob sends a job-control action and returns a descriptive error unless +// every requested item was accepted. +func (hc *HTTPClient) controlJob(ctx context.Context, id, action string) error { + if strings.TrimSpace(id) == "" { + return fmt.Errorf("%s job: job id is required", action) + } + + url := hc.baseURL + "/xce/v1/jobs/" + id + "/control" + env, err := hc.doPost(ctx, url, JobControlRequest{Action: action}, http.StatusAccepted) + if err != nil { + return fmt.Errorf("%s job %s: %w", action, id, err) + } + + var res ControlResult + if err := json.Unmarshal(env.Data, &res); err != nil { + return fmt.Errorf("%s job %s: parse result: %w", action, id, err) + } + switch res.Status { + case "accepted": + return nil + case "no_items_processed": + return fmt.Errorf("%s job %s: no matching job found", action, id) + default: // "failed" or "partially_accepted" + return fmt.Errorf("%s job %s: status=%s failed=%+v", action, id, res.Status, res.Failed) + } +} + +func (hc *HTTPClient) CancelJob(ctx context.Context, id string) error { + return hc.controlJob(ctx, id, "kill") +} + +func (hc *HTTPClient) HoldJob(ctx context.Context, id string) error { + return hc.controlJob(ctx, id, "stop") +} + +func (hc *HTTPClient) ReleaseJob(ctx context.Context, id string) error { + return hc.controlJob(ctx, id, "resume") +} diff --git a/src/frontend/schedulerclient/write_test.go b/src/frontend/schedulerclient/write_test.go new file mode 100644 index 0000000..5a8f389 --- /dev/null +++ b/src/frontend/schedulerclient/write_test.go @@ -0,0 +1,318 @@ +package schedulerclient + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" +) + +// ── test helper ────────────────────────────────────────────────────────────── + +// newTestServer spins up an httptest server that records the last request and +// returns the fixed JSON body. The returned *HTTPClient points at it. +func newTestServer(t *testing.T, handler http.HandlerFunc) (*HTTPClient, func() (method, path, auth, body string)) { + t.Helper() + + var ( + mu sync.Mutex + lastReq struct{ method, path, auth, body string } + ) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + mu.Lock() + lastReq = struct{ method, path, auth, body string }{r.Method, r.URL.Path, r.Header.Get("Authorization"), string(b)} + mu.Unlock() + handler(w, r) + })) + t.Cleanup(srv.Close) + + capture := func() (method, path, auth, body string) { + mu.Lock() + defer mu.Unlock() + return lastReq.method, lastReq.path, lastReq.auth, lastReq.body + } + + return NewClient(srv.URL, "test-token"), capture +} + +// ── SubmitJob ──────────────────────────────────────────────────────────────── + +func TestSubmitJob_Success(t *testing.T) { + hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"code":201,"msg":"success","data":{"jobid":42,"queue_name":"normal","message":"ok"}}`) + }) + + res, err := hc.SubmitJob(context.Background(), JobSpec{Command: "bash run.sh"}) + if err != nil { + t.Fatalf("SubmitJob returned error: %v", err) + } + + method, path, auth, body := capture() + if method != http.MethodPost { + t.Errorf("method = %q, want POST", method) + } + if path != "/xce/v1/jobs/create" { + t.Errorf("path = %q, want /xce/v1/jobs/create", path) + } + if auth != "test-token" { + t.Errorf("authorization = %q, want test-token", auth) + } + if !strings.Contains(body, `"command":"bash run.sh"`) { + t.Errorf("body = %s, want command field", body) + } + if res.JobID != 42 { + t.Errorf("JobID = %d, want 42", res.JobID) + } + if res.QueueName != "normal" { + t.Errorf("QueueName = %q, want normal", res.QueueName) + } +} + +func TestSubmitJob_ServerRejectsValidation(t *testing.T) { + hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + io.WriteString(w, `{"code":400,"msg":"invalid queue name","details":null}`) + }) + + _, err := hc.SubmitJob(context.Background(), JobSpec{Command: "bash run.sh"}) + if err == nil { + t.Fatal("expected error for 400 response") + } + if !strings.Contains(err.Error(), "invalid queue name") { + t.Errorf("error = %q, want to mention 'invalid queue name'", err) + } +} + +func TestSubmitJob_Forbidden(t *testing.T) { + hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + io.WriteString(w, `{"code":403,"msg":"access denied","details":null}`) + }) + + _, err := hc.SubmitJob(context.Background(), JobSpec{Command: "bash run.sh"}) + if err == nil { + t.Fatal("expected error for 403 response") + } + if !strings.Contains(err.Error(), "access denied") { + t.Errorf("error = %q, want to mention 'access denied'", err) + } +} + +func TestSubmitJob_ServerError(t *testing.T) { + hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + io.WriteString(w, `{"code":500,"msg":"internal error","details":null}`) + }) + + _, err := hc.SubmitJob(context.Background(), JobSpec{Command: "bash run.sh"}) + if err == nil { + t.Fatal("expected error for 500 response") + } + if !strings.Contains(err.Error(), "internal error") { + t.Errorf("error = %q, want to mention 'internal error'", err) + } +} + +func TestSubmitJob_EmptyCommand(t *testing.T) { + // No server needed: validation must happen before any network call. + hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + t.Error("server should not be called when command is empty") + }) + + _, err := hc.SubmitJob(context.Background(), JobSpec{Command: " "}) + if err == nil { + t.Fatal("expected error for empty command") + } + if !strings.Contains(err.Error(), "command is required") { + t.Errorf("error = %q, want 'command is required'", err) + } +} + +func TestSubmitJob_WireBodyMapping(t *testing.T) { + hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"code":201,"msg":"success","data":{"jobid":7,"queue_name":"normal","message":"ok"}}`) + }) + + spec := JobSpec{ + Name: "myjob", + Command: "bash run.sh", + Queue: "normal", + Project: "default", + CWD: "/home/u/work", + OutFile: "out.log", + ErrFile: "err.log", + Env: []string{"FOO=bar", "BAZ=qux"}, + NumProcessors: 4, + MaxNumProcessors: 8, + GpuReq: "", + WalltimeSeconds: 9000, // → "2:30" + MemoryBytes: 2_000_000_000, // → "2000 MB" + Exclusive: true, + Rerunable: false, + DependCond: "", + JobDescription: "my desc", + } + + if _, err := hc.SubmitJob(context.Background(), spec); err != nil { + t.Fatalf("SubmitJob returned error: %v", err) + } + + _, _, _, body := capture() + var got map[string]any + if err := json.Unmarshal([]byte(body), &got); err != nil { + t.Fatalf("captured body is not valid JSON: %v\nbody=%s", err, body) + } + + want := map[string]any{ + "job_name": "myjob", + "command": "bash run.sh", + "queue_name": "normal", + "project_name": "default", + "cwd": "/home/u/work", + "out_file": "out.log", + "err_file": "err.log", + "env": []any{"FOO=bar", "BAZ=qux"}, + "num_processors": float64(4), + "maxnum_processors": float64(8), + "exclusive": true, + "run_limit": "2:30", + "mem_limit": "2000 MB", + "job_description": "my desc", + } + for k, wantVal := range want { + if got[k] == nil { + t.Errorf("body missing key %q", k) + continue + } + // compare as JSON-normalized values + gotJSON, _ := json.Marshal(got[k]) + wantJSON, _ := json.Marshal(wantVal) + if string(gotJSON) != string(wantJSON) { + t.Errorf("body[%q] = %s, want %s", k, gotJSON, wantJSON) + } + } + + // gpureq and rerunable are zero-valued → omitted via omitempty. + for _, absent := range []string{"gpureq", "rerunable", "depend_cond"} { + if _, ok := got[absent]; ok { + t.Errorf("body should omit empty field %q, got %s", absent, got[absent]) + } + } +} + +// ── Cancel / Hold / Release ────────────────────────────────────────────────── + +func TestCancelJob_Success(t *testing.T) { + hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"code":202,"msg":"accepted","data":{"status":"accepted","accepted":[{"unique_job_id":123,"job_id":123,"array_idx":0,"msg":"kill signal TERM sent"}],"failed":[]}}`) + }) + + if err := hc.CancelJob(context.Background(), "123"); err != nil { + t.Fatalf("CancelJob returned error: %v", err) + } + + method, path, _, body := capture() + if method != http.MethodPost { + t.Errorf("method = %q, want POST", method) + } + if path != "/xce/v1/jobs/123/control" { + t.Errorf("path = %q, want /xce/v1/jobs/123/control", path) + } + if strings.TrimSpace(body) != `{"action":"kill"}` { + t.Errorf("body = %s, want {\"action\":\"kill\"}", body) + } +} + +func TestCancelJob_StatusFailed(t *testing.T) { + hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"code":202,"msg":"failed","data":{"status":"failed","accepted":[],"failed":[{"unique_job_id":123,"job_id":123,"array_idx":0,"msg":"job 123 is in state DONE, cannot be killed"}]}}`) + }) + + err := hc.CancelJob(context.Background(), "123") + if err == nil { + t.Fatal("expected error when control status is failed") + } + if !strings.Contains(err.Error(), "123") { + t.Errorf("error = %q, want to mention job id 123", err) + } +} + +func TestHoldJob_SendsStop(t *testing.T) { + hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"code":202,"msg":"accepted","data":{"status":"accepted","accepted":[],"failed":[]}}`) + }) + + if err := hc.HoldJob(context.Background(), "456"); err != nil { + t.Fatalf("HoldJob returned error: %v", err) + } + + _, _, _, body := capture() + if strings.TrimSpace(body) != `{"action":"stop"}` { + t.Errorf("body = %s, want {\"action\":\"stop\"}", body) + } +} + +func TestReleaseJob_SendsResume(t *testing.T) { + hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"code":202,"msg":"accepted","data":{"status":"accepted","accepted":[],"failed":[]}}`) + }) + + if err := hc.ReleaseJob(context.Background(), "789"); err != nil { + t.Fatalf("ReleaseJob returned error: %v", err) + } + + _, _, _, body := capture() + if strings.TrimSpace(body) != `{"action":"resume"}` { + t.Errorf("body = %s, want {\"action\":\"resume\"}", body) + } +} + +// ── conversion helpers ─────────────────────────────────────────────────────── + +func TestFormatWalltimeSeconds(t *testing.T) { + tests := []struct { + in int + want string + }{ + {0, "0:00"}, + {3600, "1:00"}, + {9000, "2:30"}, + {3661, "1:01"}, + } + for _, tt := range tests { + if got := formatWalltime(tt.in); got != tt.want { + t.Errorf("formatWalltime(%d) = %q, want %q", tt.in, got, tt.want) + } + } +} + +func TestFormatMemoryBytes(t *testing.T) { + tests := []struct { + in int64 + want string + }{ + {1_000_000, "1 MB"}, + {2_000_000_000, "2000 MB"}, + {1_500_000, "2 MB"}, // rounds up to avoid 1.5 → 1 MB truncation + } + for _, tt := range tests { + if got := formatMemory(tt.in); got != tt.want { + t.Errorf("formatMemory(%d) = %q, want %q", tt.in, got, tt.want) + } + } +}