feat(distributed): make MCP execution and discovery a selection

mcp.tools.execute and mcp.discovery were the only NATS subjects that
combined a queue group with a reply, and no carrier in this design
provides both. They never needed one: a queue group is a way of choosing
a subscriber, and choosing is a query.

The frontend now lists the approved, non-draining agent nodes, asks the
node_connections table in one joined statement which of those tunnels a
live replica holds, prefers one this replica holds so the call skips the
relay hop, and issues an ordinary control RPC on the path task 4 already
mounted. A peer-held tunnel is reached through the relay. That is a
choice a broker's hidden balancing could not make.

The selection reads presence and nothing else. It is filtered only on
node type and on the two statuses an operator controls, never on a health
verdict written on another clock, because refusing a worker that is
connected and answering is the same defect as picking one that is gone.
An empty fleet answers ErrNoAgentWorker, which is deliberately neither
ErrWorkerUnroutable nor anything cluster.IsWorkerAnswer accepts: nothing
was asked of any worker, so no reap guard may act on it.

A reply carrying an Error is the worker's own answer and is returned
unchanged; it is never offered to a second worker, which would turn "this
MCP server rejected your arguments" into "the fleet is broken" and could
run a tool twice. A call that never reached a worker is retried against a
different pick, at most three times, and whatever error is finally
returned is returned unwrapped so its identity survives the loop.

MCP prompts and resources now answer 501 in distributed mode instead of
an empty 200. They are served only from sessions the frontend holds, and
in distributed mode it holds none. That gap predates the removal of the
bus and is not closed by it; this only stops it being silent.

Agent workers keep every other subject, including nodes.<id>.backend.stop.
Their minted JWT loses the two MCP subjects and keeps a non-empty allow
list, because NATS reads an empty one as no restriction at all.

Assisted-by: Claude Opus 5 [claude-code]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto committed 2026-09-03 12:04:41 +00:00
1 parent 64059cd7d7
commit 5effa47527
39 files changed
+2048 -400

No files matched your search

