Skip to content

Commit 8f74f74

Browse files
owezzyowezzy
authored
feat(mcp): expose scheduling admin tools (#11228)
* feat(mcp): add scheduling client contracts Assisted-by: Hephaestus:openai/gpt-5.5 Signed-off-by: Owen Adirah <owenadira@gmail.com> * feat(mcp): add scheduling HTTP client support Assisted-by: Hephaestus:openai/gpt-5.5 Signed-off-by: Owen Adirah <owenadira@gmail.com> * feat(mcp): add in-process scheduling stubs Assisted-by: Hephaestus:openai/gpt-5.5 Signed-off-by: Owen Adirah <owenadira@gmail.com> * feat(mcp): register scheduling tools Assisted-by: Hephaestus:openai/gpt-5.5 Signed-off-by: Owen Adirah <owenadira@gmail.com> * test(mcp): map scheduling tools to REST routes Assisted-by: Hephaestus:openai/gpt-5.5 Signed-off-by: Owen Adirah <owenadira@gmail.com> * docs(mcp): document scheduling assistant tools Assisted-by: Hephaestus:openai/gpt-5.5 Signed-off-by: Owen Adirah <owenadira@gmail.com> * 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 Signed-off-by: Owen Adirah <owenadira@gmail.com> * fix(mcp): narrow scheduling dto Assisted-by: Hephaestus:openai/gpt-5.5 [opencode] Signed-off-by: Owen Adirah <owenadira@gmail.com> --------- Signed-off-by: Owen Adirah <owenadira@gmail.com>
1 parent cd62e8f commit 8f74f74

19 files changed

Lines changed: 596 additions & 3 deletions

‎core/application/application.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -553,12 +553,17 @@ func (a *Application) start() error {
553553
// once at startup and reused across chat sessions that opt in via metadata.
554554
if !a.applicationConfig.DisableLocalAIAssistant {
555555
holder := mcpTools.NewLocalAIAssistantHolder()
556+
var nodeRegistry *nodes.NodeRegistry
557+
if a.distributed != nil {
558+
nodeRegistry = a.distributed.Registry
559+
}
556560
assistantClient := localaiInproc.New(
557561
a.applicationConfig,
558562
a.applicationConfig.SystemState,
559563
a.backendLoader,
560564
a.modelLoader,
561565
a.galleryService,
566+
nodeRegistry,
562567
)
563568
// Wire usage tracking so the assistant's get_usage_stats tool
564569
// returns real data; nil values keep the tool returning a clear

‎pkg/mcp/localaitools/client.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,10 @@ type LocalAIClient interface {
6464
// ---- System ----
6565
SystemInfo(ctx context.Context) (*SystemInfo, error)
6666
ListNodes(ctx context.Context) ([]Node, error)
67+
ListScheduling(ctx context.Context) ([]ModelSchedulingConfig, error)
68+
GetScheduling(ctx context.Context, modelName string) (*ModelSchedulingConfig, error)
69+
SetScheduling(ctx context.Context, req SetSchedulingRequest) (*ModelSchedulingConfig, error)
70+
DeleteScheduling(ctx context.Context, modelName string) error
6771
// SetNodeVRAMBudget sets (or, with an empty budget, clears) a federated
6872
// node's VRAM allocation cap as a sticky admin override. Only meaningful
6973
// in distributed mode; single-process clients report it as unavailable.

‎pkg/mcp/localaitools/coverage_test.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ var toolToHTTPRoute = map[string]string{
3535
ToolListKnownBackends: "GET /backends/known",
3636
ToolSystemInfo: "GET / (welcome JSON)",
3737
ToolListNodes: "GET /api/nodes",
38+
ToolListScheduling: "GET /api/nodes/scheduling",
39+
ToolGetScheduling: "GET /api/nodes/scheduling/:model",
3840
ToolVRAMEstimate: "POST /api/models/vram-estimate",
3941
ToolGetBranding: "GET /api/branding",
4042
ToolGetUsageStats: "GET /api/usage (or /api/usage/all when all=true)",
@@ -60,6 +62,8 @@ var toolToHTTPRoute = map[string]string{
6062
ToolCreateVoiceProfile: "POST /api/voice-profiles",
6163
ToolDeleteVoiceProfile: "DELETE /api/voice-profiles/:id",
6264
ToolSetNodeVRAMBudget: "PUT /api/nodes/:id/vram-budget",
65+
ToolSetScheduling: "POST /api/nodes/scheduling",
66+
ToolDeleteScheduling: "DELETE /api/nodes/scheduling/:model",
6367
}
6468

6569
// allKnownTools is the union of expectedFullCatalog (defined in

‎pkg/mcp/localaitools/dto.go‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,41 @@ type SetNodeVRAMBudgetRequest struct {
115115
Budget string `json:"budget,omitempty" jsonschema:"VRAM allocation cap as a percentage (e.g. 80%) or absolute amount (e.g. 12GB). Empty string clears the override."`
116116
}
117117

118+
// ModelSchedulingConfig is the MCP wire shape for one per-model distributed
119+
// scheduling rule. Keep this DTO explicit instead of aliasing the node-registry
120+
// model so the MCP contract only exposes operator-facing scheduling fields.
121+
type ModelSchedulingConfig struct {
122+
ModelName string `json:"model_name"`
123+
NodeSelector string `json:"node_selector,omitempty"`
124+
MinReplicas int `json:"min_replicas"`
125+
MaxReplicas int `json:"max_replicas"`
126+
SpreadAll bool `json:"spread_all,omitempty"`
127+
RoutePolicy string `json:"route_policy,omitempty"`
128+
BalanceAbsThreshold int `json:"balance_abs_threshold,omitempty"`
129+
BalanceRelThreshold float64 `json:"balance_rel_threshold,omitempty"`
130+
MinPrefixMatch float64 `json:"min_prefix_match,omitempty"`
131+
}
132+
133+
// SetSchedulingRequest is the input for set_scheduling. It mirrors
134+
// /api/nodes/scheduling so standalone MCP and REST callers preserve the same
135+
// PATCH-style semantics for the optional prefix-cache routing fields.
136+
type SetSchedulingRequest struct {
137+
ModelName string `json:"model_name" jsonschema:"Installed model name whose distributed scheduling rule should be created or updated."`
138+
NodeSelector map[string]string `json:"node_selector,omitempty" jsonschema:"Optional node-label selector. Empty means any healthy backend node."`
139+
MinReplicas int `json:"min_replicas" jsonschema:"Minimum desired replicas. Mutually exclusive with spread_all."`
140+
MaxReplicas int `json:"max_replicas" jsonschema:"Maximum desired replicas. Must be >= min_replicas when non-zero. Mutually exclusive with spread_all."`
141+
SpreadAll bool `json:"spread_all,omitempty" jsonschema:"When true, keep one replica on every matching node. Mutually exclusive with min_replicas/max_replicas."`
142+
RoutePolicy *string `json:"route_policy,omitempty" jsonschema:"Optional prefix-cache route policy override. Omit to preserve the existing value on updates."`
143+
BalanceAbsThreshold *int `json:"balance_abs_threshold,omitempty" jsonschema:"Optional absolute imbalance threshold override. Omit to preserve the existing value on updates."`
144+
BalanceRelThreshold *float64 `json:"balance_rel_threshold,omitempty" jsonschema:"Optional relative imbalance threshold override. Omit to preserve the existing value on updates."`
145+
MinPrefixMatch *float64 `json:"min_prefix_match,omitempty" jsonschema:"Optional minimum prefix match threshold override. Omit to preserve the existing value on updates."`
146+
}
147+
148+
// DeleteSchedulingRequest identifies the model scheduling rule to remove.
149+
type DeleteSchedulingRequest struct {
150+
ModelName string `json:"model_name" jsonschema:"Installed model name whose scheduling config should be removed."`
151+
}
152+
118153
// ImportModelURIRequest is the input for import_model_uri. It mirrors the
119154
// REST surface (`/models/import-uri`) closely so both clients can produce
120155
// identical responses; the BackendPreference is a flat field rather than the

‎pkg/mcp/localaitools/fakes_test.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,10 @@ type fakeClient struct {
4242
upgradeBackend func(string) (string, error)
4343
systemInfo func() (*SystemInfo, error)
4444
listNodes func() ([]Node, error)
45+
listScheduling func() ([]ModelSchedulingConfig, error)
46+
getScheduling func(string) (*ModelSchedulingConfig, error)
47+
setScheduling func(SetSchedulingRequest) (*ModelSchedulingConfig, error)
48+
deleteScheduling func(string) error
4549
setNodeVRAMBudget func(string, string) error
4650
vramEstimate func(VRAMEstimateRequest) (*vram.EstimateResult, error)
4751
toggleModelState func(string, modeladmin.Action) error
@@ -230,6 +234,38 @@ func (f *fakeClient) ListNodes(_ context.Context) ([]Node, error) {
230234
return nil, nil
231235
}
232236

237+
func (f *fakeClient) ListScheduling(_ context.Context) ([]ModelSchedulingConfig, error) {
238+
f.record("ListScheduling", nil)
239+
if f.listScheduling != nil {
240+
return f.listScheduling()
241+
}
242+
return []ModelSchedulingConfig{}, nil
243+
}
244+
245+
func (f *fakeClient) GetScheduling(_ context.Context, modelName string) (*ModelSchedulingConfig, error) {
246+
f.record("GetScheduling", modelName)
247+
if f.getScheduling != nil {
248+
return f.getScheduling(modelName)
249+
}
250+
return &ModelSchedulingConfig{ModelName: modelName}, nil
251+
}
252+
253+
func (f *fakeClient) SetScheduling(_ context.Context, req SetSchedulingRequest) (*ModelSchedulingConfig, error) {
254+
f.record("SetScheduling", req)
255+
if f.setScheduling != nil {
256+
return f.setScheduling(req)
257+
}
258+
return &ModelSchedulingConfig{ModelName: req.ModelName, MinReplicas: req.MinReplicas, MaxReplicas: req.MaxReplicas, SpreadAll: req.SpreadAll}, nil
259+
}
260+
261+
func (f *fakeClient) DeleteScheduling(_ context.Context, modelName string) error {
262+
f.record("DeleteScheduling", modelName)
263+
if f.deleteScheduling != nil {
264+
return f.deleteScheduling(modelName)
265+
}
266+
return nil
267+
}
268+
233269
func (f *fakeClient) SetNodeVRAMBudget(_ context.Context, nodeID, budget string) error {
234270
f.record("SetNodeVRAMBudget", []any{nodeID, budget})
235271
if f.setNodeVRAMBudget != nil {

‎pkg/mcp/localaitools/httpapi/client.go‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -483,6 +483,49 @@ func (c *Client) ListNodes(ctx context.Context) ([]localaitools.Node, error) {
483483
return out, nil
484484
}
485485

486+
func (c *Client) ListScheduling(ctx context.Context) ([]localaitools.ModelSchedulingConfig, error) {
487+
var out []localaitools.ModelSchedulingConfig
488+
if err := c.do(ctx, http.MethodGet, routeScheduling, nil, &out); err != nil {
489+
if errors.Is(err, ErrHTTPNotFound) {
490+
return []localaitools.ModelSchedulingConfig{}, nil
491+
}
492+
return nil, err
493+
}
494+
return out, nil
495+
}
496+
497+
func (c *Client) GetScheduling(ctx context.Context, modelName string) (*localaitools.ModelSchedulingConfig, error) {
498+
if modelName == "" {
499+
return nil, errors.New("model_name is required")
500+
}
501+
var out localaitools.ModelSchedulingConfig
502+
if err := c.do(ctx, http.MethodGet, routeModelScheduling(modelName), nil, &out); err != nil {
503+
if errors.Is(err, ErrHTTPNotFound) {
504+
return nil, nil
505+
}
506+
return nil, err
507+
}
508+
return &out, nil
509+
}
510+
511+
func (c *Client) SetScheduling(ctx context.Context, req localaitools.SetSchedulingRequest) (*localaitools.ModelSchedulingConfig, error) {
512+
if req.ModelName == "" {
513+
return nil, errors.New("model_name is required")
514+
}
515+
var out localaitools.ModelSchedulingConfig
516+
if err := c.do(ctx, http.MethodPost, routeScheduling, req, &out); err != nil {
517+
return nil, err
518+
}
519+
return &out, nil
520+
}
521+
522+
func (c *Client) DeleteScheduling(ctx context.Context, modelName string) error {
523+
if modelName == "" {
524+
return errors.New("model_name is required")
525+
}
526+
return c.do(ctx, http.MethodDelete, routeModelScheduling(modelName), nil, nil)
527+
}
528+
486529
func (c *Client) SetNodeVRAMBudget(ctx context.Context, nodeID, budget string) error {
487530
// PUT with an empty value clears the override server-side (Task 9), so we
488531
// use PUT uniformly rather than switching to DELETE for the clear case.

‎pkg/mcp/localaitools/httpapi/client_test.go‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,43 @@ func fakeLocalAI() *httptest.Server {
8484
})
8585
})
8686

87+
mux.HandleFunc("/api/nodes/scheduling", func(w http.ResponseWriter, r *http.Request) {
88+
switch r.Method {
89+
case http.MethodGet:
90+
_ = json.NewEncoder(w).Encode([]map[string]any{{
91+
"id": "sched-1",
92+
"model_name": "qwen",
93+
"min_replicas": 1,
94+
"max_replicas": 2,
95+
"unsatisfiable_ticks": 1,
96+
"unsatisfiable_until": "2026-01-01T00:00:00Z",
97+
"created_at": "2026-01-01T00:00:00Z",
98+
"updated_at": "2026-01-01T00:00:00Z",
99+
}})
100+
case http.MethodPost:
101+
var body map[string]any
102+
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
103+
http.Error(w, err.Error(), http.StatusBadRequest)
104+
return
105+
}
106+
body["id"] = "sched-1"
107+
_ = json.NewEncoder(w).Encode(body)
108+
default:
109+
http.Error(w, "method", http.StatusMethodNotAllowed)
110+
}
111+
})
112+
113+
mux.HandleFunc("/api/nodes/scheduling/qwen", func(w http.ResponseWriter, r *http.Request) {
114+
switch r.Method {
115+
case http.MethodGet:
116+
_ = json.NewEncoder(w).Encode(map[string]any{"model_name": "qwen", "spread_all": true})
117+
case http.MethodDelete:
118+
w.WriteHeader(http.StatusNoContent)
119+
default:
120+
http.Error(w, "method", http.StatusMethodNotAllowed)
121+
}
122+
})
123+
87124
return httptest.NewServer(mux)
88125
}
89126

@@ -197,8 +234,51 @@ var _ = Describe("httpapi.Client against the LocalAI admin REST surface", func()
197234
Expect(bs[0].Installed).To(BeTrue())
198235
})
199236
})
237+
238+
Describe("Scheduling", func() {
239+
It("lists scheduling configs", func() {
240+
out, err := c.ListScheduling(ctx)
241+
Expect(err).ToNot(HaveOccurred())
242+
Expect(out).To(HaveLen(1))
243+
Expect(out[0].ModelName).To(Equal("qwen"))
244+
Expect(out[0].MinReplicas).To(Equal(1))
245+
Expect(schedulingJSONKeys(&out[0])).ToNot(Or(
246+
HaveKey("id"),
247+
HaveKey("unsatisfiable_until"),
248+
HaveKey("unsatisfiable_ticks"),
249+
HaveKey("created_at"),
250+
HaveKey("updated_at"),
251+
))
252+
})
253+
254+
It("gets one scheduling config", func() {
255+
out, err := c.GetScheduling(ctx, "qwen")
256+
Expect(err).ToNot(HaveOccurred())
257+
Expect(out.ModelName).To(Equal("qwen"))
258+
Expect(out.SpreadAll).To(BeTrue())
259+
})
260+
261+
It("sets a scheduling config", func() {
262+
out, err := c.SetScheduling(ctx, localaitools.SetSchedulingRequest{ModelName: "qwen", MinReplicas: 1, MaxReplicas: 2})
263+
Expect(err).ToNot(HaveOccurred())
264+
Expect(out.ModelName).To(Equal("qwen"))
265+
Expect(out.MaxReplicas).To(Equal(2))
266+
})
267+
268+
It("deletes a scheduling config", func() {
269+
Expect(c.DeleteScheduling(ctx, "qwen")).To(Succeed())
270+
})
271+
})
200272
})
201273

