feat(mcp): register scheduling tools

Assisted-by: Hephaestus:openai/gpt-5.5
Signed-off-by: Owen Adirah <owenadira@gmail.com>

Author: owezzy

pkg/mcp/localaitools/fakes_test.go

@@ -42,6 +42,10 @@ type fakeClient struct {
 	upgradeBackend      func(string) (string, error)
 	systemInfo          func() (*SystemInfo, error)
 	listNodes           func() ([]Node, error)
+	listScheduling      func() ([]ModelSchedulingConfig, error)
+	getScheduling       func(string) (*ModelSchedulingConfig, error)
+	setScheduling       func(SetSchedulingRequest) (*ModelSchedulingConfig, error)
+	deleteScheduling    func(string) error
 	setNodeVRAMBudget   func(string, string) error
 	vramEstimate        func(VRAMEstimateRequest) (*vram.EstimateResult, error)
 	toggleModelState    func(string, modeladmin.Action) error
@@ -230,6 +234,38 @@ func (f *fakeClient) ListNodes(_ context.Context) ([]Node, error) {
 	return nil, nil
 }
 
+func (f *fakeClient) ListScheduling(_ context.Context) ([]ModelSchedulingConfig, error) {
+	f.record("ListScheduling", nil)
+	if f.listScheduling != nil {
+		return f.listScheduling()
+	}
+	return []ModelSchedulingConfig{}, nil
+}
+
+func (f *fakeClient) GetScheduling(_ context.Context, modelName string) (*ModelSchedulingConfig, error) {
+	f.record("GetScheduling", modelName)
+	if f.getScheduling != nil {
+		return f.getScheduling(modelName)
+	}
+	return &ModelSchedulingConfig{ModelName: modelName}, nil
+}
+
+func (f *fakeClient) SetScheduling(_ context.Context, req SetSchedulingRequest) (*ModelSchedulingConfig, error) {
+	f.record("SetScheduling", req)
+	if f.setScheduling != nil {
+		return f.setScheduling(req)
+	}
+	return &ModelSchedulingConfig{ModelName: req.ModelName, MinReplicas: req.MinReplicas, MaxReplicas: req.MaxReplicas, SpreadAll: req.SpreadAll}, nil
+}
+
+func (f *fakeClient) DeleteScheduling(_ context.Context, modelName string) error {
+	f.record("DeleteScheduling", modelName)
+	if f.deleteScheduling != nil {
+		return f.deleteScheduling(modelName)
+	}
+	return nil
+}
+
 func (f *fakeClient) SetNodeVRAMBudget(_ context.Context, nodeID, budget string) error {
 	f.record("SetNodeVRAMBudget", []any{nodeID, budget})
 	if f.setNodeVRAMBudget != nil {

pkg/mcp/localaitools/server.go

@@ -47,6 +47,7 @@ func NewServer(client LocalAIClient, opts Options) *mcp.Server {
 	registerBackendTools(srv, client, opts)
 	registerConfigTools(srv, client, opts)
 	registerSystemTools(srv, client, opts)
+	registerSchedulingTools(srv, client, opts)
 	registerStateTools(srv, client, opts)
 	registerBrandingTools(srv, client, opts)
 	registerVoiceProfileTools(srv, client, opts)

pkg/mcp/localaitools/server_test.go

@@ -92,6 +92,7 @@ var expectedFullCatalog = sortedStrings(
 	ToolListInstalledModels,
 	ToolListKnownBackends,
 	ToolListNodes,
+	ToolListScheduling,
 	ToolListVoiceProfiles,
 	ToolLoadModel,
 	ToolReloadModels,
@@ -105,6 +106,9 @@ var expectedFullCatalog = sortedStrings(
 	ToolCreateVoiceProfile,
 	ToolDeleteVoiceProfile,
 	ToolSetNodeVRAMBudget,
+	ToolSetScheduling,
+	ToolDeleteScheduling,
+	ToolGetScheduling,
 )
 
 // expectedReadOnlyCatalog is the tool set when DisableMutating=true. Sorted.
@@ -123,6 +127,8 @@ var expectedReadOnlyCatalog = sortedStrings(
 	ToolListInstalledModels,
 	ToolListKnownBackends,
 	ToolListNodes,
+	ToolListScheduling,
+	ToolGetScheduling,
 	ToolListVoiceProfiles,
 	ToolSystemInfo,
 	ToolVRAMEstimate,
@@ -165,6 +171,10 @@ var _ = Describe("Tool dispatch", func() {
 		{ToolListKnownBackends, struct{}{}, "ListKnownBackends"},
 		{ToolSystemInfo, struct{}{}, "SystemInfo"},
 		{ToolListNodes, struct{}{}, "ListNodes"},
+		{ToolListScheduling, struct{}{}, "ListScheduling"},
+		{ToolGetScheduling, DeleteSchedulingRequest{ModelName: "qwen"}, "GetScheduling"},
+		{ToolSetScheduling, SetSchedulingRequest{ModelName: "qwen", MinReplicas: 1, MaxReplicas: 2}, "SetScheduling"},
+		{ToolDeleteScheduling, DeleteSchedulingRequest{ModelName: "qwen"}, "DeleteScheduling"},
 		{ToolListVoiceProfiles, struct{}{}, "ListVoiceProfiles"},
 		{ToolInstallModel, InstallModelRequest{ModelName: "test/foo"}, "InstallModel"},
 		{ToolImportModelURI, ImportModelURIRequest{URI: "Qwen/Qwen3-4B-GGUF"}, "ImportModelURI"},
@@ -219,6 +229,40 @@ var _ = Describe("Tool error surfacing", func() {
 	})
 })
 
+var _ = Describe("Scheduling tool behavior", func() {
+	It("round-trips list/get/set/delete over the MCP surface", func() {
+		fc := &fakeClient{
+			listScheduling: func() ([]ModelSchedulingConfig, error) {
+				return []ModelSchedulingConfig{{ModelName: "qwen", MinReplicas: 1, MaxReplicas: 2}}, nil
+			},
+			getScheduling: func(modelName string) (*ModelSchedulingConfig, error) {
+				return &ModelSchedulingConfig{ModelName: modelName, SpreadAll: true}, nil
+			},
+			setScheduling: func(req SetSchedulingRequest) (*ModelSchedulingConfig, error) {
+				return &ModelSchedulingConfig{ModelName: req.ModelName, MinReplicas: req.MinReplicas, MaxReplicas: req.MaxReplicas}, nil
+			},
+		}
+		ctx, sess, done := connectInMemory(fc, Options{})
+		DeferCleanup(done)
+
+		listRes := callTool(ctx, sess, ToolListScheduling, struct{}{})
+		Expect(listRes.IsError).To(BeFalse(), resultText(listRes))
+		Expect(resultText(listRes)).To(ContainSubstring(`"model_name": "qwen"`))
+
+		getRes := callTool(ctx, sess, ToolGetScheduling, DeleteSchedulingRequest{ModelName: "qwen"})
+		Expect(getRes.IsError).To(BeFalse(), resultText(getRes))
+		Expect(resultText(getRes)).To(ContainSubstring(`"spread_all": true`))
+
+		setRes := callTool(ctx, sess, ToolSetScheduling, SetSchedulingRequest{ModelName: "qwen", MinReplicas: 1, MaxReplicas: 2})
+		Expect(setRes.IsError).To(BeFalse(), resultText(setRes))
+		Expect(resultText(setRes)).To(ContainSubstring(`"max_replicas": 2`))
+
+		deleteRes := callTool(ctx, sess, ToolDeleteScheduling, DeleteSchedulingRequest{ModelName: "qwen"})
+		Expect(deleteRes.IsError).To(BeFalse(), resultText(deleteRes))
+		Expect(resultText(deleteRes)).To(ContainSubstring(`"deleted": "qwen"`))
+	})
+})
+
 var _ = Describe("Argument validation", func() {
 	type validationCase struct {
 		desc string
@@ -235,6 +279,9 @@ var _ = Describe("Argument validation", func() {
 		{"toggle_model_state rejects unknown action", ToolToggleModelState, map[string]any{"name": "foo", "action": "noop"}, "action must be one of"},
 		{"edit_model_config rejects empty patch", ToolEditModelConfig, map[string]any{"name": "foo", "patch": map[string]any{}}, "patch is required"},
 		{"create_voice_profile requires consent", ToolCreateVoiceProfile, CreateVoiceProfileRequest{Name: "Voice", Transcript: "words", AudioBase64: "UklGRg=="}, "consent_confirmed must be true"},
+		{"set_scheduling requires model_name", ToolSetScheduling, SetSchedulingRequest{}, "model_name is required"},
+		{"set_scheduling rejects invalid replica range", ToolSetScheduling, SetSchedulingRequest{ModelName: "qwen", MinReplicas: 3, MaxReplicas: 1}, "min_replicas must be <= max_replicas"},
+		{"delete_scheduling requires model_name", ToolDeleteScheduling, DeleteSchedulingRequest{}, "model_name is required"},
 	}
 
 	for _, c := range cases {

pkg/mcp/localaitools/tools_scheduling.go

@@ -0,0 +1,77 @@
+package localaitools
+
+import (
+	"context"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+func registerSchedulingTools(s *mcp.Server, client LocalAIClient, opts Options) {
+	mcp.AddTool(s, &mcp.Tool{
+		Name:        ToolListScheduling,
+		Description: "List distributed per-model scheduling configs (only meaningful in distributed mode).",
+	}, func(ctx context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) {
+		configs, err := client.ListScheduling(ctx)
+		if err != nil {
+			return errorResult(err), nil, nil
+		}
+		return jsonResult(configs), nil, nil
+	})
+
+	mcp.AddTool(s, &mcp.Tool{
+		Name:        ToolGetScheduling,
+		Description: "Get the distributed scheduling config for one model, or null when none is configured.",
+	}, func(ctx context.Context, _ *mcp.CallToolRequest, args DeleteSchedulingRequest) (*mcp.CallToolResult, any, error) {
+		if args.ModelName == "" {
+			return errorResultf("model_name is required"), nil, nil
+		}
+		config, err := client.GetScheduling(ctx, args.ModelName)
+		if err != nil {
+			return errorResult(err), nil, nil
+		}
+		return jsonResult(config), nil, nil
+	})
+
+	if opts.DisableMutating {
+		return
+	}
+
+	mcp.AddTool(s, &mcp.Tool{
+		Name:        ToolSetScheduling,
+		Description: "Create or update a distributed per-model scheduling config. Requires user confirmation per safety rule 1.",
+	}, func(ctx context.Context, _ *mcp.CallToolRequest, args SetSchedulingRequest) (*mcp.CallToolResult, any, error) {
+		if args.ModelName == "" {
+			return errorResultf("model_name is required"), nil, nil
+		}
+		if args.SpreadAll && (args.MinReplicas != 0 || args.MaxReplicas != 0) {
+			return errorResultf("spread_all and min_replicas/max_replicas are mutually exclusive"), nil, nil
+		}
+		if args.MinReplicas < 0 {
+			return errorResultf("min_replicas must be >= 0"), nil, nil
+		}
+		if args.MaxReplicas < 0 {
+			return errorResultf("max_replicas must be >= 0"), nil, nil
+		}
+		if args.MaxReplicas > 0 && args.MinReplicas > args.MaxReplicas {
+			return errorResultf("min_replicas must be <= max_replicas"), nil, nil
+		}
+		config, err := client.SetScheduling(ctx, args)
+		if err != nil {
+			return errorResult(err), nil, nil
+		}
+		return jsonResult(config), nil, nil
+	})
+
+	mcp.AddTool(s, &mcp.Tool{
+		Name:        ToolDeleteScheduling,
+		Description: "Delete a distributed per-model scheduling config. Requires user confirmation per safety rule 1.",
+	}, func(ctx context.Context, _ *mcp.CallToolRequest, args DeleteSchedulingRequest) (*mcp.CallToolResult, any, error) {
+		if args.ModelName == "" {
+			return errorResultf("model_name is required"), nil, nil
+		}
+		if err := client.DeleteScheduling(ctx, args.ModelName); err != nil {
+			return errorResult(err), nil, nil
+		}
+		return jsonResult(map[string]string{"deleted": args.ModelName}), nil, nil
+	})
+}