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 2026-07-30 11:16:36 +08:00
parent 98066db68d
commit e4af0cac73
27 changed files with 1760 additions and 7 deletions

View file

@ -8,6 +8,8 @@ Go-based REST API server (Gin) and CLI tool (Cobra) that expose cluster manageme
|--------|------|-------------|
| `apiserver` | `src/frontend/apiserver/` | `cmd/main.go` |
| `xcecli` | `src/frontend/xcecli/` | `cmd/main.go` |
| `schedulerclient` (Phase 0, ✅ done) | `src/frontend/schedulerclient/` | library, no entry point — read-only Go client for the scheduler agent's MCP Server, calls `apiserver`'s REST API. `go vet`/`go build`/`go test` clean (33 tests). See module source: `types.go` (URM types), `client.go` (interface), `httpclient.go` (resty impl), `statusmap.go`+`_test.go` (TDD, status mapping). |
| `mcp-server` (Phase 0, ✅ done) | `src/frontend/mcp-server/` | `cmd/main.go` — read-only MCP Server exposing 6 tools (`scheduler_list_jobs`/`get`/`list_queues`/`list_nodes`/`job_logs`/`cluster_metrics`); pure Go, `CGO_ENABLED=0`, uses `schedulerclient` as API client, authenticates via `XCE_API_KEY` env var. Uses `github.com/mark3labs/mcp-go@v0.57.0` for MCP stdio transport. `go build`/`go vet` clean, binary ~11MB. |
## Build & Test
```bash

View file

@ -144,8 +144,31 @@ roles:
- type: queue
values: [normal]
- action: job:control
domain:
domain:
- type: queue
values: [normal]
# mcp_readonly: service-account role for the scheduler-agent MCP server
# (Phase 0, read-only tools only). Deliberately excludes job:control,
# queue:control, host:control, cluster:control, rbac:*, file:*, user:* —
# an MCP server credential should never be able to mutate cluster state
# or manage other users/roles.
mcp_readonly:
permissions:
- action: job:view
domain:
- type: all
- action: queue:view
domain:
- type: all
- action: host:view
domain:
- type: all
- action: resource:view
domain:
- type: all
- action: cluster:view
domain:
- type: all

View file

@ -7,6 +7,7 @@ import (
"apiserver/common"
"apiserver/models"
"apiserver/store"
"github.com/gin-gonic/gin"
)
@ -16,7 +17,7 @@ func generateAPIKey() (string, error) {
if _, err := rand.Read(bytes); err != nil {
return "", err
}
return hex.EncodeToString(bytes), nil
return models.APIKeyPrefix + hex.EncodeToString(bytes), nil
}
type APIKeyController struct{}
@ -49,13 +50,34 @@ func (a *APIKeyController) GenerateAPIKey(ctx *gin.Context) {
expiryDays = req.ExpiryDays
}
createTime := time.Now()
expiryTime := createTime.AddDate(0, 0, expiryDays)
// Persist the key (by hash) so it can actually be used to authenticate.
store.AddAPIKey(apiKey, store.APIKeyEntry{
Username: req.Username,
RoleName: req.RoleName,
Description: req.Description,
CreateTime: createTime,
ExpiryTime: expiryTime,
})
// RequirePermission resolves permissions via store.GetUserRoleNames, not
// via the APIKeyEntry directly — register the RBAC user→role mapping too,
// or the issued key would authenticate but be denied every permission
// check.
if req.RoleName != "" {
store.AddRoleForUser(req.Username, req.RoleName)
}
// Create response
response := models.APIKeyResponse{
APIKey: apiKey,
Username: req.Username,
CreateTime: time.Now(),
ExpiryTime: time.Now().AddDate(0, 0, expiryDays),
RoleName: req.RoleName,
CreateTime: createTime,
ExpiryTime: expiryTime,
}
ReturnSuccess(ctx, 0, "API key generated successfully", response)
ReturnSuccess(ctx, common.SUCCESS_CODE, "API key generated successfully", response)
}

View file

@ -0,0 +1,59 @@
package controllers
import (
"net/http"
"testing"
"time"
"apiserver/models"
"apiserver/store"
"github.com/gin-gonic/gin"
)
// TestIntegration_APIKey_MCPReadonly_FullMiddlewareChain verifies the
// end-to-end story for the scheduler-agent MCP server's service account:
// an API-key-authenticated request with the mcp_readonly role can reach a
// job:view route but is denied a job:control route — through the exact same
// RequirePermission middleware JWT-authenticated users go through.
func TestIntegration_APIKey_MCPReadonly_FullMiddlewareChain(t *testing.T) {
yamlContent := `
permissions:
- job:view
- job:control
roles:
mcp_readonly:
permissions:
- action: job:view
domain:
- type: all
`
setupRBACTest(t, yamlContent)
raw := models.APIKeyPrefix + "mcp-integration-test"
store.AddAPIKey(raw, store.APIKeyEntry{
Username: "mcp-bot",
RoleName: "mcp_readonly",
ExpiryTime: time.Now().Add(time.Hour),
})
store.AddRoleForUser("mcp-bot", "mcp_readonly")
router := newTestRouter()
jobsGroup := router.Group("/xce/v1/jobs")
jobsGroup.GET("", RequirePermission("job:view"), func(ctx *gin.Context) {
ctx.JSON(http.StatusOK, gin.H{"status": "ok"})
})
jobsGroup.POST("/control", RequirePermission("job:control"), func(ctx *gin.Context) {
ctx.JSON(http.StatusAccepted, gin.H{"status": "accepted"})
})
w := doRequest(router, "GET", "/xce/v1/jobs", raw)
if w.Code != http.StatusOK {
t.Errorf("mcp_readonly on job:view: got %d, want 200, body=%s", w.Code, w.Body.String())
}
w = doRequest(router, "POST", "/xce/v1/jobs/control", raw)
if w.Code != http.StatusForbidden {
t.Errorf("mcp_readonly on job:control: got %d, want 403", w.Code)
}
}

View file

@ -0,0 +1,176 @@
package controllers
import (
"apiserver/models"
"apiserver/store"
"apiserver/utils"
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func resetAPIKeyControllerState() {
store.InitStores()
utils.Log = zap.NewNop()
}
// TestGetUsernameFromToken_APIKey verifies that a valid, unexpired API key
// resolves to its associated username via the same entry point every other
// controller already calls for JWTs — so no call site needs to change.
func TestGetUsernameFromToken_APIKey(t *testing.T) {
resetAPIKeyControllerState()
raw := models.APIKeyPrefix + "abc123"
store.AddAPIKey(raw, store.APIKeyEntry{
Username: "mcp-bot",
RoleName: "mcp_readonly",
ExpiryTime: time.Now().Add(time.Hour),
})
username, err := GetUsernameFromToken(raw)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if username != "mcp-bot" {
t.Errorf("username = %q, want mcp-bot", username)
}
}
// TestGetUsernameFromToken_APIKey_BearerPrefix verifies the "Bearer <key>"
// header form (as sent by real HTTP clients) also resolves correctly.
func TestGetUsernameFromToken_APIKey_BearerPrefix(t *testing.T) {
resetAPIKeyControllerState()
raw := models.APIKeyPrefix + "with-bearer"
store.AddAPIKey(raw, store.APIKeyEntry{
Username: "mcp-bot",
ExpiryTime: time.Now().Add(time.Hour),
})
username, err := GetUsernameFromToken("Bearer " + raw)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if username != "mcp-bot" {
t.Errorf("username = %q, want mcp-bot", username)
}
}
// TestGetUsernameFromToken_APIKey_Unknown verifies an API-key-shaped token
// that was never issued is rejected rather than silently falling through to
// JWT parsing (which would also fail, but with a confusing error message).
func TestGetUsernameFromToken_APIKey_Unknown(t *testing.T) {
resetAPIKeyControllerState()
_, err := GetUsernameFromToken(models.APIKeyPrefix + "never-issued")
if err == nil {
t.Error("expected error for unknown API key")
}
}
// TestGetUsernameFromToken_APIKey_Expired verifies an expired API key is
// rejected.
func TestGetUsernameFromToken_APIKey_Expired(t *testing.T) {
resetAPIKeyControllerState()
raw := models.APIKeyPrefix + "expired"
store.AddAPIKey(raw, store.APIKeyEntry{
Username: "mcp-bot",
ExpiryTime: time.Now().Add(-time.Hour),
})
_, err := GetUsernameFromToken(raw)
if err == nil {
t.Error("expected error for expired API key")
}
}
// TestGenerateAPIKey_PersistsToStore verifies the controller actually writes
// the generated key into the store (previously it only echoed a random hex
// string back to the caller without persisting anything, so no request
// authenticated with the returned key could ever succeed).
func TestGenerateAPIKey_PersistsToStore(t *testing.T) {
resetAPIKeyControllerState()
gin.SetMode(gin.TestMode)
router := gin.New()
handler := &APIKeyController{}
router.POST("/xce/v1/apikeys", handler.GenerateAPIKey)
body := `{"username":"mcp-bot","role_name":"mcp_readonly","expiry_days":7}`
req := httptest.NewRequest(http.MethodPost, "/xce/v1/apikeys", bytes.NewBufferString(body))
req.Header.Set("Authorization", "Bearer some-admin-token")
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", w.Code, w.Body.String())
}
var resp struct {
Data models.APIKeyResponse `json:"data"`
}
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("unmarshal response: %v", err)
}
if resp.Data.APIKey == "" {
t.Fatal("response did not include an api_key")
}
// The returned key must actually be usable for authentication —
// i.e. it must have been persisted to the store, not merely echoed.
entry, ok := store.ValidateAPIKey(resp.Data.APIKey)
if !ok {
t.Fatal("generated API key was not persisted to the store")
}
if entry.Username != "mcp-bot" {
t.Errorf("persisted username = %q, want mcp-bot", entry.Username)
}
if entry.RoleName != "mcp_readonly" {
t.Errorf("persisted role = %q, want mcp_readonly", entry.RoleName)
}
// RequirePermission/GetUserAllPermissionsWithType resolve permissions via
// store.GetUserRoleNames(username), not via the APIKeyEntry.RoleName field
// directly — so generating a key must also register the RBAC user→role
// mapping, otherwise the issued key would authenticate but be denied every
// permission check.
roles := store.GetUserRoleNames("mcp-bot")
found := false
for _, r := range roles {
if r == "mcp_readonly" {
found = true
}
}
if !found {
t.Errorf("GetUserRoleNames(mcp-bot) = %v, want to include mcp_readonly", roles)
}
}
// TestGenerateAPIKey_RequiresAuthorizationHeader verifies the endpoint still
// rejects requests with no Authorization header at all.
func TestGenerateAPIKey_RequiresAuthorizationHeader(t *testing.T) {
resetAPIKeyControllerState()
gin.SetMode(gin.TestMode)
router := gin.New()
handler := &APIKeyController{}
router.POST("/xce/v1/apikeys", handler.GenerateAPIKey)
body := `{"username":"mcp-bot","role_name":"mcp_readonly"}`
req := httptest.NewRequest(http.MethodPost, "/xce/v1/apikeys", bytes.NewBufferString(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Errorf("status = %d, want 401", w.Code)
}
}

View file

@ -181,7 +181,22 @@ func (a AuthController) Logout(ctx *gin.Context) {
}
func GetUsernameFromToken(tokenString string) (string, error) {
token, err := jwt.Parse(strings.TrimPrefix(tokenString, "Bearer "), func(token *jwt.Token) (interface{}, error) {
tokenString = strings.TrimPrefix(tokenString, "Bearer ")
// Service-account API keys are a distinct credential type from JWTs
// (recognized by a fixed prefix, see models.APIKeyPrefix) — resolve them
// against the API key store instead of attempting a JWT parse, so every
// existing GetUsernameFromToken call site works transparently for both
// credential types.
if strings.HasPrefix(tokenString, models.APIKeyPrefix) {
entry, ok := store.ValidateAPIKey(tokenString)
if !ok {
return "", fmt.Errorf("invalid or expired API key")
}
return entry.Username, nil
}
token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
// Validate the algorithm
if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])

View file

@ -0,0 +1,52 @@
package controllers
import (
"apiserver/store"
"os"
"path/filepath"
"testing"
)
// TestRBACYAML_McpReadonlyRole_LoadsReadOnlyPermissions verifies the
// mcp_readonly role (added for the scheduler-agent MCP server's service
// account, Phase 0) is present in the shipped config/rbac.yaml and grants
// only *:view permissions — never any *:control/*:submit permission, so an
// MCP server credential can never mutate cluster state.
func TestRBACYAML_McpReadonlyRole_LoadsReadOnlyPermissions(t *testing.T) {
resetRBACState()
path, err := filepath.Abs("../config/rbac.yaml")
if err != nil {
t.Fatalf("resolve rbac.yaml path: %v", err)
}
if _, err := os.Stat(path); err != nil {
t.Skipf("config/rbac.yaml not found at %s: %v", path, err)
}
if err := LoadRBACConfig(path); err != nil {
t.Fatalf("LoadRBACConfig: %v", err)
}
if !store.Roles.Has("mcp_readonly") {
t.Fatal("expected mcp_readonly role to be loaded from config/rbac.yaml")
}
perms := store.GetRolePermissions("mcp_readonly")
if len(perms) == 0 {
t.Fatal("mcp_readonly should have at least one permission")
}
for _, p := range perms {
if !hasViewSuffix(p.Action) {
t.Errorf("mcp_readonly grants non-view permission %q — MCP service accounts must be read-only", p.Action)
}
}
}
func hasViewSuffix(action string) bool {
const suffix = ":view"
if len(action) < len(suffix) {
return false
}
return action[len(action)-len(suffix):] == suffix
}

View file

@ -4,6 +4,7 @@ import (
"apiserver/common"
"apiserver/controllers"
"apiserver/models"
"apiserver/store"
"apiserver/utils"
logger "apiserver/utils"
"bytes"
@ -99,6 +100,14 @@ func validateHeader() gin.HandlerFunc {
}
func isValidToken(tokenString string, IP string) (bool, error) {
// Service-account API keys are a distinct credential type from JWTs and
// carry no IP binding (they're meant for long-lived automated clients
// like the scheduler-agent MCP server, which may not have a stable IP).
if strings.HasPrefix(tokenString, models.APIKeyPrefix) {
_, ok := store.ValidateAPIKey(tokenString)
return ok, nil
}
token, err := jwt.ParseWithClaims(tokenString, &models.JWTClaims{}, func(token *jwt.Token) (interface{}, error) {
return []byte(utils.AuthenticationYamlConfig.Token.SecretKey), nil
})

View file

@ -0,0 +1,109 @@
package middleware
import (
"net/http"
"net/http/httptest"
"testing"
"time"
"apiserver/models"
"apiserver/store"
"apiserver/utils"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func resetValidateHandlerState() {
store.InitStores()
utils.Log = zap.NewNop()
utils.AuthenticationYamlConfig = &utils.Authentication{
Token: utils.TokenConfig{
SecretKey: "test-secret",
Type: "HS256",
TokenExpireTime: 3600,
},
}
}
func newValidateHandlerRouter() *gin.Engine {
gin.SetMode(gin.TestMode)
r := gin.New()
r.Use(ValidateHandler())
r.GET("/xce/v1/jobs", func(ctx *gin.Context) { ctx.Status(http.StatusOK) })
return r
}
// TestValidateHeader_APIKey_Valid verifies a valid, unexpired API key in the
// Authorization header passes the middleware (no JWT involved at all).
func TestValidateHeader_APIKey_Valid(t *testing.T) {
resetValidateHandlerState()
raw := models.APIKeyPrefix + "valid-key"
store.AddAPIKey(raw, store.APIKeyEntry{
Username: "mcp-bot",
ExpiryTime: time.Now().Add(time.Hour),
})
router := newValidateHandlerRouter()
req := httptest.NewRequest(http.MethodGet, "/xce/v1/jobs", nil)
req.Header.Set("Authorization", "Bearer "+raw)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("status = %d, want 200, body=%s", w.Code, w.Body.String())
}
}
// TestValidateHeader_APIKey_Unknown verifies an API-key-shaped token that was
// never issued is rejected with 401, not a confusing JWT parse error.
func TestValidateHeader_APIKey_Unknown(t *testing.T) {
resetValidateHandlerState()
router := newValidateHandlerRouter()
req := httptest.NewRequest(http.MethodGet, "/xce/v1/jobs", nil)
req.Header.Set("Authorization", "Bearer "+models.APIKeyPrefix+"never-issued")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Errorf("status = %d, want 401", w.Code)
}
}
// TestValidateHeader_APIKey_Expired verifies an expired API key is rejected.
func TestValidateHeader_APIKey_Expired(t *testing.T) {
resetValidateHandlerState()
raw := models.APIKeyPrefix + "expired-key"
store.AddAPIKey(raw, store.APIKeyEntry{
Username: "mcp-bot",
ExpiryTime: time.Now().Add(-time.Hour),
})
router := newValidateHandlerRouter()
req := httptest.NewRequest(http.MethodGet, "/xce/v1/jobs", nil)
req.Header.Set("Authorization", "Bearer "+raw)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Errorf("status = %d, want 401", w.Code)
}
}
// TestValidateHeader_NoToken_StillRejected is a regression guard: adding the
// API key branch must not weaken the existing "no token" rejection.
func TestValidateHeader_NoToken_StillRejected(t *testing.T) {
resetValidateHandlerState()
router := newValidateHandlerRouter()
req := httptest.NewRequest(http.MethodGet, "/xce/v1/jobs", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusUnauthorized {
t.Errorf("status = %d, want 401", w.Code)
}
}

View file

@ -2,10 +2,14 @@ package models
import "time"
type APIKeyController struct{}
// APIKeyPrefix marks an Authorization header value as a service-account API
// key rather than a JWT, so validateHeader() can branch without attempting
// (and failing) a JWT parse first.
const APIKeyPrefix = "xce_ak_"
type APIKeyRequest struct {
Username string `json:"username" binding:"required"`
RoleName string `json:"role_name" binding:"required"`
Description string `json:"description,omitempty"`
ExpiryDays int `json:"expiry_days,omitempty"`
}
@ -13,6 +17,7 @@ type APIKeyRequest struct {
type APIKeyResponse struct {
APIKey string `json:"api_key"`
Username string `json:"username"`
RoleName string `json:"role_name"`
CreateTime time.Time `json:"create_time"`
ExpiryTime time.Time `json:"expiry_time"`
}

View file

@ -73,6 +73,14 @@ func Router() *gin.Engine {
permHandler := &controllers.PermController{}
permsApi.GET("", controllers.RequirePermission("rbac:view"), permHandler.GetPermsInfo)
}
// ################ case: apikeys route (service-account credentials, e.g. for the
// scheduler-agent MCP server; generation is admin-only, gated the same way role
// management is)
apikeysApi := r.Group("/xce/v1/apikeys")
{
apikeyHandler := &controllers.APIKeyController{}
apikeysApi.POST("", controllers.RequirePermission("rbac:control"), apikeyHandler.GenerateAPIKey)
}
// ################ case: roles route
rolesApi := r.Group("/xce/v1/roles")
{

View file

@ -0,0 +1,47 @@
package store
import (
"crypto/sha256"
"encoding/hex"
"time"
)
// APIKeys maps a hashed API key (hex-encoded SHA-256) to its metadata.
// The raw key is never stored — only its hash — so a store/config dump
// cannot leak usable credentials.
var APIKeys *DataStore[string, APIKeyEntry]
// APIKeyEntry holds the metadata associated with a service-account API key.
type APIKeyEntry struct {
Username string
RoleName string
Description string
CreateTime time.Time
ExpiryTime time.Time
}
// HashAPIKey returns the hex-encoded SHA-256 digest of a raw API key.
// Used both when persisting a newly generated key and when validating one
// presented by a client — the raw value is never compared or stored directly.
func HashAPIKey(rawKey string) string {
sum := sha256.Sum256([]byte(rawKey))
return hex.EncodeToString(sum[:])
}
// AddAPIKey persists a new API key by its hash.
func AddAPIKey(rawKey string, entry APIKeyEntry) {
APIKeys.Put(HashAPIKey(rawKey), entry)
}
// ValidateAPIKey looks up a raw API key by its hash and returns its metadata
// if present and not expired. A zero ExpiryTime never expires.
func ValidateAPIKey(rawKey string) (APIKeyEntry, bool) {
entry, ok := APIKeys.Get(HashAPIKey(rawKey))
if !ok {
return APIKeyEntry{}, false
}
if !entry.ExpiryTime.IsZero() && time.Now().After(entry.ExpiryTime) {
return APIKeyEntry{}, false
}
return entry, true
}

View file

@ -0,0 +1,94 @@
package store
import (
"testing"
"time"
)
func resetAPIKeyState() {
InitStores()
}
func TestHashAPIKey_Deterministic(t *testing.T) {
h1 := HashAPIKey("xce_ak_abc123")
h2 := HashAPIKey("xce_ak_abc123")
if h1 != h2 {
t.Errorf("HashAPIKey should be deterministic: %q != %q", h1, h2)
}
h3 := HashAPIKey("xce_ak_different")
if h1 == h3 {
t.Error("HashAPIKey should differ for different inputs")
}
}
func TestHashAPIKey_NeverPlaintext(t *testing.T) {
raw := "xce_ak_secret-value"
hashed := HashAPIKey(raw)
if hashed == raw {
t.Error("HashAPIKey must not return the raw key unchanged")
}
}
func TestAddAPIKey_ValidateAPIKey_RoundTrip(t *testing.T) {
resetAPIKeyState()
raw := "xce_ak_roundtrip"
entry := APIKeyEntry{
Username: "mcp-bot",
RoleName: "mcp_readonly",
CreateTime: time.Now(),
ExpiryTime: time.Now().Add(24 * time.Hour),
}
AddAPIKey(raw, entry)
got, ok := ValidateAPIKey(raw)
if !ok {
t.Fatal("expected ValidateAPIKey to succeed for a freshly-added key")
}
if got.Username != "mcp-bot" {
t.Errorf("Username = %q, want mcp-bot", got.Username)
}
if got.RoleName != "mcp_readonly" {
t.Errorf("RoleName = %q, want mcp_readonly", got.RoleName)
}
}
func TestValidateAPIKey_Unknown(t *testing.T) {
resetAPIKeyState()
_, ok := ValidateAPIKey("xce_ak_never-issued")
if ok {
t.Error("expected ValidateAPIKey to fail for a key that was never added")
}
}
func TestValidateAPIKey_Expired(t *testing.T) {
resetAPIKeyState()
raw := "xce_ak_expired"
AddAPIKey(raw, APIKeyEntry{
Username: "mcp-bot",
RoleName: "mcp_readonly",
CreateTime: time.Now().Add(-48 * time.Hour),
ExpiryTime: time.Now().Add(-1 * time.Hour), // already expired
})
_, ok := ValidateAPIKey(raw)
if ok {
t.Error("expected ValidateAPIKey to fail for an expired key")
}
}
func TestValidateAPIKey_StoredByHashNotPlaintext(t *testing.T) {
resetAPIKeyState()
raw := "xce_ak_hashed-storage"
AddAPIKey(raw, APIKeyEntry{Username: "mcp-bot", ExpiryTime: time.Now().Add(time.Hour)})
if APIKeys.Has(raw) {
t.Error("store must not contain the raw key as a map key")
}
if !APIKeys.Has(HashAPIKey(raw)) {
t.Error("store should contain the hashed key")
}
}

View file

@ -16,6 +16,8 @@ var (
UserRoles *DataStore[string, []string] // user_name → role_names
)
// APIKeys is declared in apikey.go (keeps API-key store logic self-contained).
// RoleEntry holds a role's permissions.
type RoleEntry struct {
RoleName string
@ -41,6 +43,7 @@ func InitStores() {
JobExecHosts = NewDataStore[int64, models.JobExecHosts]()
Roles = NewDataStore[string, RoleEntry]()
UserRoles = NewDataStore[string, []string]()
APIKeys = NewDataStore[string, APIKeyEntry]()
}
// Helper functions for RBAC.

View 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)
}
}

View 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
)

View 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=

View 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
}

View file

@ -0,0 +1,39 @@
package schedulerclient
import "context"
// ── Client interface ────────────────────────────────────────────────────────
// Client provides read-only access to the scheduler cluster through a
// scheduler-agnostic unified resource model (URM).
//
// Implementations talk to the apiserver's REST API and authenticate via
// either a Bearer JWT token or an XCE API Key (xce_ak_ prefix).
type Client interface {
// ── Jobs ──
// ListJobs returns jobs matching the given filter.
// An empty filter returns all visible jobs.
ListJobs(ctx context.Context, filter JobFilter) ([]*URMJob, *Page, error)
// GetJob returns a single job by its ID.
GetJob(ctx context.Context, id string) (*URMJob, error)
// GetJobLogs retrieves output logs for a job.
GetJobLogs(ctx context.Context, id string, opts LogOpts) ([]string, error)
// ── Queues ──
// ListQueues returns all visible queues with their current state.
ListQueues(ctx context.Context) ([]*URMQueue, error)
// ── Nodes ──
// ListNodes returns all visible compute nodes with resource usage.
ListNodes(ctx context.Context) ([]*URMNode, error)
// ── Metrics ──
// GetClusterMetrics returns a snapshot of overall cluster resource usage.
GetClusterMetrics(ctx context.Context) (*Metrics, error)
}

View file

@ -0,0 +1,7 @@
module schedulerclient
go 1.23.4
require resty.dev/v3 v3.0.0-beta.2
require golang.org/x/net v0.33.0 // indirect

View file

@ -0,0 +1,4 @@
golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
resty.dev/v3 v3.0.0-beta.2 h1:xu4mGAdbCLuc3kbk7eddWfWm4JfhwDtdapwss5nCjnQ=
resty.dev/v3 v3.0.0-beta.2/go.mod h1:OgkqiPvTDtOuV4MGZuUDhwOpkY8enjOsjjMzeOHefy4=

View file

@ -0,0 +1,439 @@
package schedulerclient
import (
"context"
"crypto/tls"
"encoding/json"
"fmt"
"math"
"net/http"
"strconv"
"strings"
"time"
"resty.dev/v3"
)
// ── Client construction ─────────────────────────────────────────────────────
// HTTPClient implements Client by calling the apiserver's REST API.
// Authenticates via either a JWT Bearer token or an XCE API Key (xce_ak_ prefix).
type HTTPClient struct {
rc *resty.Client
baseURL string
token string
}
// NewClient creates a new HTTPClient. The token is sent as the Authorization
// header value exactly as given (no "Bearer " prefix added — the apiserver
// strips it via its own TrimPrefix logic, so both forms work).
func NewClient(baseURL, token string) *HTTPClient {
c := resty.New()
c.SetTimeout(10 * time.Second)
c.SetTLSClientConfig(&tls.Config{InsecureSkipVerify: true})
return &HTTPClient{
rc: c,
baseURL: strings.TrimRight(baseURL, "/"),
token: token,
}
}
// ── Client interface implementation ─────────────────────────────────────────
func (hc *HTTPClient) ListJobs(ctx context.Context, filter JobFilter) ([]*URMJob, *Page, error) {
url := hc.baseURL + "/xce/v1/jobs"
req := hc.rc.R().SetContext(ctx).SetHeader("Authorization", hc.token)
if filter.Queue != "" {
req.SetQueryParam("queue_name", filter.Queue)
}
if filter.Owner != "" {
req.SetQueryParam("submit_user", filter.Owner)
}
if filter.Limit > 0 {
req.SetQueryParam("limit", strconv.Itoa(filter.Limit))
}
if filter.PageToken != "" {
req.SetQueryParam("page_token", filter.PageToken)
}
// Note: status filter on URM statuses is applied in post-processing
// because the apiserver works with raw status strings, not URM statuses.
return hc.getJobs(req, url, filter)
}
func (hc *HTTPClient) GetJob(ctx context.Context, id string) (*URMJob, error) {
url := hc.baseURL + "/xce/v1/jobs/" + id
req := hc.rc.R().SetContext(ctx).SetHeader("Authorization", hc.token)
jobs, _, err := hc.getJobs(req, url, JobFilter{})
if err != nil {
return nil, err
}
if len(jobs) == 0 {
return nil, fmt.Errorf("job %s not found", id)
}
return jobs[0], nil
}
func (hc *HTTPClient) GetJobLogs(ctx context.Context, id string, opts LogOpts) ([]string, error) {
url := hc.baseURL + "/xce/v1/jobs/" + id + "/output"
req := hc.rc.R().SetContext(ctx).SetHeader("Authorization", hc.token)
if opts.Tail > 0 {
req.SetQueryParam("tail", strconv.Itoa(opts.Tail))
}
if opts.OffsetBytes > 0 {
req.SetQueryParam("offset", strconv.FormatInt(opts.OffsetBytes, 10))
}
var env apiEnvelope
resp, err := req.SetResult(&env).Get(url)
if err != nil {
return nil, fmt.Errorf("get job logs: %w", err)
}
if resp.StatusCode() != http.StatusOK {
return nil, fmt.Errorf("get job logs: HTTP %d: %s", resp.StatusCode(), env.Msg)
}
if env.Code != http.StatusOK {
return nil, fmt.Errorf("get job logs: api error %d: %s", env.Code, env.Msg)
}
var out struct {
Content string `json:"content"`
JobStatus string `json:"job_status"`
}
if err := json.Unmarshal(env.Data, &out); err != nil {
return nil, fmt.Errorf("parse job logs: %w", err)
}
if out.Content == "" {
return []string{}, nil
}
return strings.Split(strings.TrimSuffix(out.Content, "\n"), "\n"), nil
}
func (hc *HTTPClient) ListQueues(ctx context.Context) ([]*URMQueue, error) {
url := hc.baseURL + "/xce/v1/queues"
req := hc.rc.R().SetContext(ctx).SetHeader("Authorization", hc.token)
var env apiEnvelope
resp, err := req.SetResult(&env).Get(url)
if err != nil {
return nil, fmt.Errorf("list queues: %w", err)
}
if resp.StatusCode() != http.StatusOK || env.Code != http.StatusOK {
return nil, fmt.Errorf("list queues: HTTP %d: %s", resp.StatusCode(), env.Msg)
}
var rows struct {
Count int `json:"count"`
Rows []apiQueue `json:"rows"`
}
if err := json.Unmarshal(env.Data, &rows); err != nil {
return nil, fmt.Errorf("parse queues: %w", err)
}
queues := make([]*URMQueue, 0, len(rows.Rows))
for i := range rows.Rows {
q := rows.Rows[i]
queues = append(queues, &URMQueue{
Name: q.QueueName,
Status: QueueStatusFromString(q.QueueStatus),
RunningJobs: q.NumRun,
PendingJobs: q.NumPend,
SuspendedJobs: q.NumSSusp + q.NumUSusp,
NodeCount: 0, // apiserver does not expose queue→node mapping directly
Extensions: map[string]any{
"raw_status": q.QueueStatus,
"priority": q.Priority,
"num_ssusp": q.NumSSusp,
"num_ususp": q.NumUSusp,
"num_done": q.NumDone,
},
})
}
return queues, nil
}
func (hc *HTTPClient) ListNodes(ctx context.Context) ([]*URMNode, error) {
url := hc.baseURL + "/xce/v1/hosts"
req := hc.rc.R().SetContext(ctx).SetHeader("Authorization", hc.token)
var env apiEnvelope
resp, err := req.SetResult(&env).Get(url)
if err != nil {
return nil, fmt.Errorf("list nodes: %w", err)
}
if resp.StatusCode() != http.StatusOK || env.Code != http.StatusOK {
return nil, fmt.Errorf("list nodes: HTTP %d: %s", resp.StatusCode(), env.Msg)
}
var rows struct {
Count int `json:"count"`
Rows []apiHost `json:"rows"`
}
if err := json.Unmarshal(env.Data, &rows); err != nil {
return nil, fmt.Errorf("parse nodes: %w", err)
}
nodes := make([]*URMNode, 0, len(rows.Rows))
for i := range rows.Rows {
h := rows.Rows[i]
nodes = append(nodes, &URMNode{
Name: h.HostName,
Status: NodeStatusFromString(h.HostStatus),
CoresTotal: h.NumCPUs,
CoresUsed: h.NumRun, // running jobs ≈ cores used
MemoryBytesTotal: int64(h.MaxMem) * mbToBytes,
MemoryBytesUsed: 0, // apiserver does not expose live used-memory for hosts
GPUCount: 0, // TODO: parse from extensions/topology when available
Features: nil,
Extensions: map[string]any{
"raw_status": h.HostStatus,
"host_type": h.HostType,
"host_model": h.HostModel,
"cpu_factor": h.CPUFactor,
"max_slot": h.MaxSlot,
"max_swap_mb": h.MaxSwap,
"max_tmp_mb": h.MaxTmp,
"num_run": h.NumRun,
"num_ssusp": h.NumSSusp,
"num_ususp": h.NumUSusp,
},
})
}
return nodes, nil
}
func (hc *HTTPClient) GetClusterMetrics(ctx context.Context) (*Metrics, error) {
// Fetch hosts to compute total/used resources.
nodes, err := hc.ListNodes(ctx)
if err != nil {
return nil, fmt.Errorf("cluster metrics: %w", err)
}
var (
totalCores, usedCores int
totalMem, usedMem int64
totalGPUs int
pendingJobs, runningJobs int
)
for _, n := range nodes {
totalCores += n.CoresTotal
usedCores += n.CoresUsed
totalMem += n.MemoryBytesTotal
// usedMem not directly available; try from extensions
}
// Also fetch jobs for running/pending counts — we do this as a cheap
// parallel operation to get job counts across all queues.
url := hc.baseURL + "/xce/v1/jobs"
var env apiEnvelope
req := hc.rc.R().SetContext(ctx).SetHeader("Authorization", hc.token)
resp, err := req.SetResult(&env).Get(url)
if err == nil && resp.StatusCode() == http.StatusOK && env.Code == http.StatusOK {
var r struct {
Rows []json.RawMessage `json:"rows"`
}
if json.Unmarshal(env.Data, &r) == nil {
for _, raw := range r.Rows {
var stub struct {
JobStatus string `json:"job_status"`
}
if json.Unmarshal(raw, &stub) == nil {
switch JobStatusFromString(stub.JobStatus) {
case JobStatusRunning:
runningJobs++
case JobStatusPending:
pendingJobs++
}
}
}
}
}
util := 0.0
if totalCores > 0 {
util = math.Round(float64(usedCores)/float64(totalCores)*100) / 100
}
return &Metrics{
TotalCores: totalCores,
UsedCores: usedCores,
TotalMemory: totalMem,
UsedMemory: usedMem,
TotalGPUs: totalGPUs,
UsedGPUs: 0,
RunningJobs: runningJobs,
PendingJobs: pendingJobs,
Utilization: util,
}, nil
}
// ── shared helper ───────────────────────────────────────────────────────────
// getJobs makes an HTTP GET and parses the job list response.
// When filter.Statuses is non-empty, results are filtered in Go after decoding.
func (hc *HTTPClient) getJobs(req *resty.Request, url string, filter JobFilter) ([]*URMJob, *Page, error) {
var env apiEnvelope
resp, err := req.SetResult(&env).Get(url)
if err != nil {
return nil, nil, fmt.Errorf("list jobs: %w", err)
}
if resp.StatusCode() != http.StatusOK || env.Code != http.StatusOK {
return nil, nil, fmt.Errorf("list jobs: HTTP %d: %s", resp.StatusCode(), env.Msg)
}
var rows struct {
Count int `json:"count"`
Rows []apiJob `json:"rows"`
PageToken string `json:"page_token"`
}
if err := json.Unmarshal(env.Data, &rows); err != nil {
return nil, nil, fmt.Errorf("parse jobs: %w", err)
}
jobs := make([]*URMJob, 0, len(rows.Rows))
for i := range rows.Rows {
j := rows.Rows[i]
urmStatus := JobStatusFromString(j.JobStatus)
// Apply URM status filter in post-processing
if len(filter.Statuses) > 0 {
matched := false
for _, s := range filter.Statuses {
if urmStatus == s {
matched = true
break
}
}
if !matched {
continue
}
}
// Parse exit code
var exitCode *int
if j.ExitStatus != "" {
if code, err := strconv.Atoi(j.ExitStatus); err == nil {
exitCode = &code
}
}
// Parse timestamps
var submitTime time.Time
if j.SubmitTime > 0 {
submitTime = time.Unix(j.SubmitTime, 0)
}
var startTime *time.Time
if j.StartTime > 0 {
t := time.Unix(j.StartTime, 0)
startTime = &t
}
// Walltime ≈ total CPU time available. For XCE jobs, run_limit
// is available as a string ("HH:MM") and MaxNumProcessors gives
// slot count. We store what we can.
walltime := 0 // TODO: parse run_limit string when available
// Memory: MaxMem from apiserver is in MB (standard HPC convention).
memBytes := int64(math.Round(j.MaxMem * mbToBytes))
urmJob := &URMJob{
ID: strconv.FormatInt(j.JobID, 10),
Name: j.JobName,
Status: urmStatus,
RawStatus: j.JobStatus,
Queue: j.QueueName,
Nodes: j.NumHosts,
Cores: j.NumProcessors,
GPUs: 0, // TODO: parse from resreq/extensions
MemoryBytes: memBytes,
WalltimeSeconds: walltime,
SubmissionTime: submitTime,
StartTime: startTime,
ExitCode: exitCode,
Owner: j.SubmitUser,
Extensions: map[string]any{
"command": j.Command,
"project_name": j.ProjectName,
"pending_reason": j.PendingReason,
"resreq": j.Resreq,
"exit_status_raw": j.ExitStatus,
"num_processors": j.NumProcessors,
"submit_time": j.SubmitTime,
"start_time": j.StartTime,
"end_time": j.EndTime,
},
}
jobs = append(jobs, urmJob)
}
page := &Page{
Total: rows.Count,
NextPageToken: rows.PageToken,
}
return jobs, page, nil
}
// ── internal types for JSON decoding (apiserver response shapes) ────────────
const mbToBytes = 1_000_000 // MaxMem from apiserver is in MB → stored as bytes in URM
// apiEnvelope is the standard top-level response: {"code":200,"msg":"ok","data":...}
type apiEnvelope struct {
Code int `json:"code"`
Msg string `json:"msg"`
Data json.RawMessage `json:"data"`
}
// apiJob mirrors the fields of apiserver's QueryJob that we need for URM conversion.
type apiJob struct {
JobID int64 `json:"job_id"`
SubmitUser string `json:"submit_user"`
JobStatus string `json:"job_status"`
JobName string `json:"job_name"`
Command string `json:"command"`
QueueName string `json:"queue_name"`
ProjectName string `json:"project_name"`
NumProcessors int `json:"num_processors"`
NumHosts int `json:"num_hosts"`
SubmitTime int64 `json:"submit_time"`
StartTime int64 `json:"start_time"`
EndTime int64 `json:"end_time"`
MaxMem float64 `json:"max_mem"`
ExitStatus string `json:"exit_status"`
PendingReason string `json:"pending_reason"`
Resreq string `json:"resreq"`
}
// apiQueue mirrors fields of DBQueueInfo needed for URM conversion.
type apiQueue struct {
QueueName string `json:"queue_name"`
QueueStatus string `json:"queue_status"`
Priority int `json:"priority"`
NumPend int `json:"num_pend"`
NumRun int `json:"num_run"`
NumSSusp int `json:"num_ssusp"`
NumUSusp int `json:"num_ususp"`
NumDone int `json:"num_done"`
}
// apiHost mirrors fields of XCEHost needed for URM conversion.
type apiHost struct {
HostName string `json:"host_name"`
HostStatus string `json:"host_status"`
HostType string `json:"host_type"`
HostModel string `json:"host_model"`
CPUFactor float64 `json:"cpu_factor"`
NumCPUs int `json:"num_cpus"`
MaxSlot int `json:"max_slot"`
MaxMem int `json:"max_mem"`
MaxSwap int `json:"max_swap"`
MaxTmp int `json:"max_tmp"`
NumRun int `json:"num_run"`
NumSSusp int `json:"num_ssusp"`
NumUSusp int `json:"num_ususp"`
NumDone int `json:"num_done"`
}

View file

@ -0,0 +1,93 @@
package schedulerclient
import "strings"
// ── JobStatus mapping ───────────────────────────────────────────────────────
// JobStatusFromString maps the apiserver's job status string to a URM JobStatus.
//
// The input is the output of apiserver models.JobStateStr() — one of:
//
// PEND, PSUSP, RUN, USUSP, SSUSP, RUNKWN, EXIT, DONE, UNKWN, ERR, TRANSIENT
func JobStatusFromString(s string) JobStatus {
switch s {
case "PEND", "PSUSP":
return JobStatusPending
case "RUN", "RUNKWN":
return JobStatusRunning
case "USUSP", "SSUSP":
return JobStatusSuspended
case "EXIT":
return JobStatusFailed
case "DONE":
return JobStatusSucceeded
default:
// UNKWN, ERR, TRANSIENT, empty, or unknown values → unknown
return JobStatusUnknown
}
}
// ── QueueStatus mapping ─────────────────────────────────────────────────────
// QueueStatusFromString maps the apiserver's composite queue status string
// (e.g. "open:active", "closed:inactive_adm") to a URM status:
//
// active — queue is open and accepting jobs
// inactive — queue is administratively or window-closed
// draining — open but has an inactive-time window active (will close)
// unknown — unrecognised format
func QueueStatusFromString(s string) string {
if s == "" {
return "unknown"
}
// Format is "<open|closed>:<active|inactive_adm|inactive_win>"
parts := strings.SplitN(s, ":", 2)
if len(parts) < 1 {
return "unknown"
}
openClosed := parts[0]
var inactiveKind string
if len(parts) > 1 {
inactiveKind = parts[1]
}
switch openClosed {
case "open":
if inactiveKind == "inactive_adm" || inactiveKind == "inactive_win" {
return "draining"
}
return "active"
case "closed":
return "inactive"
default:
return "unknown"
}
}
// ── NodeStatus mapping ─────────────────────────────────────────────────────
// NodeStatusFromString maps the apiserver's host status string (output of
// models.HostStateStr()) to a URM node status:
//
// idle — OK, accepting work
// busy — running jobs
// down — CLOSED, CLOSED_BY_ADMIN
// drain — LOCKU, LOCKW (not accepting new work, existing jobs finish)
// unknown — UNAVAIL or unrecognised
func NodeStatusFromString(s string) string {
switch s {
case "OK":
return "idle"
case "BUSY":
return "busy"
case "CLOSED", "CLOSED_BY_ADMIN":
return "down"
case "LOCKU", "LOCKW":
return "drain"
default:
// UNAVAIL, empty, or unknown → unknown
return "unknown"
}
}

View file

@ -0,0 +1,94 @@
package schedulerclient
import "testing"
// ── JobStatus mapping tests ─────────────────────────────────────────────────
func TestJobStatusFromString(t *testing.T) {
tests := []struct {
input string
want JobStatus
}{
// PEND → pending
{"PEND", JobStatusPending},
{"PSUSP", JobStatusPending}, // user-suspended before start → still pending
// RUN → running
{"RUN", JobStatusRunning},
{"RUNKWN", JobStatusRunning}, // running but state uncertain → treat as running
// suspended
{"USUSP", JobStatusSuspended}, // user-suspended while running
{"SSUSP", JobStatusSuspended}, // system-suspended while running
// terminal
{"EXIT", JobStatusFailed},
{"DONE", JobStatusSucceeded},
// unknown / error
{"UNKWN", JobStatusUnknown},
{"ERR", JobStatusUnknown},
// edge cases
{"TRANSIENT", JobStatusUnknown},
{"", JobStatusUnknown},
{"NONEXISTENT", JobStatusUnknown},
}
for _, tt := range tests {
t.Run(tt.input, func(t *testing.T) {
got := JobStatusFromString(tt.input)
if got != tt.want {
t.Errorf("JobStatusFromString(%q) = %q, want %q", tt.input, got, tt.want)
}
})
}
}
// ── QueueStatus mapping tests ────────────────────────────────────────────────
func TestQueueStatusFromString(t *testing.T) {
tests := []struct {
input string
want string
}{
{"open:active", "active"},
{"closed:active", "inactive"},
{"open:inactive_adm", "draining"},
{"open:inactive_win", "draining"},
{"closed:inactive_adm", "inactive"},
{"closed:inactive_win", "inactive"},
{"open", "active"},
{"closed", "inactive"},
{"", "unknown"},
}
for _, tt := range tests {
t.Run(tt.input, func(t *testing.T) {
got := QueueStatusFromString(tt.input)
if got != tt.want {
t.Errorf("QueueStatusFromString(%q) = %q, want %q", tt.input, got, tt.want)
}
})
}
}
// ── NodeStatus mapping tests ────────────────────────────────────────────────
func TestNodeStatusFromString(t *testing.T) {
tests := []struct {
input string
want string
}{
{"OK", "idle"},
{"BUSY", "busy"},
{"CLOSED", "down"},
{"CLOSED_BY_ADMIN", "down"},
{"LOCKU", "drain"},
{"LOCKW", "drain"},
{"UNAVAIL", "unknown"},
{"", "unknown"},
{"NONEXISTENT", "unknown"},
}
for _, tt := range tests {
t.Run(tt.input, func(t *testing.T) {
got := NodeStatusFromString(tt.input)
if got != tt.want {
t.Errorf("NodeStatusFromString(%q) = %q, want %q", tt.input, got, tt.want)
}
})
}
}

View file

@ -0,0 +1,104 @@
// Package schedulerclient provides a read-only 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.
package schedulerclient
import "time"
// ── Unified Resource Model (URM) types ──────────────────────────────────────
// JobStatus is the canonical status of a job across all scheduler backends.
type JobStatus string
const (
JobStatusPending JobStatus = "pending"
JobStatusRunning JobStatus = "running"
JobStatusSucceeded JobStatus = "succeeded"
JobStatusFailed JobStatus = "failed"
JobStatusCancelled JobStatus = "cancelled"
JobStatusSuspended JobStatus = "suspended"
JobStatusUnknown JobStatus = "unknown"
)
// URMJob is the unified representation of a scheduler job.
type URMJob struct {
ID string `json:"id"`
Name string `json:"name"`
Status JobStatus `json:"status"`
RawStatus string `json:"raw_status"` // backend-native status string
Queue string `json:"queue"`
Nodes int `json:"nodes"`
Cores int `json:"cores"`
GPUs int `json:"gpus"`
MemoryBytes int64 `json:"memory_bytes"`
WalltimeSeconds int `json:"walltime_seconds"`
SubmissionTime time.Time `json:"submission_time"`
StartTime *time.Time `json:"start_time,omitempty"`
ExitCode *int `json:"exit_code,omitempty"`
Owner string `json:"owner"`
LogsURL string `json:"logs_url,omitempty"`
Extensions map[string]any `json:"extensions,omitempty"`
}
// URMQueue is the unified representation of a scheduler queue/partition.
type URMQueue struct {
Name string `json:"name"`
Status string `json:"status"` // active | inactive | draining
RunningJobs int `json:"running_jobs"`
PendingJobs int `json:"pending_jobs"`
SuspendedJobs int `json:"suspended_jobs"`
NodeCount int `json:"node_count"`
MaxWalltimeSeconds int `json:"max_walltime_seconds"`
Extensions map[string]any `json:"extensions,omitempty"`
}
// URMNode is the unified representation of a compute node.
type URMNode struct {
Name string `json:"name"`
Status string `json:"status"` // idle | busy | down | drain | unknown
CoresTotal int `json:"cores_total"`
CoresUsed int `json:"cores_used"`
MemoryBytesTotal int64 `json:"memory_bytes_total"`
MemoryBytesUsed int64 `json:"memory_bytes_used"`
GPUCount int `json:"gpu_count"`
Features []string `json:"features"`
Extensions map[string]any `json:"extensions,omitempty"`
}
// ── Filter / input types ────────────────────────────────────────────────────
// JobFilter specifies criteria for listing jobs.
type JobFilter struct {
Statuses []JobStatus // filter by URM status(es)
Queue string // filter by queue name
Owner string // filter by owner username
Limit int // max results (0 = default)
PageToken string // pagination cursor from a prior response
}
// LogOpts specifies how to retrieve job logs.
type LogOpts struct {
Tail int // last N lines (0 = all)
OffsetBytes int64 // byte offset to start from
}
// Page carries pagination metadata returned with list operations.
type Page struct {
Total int `json:"total"`
NextPageToken string `json:"next_page_token,omitempty"`
}
// Metrics is a snapshot of overall cluster resource usage.
type Metrics struct {
TotalCores int `json:"total_cores"`
UsedCores int `json:"used_cores"`
TotalMemory int64 `json:"total_memory_bytes"`
UsedMemory int64 `json:"used_memory_bytes"`
TotalGPUs int `json:"total_gpus"`
UsedGPUs int `json:"used_gpus"`
RunningJobs int `json:"running_jobs"`
PendingJobs int `json:"pending_jobs"`
Utilization float64 `json:"utilization"` // 0.0 – 1.0
}