myxce/src/frontend/mcp-server/internal/templates.go

package internal

import (
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"strings"

	sclient "schedulerclient"
)

// ── job templates ────────────────────────────────────────────────────────────

// TemplateGeneralSerial is the "通用串行脚本" template — submit a shell
// script/command to a queue with optional resource limits.
const TemplateGeneralSerial = "general_serial"

// JobTemplate resolves raw LLM tool arguments into a validated JobSpec.
// Implementations read only the keys they need, so future templates
// (GPU training, job arrays, dependency chains) add new params without
// touching tool registration in cmd/main.go — they register here instead.
type JobTemplate interface {
	Name() string
	Description() string
	Resolve(args map[string]any) (*sclient.JobSpec, error)
}

// TemplateRegistry holds all registered job templates.
type TemplateRegistry struct {
	m map[string]JobTemplate
}

// NewTemplateRegistry returns a registry seeded with the built-in templates.
func NewTemplateRegistry() *TemplateRegistry {
	r := &TemplateRegistry{m: make(map[string]JobTemplate)}
	r.Register(generalSerialTemplate{})
	return r
}

// Register adds a template, overwriting any existing template with the same name.
func (r *TemplateRegistry) Register(t JobTemplate) {
	r.m[t.Name()] = t
}

// Get returns the template with the given name.
func (r *TemplateRegistry) Get(name string) (JobTemplate, bool) {
	t, ok := r.m[name]
	return t, ok
}

// Has reports whether a template with the given name is registered.
func (r *TemplateRegistry) Has(name string) bool {
	_, ok := r.m[name]
	return ok
}

// Names returns the sorted list of registered template names.
func (r *TemplateRegistry) Names() []string {
	names := make([]string, 0, len(r.m))
	for n := range r.m {
		names = append(names, n)
	}
	// deterministic order for error messages and tool enums
	sortStrings(names)
	return names
}

func sortStrings(s []string) {
	for i := 1; i < len(s); i++ {
		for j := i; j > 0 && s[j-1] > s[j]; j-- {
			s[j-1], s[j] = s[j], s[j-1]
		}
	}
}

// ── general_serial ───────────────────────────────────────────────────────────

// generalSerialTemplate implements the "通用串行脚本" template: a single shell
// command/script run on one or more processors in a queue.
type generalSerialTemplate struct{}

func (generalSerialTemplate) Name() string { return TemplateGeneralSerial }

func (generalSerialTemplate) Description() string {
	return "提交一个 shell 脚本/命令（通用串行作业）：command 必填，可选 queue/cores/walltime/memory/cwd/env 等"
}

func (generalSerialTemplate) Resolve(args map[string]any) (*sclient.JobSpec, error) {
	command := argString(args, "command")
	if strings.TrimSpace(command) == "" {
		return nil, errors.New("scheduler_submit: command is required")
	}

	spec := &sclient.JobSpec{
		Name:           argString(args, "name"),
		Command:        command,
		Queue:          firstNonEmpty(argString(args, "queue"), "normal"),
		Project:        argString(args, "project"),
		CWD:            argString(args, "cwd"),
		OutFile:        argString(args, "out_file"),
		ErrFile:        argString(args, "err_file"),
		JobDescription: argString(args, "description"),
	}

	if cores := argNumber(args, "cores"); cores > 0 {
		spec.NumProcessors = int(cores)
	}

	if mem := argString(args, "memory"); mem != "" {
		bytes, err := parseMemoryToBytes(mem)
		if err != nil {
			return nil, fmt.Errorf("scheduler_submit: invalid memory %q: %w", mem, err)
		}
		spec.MemoryBytes = bytes
	}

	if wt := argString(args, "walltime"); wt != "" {
		secs, err := parseWalltimeSeconds(wt)
		if err != nil {
			return nil, fmt.Errorf("scheduler_submit: invalid walltime %q: %w", wt, err)
		}
		spec.WalltimeSeconds = secs
	}

	if env := argString(args, "env"); env != "" {
		for _, kv := range strings.Split(env, ",") {
			kv = strings.TrimSpace(kv)
			if kv != "" {
				spec.Env = append(spec.Env, kv)
			}
		}
	}

	return spec, nil
}

// ── argument helpers ─────────────────────────────────────────────────────────

func argString(args map[string]any, key string) string {
	raw, ok := args[key]
	if !ok {
		return ""
	}
	s, _ := raw.(string)
	return s
}

func argNumber(args map[string]any, key string) float64 {
	raw, ok := args[key]
	if !ok {
		return 0
	}
	n, _ := raw.(float64)
	return n
}

func firstNonEmpty(vals ...string) string {
	for _, v := range vals {
		if v != "" {
			return v
		}
	}
	return ""
}

// ── unit parsers (pure functions, unit-testable) ─────────────────────────────

// memoryRe matches "<number> <unit>", "<number><unit>", and "<number> <unit>"
// with optional whitespace between value and unit.
var memoryRe = regexp.MustCompile(`^([0-9]*\.?[0-9]+)\s*([A-Za-z]+)?$`)

// parseMemoryToBytes parses "512 MB" / "512M" / "4 GB" / "1.5 G" into bytes
// using the package's decimal convention (1 MB = 1_000_000 bytes, matching
// schedulerclient.mbToBytes) so "4 GB" round-trips to "4000 MB" on the wire.
// A bare number without a unit is rejected (unit is required).
func parseMemoryToBytes(s string) (int64, error) {
	s = strings.TrimSpace(s)
	if s == "" {
		return 0, errors.New("empty memory value")
	}
	m := memoryRe.FindStringSubmatch(s)
	if m == nil {
		return 0, fmt.Errorf("expected '<number> <unit>'")
	}
	value, err := strconv.ParseFloat(m[1], 64)
	if err != nil {
		return 0, fmt.Errorf("invalid number %q", m[1])
	}

	unit := strings.ToUpper(m[2])
	var mult float64
	switch unit {
	case "MB", "M":
		mult = 1
	case "GB", "G":
		mult = 1_000
	case "TB", "T":
		mult = 1_000_000
	default:
		return 0, fmt.Errorf("unsupported unit %q", m[2])
	}
	return int64(value * mult * 1_000_000), nil
}

// parseWalltimeSeconds parses "HH:MM" or "H:MM" into seconds. Bare integers
// are rejected — the apiserver's run_limit treats bare ints as MINUTES, which
// would silently change semantics, so the template requires explicit H:MM.
func parseWalltimeSeconds(s string) (int, error) {
	s = strings.TrimSpace(s)
	parts := strings.Split(s, ":")
	if len(parts) != 2 {
		return 0, fmt.Errorf("expected 'HH:MM', got %q", s)
	}
	hours, err := strconv.Atoi(parts[0])
	if err != nil || hours < 0 || hours > 24 {
		return 0, fmt.Errorf("hour must be between 0 and 24, got %q", parts[0])
	}
	minutes, err := strconv.Atoi(parts[1])
	if err != nil || minutes < 0 || minutes > 59 {
		return 0, fmt.Errorf("minute must be between 0 and 59, got %q", parts[1])
	}
	if hours == 24 && minutes != 0 {
		return 0, fmt.Errorf("24:00 is the max; got %q", s)
	}
	return hours*3600 + minutes*60, nil
}