fix(mcp): wire in-process scheduling

Use an explicit MCP scheduling DTO and route in-process scheduling calls through the distributed node registry so the embedded assistant matches the REST scheduling surface.

Assisted-by: Hephaestus:openai/gpt-5.5

Author: owezzy

core/application/application.go

@@ -553,12 +553,17 @@ func (a *Application) start() error {
 	// once at startup and reused across chat sessions that opt in via metadata.
 	if !a.applicationConfig.DisableLocalAIAssistant {
 		holder := mcpTools.NewLocalAIAssistantHolder()
+		var nodeRegistry *nodes.NodeRegistry
+		if a.distributed != nil {
+			nodeRegistry = a.distributed.Registry
+		}
 		assistantClient := localaiInproc.New(
 			a.applicationConfig,
 			a.applicationConfig.SystemState,
 			a.backendLoader,
 			a.modelLoader,
 			a.galleryService,
+			nodeRegistry,
 		)
 		// Wire usage tracking so the assistant's get_usage_stats tool
 		// returns real data; nil values keep the tool returning a clear

pkg/mcp/localaitools/dto.go

@@ -1,7 +1,8 @@
 package localaitools
 
 import (
-	"github.com/mudler/LocalAI/core/services/nodes"
+	"time"
+
 	"github.com/mudler/LocalAI/core/services/voiceprofile"
 )
 
@@ -119,9 +120,25 @@ type SetNodeVRAMBudgetRequest struct {
 }
 
 // ModelSchedulingConfig is the REST/MCP wire shape for one per-model
-// distributed scheduling rule. The NodeSelector field is a JSON object encoded
-// as a string, matching /api/nodes/scheduling responses.
-type ModelSchedulingConfig = nodes.ModelSchedulingConfig
+// distributed scheduling rule. Keep this DTO explicit instead of aliasing the
+// node-registry model so MCP contracts remain owned by localaitools while still
+// matching /api/nodes/scheduling JSON.
+type ModelSchedulingConfig struct {
+	ID                  string     `json:"id"`
+	ModelName           string     `json:"model_name"`
+	NodeSelector        string     `json:"node_selector,omitempty"`
+	MinReplicas         int        `json:"min_replicas"`
+	MaxReplicas         int        `json:"max_replicas"`
+	SpreadAll           bool       `json:"spread_all,omitempty"`
+	RoutePolicy         string     `json:"route_policy,omitempty"`
+	BalanceAbsThreshold int        `json:"balance_abs_threshold,omitempty"`
+	BalanceRelThreshold float64    `json:"balance_rel_threshold,omitempty"`
+	MinPrefixMatch      float64    `json:"min_prefix_match,omitempty"`
+	UnsatisfiableUntil  *time.Time `json:"unsatisfiable_until,omitempty"`
+	UnsatisfiableTicks  int        `json:"unsatisfiable_ticks"`
+	CreatedAt           time.Time  `json:"created_at"`
+	UpdatedAt           time.Time  `json:"updated_at"`
+}
 
 // SetSchedulingRequest is the input for set_scheduling. It mirrors
 // /api/nodes/scheduling so standalone MCP and REST callers preserve the same

pkg/mcp/localaitools/inproc/client.go

@@ -23,6 +23,8 @@ import (
 	"github.com/mudler/LocalAI/core/schema"
 	"github.com/mudler/LocalAI/core/services/galleryop"
 	"github.com/mudler/LocalAI/core/services/modeladmin"
+	"github.com/mudler/LocalAI/core/services/nodes"
+	"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
 	"github.com/mudler/LocalAI/core/services/routing/billing"
 	"github.com/mudler/LocalAI/core/services/routing/pii"
 	"github.com/mudler/LocalAI/core/services/routing/router"
@@ -47,6 +49,7 @@ type Client struct {
 	ConfigLoader  *config.ModelConfigLoader
 	ModelLoader   *model.ModelLoader
 	Gallery       *galleryop.GalleryService
+	NodeRegistry  *nodes.NodeRegistry
 	VoiceProfiles *voiceprofile.Store
 
 	// StatsRecorder and FallbackUser are optional — they back the
@@ -75,13 +78,18 @@ type Client struct {
 // except ModelLoader (used only for SystemInfo's loaded-models report and
 // best-effort ShutdownModel calls during config edits) and the stats
 // fields (StatsRecorder, FallbackUser) which gate get_usage_stats.
-func New(appConfig *config.ApplicationConfig, systemState *system.SystemState, cl *config.ModelConfigLoader, ml *model.ModelLoader, gs *galleryop.GalleryService) *Client {
+func New(appConfig *config.ApplicationConfig, systemState *system.SystemState, cl *config.ModelConfigLoader, ml *model.ModelLoader, gs *galleryop.GalleryService, registries ...*nodes.NodeRegistry) *Client {
+	var registry *nodes.NodeRegistry
+	if len(registries) > 0 {
+		registry = registries[0]
+	}
 	return &Client{
 		AppConfig:     appConfig,
 		SystemState:   systemState,
 		ConfigLoader:  cl,
 		ModelLoader:   ml,
 		Gallery:       gs,
+		NodeRegistry:  registry,
 		VoiceProfiles: voiceprofile.NewStore(appConfig.DataPath),
 		modelAdmin:    modeladmin.NewConfigService(cl, appConfig),
 	}
@@ -501,20 +509,119 @@ func (c *Client) ListNodes(_ context.Context) ([]localaitools.Node, error) {
 	return []localaitools.Node{}, nil
 }
 
-func (c *Client) ListScheduling(_ context.Context) ([]localaitools.ModelSchedulingConfig, error) {
-	return []localaitools.ModelSchedulingConfig{}, nil
+func (c *Client) ListScheduling(ctx context.Context) ([]localaitools.ModelSchedulingConfig, error) {
+	if c.NodeRegistry == nil {
+		return []localaitools.ModelSchedulingConfig{}, nil
+	}
+	configs, err := c.NodeRegistry.ListModelSchedulings(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return localaitools.SchedulingConfigsFromNodes(configs), nil
 }
 
-func (c *Client) GetScheduling(_ context.Context, _ string) (*localaitools.ModelSchedulingConfig, error) {
-	return nil, errors.New("model scheduling is only available in distributed mode")
+func (c *Client) GetScheduling(ctx context.Context, modelName string) (*localaitools.ModelSchedulingConfig, error) {
+	if modelName == "" {
+		return nil, errors.New("model_name is required")
+	}
+	if c.NodeRegistry == nil {
+		return nil, errors.New("model scheduling is only available in distributed mode")
+	}
+	config, err := c.NodeRegistry.GetModelScheduling(ctx, modelName)
+	if err != nil {
+		return nil, err
+	}
+	if config == nil {
+		return nil, nil
+	}
+	out := localaitools.SchedulingConfigFromNode(*config)
+	return &out, nil
 }
 
-func (c *Client) SetScheduling(_ context.Context, _ localaitools.SetSchedulingRequest) (*localaitools.ModelSchedulingConfig, error) {
-	return nil, errors.New("model scheduling is only available in distributed mode")
+func (c *Client) SetScheduling(ctx context.Context, req localaitools.SetSchedulingRequest) (*localaitools.ModelSchedulingConfig, error) {
+	if req.ModelName == "" {
+		return nil, errors.New("model_name is required")
+	}
+	if c.NodeRegistry == nil {
+		return nil, errors.New("model scheduling is only available in distributed mode")
+	}
+
+	existing, err := c.NodeRegistry.GetModelScheduling(ctx, req.ModelName)
+	if err != nil {
+		return nil, fmt.Errorf("load existing scheduling config: %w", err)
+	}
+	routePolicy := ""
+	absThr := 0
+	relThr := 0.0
+	minMatch := 0.0
+	if existing != nil {
+		routePolicy = existing.RoutePolicy
+		absThr = existing.BalanceAbsThreshold
+		relThr = existing.BalanceRelThreshold
+		minMatch = existing.MinPrefixMatch
+	}
+	if req.RoutePolicy != nil {
+		routePolicy = *req.RoutePolicy
+	}
+	if req.BalanceAbsThreshold != nil {
+		absThr = *req.BalanceAbsThreshold
+	}
+	if req.BalanceRelThreshold != nil {
+		relThr = *req.BalanceRelThreshold
+	}
+	if req.MinPrefixMatch != nil {
+		minMatch = *req.MinPrefixMatch
+	}
+	if req.SpreadAll && (req.MinReplicas != 0 || req.MaxReplicas != 0) {
+		return nil, errors.New("spread_all and min_replicas/max_replicas are mutually exclusive")
+	}
+	if req.MinReplicas < 0 {
+		return nil, errors.New("min_replicas must be >= 0")
+	}
+	if req.MaxReplicas < 0 {
+		return nil, errors.New("max_replicas must be >= 0")
+	}
+	if req.MaxReplicas > 0 && req.MinReplicas > req.MaxReplicas {
+		return nil, errors.New("min_replicas must be <= max_replicas")
+	}
+	if err := prefixcache.ValidateThresholds(routePolicy, absThr, relThr, minMatch); err != nil {
+		return nil, err
+	}
+
+	var selectorJSON string
+	if len(req.NodeSelector) > 0 {
+		b, err := json.Marshal(req.NodeSelector)
+		if err != nil {
+			return nil, fmt.Errorf("invalid node_selector: %w", err)
+		}
+		selectorJSON = string(b)
+	}
+	config := &nodes.ModelSchedulingConfig{
+		ModelName:           req.ModelName,
+		NodeSelector:        selectorJSON,
+		MinReplicas:         req.MinReplicas,
+		MaxReplicas:         req.MaxReplicas,
+		SpreadAll:           req.SpreadAll,
+		RoutePolicy:         routePolicy,
+		BalanceAbsThreshold: absThr,
+		BalanceRelThreshold: relThr,
+		MinPrefixMatch:      minMatch,
+	}
+	if err := c.NodeRegistry.SetModelScheduling(ctx, config); err != nil {
+		return nil, err
+	}
+	out := localaitools.SchedulingConfigFromNode(*config)
+	return &out, nil
 }
 
-func (c *Client) DeleteScheduling(_ context.Context, _ string) error {
-	return errors.New("model scheduling is only available in distributed mode")
+func (c *Client) DeleteScheduling(ctx context.Context, modelName string) error {
+	if modelName == "" {
+		return errors.New("model_name is required")
+	}
+	if c.NodeRegistry == nil {
+		return errors.New("model scheduling is only available in distributed mode")
+	}
+	return c.NodeRegistry.DeleteModelScheduling(ctx, modelName)
 }
 
 func (c *Client) SetNodeVRAMBudget(_ context.Context, _, _ string) error {

pkg/mcp/localaitools/inproc/client_test.go

@@ -13,8 +13,11 @@ import (
 	"github.com/mudler/LocalAI/core/config"
 	"github.com/mudler/LocalAI/core/gallery"
 	"github.com/mudler/LocalAI/core/services/galleryop"
+	"github.com/mudler/LocalAI/core/services/nodes"
 	localaitools "github.com/mudler/LocalAI/pkg/mcp/localaitools"
 	"github.com/mudler/LocalAI/pkg/system"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
 )
 
 // Regression spec for the bug we fixed when channel sends were
@@ -124,3 +127,69 @@ var _ = Describe("inproc.Client model aliases", func() {
 		})
 	})
 })
+
+var _ = Describe("inproc.Client model scheduling", func() {
+	var (
+		ctx      context.Context
+		registry *nodes.NodeRegistry
+		c        *Client
+		stringp  = func(s string) *string { return &s }
+		floatp   = func(f float64) *float64 { return &f }
+	)
+
+	BeforeEach(func() {
+		ctx = context.Background()
+		db, err := gorm.Open(sqlite.Open(filepath.Join(GinkgoT().TempDir(), "nodes.db")), &gorm.Config{})
+		Expect(err).ToNot(HaveOccurred())
+		registry, err = nodes.NewNodeRegistry(db)
+		Expect(err).ToNot(HaveOccurred())
+
+		tempDir := GinkgoT().TempDir()
+		systemState, err := system.GetSystemState(system.WithModelPath(tempDir))
+		Expect(err).ToNot(HaveOccurred())
+		appConfig := config.NewApplicationConfig(config.WithSystemState(systemState))
+		c = New(appConfig, systemState, config.NewModelConfigLoader(tempDir), nil, nil, registry)
+	})
+
+	It("sets, lists, gets, merges, and deletes scheduling configs through the registry", func() {
+		created, err := c.SetScheduling(ctx, localaitools.SetSchedulingRequest{
+			ModelName:      "qwen",
+			NodeSelector:   map[string]string{"gpu": "nvidia"},
+			MinReplicas:    1,
+			MaxReplicas:    2,
+			RoutePolicy:    stringp("prefix_cache"),
+			MinPrefixMatch: floatp(0.4),
+		})
+		Expect(err).ToNot(HaveOccurred())
+		Expect(created.ModelName).To(Equal("qwen"))
+		Expect(created.NodeSelector).To(Equal(`{"gpu":"nvidia"}`))
+		Expect(created.RoutePolicy).To(Equal("prefix_cache"))
+		Expect(created.MinPrefixMatch).To(Equal(0.4))
+
+		listed, err := c.ListScheduling(ctx)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(listed).To(HaveLen(1))
+		Expect(listed[0].ModelName).To(Equal("qwen"))
+
+		updated, err := c.SetScheduling(ctx, localaitools.SetSchedulingRequest{
+			ModelName:   "qwen",
+			MinReplicas: 2,
+			MaxReplicas: 3,
+		})
+		Expect(err).ToNot(HaveOccurred())
+		Expect(updated.MinReplicas).To(Equal(2))
+		Expect(updated.MaxReplicas).To(Equal(3))
+		Expect(updated.RoutePolicy).To(Equal("prefix_cache"))
+		Expect(updated.MinPrefixMatch).To(Equal(0.4))
+
+		got, err := c.GetScheduling(ctx, "qwen")
+		Expect(err).ToNot(HaveOccurred())
+		Expect(got).ToNot(BeNil())
+		Expect(got.MaxReplicas).To(Equal(3))
+
+		Expect(c.DeleteScheduling(ctx, "qwen")).To(Succeed())
+		missing, err := c.GetScheduling(ctx, "qwen")
+		Expect(err).ToNot(HaveOccurred())
+		Expect(missing).To(BeNil())
+	})
+})

pkg/mcp/localaitools/scheduling.go

@@ -0,0 +1,30 @@
+package localaitools
+
+import "github.com/mudler/LocalAI/core/services/nodes"
+
+func SchedulingConfigFromNode(config nodes.ModelSchedulingConfig) ModelSchedulingConfig {
+	return ModelSchedulingConfig{
+		ID:                  config.ID,
+		ModelName:           config.ModelName,
+		NodeSelector:        config.NodeSelector,
+		MinReplicas:         config.MinReplicas,
+		MaxReplicas:         config.MaxReplicas,
+		SpreadAll:           config.SpreadAll,
+		RoutePolicy:         config.RoutePolicy,
+		BalanceAbsThreshold: config.BalanceAbsThreshold,
+		BalanceRelThreshold: config.BalanceRelThreshold,
+		MinPrefixMatch:      config.MinPrefixMatch,
+		UnsatisfiableUntil:  config.UnsatisfiableUntil,
+		UnsatisfiableTicks:  config.UnsatisfiableTicks,
+		CreatedAt:           config.CreatedAt,
+		UpdatedAt:           config.UpdatedAt,
+	}
+}
+
+func SchedulingConfigsFromNodes(configs []nodes.ModelSchedulingConfig) []ModelSchedulingConfig {
+	out := make([]ModelSchedulingConfig, 0, len(configs))
+	for _, config := range configs {
+		out = append(out, SchedulingConfigFromNode(config))
+	}
+	return out
+}