274+
func schedulingJSONKeys(config *localaitools.ModelSchedulingConfig) map[string]any {
275+
var out map[string]any
276+
b, err := json.Marshal(config)
277+
Expect(err).ToNot(HaveOccurred())
278+
Expect(json.Unmarshal(b, &out)).To(Succeed())
279+
return out
280+
}
281+
202282
var _ = Describe("Model aliases", func() {
203283
Describe("ListAliases", func() {
204284
It("passes the GET /api/aliases payload through unchanged", func() {

‎pkg/mcp/localaitools/httpapi/routes.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ const (
2424
routeBackendsKnown = "/backends/known"
2525
routeBackendsApply = "/backends/apply"
2626
routeNodes = "/api/nodes"
27+
routeScheduling = "/api/nodes/scheduling"
2728
routeVRAMEstimate = "/api/models/vram-estimate"
2829
routeBranding = "/api/branding"
2930
routeSettings = "/api/settings"
@@ -66,3 +67,7 @@ func routeVoiceProfileDelete(id string) string {
6667
func routeNodeVRAMBudget(id string) string {
6768
return "/api/nodes/" + url.PathEscape(id) + "/vram-budget"
6869
}
70+
71+
func routeModelScheduling(modelName string) string {
72+
return "/api/nodes/scheduling/" + url.PathEscape(modelName)
73+
}

0 commit comments

Comments
 (0)