+41
View File
@@ -0,0 +1,41 @@
// SPDX-License-Identifier: MIT
package application
import (
"fmt"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/services/nodes"
)
// newAgentControl builds the frontend's agent control client: the SELECTION
// that decides which agent worker answers a verb, and the control client that
// carries the verb to it.
//
// A named function rather than two more lines in initDistributed, for the
// reason distributedSchedulerOptions is one: initDistributed opens a database
// and a bus, so no unit spec reaches it, and the argument that matters here has
// no symptom when it is wrong. An empty instance id makes every connection read
// report nothing as held by this replica, so every MCP call takes a relay hop
// through a peer even when this replica holds the worker's tunnel itself, and
// nothing anywhere says so: the calls all succeed, just through one more
// process than they need. It is refused here instead of shipped as latency.
//
// The other two refusals are the ordinary kind. A selector with no registry has
// nothing to select from and a client with no transport reaches nobody, and
// both would present as MCP being quietly unavailable in a deployment that
// looks healthy.
func newAgentControl(cfg config.DistributedConfig, registry *nodes.NodeRegistry,
conns nodes.AgentConnectionReader, control *nodes.ControlClient) (*nodes.AgentControlClient, error) {
if cfg.InstanceID == "" {
return nil, fmt.Errorf("the agent control client was built with no instance id: every MCP call would relay through a peer even for a worker whose tunnel this replica holds")
}
if registry == nil || conns == nil {
return nil, fmt.Errorf("the agent control client was built with no way to find a connected agent worker")
}
if control == nil {
return nil, fmt.Errorf("the agent control client was built with no control transport to reach an agent worker over")
}
return nodes.NewAgentControlClient(nodes.NewAgentSelector(registry, conns, cfg.InstanceID), control), nil
}
@@ -0,0 +1,89 @@
// SPDX-License-Identifier: MIT
package application
import (
"context"
"runtime"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/config"
mcpremote "github.com/mudler/LocalAI/core/services/mcp"
"github.com/mudler/LocalAI/core/services/nodes"
"github.com/mudler/LocalAI/core/services/testutil"
)
// recordingConnections captures the owner id the selection was made with. It is
// how this spec sees the one argument whose loss has no other symptom.
type recordingConnections struct{ owners []string }
func (r *recordingConnections) ConnectedAmong(_ context.Context, _ []string, owner string) ([]string, []string, error) {
r.owners = append(r.owners, owner)
return nil, nil, nil
}
// The wiring that connects MCP to the agent workers, guarded the way the
// absence wiring is and for the same reason: initDistributed opens a database
// and a bus, so no unit spec reaches the construction literal, and one of these
// arguments is silent when it is wrong.
var _ = Describe("building the frontend's agent control client", func() {
var registry *nodes.NodeRegistry
var ctx context.Context
BeforeEach(func() {
if runtime.GOOS == "darwin" {
Skip("testcontainers requires Docker, not available on macOS CI")
}
ctx = context.Background()
var err error
registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB())
Expect(err).ToNot(HaveOccurred())
})
// Every refusal below is given a valid value for everything except the one
// argument it is about, so no assertion can be satisfied by a guard that
// fires for the wrong reason.
It("makes the selection with THIS replica's instance id", func() {
// The silent one. With an empty id every connection read reports
// nothing as held here, so every MCP call relays through a peer even
// for a worker whose tunnel this replica holds: correct answers, one
// extra hop, and no log line anywhere.
Expect(registry.Register(ctx, &nodes.BackendNode{
Name: "agent-1", NodeType: nodes.NodeTypeAgent, Address: "a:50051",
}, true)).To(Succeed())
conns := &recordingConnections{}
client, err := newAgentControl(
config.DistributedConfig{InstanceID: "replica-7"}, registry, conns,
nodes.NewControlClient(nil, "token"))
Expect(err).ToNot(HaveOccurred())
// Driven through a real call rather than read off a field: what has to
// be true is that the id reaches the SELECTION, not that it was stored.
_, _ = client.ExecuteMCPTool(ctx, mcpremote.MCPToolRequest{ModelName: "m"})
Expect(conns.owners).To(ConsistOf("replica-7"))
})
It("refuses to build with no instance id", func() {
_, err := newAgentControl(config.DistributedConfig{}, registry, &recordingConnections{},
nodes.NewControlClient(nil, "token"))
Expect(err).To(HaveOccurred())
Expect(err.Error()).To(ContainSubstring("instance id"))
})
It("refuses to build with nothing to read connections through", func() {
_, err := newAgentControl(config.DistributedConfig{InstanceID: "replica-7"}, registry, nil,
nodes.NewControlClient(nil, "token"))
Expect(err).To(HaveOccurred())
Expect(err.Error()).To(ContainSubstring("connected agent worker"))
})
It("refuses to build with no control transport", func() {
_, err := newAgentControl(config.DistributedConfig{InstanceID: "replica-7"}, registry,
&recordingConnections{}, nil)
Expect(err).To(HaveOccurred())
Expect(err.Error()).To(ContainSubstring("control transport"))
})
})
+18
View File
@@ -83,6 +83,13 @@ type DistributedServices struct {
// WorkerDialer. Exposed so the model store built in startup.go reaches
// remote models the same way every other caller does.
BackendClients nodes.BackendClientFactory
// AgentControl carries the frontend's MCP verbs to whichever agent worker
// holds a tunnel this deployment can reach. It is what the chat, responses,
// messages and MCP endpoints reach an agent worker through; a nil one means
// this frontend cannot run MCP at all, which is why initDistributed refuses
// to come up without it rather than leaving the endpoints to discover it
// one request at a time.
AgentControl *nodes.AgentControlClient
shutdownOnce sync.Once
}
@@ -438,6 +445,16 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
// second client would open its own and the two would never share one.
controlClient := nodes.NewControlClient(workerHTTPDialer, cfg.Distributed.RegistrationToken)
// The caller the agent worker's control plane has been waiting for. MCP
// execution and discovery used to be a NATS request onto a queue group,
// where the bus chose the worker and neither side could say which one had
// answered; they are now a query against the connection rows plus an
// ordinary control RPC over the chosen worker's tunnel.
agentControl, err := newAgentControl(cfg.Distributed, registry, clusterRegistry, controlClient)
if err != nil {
return nil, fmt.Errorf("wiring the agent control client: %w", err)
}
// Create FileStager for distributed file transfer
var fileStager nodes.FileStager
if cfg.Distributed.StorageURL != "" {
@@ -666,6 +683,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
Tunnels: tunnels,
WorkerDialer: workerDialer,
BackendClients: backendClients,
AgentControl: agentControl,
Bus: bus,
}
// Checked once, here, on the assembled struct. See requireBroadcastCarrier.