forked from lei.chen/myxce
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:
parent
98066db68d
commit
e4af0cac73
27 changed files with 1760 additions and 7 deletions
|
|
@ -6082,3 +6082,5 @@
|
|||
| 01:56 | Edited ../../../../../claude/.claude/projects/-home-rshen-projects-xce-xce/memory/MEMORY.md | 1→2 lines | ~142 |
|
||||
| 01:56 | Session end: 33 writes across 20 files (cluster.c, sched.c, limits.yaml, project_bug1462_config_replication_fixed.md, MEMORY.md) | 39 reads | ~71360 tok |
|
||||
| 02:11 | Session end: 33 writes across 20 files (cluster.c, sched.c, limits.yaml, project_bug1462_config_replication_fixed.md, MEMORY.md) | 39 reads | ~71360 tok |
|
||||
| 12:10 | Updated root CLAUDE.md: corrected Native Adapter Go-import assumption (myxce kernel is C/cgo, Go only at apiserver/xcecli edges), logged API Key auth is half-built, confirmed 2 decisions (code location, service-account API Key auth) | /Users/dankeer/Projects/chatbot/CLAUDE.md | done | ~600 |
|
||||
| 12:10 | Updated myxce/CLAUDE.md and src/frontend/CLAUDE.md: registered planned schedulerclient/ and mcp-server/ modules (Phase 0, in progress) pointing to root CLAUDE.md for overall plan | myxce/CLAUDE.md, myxce/src/frontend/CLAUDE.md | done | ~300 |
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ For details on any topic below, read the corresponding steering file in `.kiro/s
|
|||
- `src/worker/` — See `src/worker/CLAUDE.md`
|
||||
- `src/scheduler/` — See `src/scheduler/CLAUDE.md`
|
||||
- `src/lib/` — See `src/lib/CLAUDE.md`
|
||||
- `src/frontend/schedulerclient/`, `src/frontend/mcp-server/` — Phase 0 已完成(2026-07-26)。`schedulerclient` 是独立 Go library 提供只读 Client interface + resty HTTP 实现(33 单测)。`mcp-server` 是 MCP stdio 二进制(`github.com/mark3labs/mcp-go@v0.57.0`),通过 `schedulerclient` 调用 `apiserver` REST API,暴露 6 个只读 MCP Tool。Pure Go, `CGO_ENABLED=0`。总体架构与分阶段计划见根目录 `CLAUDE.md`("Unified Scheduler Agent Framework")。
|
||||
|
||||
# Agentic SDLC and Spec-Driven Development
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -148,4 +148,27 @@ roles:
|
|||
- 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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
176
src/frontend/apiserver/controllers/apikey_test.go
Normal file
176
src/frontend/apiserver/controllers/apikey_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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"])
|
||||
|
|
|
|||
52
src/frontend/apiserver/controllers/rbac_mcp_readonly_test.go
Normal file
52
src/frontend/apiserver/controllers/rbac_mcp_readonly_test.go
Normal 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
|
||||
}
|
||||
|
|
@ -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
|
||||
})
|
||||
|
|
|
|||
109
src/frontend/apiserver/middleware/validatehandler_test.go
Normal file
109
src/frontend/apiserver/middleware/validatehandler_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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"`
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
{
|
||||
|
|
|
|||
47
src/frontend/apiserver/store/apikey.go
Normal file
47
src/frontend/apiserver/store/apikey.go
Normal 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
|
||||
}
|
||||
94
src/frontend/apiserver/store/apikey_test.go
Normal file
94
src/frontend/apiserver/store/apikey_test.go
Normal 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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
86
src/frontend/mcp-server/cmd/main.go
Normal file
86
src/frontend/mcp-server/cmd/main.go
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"log"
|
||||
"os"
|
||||
|
||||
mcps "github.com/mark3labs/mcp-go/server"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
"mcp-server/internal"
|
||||
)
|
||||
|
||||
func main() {
|
||||
baseURL := os.Getenv("XCE_API_URL")
|
||||
token := os.Getenv("XCE_API_KEY")
|
||||
|
||||
if baseURL == "" {
|
||||
baseURL = "https://localhost:8080"
|
||||
log.Printf("XCE_API_URL not set, defaulting to %s", baseURL)
|
||||
}
|
||||
if token == "" {
|
||||
log.Fatal("XCE_API_KEY environment variable is required")
|
||||
}
|
||||
|
||||
app, err := internal.NewApp(baseURL, token)
|
||||
if err != nil {
|
||||
log.Fatalf("failed to initialize: %v", err)
|
||||
}
|
||||
|
||||
s := mcps.NewMCPServer(
|
||||
"scheduler-agent",
|
||||
"0.1.0",
|
||||
mcps.WithResourceCapabilities(true, true),
|
||||
mcps.WithLogging(),
|
||||
)
|
||||
|
||||
// ── scheduler_list_jobs ──
|
||||
s.AddTool(mcp.NewTool("scheduler_list_jobs",
|
||||
mcp.WithDescription("List jobs in the cluster. Supports filtering by queue, owner, status, and pagination."),
|
||||
mcp.WithString("queue", mcp.Description("Filter by queue name")),
|
||||
mcp.WithString("owner", mcp.Description("Filter by job owner username")),
|
||||
mcp.WithString("status", mcp.Description("Filter by status: pending|running|succeeded|failed|cancelled|suspended|unknown")),
|
||||
mcp.WithNumber("limit", mcp.Description("Max results to return")),
|
||||
mcp.WithString("page_token", mcp.Description("Pagination token from a prior response")),
|
||||
), app.HandleListJobs)
|
||||
|
||||
// ── scheduler_get_job ──
|
||||
s.AddTool(mcp.NewTool("scheduler_get_job",
|
||||
mcp.WithDescription("Get detailed information about a specific job by its ID."),
|
||||
mcp.WithString("id",
|
||||
mcp.Description("Job ID"),
|
||||
mcp.Required(),
|
||||
),
|
||||
), app.HandleGetJob)
|
||||
|
||||
// ── scheduler_list_queues ──
|
||||
s.AddTool(mcp.NewTool("scheduler_list_queues",
|
||||
mcp.WithDescription("List all queues/partitions and their current state (active, inactive, draining)."),
|
||||
), app.HandleListQueues)
|
||||
|
||||
// ── scheduler_list_nodes ──
|
||||
s.AddTool(mcp.NewTool("scheduler_list_nodes",
|
||||
mcp.WithDescription("List all compute nodes with their resource usage and status."),
|
||||
), app.HandleListNodes)
|
||||
|
||||
// ── scheduler_job_logs ──
|
||||
s.AddTool(mcp.NewTool("scheduler_job_logs",
|
||||
mcp.WithDescription("Retrieve the output logs of a job."),
|
||||
mcp.WithString("id",
|
||||
mcp.Description("Job ID"),
|
||||
mcp.Required(),
|
||||
),
|
||||
mcp.WithNumber("tail", mcp.Description("Number of lines from the end to return (0 = all)")),
|
||||
mcp.WithNumber("offset_bytes", mcp.Description("Byte offset to start reading from")),
|
||||
), app.HandleJobLogs)
|
||||
|
||||
// ── scheduler_cluster_metrics ──
|
||||
s.AddTool(mcp.NewTool("scheduler_cluster_metrics",
|
||||
mcp.WithDescription("Get a snapshot of overall cluster resource usage and utilization."),
|
||||
), app.HandleClusterMetrics)
|
||||
|
||||
log.Printf("Starting scheduler-agent MCP server (base-url: %s)", baseURL)
|
||||
if err := mcps.ServeStdio(s); err != nil {
|
||||
log.Fatalf("server error: %v", err)
|
||||
}
|
||||
}
|
||||
21
src/frontend/mcp-server/go.mod
Normal file
21
src/frontend/mcp-server/go.mod
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
module mcp-server
|
||||
|
||||
go 1.25.5
|
||||
|
||||
replace schedulerclient => ../schedulerclient
|
||||
|
||||
require (
|
||||
github.com/mark3labs/mcp-go v0.57.0
|
||||
schedulerclient v0.0.0-00010101000000-000000000000
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/google/jsonschema-go v0.4.2 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect
|
||||
github.com/spf13/cast v1.7.1 // indirect
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
|
||||
golang.org/x/net v0.33.0 // indirect
|
||||
golang.org/x/text v0.21.0 // indirect
|
||||
resty.dev/v3 v3.0.0-beta.2 // indirect
|
||||
)
|
||||
38
src/frontend/mcp-server/go.sum
Normal file
38
src/frontend/mcp-server/go.sum
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI=
|
||||
github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
|
||||
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
|
||||
github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8=
|
||||
github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
||||
github.com/mark3labs/mcp-go v0.57.0 h1:jzWKyCzdWnwnZt05cvcQQ+ngiUl2RnixXJa7Kj4qP1E=
|
||||
github.com/mark3labs/mcp-go v0.57.0/go.mod h1:+8WclSK1ZUweCP3hvktSji8n8ABG/95QaEkeVE/Uwas=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
||||
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
|
||||
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEVZGK7IN2kJkjTuQ=
|
||||
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU=
|
||||
github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
|
||||
github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
|
||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
|
||||
golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
|
||||
golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
|
||||
golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
|
||||
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
resty.dev/v3 v3.0.0-beta.2 h1:xu4mGAdbCLuc3kbk7eddWfWm4JfhwDtdapwss5nCjnQ=
|
||||
resty.dev/v3 v3.0.0-beta.2/go.mod h1:OgkqiPvTDtOuV4MGZuUDhwOpkY8enjOsjjMzeOHefy4=
|
||||
201
src/frontend/mcp-server/internal/app.go
Normal file
201
src/frontend/mcp-server/internal/app.go
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
// Package internal provides the scheduler-agent MCP tool handlers.
|
||||
package internal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
sclient "schedulerclient"
|
||||
)
|
||||
|
||||
// App wires the schedulerclient to the MCP tool handlers.
|
||||
type App struct {
|
||||
client sclient.Client
|
||||
}
|
||||
|
||||
// NewApp creates a new App with a schedulerclient connected to the given base URL.
|
||||
func NewApp(baseURL, token string) (*App, error) {
|
||||
hc := sclient.NewClient(baseURL, token)
|
||||
return &App{client: hc}, nil
|
||||
}
|
||||
|
||||
// ── helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
// textResult returns an MCP tool result containing the JSON representation of v.
|
||||
func textResult(v any) (*mcp.CallToolResult, error) {
|
||||
b, err := json.MarshalIndent(v, "", " ")
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("marshal error: %v", err)), nil
|
||||
}
|
||||
return mcp.NewToolResultText(string(b)), nil
|
||||
}
|
||||
|
||||
// textResultStr returns an MCP tool result wrapping a plain string.
|
||||
func textResultStr(s string) *mcp.CallToolResult {
|
||||
return mcp.NewToolResultText(s)
|
||||
}
|
||||
|
||||
// getString extracts a string parameter from an MCP tool request.
|
||||
func getString(req mcp.CallToolRequest, key string) string {
|
||||
args := req.GetArguments()
|
||||
raw, ok := args[key]
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
s, _ := raw.(string)
|
||||
return s
|
||||
}
|
||||
|
||||
// getNumber extracts a float64 parameter from an MCP tool request.
|
||||
func getNumber(req mcp.CallToolRequest, key string) float64 {
|
||||
args := req.GetArguments()
|
||||
raw, ok := args[key]
|
||||
if !ok {
|
||||
return 0
|
||||
}
|
||||
n, _ := raw.(float64)
|
||||
return n
|
||||
}
|
||||
|
||||
// ── tool handlers ───────────────────────────────────────────────────────────
|
||||
|
||||
func (a *App) HandleListJobs(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
filter := sclient.JobFilter{
|
||||
Queue: getString(req, "queue"),
|
||||
Owner: getString(req, "owner"),
|
||||
PageToken: getString(req, "page_token"),
|
||||
}
|
||||
if n := getNumber(req, "limit"); n > 0 {
|
||||
filter.Limit = int(n)
|
||||
}
|
||||
if s := getString(req, "status"); s != "" {
|
||||
filter.Statuses = parseStatuses(s)
|
||||
}
|
||||
|
||||
jobs, page, err := a.client.ListJobs(ctx, filter)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("list jobs failed: %v", err)), nil
|
||||
}
|
||||
|
||||
resp := map[string]any{
|
||||
"jobs": jobs,
|
||||
"page": page,
|
||||
}
|
||||
return textResult(resp)
|
||||
}
|
||||
|
||||
func (a *App) HandleGetJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
|
||||
job, err := a.client.GetJob(ctx, id)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("get job failed: %v", err)), nil
|
||||
}
|
||||
return textResult(job)
|
||||
}
|
||||
|
||||
func (a *App) HandleListQueues(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
queues, err := a.client.ListQueues(ctx)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("list queues failed: %v", err)), nil
|
||||
}
|
||||
return textResult(queues)
|
||||
}
|
||||
|
||||
func (a *App) HandleListNodes(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
nodes, err := a.client.ListNodes(ctx)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("list nodes failed: %v", err)), nil
|
||||
}
|
||||
return textResult(nodes)
|
||||
}
|
||||
|
||||
func (a *App) HandleJobLogs(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
opts := sclient.LogOpts{
|
||||
Tail: int(getNumber(req, "tail")),
|
||||
OffsetBytes: int64(getNumber(req, "offset_bytes")),
|
||||
}
|
||||
|
||||
lines, err := a.client.GetJobLogs(ctx, id, opts)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("get job logs failed: %v", err)), nil
|
||||
}
|
||||
return textResultStr(strconv.Itoa(len(lines)) + " lines\n" + joinLines(lines)), nil
|
||||
}
|
||||
|
||||
func (a *App) HandleClusterMetrics(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
metrics, err := a.client.GetClusterMetrics(ctx)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("get cluster metrics failed: %v", err)), nil
|
||||
}
|
||||
return textResult(metrics)
|
||||
}
|
||||
|
||||
// ── helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
// parseStatuses parses a comma-separated list of URM status strings.
|
||||
func parseStatuses(s string) []sclient.JobStatus {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
parts := splitComma(s)
|
||||
statuses := make([]sclient.JobStatus, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
statuses = append(statuses, sclient.JobStatus(p))
|
||||
}
|
||||
return statuses
|
||||
}
|
||||
|
||||
// joinLines joins log lines into a single string.
|
||||
func joinLines(lines []string) string {
|
||||
if len(lines) == 0 {
|
||||
return ""
|
||||
}
|
||||
if len(lines) == 1 {
|
||||
return lines[0]
|
||||
}
|
||||
// Estimate total length for capacity
|
||||
total := 0
|
||||
for _, l := range lines {
|
||||
total += len(l) + 1
|
||||
}
|
||||
b := make([]byte, 0, total)
|
||||
for i, l := range lines {
|
||||
if i > 0 {
|
||||
b = append(b, '\n')
|
||||
}
|
||||
b = append(b, []byte(l)...)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// splitComma splits a comma-separated string, trimming whitespace.
|
||||
func splitComma(s string) []string {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
var result []string
|
||||
start := 0
|
||||
for i := 0; i < len(s); i++ {
|
||||
if s[i] == ',' {
|
||||
if i > start {
|
||||
result = append(result, s[start:i])
|
||||
}
|
||||
start = i + 1
|
||||
}
|
||||
}
|
||||
if start < len(s) {
|
||||
result = append(result, s[start:])
|
||||
}
|
||||
return result
|
||||
}
|
||||
39
src/frontend/schedulerclient/client.go
Normal file
39
src/frontend/schedulerclient/client.go
Normal 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)
|
||||
}
|
||||
7
src/frontend/schedulerclient/go.mod
Normal file
7
src/frontend/schedulerclient/go.mod
Normal 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
|
||||
4
src/frontend/schedulerclient/go.sum
Normal file
4
src/frontend/schedulerclient/go.sum
Normal 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=
|
||||
439
src/frontend/schedulerclient/httpclient.go
Normal file
439
src/frontend/schedulerclient/httpclient.go
Normal 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"`
|
||||
}
|
||||
93
src/frontend/schedulerclient/statusmap.go
Normal file
93
src/frontend/schedulerclient/statusmap.go
Normal 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"
|
||||
}
|
||||
}
|
||||
94
src/frontend/schedulerclient/statusmap_test.go
Normal file
94
src/frontend/schedulerclient/statusmap_test.go
Normal 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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
104
src/frontend/schedulerclient/types.go
Normal file
104
src/frontend/schedulerclient/types.go
Normal 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
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue