From 5effa4752789f2bbe4e17029ee00ff731d8afffa Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 3 Sep 2026 12:04:41 +0000 Subject: [PATCH] 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..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 --- core/application/agent_control_wiring.go | 41 ++ core/application/agent_control_wiring_test.go | 89 +++++ core/application/distributed.go | 18 + core/cli/agent_worker.go | 49 +-- core/cli/agent_worker_test.go | 52 +-- core/http/endpoints/anthropic/messages.go | 4 +- core/http/endpoints/localai/mcp.go | 4 +- core/http/endpoints/localai/mcp_local_only.go | 41 ++ .../endpoints/localai/mcp_local_only_test.go | 114 ++++++ core/http/endpoints/localai/mcp_prompts.go | 11 + core/http/endpoints/localai/mcp_resources.go | 10 + core/http/endpoints/localai/mcp_tools.go | 16 +- core/http/endpoints/mcp/executor.go | 45 ++- core/http/endpoints/mcp/remote_test.go | 145 +++++++ core/http/endpoints/mcp/tools.go | 96 +++-- core/http/endpoints/openai/chat.go | 4 +- .../http/endpoints/openresponses/responses.go | 4 +- core/http/routes/agent_control.go | 38 ++ .../routes/agent_control_internal_test.go | 36 ++ core/http/routes/anthropic.go | 8 +- core/http/routes/localai.go | 10 +- core/http/routes/openai.go | 10 +- core/http/routes/openresponses.go | 9 +- core/services/agentworker/handlers.go | 2 + core/services/cluster/connected_among_test.go | 173 ++++++++ core/services/cluster/ownership.go | 92 +++++ core/services/cluster/ownership_test.go | 8 + core/services/mcp/remote.go | 26 +- core/services/messaging/subjects.go | 7 - core/services/nodes/agent_control.go | 137 +++++++ core/services/nodes/agent_control_test.go | 375 ++++++++++++++++++ core/services/nodes/agent_selector.go | 171 ++++++++ core/services/nodes/agent_selector_test.go | 271 +++++++++++++ docs/content/features/distributed-mode.md | 33 +- pkg/natsauth/permissions.go | 10 +- pkg/natsauth/permissions_coverage_test.go | 24 +- tests/e2e/distributed/mcp_nats_test.go | 177 --------- tests/e2e/distributed/nats_jwt_test.go | 11 +- .../e2e/distributed/nats_queue_reply_test.go | 77 ++++ 39 files changed, 2048 insertions(+), 400 deletions(-) create mode 100644 core/application/agent_control_wiring.go create mode 100644 core/application/agent_control_wiring_test.go create mode 100644 core/http/endpoints/localai/mcp_local_only.go create mode 100644 core/http/endpoints/localai/mcp_local_only_test.go create mode 100644 core/http/endpoints/mcp/remote_test.go create mode 100644 core/http/routes/agent_control.go create mode 100644 core/http/routes/agent_control_internal_test.go create mode 100644 core/services/cluster/connected_among_test.go create mode 100644 core/services/nodes/agent_control.go create mode 100644 core/services/nodes/agent_control_test.go create mode 100644 core/services/nodes/agent_selector.go create mode 100644 core/services/nodes/agent_selector_test.go delete mode 100644 tests/e2e/distributed/mcp_nats_test.go create mode 100644 tests/e2e/distributed/nats_queue_reply_test.go diff --git a/core/application/agent_control_wiring.go b/core/application/agent_control_wiring.go new file mode 100644 index 000000000..f991ee875 --- /dev/null +++ b/core/application/agent_control_wiring.go @@ -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 +} diff --git a/core/application/agent_control_wiring_test.go b/core/application/agent_control_wiring_test.go new file mode 100644 index 000000000..d8c1a0226 --- /dev/null +++ b/core/application/agent_control_wiring_test.go @@ -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")) + }) +}) diff --git a/core/application/distributed.go b/core/application/distributed.go index c782fa654..6fd3a826a 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -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. diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index fa49f606c..0b2ee70e2 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -178,11 +178,14 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { // The tunnel, and the loopback control plane behind it. // - // ADDED to this worker rather than swapping anything out: every verb below - // still arrives on NATS, and will until the tasks that move them land. What - // this buys today is that the frontend can reach an agent worker by RPC at - // all, on the same carrier and with the same failure vocabulary a backend - // worker already uses, without the agent worker opening an inbound port. + // This is now the ONLY way MCP tool execution and discovery reach this + // worker: their queue-group subjects are gone, because a queue group was + // only ever a way of SELECTING a worker and the frontend makes that + // selection itself (nodes.AgentSelector). The remaining verbs below still + // arrive on NATS and will until the tasks that move them land. + // + // The worker opens no inbound port for any of it: it dials out and the + // control plane rides the tunnel it holds. // // The credential is read through credMgr rather than captured from res, // because every re-registration the manager performs ROTATES it and a @@ -191,9 +194,9 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { // // It is started AFTER registration, which is what supplies both the node // identity the dial names and the credential it presents, and BEFORE the - // NATS subscriptions, so that a frontend that reaches this worker over the - // tunnel finds its verbs mounted rather than a 404 it would read as a - // version skew. + // remaining NATS subscriptions, so that a frontend that reaches this worker + // over the tunnel finds its verbs mounted rather than a 404 it would read + // as a version skew. agentCtl, err := agentworker.Start(shutdownCtx, agentworker.Options{ FrontendURL: cmd.RegisterTo, NodeID: nodeID, @@ -244,19 +247,6 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { return fmt.Errorf("starting dispatcher: %w", err) } - // Subscribe to MCP tool execution requests (load-balanced across workers). - // The frontend routes model-level MCP tool calls here via NATS request-reply. - if _, err := natsClient.QueueSubscribeReply(messaging.SubjectMCPToolExecute, messaging.QueueAgentWorkers, - replyOverNATS(messaging.SubjectMCPToolExecute, serveMCPToolRequest)); err != nil { - return fmt.Errorf("subscribing to %s: %w", messaging.SubjectMCPToolExecute, err) - } - - // Subscribe to MCP discovery requests (load-balanced across workers). - if _, err := natsClient.QueueSubscribeReply(messaging.SubjectMCPDiscovery, messaging.QueueAgentWorkers, - replyOverNATS(messaging.SubjectMCPDiscovery, serveMCPDiscoveryRequest)); err != nil { - return fmt.Errorf("subscribing to %s: %w", messaging.SubjectMCPDiscovery, err) - } - // Subscribe to MCP CI job execution (load-balanced across agent workers). // In distributed mode, MCP CI jobs are routed here because the frontend // cannot create MCP sessions (e.g., stdio servers using docker). @@ -456,23 +446,6 @@ func encodeMCPReply(resp any) (json.RawMessage, error) { return out, nil } -// replyOverNATS adapts one of the serve* functions to a NATS request-reply -// subscription, so the bus carries exactly the bytes the tunnel does. -func replyOverNATS(subject string, serve func(context.Context, json.RawMessage) (json.RawMessage, error)) func([]byte, func([]byte)) { - return func(data []byte, reply func([]byte)) { - out, err := serve(context.Background(), data) - if err != nil { - // Nothing is sent. A requester on the bus reads that as a timeout, - // which is the closest the carrier has to "this worker did not - // answer"; inventing a reply body here would put a failure to serve - // into the bucket reserved for the worker's own verdict. - xlog.Error("Agent worker could not serve a bus request", "subject", subject, "error", err) - return - } - reply(out) - } -} - // handleMCPCIJob processes an MCP CI job on the agent worker. // The agent worker can create MCP sessions (has docker) and call the LocalAI API for inference. func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken string, natsClient messaging.MessagingClient, jobTimeout time.Duration) { diff --git a/core/cli/agent_worker_test.go b/core/cli/agent_worker_test.go index 7fad5a0b4..1e8e1d7ec 100644 --- a/core/cli/agent_worker_test.go +++ b/core/cli/agent_worker_test.go @@ -11,20 +11,16 @@ import ( "github.com/mudler/LocalAI/core/services/messaging" ) -// The agent worker answers its MCP verbs on two carriers at once: the NATS -// subject it has always answered, and the control route on the tunnel it now -// holds. These specs pin the two properties that stops those drifting apart. +// The agent worker answers its MCP verbs on ONE carrier now: the control route +// on the tunnel it holds. The queue-group subjects these used to arrive on are +// gone, because a queue group was only ever a way of SELECTING a worker, and +// the frontend now makes that selection itself (nodes.AgentSelector). // -// The first is that there is ONE implementation. A frontend that reaches a -// worker over the bus and one that reaches the same worker over its tunnel must -// get the same bytes, because during this migration both are live and which one -// is used is not a decision anybody makes deliberately. -// -// The second is the split between an ANSWER and a FAILURE TO SERVE. A tool that -// ran and failed is the worker's own verdict and travels inside the reply; a -// verb this worker could not serve at all is a Go error, which becomes a -// non-2xx over the tunnel and silence on the bus, and which nothing may read as -// evidence about anything. +// What these specs pin is the split between an ANSWER and a FAILURE TO SERVE. A +// tool that ran and failed is the worker's own verdict and travels inside the +// reply, on a 200; a verb this worker could not serve at all is a Go error, +// which becomes a non-2xx, and which nothing may read as evidence about +// anything. var _ = Describe("The agent worker's MCP verbs", func() { It("answers a tool request it could not decode, rather than failing to serve it", func() { // The decode happened on this worker and its outcome is something the @@ -48,36 +44,6 @@ var _ = Describe("The agent worker's MCP verbs", func() { Expect(resp.Error).To(ContainSubstring("unmarshal error")) }) - It("puts on the bus exactly the bytes the tunnel route returns", func() { - // The one property that keeps the two carriers honest. A second - // implementation for the bus is how a deployment ends up behaving - // differently depending on which one a frontend happened to pick. - request := json.RawMessage(`{"tool_name":`) - overTunnel, err := serveMCPToolRequest(context.Background(), request) - Expect(err).ToNot(HaveOccurred()) - - sent := make(chan []byte, 1) - replyOverNATS("mcp.tools.execute", serveMCPToolRequest)(request, func(b []byte) { sent <- b }) - - var overBus []byte - Eventually(sent).Should(Receive(&overBus)) - Expect(string(overBus)).To(Equal(string(overTunnel))) - }) - - It("sends nothing on the bus when the verb could not be served", func() { - // A requester reads the silence as a timeout, which is the closest the - // bus has to "this worker did not answer". Inventing a reply body would - // put a failure to serve into the bucket reserved for the worker's own - // verdict, which is the collapse this whole phase exists to prevent. - sent := make(chan []byte, 1) - failing := func(context.Context, json.RawMessage) (json.RawMessage, error) { - return nil, context.DeadlineExceeded - } - - replyOverNATS("mcp.tools.execute", failing)(json.RawMessage(`{}`), func(b []byte) { sent <- b }) - - Expect(sent).ToNot(Receive()) - }) }) var _ = Describe("The agent worker's backend stop", func() { diff --git a/core/http/endpoints/anthropic/messages.go b/core/http/endpoints/anthropic/messages.go index 9310cabd9..231444ded 100644 --- a/core/http/endpoints/anthropic/messages.go +++ b/core/http/endpoints/anthropic/messages.go @@ -29,7 +29,7 @@ import ( // @Param request body schema.AnthropicRequest true "query params" // @Success 200 {object} schema.AnthropicResponse "Response" // @Router /v1/messages [post] -func MessagesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient) echo.HandlerFunc { +func MessagesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, agentControl mcpTools.AgentControl) echo.HandlerFunc { return func(c echo.Context) error { id := uuid.New().String() @@ -70,7 +70,7 @@ func MessagesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evalu if (len(mcpServers) > 0 || mcpPromptName != "" || len(mcpResourceURIs) > 0) && (cfg.MCP.Servers != "" || cfg.MCP.Stdio != "") { remote, stdio, mcpErr := cfg.MCP.MCPConfigFromYAML() if mcpErr == nil { - mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), natsClient, cfg.Name, remote, stdio, mcpServers) + mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), agentControl, cfg.Name, remote, stdio, mcpServers) // Prompt and resource injection (pre-processing step — resolves locally regardless of distributed mode) namedSessions, sessErr := mcpTools.NamedSessionsFromMCPConfig(cfg.Name, remote, stdio, mcpServers) diff --git a/core/http/endpoints/localai/mcp.go b/core/http/endpoints/localai/mcp.go index f3905442d..9687afac6 100644 --- a/core/http/endpoints/localai/mcp.go +++ b/core/http/endpoints/localai/mcp.go @@ -57,7 +57,7 @@ type MCPErrorEvent struct { // @Param request body schema.OpenAIRequest true "query params" // @Success 200 {object} schema.OpenAIResponse "Response" // @Router /v1/mcp/chat/completions [post] -func MCPEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient, compressor middleware.ChatCompressor) echo.HandlerFunc { +func MCPEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, agentControl mcpTools.AgentControl, compressor middleware.ChatCompressor) echo.HandlerFunc { // The legacy /v1/mcp/chat/completions endpoint never opts into the // in-process LocalAI Assistant tool surface — pass nil holder so the // assistant branch in chat.go is unreachable from this code path. @@ -65,7 +65,7 @@ func MCPEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator // the per-model PII config and is kept for backward compatibility. // The request-side middleware on the main chat route handles // filtering for the standard /v1/chat/completions path. - chatHandler := openai.ChatEndpoint(cl, ml, evaluator, appConfig, natsClient, nil, compressor) + chatHandler := openai.ChatEndpoint(cl, ml, evaluator, appConfig, agentControl, nil, compressor) return func(c echo.Context) error { input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest) diff --git a/core/http/endpoints/localai/mcp_local_only.go b/core/http/endpoints/localai/mcp_local_only.go new file mode 100644 index 000000000..281b0f456 --- /dev/null +++ b/core/http/endpoints/localai/mcp_local_only.go @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: MIT + +package localai + +import ( + "net/http" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/config" +) + +// mcpLocalSessionsOnly answers the request itself when this frontend cannot +// serve it, and reports whether it did. +// +// MCP prompts and resources are served ONLY from MCP sessions this process +// holds. In distributed mode the frontend holds none: creating them is what an +// agent worker is for, because a stdio server usually means running docker, and +// nothing in this programme carries prompts or resources to one. That is a hole +// that PREDATES the removal of the message bus and is not closed by it; tools +// and discovery had a carrier and these two never did. +// +// Until this task the endpoints did not say so. A model with MCP servers +// configured answered 200 with an empty list, which reads as "this model has no +// prompts" and is a different statement from "this deployment cannot tell you". +// A client cannot act on the first and can act on the second, and an operator +// reading an empty list has no reason to look further. 501 with the reason in +// the body is the smallest honest answer, and it is deliberately NOT a 404 or +// an empty 200: the resource may well exist, on a worker this frontend has no +// verb to ask. +// +// One function for four endpoints, so the four cannot drift into four different +// stories about the same limitation. Each endpoint still has to CALL it, and +// each call is pinned by its own spec. +func mcpLocalSessionsOnly(c echo.Context, appConfig *config.ApplicationConfig, surface string) (bool, error) { + if appConfig == nil || !appConfig.Distributed.Enabled { + return false, nil + } + return true, c.JSON(http.StatusNotImplemented, map[string]string{ + "error": "MCP " + surface + " are not available in distributed mode: they are served only from MCP sessions held by this frontend, and in distributed mode the sessions live on an agent worker, which serves no " + surface + " verb", + }) +} diff --git a/core/http/endpoints/localai/mcp_local_only_test.go b/core/http/endpoints/localai/mcp_local_only_test.go new file mode 100644 index 000000000..89c91f116 --- /dev/null +++ b/core/http/endpoints/localai/mcp_local_only_test.go @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: MIT + +package localai + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + + "github.com/labstack/echo/v4" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" +) + +// The MCP prompts and resources endpoints work only against sessions THIS +// process holds, and in distributed mode it holds none. That is a hole older +// than the removal of the message bus: tools and discovery had a carrier to an +// agent worker and these two never did, and nothing in this programme gives +// them one. +// +// What these specs pin is that the hole is now HONEST. The old answer was 200 +// with an empty list, which reads as "this model has no prompts": a client +// cannot act on it, and an operator reading it has no reason to look further. +// +// Each endpoint is asserted separately on purpose. The refusal is one function +// but four CALLS, and a call site that lost its call would still leave the +// other three green. +var _ = Describe("MCP prompts and resources in distributed mode", func() { + distributed := &config.ApplicationConfig{} + distributed.Distributed.Enabled = true + + standalone := &config.ApplicationConfig{} + + // loader with no models at all. Every endpoint below looks a model up, and + // the refusal has to come FIRST: a 404 for an unknown model would be a + // different, and wrong, story about the same request. + emptyLoader := config.NewModelConfigLoader("") + + call := func(h echo.HandlerFunc, method, path, body string, params map[string]string) *httptest.ResponseRecorder { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(method, path, strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + names := make([]string, 0, len(params)) + values := make([]string, 0, len(params)) + for k, v := range params { + names = append(names, k) + values = append(values, v) + } + c.SetParamNames(names...) + c.SetParamValues(values...) + Expect(h(c)).To(Succeed()) + return rec + } + + // assertRefused checks the STATUS and the BODY. A status-only assertion + // cannot tell a 501 with an empty body from the empty 200 this replaces, + // and the body is the whole reason the answer is worth giving. + assertRefused := func(rec *httptest.ResponseRecorder, surface string) { + GinkgoHelper() + Expect(rec.Code).To(Equal(http.StatusNotImplemented)) + var body map[string]string + Expect(json.Unmarshal(rec.Body.Bytes(), &body)).To(Succeed()) + Expect(body["error"]).To(ContainSubstring(surface)) + Expect(body["error"]).To(ContainSubstring("distributed mode")) + Expect(body["error"]).To(ContainSubstring("agent worker")) + } + + It("refuses to list prompts rather than answering with an empty list", func() { + assertRefused(call(MCPPromptsEndpoint(emptyLoader, distributed), + http.MethodGet, "/v1/mcp/prompts/m", "", map[string]string{"model": "m"}), "prompts") + }) + + It("refuses to expand a prompt", func() { + assertRefused(call(MCPGetPromptEndpoint(emptyLoader, distributed), + http.MethodPost, "/v1/mcp/prompts/m/p", `{"arguments":{}}`, + map[string]string{"model": "m", "prompt": "p"}), "prompts") + }) + + It("refuses to list resources rather than answering with an empty list", func() { + assertRefused(call(MCPResourcesEndpoint(emptyLoader, distributed), + http.MethodGet, "/v1/mcp/resources/m", "", map[string]string{"model": "m"}), "resources") + }) + + It("refuses to read a resource", func() { + assertRefused(call(MCPReadResourceEndpoint(emptyLoader, distributed), + http.MethodPost, "/v1/mcp/resources/m/read", `{"uri":"file:///x"}`, + map[string]string{"model": "m"}), "resources") + }) + + It("does not refuse in a standalone deployment, where the sessions are here", func() { + // The negative control. A guard that fired unconditionally would pass + // every assertion above and break MCP prompts for every single-binary + // deployment, which is where they actually work. + e := echo.New() + req := httptest.NewRequest(http.MethodGet, "/v1/mcp/prompts/m", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("model") + c.SetParamValues("m") + + err := MCPPromptsEndpoint(emptyLoader, standalone)(c) + // The model does not exist, so this fails on the lookup. What matters + // is that it got PAST the guard: a 501 here would mean the guard fired + // with distributed mode off. + Expect(err).To(HaveOccurred()) + Expect(rec.Code).ToNot(Equal(http.StatusNotImplemented)) + }) +}) diff --git a/core/http/endpoints/localai/mcp_prompts.go b/core/http/endpoints/localai/mcp_prompts.go index 8f04ee6c8..b9cc351bd 100644 --- a/core/http/endpoints/localai/mcp_prompts.go +++ b/core/http/endpoints/localai/mcp_prompts.go @@ -17,6 +17,13 @@ func MCPPromptsEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applicat return echo.ErrBadRequest } + // Before the model lookup and before the empty-config shortcut: both + // answer 200 with an empty list, which is the very answer this refuses + // to give in a deployment that cannot look. + if refused, err := mcpLocalSessionsOnly(c, appConfig, "prompts"); refused { + return err + } + cfg, exists := cl.GetModelConfig(modelName) if !exists { return fmt.Errorf("model %q not found", modelName) @@ -86,6 +93,10 @@ func MCPGetPromptEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applic return echo.ErrBadRequest } + if refused, err := mcpLocalSessionsOnly(c, appConfig, "prompts"); refused { + return err + } + cfg, exists := cl.GetModelConfig(modelName) if !exists { return fmt.Errorf("model %q not found", modelName) diff --git a/core/http/endpoints/localai/mcp_resources.go b/core/http/endpoints/localai/mcp_resources.go index 0cacec3c1..0683362a3 100644 --- a/core/http/endpoints/localai/mcp_resources.go +++ b/core/http/endpoints/localai/mcp_resources.go @@ -17,6 +17,12 @@ func MCPResourcesEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applic return echo.ErrBadRequest } + // See the prompts endpoint: this has to come before the empty-config + // shortcut, which answers with the empty list this refuses to invent. + if refused, err := mcpLocalSessionsOnly(c, appConfig, "resources"); refused { + return err + } + cfg, exists := cl.GetModelConfig(modelName) if !exists { return fmt.Errorf("model %q not found", modelName) @@ -73,6 +79,10 @@ func MCPReadResourceEndpoint(cl *config.ModelConfigLoader, appConfig *config.App return echo.ErrBadRequest } + if refused, err := mcpLocalSessionsOnly(c, appConfig, "resources"); refused { + return err + } + cfg, exists := cl.GetModelConfig(modelName) if !exists { return fmt.Errorf("model %q not found", modelName) diff --git a/core/http/endpoints/localai/mcp_tools.go b/core/http/endpoints/localai/mcp_tools.go index f5db27bd7..86ca3bdc6 100644 --- a/core/http/endpoints/localai/mcp_tools.go +++ b/core/http/endpoints/localai/mcp_tools.go @@ -13,7 +13,7 @@ import ( // MCPServersEndpoint returns the list of MCP servers and their tools for a given model. // GET /v1/mcp/servers/:model -func MCPServersEndpoint(cl *config.ModelConfigLoader, appConfig *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient) echo.HandlerFunc { +func MCPServersEndpoint(cl *config.ModelConfigLoader, appConfig *config.ApplicationConfig, agentControl mcpTools.AgentControl) echo.HandlerFunc { return func(c echo.Context) error { modelName := c.Param("model") if modelName == "" { @@ -45,10 +45,10 @@ func MCPServersEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applicat }) } - // In distributed mode, route discovery through NATS to an agent worker + // In distributed mode, ask an agent worker over the tunnel it holds // that can actually connect to the MCP servers. - if natsClient != nil { - resp, err := mcpTools.DiscoverMCPToolsRemote(c.Request().Context(), natsClient, cfg.Name, remote, stdio) + if agentControl != nil { + resp, err := mcpTools.DiscoverMCPToolsRemote(c.Request().Context(), agentControl, cfg.Name, remote, stdio) if err != nil { return c.JSON(http.StatusOK, map[string]any{ "model": modelName, @@ -80,7 +80,7 @@ func MCPServersEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applicat // MCPServersEndpointFromMiddleware is a version that uses the middleware-resolved model config. // This allows it to use the same middleware chain as other endpoints. -func MCPServersEndpointFromMiddleware(natsClient mcpTools.MCPNATSClient) echo.HandlerFunc { +func MCPServersEndpointFromMiddleware(agentControl mcpTools.AgentControl) echo.HandlerFunc { return func(c echo.Context) error { cfg, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) if !ok || cfg == nil { @@ -103,9 +103,9 @@ func MCPServersEndpointFromMiddleware(natsClient mcpTools.MCPNATSClient) echo.Ha }) } - // In distributed mode, route discovery through NATS to an agent worker. - if natsClient != nil { - resp, err := mcpTools.DiscoverMCPToolsRemote(c.Request().Context(), natsClient, cfg.Name, remote, stdio) + // In distributed mode, ask an agent worker over the tunnel it holds. + if agentControl != nil { + resp, err := mcpTools.DiscoverMCPToolsRemote(c.Request().Context(), agentControl, cfg.Name, remote, stdio) if err != nil { return c.JSON(http.StatusOK, map[string]any{ "model": cfg.Name, diff --git a/core/http/endpoints/mcp/executor.go b/core/http/endpoints/mcp/executor.go index 9f9b279d6..11dcd4574 100644 --- a/core/http/endpoints/mcp/executor.go +++ b/core/http/endpoints/mcp/executor.go @@ -11,7 +11,8 @@ import ( ) // ToolExecutor abstracts MCP tool discovery and execution. -// Implementations handle local (in-process sessions) vs distributed (NATS) modes. +// Implementations handle local (in-process sessions) vs distributed (an agent +// worker reached over its tunnel) modes. type ToolExecutor interface { // DiscoverTools returns the tool function schemas available from MCP servers. DiscoverTools(ctx context.Context) ([]functions.Function, error) @@ -58,28 +59,29 @@ func (e *LocalToolExecutor) HasTools() bool { return len(e.tools) > 0 } -// DistributedToolExecutor routes tool operations through NATS to agent workers. +// DistributedToolExecutor routes tool operations to agent workers over the +// tunnels they hold. type DistributedToolExecutor struct { - natsClient MCPNATSClient - modelName string - remote config.MCPGenericConfig[config.MCPRemoteServers] - stdio config.MCPGenericConfig[config.MCPSTDIOServers] - toolDefs []mcpRemote.MCPToolDef + agent AgentControl + modelName string + remote config.MCPGenericConfig[config.MCPRemoteServers] + stdio config.MCPGenericConfig[config.MCPSTDIOServers] + toolDefs []mcpRemote.MCPToolDef } -// NewDistributedToolExecutor creates a ToolExecutor that routes through NATS. -// It discovers tools immediately via a NATS request-reply to an agent worker. -func NewDistributedToolExecutor(ctx context.Context, natsClient MCPNATSClient, modelName string, +// NewDistributedToolExecutor creates a ToolExecutor that routes to agent +// workers. It discovers tools immediately with a control RPC to one of them. +func NewDistributedToolExecutor(ctx context.Context, agent AgentControl, modelName string, remote config.MCPGenericConfig[config.MCPRemoteServers], stdio config.MCPGenericConfig[config.MCPSTDIOServers], ) *DistributedToolExecutor { e := &DistributedToolExecutor{ - natsClient: natsClient, - modelName: modelName, - remote: remote, - stdio: stdio, + agent: agent, + modelName: modelName, + remote: remote, + stdio: stdio, } - resp, err := DiscoverMCPToolsRemote(ctx, natsClient, modelName, remote, stdio) + resp, err := DiscoverMCPToolsRemote(ctx, agent, modelName, remote, stdio) if err != nil { xlog.Error("Failed to discover MCP tools (distributed)", "error", err) } else if resp != nil { @@ -103,7 +105,7 @@ func (e *DistributedToolExecutor) IsTool(name string) bool { } func (e *DistributedToolExecutor) ExecuteTool(ctx context.Context, toolName, arguments string) (string, error) { - return ExecuteMCPToolCallRemote(ctx, e.natsClient, e.modelName, e.remote, e.stdio, toolName, arguments) + return ExecuteMCPToolCallRemote(ctx, e.agent, e.modelName, e.remote, e.stdio, toolName, arguments) } func (e *DistributedToolExecutor) HasTools() bool { @@ -111,15 +113,16 @@ func (e *DistributedToolExecutor) HasTools() bool { } // NewToolExecutor creates the appropriate ToolExecutor based on the current mode. -// When natsClient is non-nil, returns a DistributedToolExecutor that routes through NATS. -// When natsClient is nil, creates local sessions and returns a LocalToolExecutor. -func NewToolExecutor(ctx context.Context, natsClient MCPNATSClient, modelName string, +// When agent is non-nil, returns a DistributedToolExecutor that reaches an agent +// worker over its tunnel. When agent is nil, creates local sessions and returns +// a LocalToolExecutor. +func NewToolExecutor(ctx context.Context, agent AgentControl, modelName string, remote config.MCPGenericConfig[config.MCPRemoteServers], stdio config.MCPGenericConfig[config.MCPSTDIOServers], enabledServers []string, ) ToolExecutor { - if natsClient != nil { - return NewDistributedToolExecutor(ctx, natsClient, modelName, remote, stdio) + if agent != nil { + return NewDistributedToolExecutor(ctx, agent, modelName, remote, stdio) } sessions, err := NamedSessionsFromMCPConfig(modelName, remote, stdio, enabledServers) if err != nil || len(sessions) == 0 { diff --git a/core/http/endpoints/mcp/remote_test.go b/core/http/endpoints/mcp/remote_test.go new file mode 100644 index 000000000..6a82b8f24 --- /dev/null +++ b/core/http/endpoints/mcp/remote_test.go @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: MIT + +package mcp + +import ( + "context" + "errors" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + mcpRemote "github.com/mudler/LocalAI/core/services/mcp" +) + +// recordingAgent records the context each verb was given, so a spec can read +// the budget the caller applied. +// +// The budget is invisible any other way. It travels as a context deadline, not +// as anything on the wire, so the only place it can be observed is inside the +// call; a spec that waited for it to expire would be a spec that sleeps for +// minutes, and one that read a constant would pass with the constant unused. +type recordingAgent struct { + toolDeadline time.Time + toolHasDL bool + discoDeadline time.Time + discoHasDL bool + + toolReply *mcpRemote.MCPToolResponse + discoReply *mcpRemote.MCPDiscoveryResponse + err error +} + +func (a *recordingAgent) ExecuteMCPTool(ctx context.Context, _ mcpRemote.MCPToolRequest) (*mcpRemote.MCPToolResponse, error) { + a.toolDeadline, a.toolHasDL = ctx.Deadline() + if a.err != nil { + return nil, a.err + } + return a.toolReply, nil +} + +func (a *recordingAgent) DiscoverMCPTools(ctx context.Context, _ mcpRemote.MCPDiscoveryRequest) (*mcpRemote.MCPDiscoveryResponse, error) { + a.discoDeadline, a.discoHasDL = ctx.Deadline() + if a.err != nil { + return nil, a.err + } + return a.discoReply, nil +} + +var ( + noRemote = config.MCPGenericConfig[config.MCPRemoteServers]{} + noStdio = config.MCPGenericConfig[config.MCPSTDIOServers]{} +) + +var _ = Describe("Routing MCP to an agent worker", func() { + ctx := context.Background() + + Describe("with no agent control client wired", func() { + // A frontend in distributed mode with nothing to reach an agent worker + // through must say so. The alternative shape, which this replaces, was + // a nil-pointer dereference inside a chat request. + It("names the missing client rather than dialling nothing", func() { + _, err := ExecuteMCPToolCallRemote(ctx, nil, "m", noRemote, noStdio, "weather", "{}") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("agent control client")) + }) + + It("names it for discovery too", func() { + _, err := DiscoverMCPToolsRemote(ctx, nil, "m", noRemote, noStdio) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("agent control client")) + }) + }) + + Describe("the budget it gives the call", func() { + // Two verbs, two constants, two assertions. The rule is written at both + // sites, so a single spec would leave whichever site it did not cover + // free to lose its deadline: a tool call whose worker went quiet would + // then hold the caller until the tunnel's keepalive noticed. + // + // Asserted as a WINDOW around the deadline rather than an equality, + // because the deadline is stamped from a clock this spec does not hold. + // The window is far tighter than the difference between the two + // constants, so a call given the wrong one still fails here. + const slack = 5 * time.Second + + It("bounds a tool call by the documented tool timeout", func() { + agent := &recordingAgent{toolReply: &mcpRemote.MCPToolResponse{Result: "ok"}} + start := time.Now() + + out, err := ExecuteMCPToolCallRemote(ctx, agent, "m", noRemote, noStdio, "weather", `{"city":"London"}`) + Expect(err).ToNot(HaveOccurred()) + Expect(out).To(Equal("ok")) + + Expect(agent.toolHasDL).To(BeTrue(), "the tool call was given no deadline at all") + Expect(agent.toolDeadline).To(BeTemporally("~", start.Add(config.DefaultMCPToolTimeout), slack)) + }) + + It("bounds discovery by the documented discovery timeout", func() { + agent := &recordingAgent{discoReply: &mcpRemote.MCPDiscoveryResponse{}} + start := time.Now() + + _, err := DiscoverMCPToolsRemote(ctx, agent, "m", noRemote, noStdio) + Expect(err).ToNot(HaveOccurred()) + + Expect(agent.discoHasDL).To(BeTrue(), "discovery was given no deadline at all") + Expect(agent.discoDeadline).To(BeTemporally("~", start.Add(config.DefaultMCPDiscoveryTimeout), slack)) + }) + + It("does not extend a caller's shorter deadline", func() { + // The caller's own budget wins. context.WithTimeout keeps the + // earlier of the two, and a hand-rolled deadline would not. + short, cancel := context.WithTimeout(ctx, time.Second) + defer cancel() + agent := &recordingAgent{toolReply: &mcpRemote.MCPToolResponse{Result: "ok"}} + + _, err := ExecuteMCPToolCallRemote(short, agent, "m", noRemote, noStdio, "weather", "{}") + Expect(err).ToNot(HaveOccurred()) + Expect(agent.toolDeadline).To(BeTemporally("<", time.Now().Add(config.DefaultMCPToolTimeout))) + }) + }) + + It("passes the agent's failure through with its identity intact", func() { + // The classification is the control client's, and re-deciding it here + // would be the same rule in two places. What this pins is that + // wrapping it for a human does not hide it from errors.Is. + boom := errors.New("the fleet said no") + agent := &recordingAgent{err: boom} + + _, err := ExecuteMCPToolCallRemote(ctx, agent, "m", noRemote, noStdio, "weather", "{}") + Expect(err).To(MatchError(boom)) + _, err = DiscoverMCPToolsRemote(ctx, agent, "m", noRemote, noStdio) + Expect(err).To(MatchError(boom)) + }) + + It("refuses tool arguments that are not JSON before reaching for a worker", func() { + agent := &recordingAgent{toolReply: &mcpRemote.MCPToolResponse{Result: "ok"}} + + _, err := ExecuteMCPToolCallRemote(ctx, agent, "m", noRemote, noStdio, "weather", "{not json") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("invalid tool arguments JSON")) + Expect(agent.toolHasDL).To(BeFalse(), "no worker should have been asked") + }) +}) diff --git a/core/http/endpoints/mcp/tools.go b/core/http/endpoints/mcp/tools.go index 0b4931a05..f6524bf1a 100644 --- a/core/http/endpoints/mcp/tools.go +++ b/core/http/endpoints/mcp/tools.go @@ -16,7 +16,6 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/schema" mcpRemote "github.com/mudler/LocalAI/core/services/mcp" - "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/pkg/functions" "github.com/mudler/LocalAI/pkg/httpclient" @@ -100,9 +99,23 @@ var ( client = mcp.NewClient(&mcp.Implementation{Name: "LocalAI", Version: "v1.0.0"}, nil) ) -// MCPNATSClient is the interface for NATS request-reply operations needed by MCP routing. -type MCPNATSClient interface { - Request(subject string, data []byte, timeout time.Duration) ([]byte, error) +// AgentControl is the frontend's port onto the MCP verbs an agent worker serves +// over the tunnel it holds. *nodes.AgentControlClient is the only production +// implementation, and it is what decides WHICH agent worker answers. +// +// The contract a caller here relies on, and the reason these functions do no +// classification of their own: an implementation returns the worker's own Error +// field AS A GO ERROR, so a nil error means the verb succeeded. Re-reading the +// Error field at every call site would be the same rule written twice, and the +// copy that gets forgotten is the one that reports a failed tool call as an +// empty success. +// +// It replaces an interface over NATS request-reply. What that carried was a +// subject and a queue group, which between them chose a worker; the choosing is +// now a query and the carrying is an ordinary control RPC. +type AgentControl interface { + ExecuteMCPTool(ctx context.Context, req mcpRemote.MCPToolRequest) (*mcpRemote.MCPToolResponse, error) + DiscoverMCPTools(ctx context.Context, req mcpRemote.MCPDiscoveryRequest) (*mcpRemote.MCPDiscoveryResponse, error) } // MetadataKeyLocalAIAssistant is the request-metadata key the chat handler @@ -510,18 +523,28 @@ func ExecuteMCPToolCall(ctx context.Context, tools []MCPToolInfo, toolName strin return string(combined), nil } -// ExecuteMCPToolCallRemote routes an MCP tool execution request to an agent worker via NATS. -// Used in distributed mode when the frontend doesn't hold MCP sessions locally. +// ExecuteMCPToolCallRemote runs one MCP tool on an agent worker. +// +// Used in distributed mode, where the frontend holds no MCP sessions of its +// own: an agent worker is what can create them (stdio servers under docker), +// so the frontend serialises the model's MCP configuration and asks one. +// +// The budget is applied HERE, as a context deadline, and that is the whole of +// the change in where it lives. On the bus it was the request-reply timeout, +// which was the only thing bounding a worker that never answered; the control +// RPC carries no deadline of its own (see nodes.ControlClient.clientFor), so +// without this a tool call whose worker went quiet holds the caller until the +// tunnel's own keepalive notices, which is far longer than any caller expects. func ExecuteMCPToolCallRemote( ctx context.Context, - natsClient MCPNATSClient, + agent AgentControl, modelName string, remote config.MCPGenericConfig[config.MCPRemoteServers], stdio config.MCPGenericConfig[config.MCPSTDIOServers], toolName, arguments string, ) (string, error) { - if natsClient == nil { - return "", fmt.Errorf("NATS client not configured for distributed MCP") + if agent == nil { + return "", fmt.Errorf("no agent control client is configured for distributed MCP: this frontend cannot reach an agent worker to run tool %q", toolName) } var args map[string]any @@ -531,63 +554,50 @@ func ExecuteMCPToolCallRemote( } } - req := mcpRemote.MCPToolRequest{ + ctx, cancel := context.WithTimeout(ctx, config.DefaultMCPToolTimeout) + defer cancel() + + resp, err := agent.ExecuteMCPTool(ctx, mcpRemote.MCPToolRequest{ ModelName: modelName, ToolName: toolName, Arguments: args, RemoteServers: remote, StdioServers: stdio, - } - reqData, _ := json.Marshal(req) - - replyData, err := natsClient.Request(messaging.SubjectMCPToolExecute, reqData, config.DefaultMCPToolTimeout) + }) if err != nil { - return "", fmt.Errorf("NATS MCP tool request failed: %w", err) - } - - var resp mcpRemote.MCPToolResponse - if err := json.Unmarshal(replyData, &resp); err != nil { - return "", fmt.Errorf("unmarshal MCP reply: %w", err) - } - if resp.Error != "" { - return "", fmt.Errorf("remote MCP tool error: %s", resp.Error) + return "", fmt.Errorf("the MCP tool call could not be run on an agent worker: %w", err) } return resp.Result, nil } -// DiscoverMCPToolsRemote routes an MCP discovery request to an agent worker via NATS. -// Returns server info and tool function schemas from the remote worker. +// DiscoverMCPToolsRemote asks an agent worker which MCP servers and tool +// schemas a model's configuration reaches. func DiscoverMCPToolsRemote( ctx context.Context, - natsClient MCPNATSClient, + agent AgentControl, modelName string, remote config.MCPGenericConfig[config.MCPRemoteServers], stdio config.MCPGenericConfig[config.MCPSTDIOServers], ) (*mcpRemote.MCPDiscoveryResponse, error) { - if natsClient == nil { - return nil, fmt.Errorf("NATS client not configured for distributed MCP") + if agent == nil { + return nil, fmt.Errorf("no agent control client is configured for distributed MCP: this frontend cannot reach an agent worker to discover the tools of model %q", modelName) } - req := mcpRemote.MCPDiscoveryRequest{ + // Its own budget, and its own constant. Discovery opens every configured + // MCP server, which a tool call on an already-open session does not, so the + // two are not the same wait and never were. + ctx, cancel := context.WithTimeout(ctx, config.DefaultMCPDiscoveryTimeout) + defer cancel() + + resp, err := agent.DiscoverMCPTools(ctx, mcpRemote.MCPDiscoveryRequest{ ModelName: modelName, RemoteServers: remote, StdioServers: stdio, - } - reqData, _ := json.Marshal(req) - - replyData, err := natsClient.Request(messaging.SubjectMCPDiscovery, reqData, config.DefaultMCPDiscoveryTimeout) + }) if err != nil { - return nil, fmt.Errorf("NATS MCP discovery request failed: %w", err) + return nil, fmt.Errorf("MCP discovery could not be run on an agent worker: %w", err) } - - var resp mcpRemote.MCPDiscoveryResponse - if err := json.Unmarshal(replyData, &resp); err != nil { - return nil, fmt.Errorf("unmarshal MCP discovery reply: %w", err) - } - if resp.Error != "" { - return nil, fmt.Errorf("remote MCP discovery error: %s", resp.Error) - } - return &resp, nil + return resp, nil } // ListMCPServers returns server info with tool, prompt, and resource names for each session. diff --git a/core/http/endpoints/openai/chat.go b/core/http/endpoints/openai/chat.go index fcec18932..72db7ff0b 100644 --- a/core/http/endpoints/openai/chat.go +++ b/core/http/endpoints/openai/chat.go @@ -129,7 +129,7 @@ func applyAutoparserOverride( // @Param request body schema.OpenAIRequest true "query params" // @Success 200 {object} schema.OpenAIResponse "Response" // @Router /v1/chat/completions [post] -func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, startupOptions *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient, assistantHolder *mcpTools.LocalAIAssistantHolder, compressor middleware.ChatCompressor) echo.HandlerFunc { +func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, startupOptions *config.ApplicationConfig, agentControl mcpTools.AgentControl, assistantHolder *mcpTools.LocalAIAssistantHolder, compressor middleware.ChatCompressor) echo.HandlerFunc { return func(c echo.Context) error { var textContentToReturn string id := uuid.New().String() @@ -218,7 +218,7 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator if (len(mcpServers) > 0 || mcpPromptName != "" || len(mcpResourceURIs) > 0) && (config.MCP.Servers != "" || config.MCP.Stdio != "") { remote, stdio, mcpErr := config.MCP.MCPConfigFromYAML() if mcpErr == nil { - mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), natsClient, config.Name, remote, stdio, mcpServers) + mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), agentControl, config.Name, remote, stdio, mcpServers) // Prompt and resource injection (pre-processing step — resolves locally regardless of distributed mode) namedSessions, sessErr := mcpTools.NamedSessionsFromMCPConfig(config.Name, remote, stdio, mcpServers) diff --git a/core/http/endpoints/openresponses/responses.go b/core/http/endpoints/openresponses/responses.go index d86780fbd..672f0025f 100644 --- a/core/http/endpoints/openresponses/responses.go +++ b/core/http/endpoints/openresponses/responses.go @@ -31,7 +31,7 @@ import ( // @Param request body schema.OpenResponsesRequest true "Request body" // @Success 200 {object} schema.ORResponseResource "Response" // @Router /v1/responses [post] -func ResponsesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient) echo.HandlerFunc { +func ResponsesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, agentControl mcpTools.AgentControl) echo.HandlerFunc { return func(c echo.Context) error { createdAt := time.Now().Unix() responseID := fmt.Sprintf("resp_%s", uuid.New().String()) @@ -124,7 +124,7 @@ func ResponsesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, eval if !hasMCPRequest { enabledServers = nil // backward compat: auto-activate all servers } - mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), natsClient, cfg.Name, remote, stdio, enabledServers) + mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), agentControl, cfg.Name, remote, stdio, enabledServers) // Prompt and resource injection (pre-processing step — resolves locally regardless of distributed mode) if hasMCPRequest { diff --git a/core/http/routes/agent_control.go b/core/http/routes/agent_control.go new file mode 100644 index 000000000..1fc23156c --- /dev/null +++ b/core/http/routes/agent_control.go @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: MIT + +package routes + +import ( + "github.com/mudler/LocalAI/core/application" + mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" +) + +// mcpAgentControl returns the port the MCP endpoints reach an agent worker +// through, or a nil interface in a deployment that has no agent workers to +// reach. +// +// One function for four route files, and it exists because of the trap rather +// than to save three lines. Every one of those endpoints decides between its +// local and its distributed path by asking whether this value is nil, and a nil +// *nodes.AgentControlClient assigned straight into the interface is NOT nil: +// the interface carries a type, so the check passes, the distributed path is +// taken, and every MCP request in a standalone deployment fails instead of +// using the in-process sessions it has. Returning the untyped nil explicitly is +// the only spelling that keeps that check meaning what its call sites think it +// means. +func mcpAgentControl(app *application.Application) mcpTools.AgentControl { + return mcpAgentControlOf(app.Distributed()) +} + +// mcpAgentControlOf is the half of the rule above that a spec can reach. +// +// It is separate because an *application.Application carries its distributed +// services in an unexported field with no way in from outside the package, so +// the typed-nil conversion could not otherwise be asserted at all, and it is +// precisely the conversion that is easy to get wrong and silent when it is. +func mcpAgentControlOf(d *application.DistributedServices) mcpTools.AgentControl { + if d == nil || d.AgentControl == nil { + return nil + } + return d.AgentControl +} diff --git a/core/http/routes/agent_control_internal_test.go b/core/http/routes/agent_control_internal_test.go new file mode 100644 index 000000000..64fb60e2a --- /dev/null +++ b/core/http/routes/agent_control_internal_test.go @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: MIT + +package routes + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/application" + "github.com/mudler/LocalAI/core/services/nodes" +) + +// Every MCP endpoint decides between its LOCAL sessions and an agent worker by +// asking whether this value is nil. That check means what its call sites think +// it means only if the nil that reaches it is an untyped one. +var _ = Describe("the MCP endpoints' agent control port", func() { + It("is nil in a deployment with no distributed services", func() { + Expect(mcpAgentControlOf(nil)).To(BeNil()) + }) + + It("is nil, and not a non-nil interface holding a nil pointer, when nothing built the client", func() { + // The trap, stated as an assertion. A nil *nodes.AgentControlClient + // assigned straight into the interface compares NON-nil, so every MCP + // request would take the distributed path and fail, in a deployment + // that has perfectly good in-process sessions. + Expect(mcpAgentControlOf(&application.DistributedServices{})).To(BeNil()) + }) + + It("carries the client through when there is one", func() { + // The negative control: a function that returned nil unconditionally + // would pass both assertions above and disable distributed MCP + // entirely, with the local path silently taken instead. + client := nodes.NewAgentControlClient(nil, nil) + Expect(mcpAgentControlOf(&application.DistributedServices{AgentControl: client})).To(BeIdenticalTo(client)) + }) +}) diff --git a/core/http/routes/anthropic.go b/core/http/routes/anthropic.go index 124557655..3c305087a 100644 --- a/core/http/routes/anthropic.go +++ b/core/http/routes/anthropic.go @@ -10,7 +10,6 @@ import ( "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/endpoints/anthropic" - mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" "github.com/mudler/LocalAI/core/http/middleware" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services/routing/pii" @@ -25,17 +24,14 @@ func RegisterAnthropicRoutes(app *echo.Echo, application *application.Application, ) { // Anthropic Messages API endpoint - var natsClient mcpTools.MCPNATSClient - if d := application.Distributed(); d != nil { - natsClient = d.Nats - } + agentControl := mcpAgentControl(application) messagesHandler := anthropic.MessagesEndpoint( application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig(), - natsClient, + agentControl, ) messagesMiddleware := []echo.MiddlewareFunc{ diff --git a/core/http/routes/localai.go b/core/http/routes/localai.go index 2120a07f1..ae5ad8667 100644 --- a/core/http/routes/localai.go +++ b/core/http/routes/localai.go @@ -6,7 +6,6 @@ import ( "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/endpoints/localai" - mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" "github.com/mudler/LocalAI/core/http/middleware" "github.com/mudler/LocalAI/core/schema" compressionservice "github.com/mudler/LocalAI/core/services/compression" @@ -455,11 +454,8 @@ func RegisterLocalAIRoutes(router *echo.Echo, compressionservice.CounterFunc(tokens.CountMessages), compressionservice.NewInferenceSummarizer(cl, ml, appConfig), ) - var mcpNATS mcpTools.MCPNATSClient - if d := app.Distributed(); d != nil { - mcpNATS = d.Nats - } - mcpStreamHandler := localai.MCPEndpoint(cl, ml, evaluator, appConfig, mcpNATS, chatCompressor) + agentControl := mcpAgentControl(app) + mcpStreamHandler := localai.MCPEndpoint(cl, ml, evaluator, appConfig, agentControl, chatCompressor) mcpStreamMiddleware := []echo.MiddlewareFunc{ requestExtractor.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)), requestExtractor.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }), @@ -478,7 +474,7 @@ func RegisterLocalAIRoutes(router *echo.Echo, router.POST("/mcp/chat/completions", mcpStreamHandler, mcpStreamMiddleware...) // MCP server listing endpoint - router.GET("/v1/mcp/servers/:model", localai.MCPServersEndpoint(cl, appConfig, mcpNATS), mcpMw) + router.GET("/v1/mcp/servers/:model", localai.MCPServersEndpoint(cl, appConfig, agentControl), mcpMw) // MCP prompts endpoints router.GET("/v1/mcp/prompts/:model", localai.MCPPromptsEndpoint(cl, appConfig), mcpMw) diff --git a/core/http/routes/openai.go b/core/http/routes/openai.go index 6a9012626..4eb1250ee 100644 --- a/core/http/routes/openai.go +++ b/core/http/routes/openai.go @@ -5,7 +5,6 @@ import ( "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/endpoints/localai" - mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" "github.com/mudler/LocalAI/core/http/endpoints/openai" "github.com/mudler/LocalAI/core/http/middleware" "github.com/mudler/LocalAI/core/schema" @@ -38,18 +37,15 @@ func RegisterOpenAIRoutes(app *echo.Echo, app.POST("/v1/realtime/transcription_session", openai.RealtimeTranscriptionSession(application), traceMiddleware) app.POST("/v1/realtime/calls", openai.RealtimeCalls(application), traceMiddleware) - // NATS client for distributed MCP tool routing (nil when not in distributed mode) - var natsClient mcpTools.MCPNATSClient - if d := application.Distributed(); d != nil { - natsClient = d.Nats - } + // How the MCP endpoints reach an agent worker; nil outside distributed mode. + agentControl := mcpAgentControl(application) // chat chatCompressor := compressionservice.New( compressionservice.CounterFunc(tokens.CountMessages), compressionservice.NewInferenceSummarizer(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()), ) - chatHandler := openai.ChatEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig(), natsClient, application.LocalAIAssistant(), chatCompressor) + chatHandler := openai.ChatEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig(), agentControl, application.LocalAIAssistant(), chatCompressor) chatMiddleware := []echo.MiddlewareFunc{ nodeHeaderMiddleware, usageMiddleware, diff --git a/core/http/routes/openresponses.go b/core/http/routes/openresponses.go index 8aff6ccf1..18a283b82 100644 --- a/core/http/routes/openresponses.go +++ b/core/http/routes/openresponses.go @@ -5,7 +5,6 @@ import ( "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/config" localai "github.com/mudler/LocalAI/core/http/endpoints/localai" - mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" "github.com/mudler/LocalAI/core/http/endpoints/openresponses" "github.com/mudler/LocalAI/core/http/middleware" "github.com/mudler/LocalAI/core/schema" @@ -16,11 +15,9 @@ func RegisterOpenResponsesRoutes(app *echo.Echo, re *middleware.RequestExtractor, application *application.Application) { - // NATS client for distributed MCP tool routing (nil when not in distributed mode) - var natsClient mcpTools.MCPNATSClient + // How the MCP endpoints reach an agent worker; nil outside distributed mode. + agentControl := mcpAgentControl(application) if d := application.Distributed(); d != nil { - natsClient = d.Nats - // Replicate response metadata across frontend replicas and subscribe to // delegated cancels. Without this a GET, a previous_response_id lookup or // a cancel that the load balancer sends to a replica other than the @@ -38,7 +35,7 @@ func RegisterOpenResponsesRoutes(app *echo.Echo, application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig(), - natsClient, + agentControl, ) responsesMiddleware := []echo.MiddlewareFunc{ diff --git a/core/services/agentworker/handlers.go b/core/services/agentworker/handlers.go index 435040edd..b933362df 100644 --- a/core/services/agentworker/handlers.go +++ b/core/services/agentworker/handlers.go @@ -11,6 +11,8 @@ // This package is the worker half of moving them onto the same HTTP control // plane a backend worker already serves. The frontend half, which chooses WHICH // agent worker to ask, is not here: this package answers, it does not select. +// That selection is nodes.AgentSelector, and it is a query against the +// connection rows rather than anything a broker does. package agentworker import ( diff --git a/core/services/cluster/connected_among_test.go b/core/services/cluster/connected_among_test.go new file mode 100644 index 000000000..77916d16a --- /dev/null +++ b/core/services/cluster/connected_among_test.go @@ -0,0 +1,173 @@ +// SPDX-License-Identifier: MIT + +package cluster_test + +import ( + "context" + "strings" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// ConnectedAmong is the read a SELECTION is built on: which of these nodes can +// a request reach right now, and which of those can it reach without a relay +// hop. Every spec below is about one of the two ways that answer can be wrong, +// and both are the invariant this phase is built around rather than tidiness: +// +// - naming a node NO live replica holds sends a request at a process that is +// gone, and the caller reads the failure as the worker's own answer; +// - dropping a node a live replica DOES hold refuses work to a fleet that is +// fine, which is what an absence read leaking into a routing read looks +// like. +var _ = Describe("ConnectedAmong", func() { + var ( + ctx context.Context + db *gorm.DB + reg *cluster.Registry + ) + + // live registers a replica that is heartbeating now. + live := func(id string) { + GinkgoHelper() + Expect(reg.Register(ctx, id, "10.0.0.1:8080", "v1")).To(Succeed()) + } + + // kill stops a replica heartbeating, WITHOUT stamping a departure on the + // rows it holds. That is exactly the state a replica that died leaves + // behind until a peer's membership sweep runs. + kill := func(id string) { + GinkgoHelper() + age(ctx, db, "instances", "last_seen", "id", id, cluster.InstanceLiveness+5*time.Second) + } + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + It("reports a node this replica holds in BOTH lists", func() { + // heldByOwner is a subset of held and not an alternative to it. A + // caller that prefers what it owns and falls back to the rest would + // otherwise have to union the two itself, and a caller that forgot + // would never fall back at all. + live("me") + _, err := reg.Claim(ctx, "agent-1", "me") + Expect(err).ToNot(HaveOccurred()) + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(ConsistOf("agent-1")) + Expect(byOwner).To(ConsistOf("agent-1")) + }) + + It("reports a node a PEER holds as held, and not as held by this replica", func() { + live("me") + live("peer") + _, err := reg.Claim(ctx, "agent-1", "peer") + Expect(err).ToNot(HaveOccurred()) + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(ConsistOf("agent-1")) + Expect(byOwner).To(BeEmpty()) + }) + + It("drops a node whose owning replica is no longer live", func() { + // The whole reason the statement joins instances. A connection row + // outlives its owner by up to a liveness window plus a heartbeat, so + // without the join this read names a dead replica for that entire + // window and every caller acting on it dials a corpse. + live("dead") + _, err := reg.Claim(ctx, "agent-1", "dead") + Expect(err).ToNot(HaveOccurred()) + kill("dead") + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(BeEmpty()) + Expect(byOwner).To(BeEmpty()) + }) + + It("drops a node whose row records a departure, however recent", func() { + // A departed row is not a candidate for ROUTING, whatever its age. How + // old the departure is decides whether the worker has GONE, which is + // Presence's question and not this one; answering it here would put two + // windows on one fact. + live("me") + epoch, err := reg.Claim(ctx, "agent-1", "me") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "agent-1", "me", epoch)).To(Succeed()) + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(BeEmpty()) + Expect(byOwner).To(BeEmpty()) + }) + + It("answers about the ids it was given and no others", func() { + live("me") + for _, id := range []string{"agent-1", "agent-2", "backend-9"} { + _, err := reg.Claim(ctx, id, "me") + Expect(err).ToNot(HaveOccurred()) + } + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1", "agent-2"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(ConsistOf("agent-1", "agent-2")) + Expect(byOwner).To(ConsistOf("agent-1", "agent-2")) + }) + + It("answers an empty candidate set without asking the database", func() { + // `IN ()` is a syntax error in PostgreSQL, so a statement issued here + // would fail rather than answer nothing. + rec := newSQLRecorder() + recording := cluster.NewRegistry(db.Session(&gorm.Session{Logger: rec})) + + held, byOwner, err := recording.ConnectedAmong(ctx, nil, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(BeEmpty()) + Expect(byOwner).To(BeEmpty()) + Expect(rec.statementCount()).To(Equal(0)) + }) + + It("answers in one joined statement measured on the database clock", func() { + live("me") + _, err := reg.Claim(ctx, "agent-1", "me") + Expect(err).ToNot(HaveOccurred()) + + rec := newSQLRecorder() + recording := cluster.NewRegistry(db.Session(&gorm.Session{Logger: rec})) + _, _, err = recording.ConnectedAmong(ctx, []string{"agent-1", "agent-2"}, "me") + Expect(err).ToNot(HaveOccurred()) + + sql := strings.ToLower(rec.only()) + // only() rules out the read-per-id shape: between two statements the + // owning replica can die, and the answer would then be assembled from + // two different snapshots of the cluster. + Expect(sql).To(ContainSubstring("join")) + Expect(sql).To(ContainSubstring("instances")) + // The liveness window is computed by the DATABASE. A Go-side cutoff + // would appear here as a bound literal and would then move with each + // replica's clock skew. The test container shares this host's clock, so + // no behavioural spec can catch that; the statement shape is the only + // place it is visible. + Expect(sql).To(ContainSubstring("make_interval")) + Expect(sql).To(ContainSubstring("now()")) + Expect(sql).ToNot(MatchRegexp(`last_seen\s*>\s*'`), + "the liveness cutoff must not be a literal timestamp from this process's clock") + // Held-ness is asked in the SQL rather than left to the join alone. An + // empty owner id matches no instance today only because no replica + // registers under one, which is an accident of who registers rather + // than a property of ownership. + Expect(strings.Join(strings.Fields(sql), " ")).To(ContainSubstring( + "node_connections.owner_instance_id <> ''")) + }) +}) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go index 1f253cd07..26ef64dde 100644 --- a/core/services/cluster/ownership.go +++ b/core/services/cluster/ownership.go @@ -396,3 +396,95 @@ func (r *Registry) PurgeDepartedBefore(ctx context.Context, olderThan time.Durat } return res.RowsAffected, nil } + +// connectedAmongQuery answers, for a set of nodes at once, which of their +// tunnels a LIVE replica holds and which of those are held by one named +// replica. +// +// ONE statement rather than a read per id, for the reason Owner is one +// statement: between two reads an owning replica can die, and the answer would +// then be assembled from two different snapshots of the cluster. A selection +// built that way can name a node whose owner the second read would have +// rejected, which is the one thing this read exists to prevent. +// +// The join is what makes a DEAD owner's row invisible here. Without it a row +// outlives its owner by up to a liveness window plus a heartbeat, and every +// caller that acted on this answer would be sent at a process that is gone. +// +// Held-ness is asked in SQL and the departure is not asked at all, and that is +// the whole difference between this read and Presence. This one answers a +// ROUTING question ("can a request get there right now"), so a row nobody holds +// is simply not a candidate, whatever its age. Deciding whether a worker has +// GONE is Presence's job and needs the grace; nothing here may be read as +// absence, and nothing here reports any. +// +// The liveness window is computed by the DATABASE, like every other window in +// this package: it is compared across replicas, so a Go-side cutoff would make +// the effective window depend on each replica's clock skew. The test container +// shares this host's clock, so no behavioural spec can see the difference and +// the statement shape is pinned instead. +// +// The predicates are the package's own rather than copies, so two spellings of +// "live" or of "held" cannot drift into a node that one query calls connected +// and another calls gone. +// +// The bind order is the order the placeholders appear in the text: the owner, +// then the liveness window, then the node ids. +const connectedAmongQuery = ` +SELECT node_connections.node_id AS node_id, + (node_connections.owner_instance_id = ?) AS by_owner +FROM node_connections +JOIN instances + ON instances.id = node_connections.owner_instance_id + AND ` + instanceIsLive + ` +WHERE ` + connectionIsHeld + ` + AND node_connections.node_id IN ?` + +// ConnectedAmong returns the subset of nodeIDs whose tunnel a LIVE replica +// currently holds, and separately those held by owner. +// +// It takes ids rather than asking which nodes are agents, because this package +// is a leaf and node type lives in core/services/nodes. +// +// heldByOwner is a SUBSET of held, not an alternative to it. A caller that +// prefers the ids it owns falls back to the rest, and expressing that as two +// disjoint lists would make the fallback a union the caller had to build. +// +// The answer is a routing fact and NEVER an absence one. An id missing from +// both lists means no live replica holds its tunnel at this instant, which a +// worker reconnecting between replicas produces routinely; nothing may read it +// as the worker having gone away. See Presence for the read that can answer +// that question, and only that one. +func (r *Registry) ConnectedAmong(ctx context.Context, nodeIDs []string, owner string) (held []string, heldByOwner []string, err error) { + if len(nodeIDs) == 0 { + // Answered without a statement, because `IN ()` is a syntax error in + // PostgreSQL and gorm's expansion of an empty slice is a filter nobody + // wrote. An empty candidate set has exactly one honest answer. + return nil, nil, nil + } + // Refused rather than attempted, for the reason Owner refuses: now() and + // make_interval are PostgreSQL, so on the single-binary SQLite path this + // would fail as a missing function and read as a missing migration. It is + // deliberately not an empty answer: a deployment with no cluster has + // nothing to say about who holds a tunnel, and an empty list would be read + // as "nobody does". + if !isPostgres(r.db) { + return nil, nil, fmt.Errorf("reading which of %d nodes are connected: connection ownership requires PostgreSQL, this deployment runs on %q", len(nodeIDs), r.db.Dialector.Name()) + } + var rows []struct { + NodeID string + ByOwner bool + } + if err := r.db.WithContext(ctx).Raw(connectedAmongQuery, + owner, InstanceLiveness.Seconds(), nodeIDs, + ).Scan(&rows).Error; err != nil { + return nil, nil, fmt.Errorf("reading which of %d nodes are connected: %w", len(nodeIDs), err) + } + for _, row := range rows { + held = append(held, row.NodeID) + if row.ByOwner { + heldByOwner = append(heldByOwner, row.NodeID) + } + } + return held, heldByOwner, nil +} diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go index 54195cfd8..d23f9d98d 100644 --- a/core/services/cluster/ownership_test.go +++ b/core/services/cluster/ownership_test.go @@ -58,6 +58,14 @@ func (r *sqlRecorder) only() string { return r.statements[0] } +// statementCount reports how many statements were recorded, for the specs whose +// claim is that a call issued NONE. +func (r *sqlRecorder) statementCount() int { + r.mu.Lock() + defer r.mu.Unlock() + return len(r.statements) +} + // writeTarget matches the table a statement writes to, anchored at the verb so // the UPDATE inside an upsert's ON CONFLICT clause cannot be mistaken for one. var writeTarget = regexp.MustCompile(`^\s*(?i:delete\s+from|update)\s+"?([a-z_]+)"?`) diff --git a/core/services/mcp/remote.go b/core/services/mcp/remote.go index 17cfc1f36..ae789a1ab 100644 --- a/core/services/mcp/remote.go +++ b/core/services/mcp/remote.go @@ -5,9 +5,13 @@ import ( "github.com/mudler/LocalAI/pkg/functions" ) -// MCPToolRequest is the NATS request-reply payload for executing an MCP tool -// on an agent worker. The frontend serializes the model's MCP server config -// so the worker can create sessions and execute the tool. +// MCPToolRequest is the control-RPC payload for executing an MCP tool on an +// agent worker. The frontend serializes the model's MCP server config so the +// worker can create sessions and execute the tool. +// +// It lives in this package rather than beside either side of the call, because +// both sides need it and they are in packages that must not import each other: +// core/http/endpoints/mcp is the caller and core/cli is the worker. type MCPToolRequest struct { ModelName string `json:"model_name"` ToolName string `json:"tool_name"` @@ -16,21 +20,25 @@ type MCPToolRequest struct { StdioServers config.MCPGenericConfig[config.MCPSTDIOServers] `json:"stdio_servers"` } -// MCPToolResponse is the NATS reply for an MCP tool execution. +// MCPToolResponse is the reply to an MCP tool execution. +// +// A set Error is the WORKER'S OWN ANSWER, carried inside a successful reply +// rather than as a transport failure, which is what lets a caller tell "this +// tool rejected your arguments" from "this worker could not be reached". type MCPToolResponse struct { Result string `json:"result,omitempty"` Error string `json:"error,omitempty"` } -// MCPDiscoveryRequest is the NATS request-reply payload for discovering -// available MCP tools, prompts, and resources from a model's MCP servers. +// MCPDiscoveryRequest is the control-RPC payload for discovering available MCP +// tools, prompts, and resources from a model's MCP servers. type MCPDiscoveryRequest struct { ModelName string `json:"model_name"` RemoteServers config.MCPGenericConfig[config.MCPRemoteServers] `json:"remote_servers"` StdioServers config.MCPGenericConfig[config.MCPSTDIOServers] `json:"stdio_servers"` } -// MCPDiscoveryResponse is the NATS reply for an MCP discovery request. +// MCPDiscoveryResponse is the reply to an MCP discovery request. type MCPDiscoveryResponse struct { Servers []MCPServerInfo `json:"servers,omitempty"` Tools []MCPToolDef `json:"tools,omitempty"` // flattened tool list with functions @@ -48,8 +56,8 @@ type MCPServerInfo struct { } // MCPToolDef is a serializable tool definition (function schema) that can -// travel over NATS. Unlike MCPToolInfo which holds a live session pointer, -// this is pure data. +// travel between processes. Unlike MCPToolInfo which holds a live session +// pointer, this is pure data. type MCPToolDef struct { ServerName string `json:"server_name"` ToolName string `json:"tool_name"` diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index 8f5429a39..657a99f92 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -47,13 +47,6 @@ func SubjectJobResult(jobID string) string { return subjectJobProgressPrefix + sanitizeSubjectToken(jobID) + ".result" } -// MCP Tool Execution (Request-Reply via NATS — load-balanced across agent workers) -const ( - SubjectMCPToolExecute = "mcp.tools.execute" - SubjectMCPDiscovery = "mcp.discovery" - QueueAgentWorkers = "agent-workers" -) - // SubjectFineTuneProgress returns the NATS subject for fine-tune progress. func SubjectFineTuneProgress(jobID string) string { return subjectFineTunePrefix + sanitizeSubjectToken(jobID) + ".progress" diff --git a/core/services/nodes/agent_control.go b/core/services/nodes/agent_control.go new file mode 100644 index 000000000..846a5ec96 --- /dev/null +++ b/core/services/nodes/agent_control.go @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: MIT + +package nodes + +import ( + "context" + "errors" + "fmt" + + mcpremote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// maxAgentPicks bounds how many agent workers one verb is offered to before it +// gives up. +// +// Three rather than "every agent in the fleet": a caller is waiting, and a +// deployment whose first three connected agents all fail to answer has a +// problem that trying a fourth does not fix. It is also not one: a worker whose +// tunnel died between the connection read and the dial is ordinary, and giving +// up on the first of those would make every re-homing agent an outage. +const maxAgentPicks = 3 + +// ErrNoAgentControl reports that this deployment has no agent control client +// wired at all. +// +// Like ErrNoAgentWorker it is neither ErrWorkerUnroutable nor anything +// cluster.IsWorkerAnswer accepts: it is a fact about this process's own wiring +// and says nothing about any worker. +var ErrNoAgentControl = errors.New("nodes: this deployment has no agent control client") + +// AgentControlClient issues the frontend's control RPCs to whichever agent +// worker the selector picks. +// +// It is the caller the agent worker's control plane has been waiting for. What +// used to be a NATS request onto a queue group is now two ordinary things: a +// SELECTION, which is a query against the connection rows, and a control RPC +// over the picked worker's own tunnel, relayed by the peer mesh when another +// replica holds it. +type AgentControlClient struct { + sel *AgentSelector + cc *ControlClient +} + +// NewAgentControlClient returns the client that carries the frontend's agent +// verbs to the workers sel picks, over cc. +func NewAgentControlClient(sel *AgentSelector, cc *ControlClient) *AgentControlClient { + return &AgentControlClient{sel: sel, cc: cc} +} + +// ExecuteMCPTool runs one MCP tool on an agent worker and returns its answer. +// +// A reply carrying an Error is the WORKER'S OWN ANSWER and comes back as an +// error naming it, never as a retry: see agentVerb. +func (a *AgentControlClient) ExecuteMCPTool(ctx context.Context, req mcpremote.MCPToolRequest) (*mcpremote.MCPToolResponse, error) { + return agentVerb(ctx, a, workerctl.PathMCPToolExecute, req, + func(r *mcpremote.MCPToolResponse) string { return r.Error }) +} + +// DiscoverMCPTools asks an agent worker which MCP servers and tools a model's +// configuration reaches. +func (a *AgentControlClient) DiscoverMCPTools(ctx context.Context, req mcpremote.MCPDiscoveryRequest) (*mcpremote.MCPDiscoveryResponse, error) { + return agentVerb(ctx, a, workerctl.PathMCPDiscovery, req, + func(r *mcpremote.MCPDiscoveryResponse) string { return r.Error }) +} + +// agentVerb is the ONE place the select-call-retry rule lives, and the one +// place the line between a retryable failure and an answer is drawn. +// +// The rule, stated as the code enforces it: +// +// - A DECODED REPLY whose Error field is set is the worker's own answer to +// this verb. It is returned as an error naming what the worker said and is +// NEVER offered to another worker: retrying it would turn "this MCP server +// rejected your arguments" into "the fleet is broken", and could run a tool +// twice. +// - Everything the RPC returns as a Go ERROR is a failure to obtain an +// answer. The request never reached a handler (a refused stream, an +// unreachable peer, a lost tunnel, a 404 from an older build) or its answer +// could not be read, so nothing ran to completion here and another worker +// may be offered the verb. +// +// This is deliberately NOT keyed on cluster.IsWorkerAnswer, and the reason is +// worth stating because the plan for this task said it should be. +// IsWorkerAnswer accepts the tunnel's stream-refusal vocabulary, which a worker +// writes BEFORE any request body reaches its control server: ErrStreamTagUnknown +// and ErrStreamTargetUnavailable both mean "I could not carry this to my own +// control plane", not "I ran your tool and here is what happened". Returning +// those unchanged without trying another agent would make one agent worker with +// a dead control server take down MCP for the whole deployment, and it is the +// case the plan's own spec list requires to be retried. +// +// The taxonomy is preserved either way: whatever error is returned is returned +// UNWRAPPED, so cluster.IsWorkerAnswer and ErrWorkerUnroutable still see +// exactly what ControlClient produced. +func agentVerb[Req any, Rep any](ctx context.Context, a *AgentControlClient, path string, req Req, + answerError func(*Rep) string) (*Rep, error) { + if a == nil || a.sel == nil || a.cc == nil { + // Guarded on the nil RECEIVER too, because the interface a caller holds + // this through is satisfied by a typed nil pointer, which is not an + // untyped nil and would otherwise panic inside a request. + return nil, fmt.Errorf("control rpc %s: %w", path, ErrNoAgentControl) + } + tried := make(map[string]bool, maxAgentPicks) + var lastErr error + for range maxAgentPicks { + nodeID, _, err := a.sel.pickConnectedExcluding(ctx, tried) + if err != nil { + if lastErr != nil && errors.Is(err, ErrNoAgentWorker) { + // Every worker that held a tunnel has now failed to answer. The + // evidence is what they said, not that the candidate list ran + // out, and reporting an empty fleet here would hide a fleet + // that is refusing. + return nil, lastErr + } + return nil, err + } + tried[nodeID] = true + + var reply Rep + if err := a.cc.Call(ctx, nodeID, path, req, &reply); err != nil { + lastErr = err + if ctx.Err() != nil { + // The caller's own budget is spent. Another pick would fail the + // same way immediately, and the error already says whose fault + // the expiry was. + return nil, lastErr + } + continue + } + if msg := answerError(&reply); msg != "" { + return nil, fmt.Errorf("agent worker %q answered %s with an error: %s", nodeID, path, msg) + } + return &reply, nil + } + return nil, lastErr +} diff --git a/core/services/nodes/agent_control_test.go b/core/services/nodes/agent_control_test.go new file mode 100644 index 000000000..aaeaa81f1 --- /dev/null +++ b/core/services/nodes/agent_control_test.go @@ -0,0 +1,375 @@ +// SPDX-License-Identifier: MIT + +package nodes_test + +import ( + "context" + "encoding/json" + "errors" + "net" + "net/http" + "net/http/httptest" + "runtime" + "sync/atomic" + + "github.com/gorilla/websocket" + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/agentworker" + "github.com/mudler/LocalAI/core/services/cluster" + mcpremote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/worker" +) + +// This file drives the WHOLE path a distributed MCP call takes: a real agent +// worker holding a real tunnel, a real connection row deciding which replica +// owns it, a real relay when the owner is a peer, and the real control client +// on top. Nothing here is a double of a transport. +// +// That matters more than usual for this task. The selection it proves is a +// query against the same rows the scheduler reads absence from, and a fake +// registry that never dials would prove nothing at all about a relayed pick: +// every interesting failure (a stream a worker refuses, a tunnel a peer holds, +// a candidate whose owner is dead) exists only on the wire. + +const agentControlToken = "agent-control-token" + +// fakeFrontend is the far side of a worker's tunnel: the real WebSocket +// upgrade and the real yamux server handshake. +type fakeFrontend struct { + srv *httptest.Server + sessions chan *yamux.Session +} + +func newFakeFrontend() *fakeFrontend { + GinkgoHelper() + f := &fakeFrontend{sessions: make(chan *yamux.Session, 4)} + upgrader := websocket.Upgrader{} + f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != cluster.ConnectPath { + w.WriteHeader(http.StatusNotFound) + return + } + ws, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + sess, err := yamux.Server(cluster.WebsocketConn(ws), nil, nil) + if err != nil { + _ = ws.Close() + return + } + select { + case f.sessions <- sess: + default: + _ = sess.Close() + } + })) + DeferCleanup(f.srv.Close) + return f +} + +// session waits for the worker to dial in and hands back its tunnel session. +// +// Waited for on a channel rather than slept on: the dial is the worker's own +// and nothing in this process orders it against the next line of the spec. +func (f *fakeFrontend) session() *yamux.Session { + GinkgoHelper() + var sess *yamux.Session + Eventually(f.sessions, "20s").Should(Receive(&sess)) + DeferCleanup(func() { _ = sess.Close() }) + return sess +} + +// yamuxPair is the peer link between two frontend replicas, over an in-process +// pipe: the client half is what the dialling replica opens relay streams on, +// the server half is what the owning replica accepts them from. +func yamuxPair() (client, server *yamux.Session) { + GinkgoHelper() + a, b := net.Pipe() + var err error + server, err = yamux.Server(a, nil, nil) + Expect(err).ToNot(HaveOccurred()) + client, err = yamux.Client(b, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { + _ = client.Close() + _ = server.Close() + }) + return client, server +} + +// peerLinks is the outbound half of the peer mesh, standing in for +// cluster.PeerPool. It opens a real yamux stream onto the owning replica's +// accepted session, which is what the relay reads its request frame from. +type peerLinks struct{ sessions map[string]*yamux.Session } + +func (p *peerLinks) Open(ctx context.Context, peerID string) (net.Conn, error) { + sess, ok := p.sessions[peerID] + if !ok { + return nil, cluster.ErrPeerUnreachable + } + return sess.OpenStream(ctx) +} + +var _ = Describe("AgentControlClient", func() { + const ( + selfInstance = "replica-me" + peerInstance = "replica-peer" + toolCallReply = "Weather in London: 15C, cloudy" + ) + + var ( + ctx context.Context + cancel context.CancelFunc + db *gorm.DB + registry *nodes.NodeRegistry + clusterReg *cluster.Registry + mine *cluster.TunnelRegistry + theirs *cluster.TunnelRegistry + peers *peerLinks + control *nodes.ControlClient + ) + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + ctx, cancel = context.WithCancel(context.Background()) + DeferCleanup(cancel) + + db = testutil.SetupTestDB() + var err error + registry, err = nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + clusterReg = cluster.NewRegistry(db) + Expect(clusterReg.Register(ctx, selfInstance, "10.0.0.1:8080", "v1")).To(Succeed()) + Expect(clusterReg.Register(ctx, peerInstance, "10.0.0.2:8080", "v1")).To(Succeed()) + + mine = cluster.NewTunnelRegistry(clusterReg, selfInstance) + theirs = cluster.NewTunnelRegistry(clusterReg, peerInstance) + + // The owning replica's inbound half, with the relay installed, and this + // replica's outbound half pointed at it. Together they are the path a + // request takes when the worker's tunnel landed somewhere else. + relayed := cluster.NewSessionStore(cluster.NewRelay(theirs).Stream) + DeferCleanup(relayed.CloseAll) + client, server := yamuxPair() + relayed.Accept(selfInstance, server) + peers = &peerLinks{sessions: map[string]*yamux.Session{peerInstance: client}} + + dialer := cluster.NewWorkerDialer(mine, peers) + control = nodes.NewControlClient(nodes.WorkerNetDialerFor(func(nodeID string) func(context.Context, string, string) (net.Conn, error) { + return dialer.DialerFor(nodeID, cluster.StreamTagHTTP) + }), agentControlToken) + }) + + // registerAgent puts an approved agent node in the registry and returns its + // assigned id, which is the id the connection row and the tunnel both use. + registerAgent := func(name string) string { + GinkgoHelper() + Expect(registry.Register(ctx, &nodes.BackendNode{ + Name: name, NodeType: nodes.NodeTypeAgent, Address: name + ":50051", + }, true)).To(Succeed()) + node, err := registry.GetByName(ctx, name) + Expect(err).ToNot(HaveOccurred()) + return node.ID + } + + // startAgent runs a REAL agent worker serving cfg, dialling a tunnel that + // lands on holder, and returns its node id. + startAgent := func(name string, holder *cluster.TunnelRegistry, cfg agentworker.Config) string { + GinkgoHelper() + nodeID := registerAgent(name) + frontend := newFakeFrontend() + rt, err := agentworker.Start(ctx, agentworker.Options{ + FrontendURL: frontend.srv.URL, + NodeID: nodeID, + TunnelToken: func() string { return "tunnel-secret" }, + ControlToken: agentControlToken, + Handlers: cfg, + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = rt.Close() }) + _, err = holder.Attach(ctx, nodeID, frontend.session()) + Expect(err).ToNot(HaveOccurred()) + return nodeID + } + + // startRefusingAgent runs a worker whose tunnel is real and whose local + // service will not open, which is what a worker with a dead control server + // looks like from the frontend: the worker itself writes the refusal, and + // the frontend reads cluster.ErrStreamTargetUnavailable off the wire. + startRefusingAgent := func(name string, holder *cluster.TunnelRegistry, refusals *atomic.Int32) string { + GinkgoHelper() + nodeID := registerAgent(name) + frontend := newFakeFrontend() + tunnel, err := worker.StartTunnel(ctx, worker.TunnelConfig{ + FrontendURL: frontend.srv.URL, + NodeID: nodeID, + Token: func() string { return "tunnel-secret" }, + Services: map[string]worker.LocalService{ + cluster.StreamTagHTTP: func(context.Context, string) (net.Conn, error) { + refusals.Add(1) + return nil, errors.New("this worker's control server is not listening") + }, + }, + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = tunnel.Close() }) + _, err = holder.Attach(ctx, nodeID, frontend.session()) + Expect(err).ToNot(HaveOccurred()) + return nodeID + } + + // recordingTool answers every tool call with reply and counts its calls. + recordingTool := func(calls *atomic.Int32, reply mcpremote.MCPToolResponse) agentworker.UnaryHandler { + return func(_ context.Context, raw json.RawMessage) (json.RawMessage, error) { + calls.Add(1) + var req mcpremote.MCPToolRequest + Expect(json.Unmarshal(raw, &req)).To(Succeed()) + out, err := json.Marshal(reply) + Expect(err).ToNot(HaveOccurred()) + return out, nil + } + } + + agentControl := func() *nodes.AgentControlClient { + return nodes.NewAgentControlClient( + nodes.NewAgentSelector(registry, clusterReg, selfInstance), control) + } + + toolRequest := mcpremote.MCPToolRequest{ + ModelName: "test-model", + ToolName: "weather", + Arguments: map[string]any{"city": "London"}, + } + + It("carries a tool call over the tunnel of an agent THIS replica holds", func() { + var calls atomic.Int32 + startAgent("agent-local", mine, agentworker.Config{ + MCPTool: recordingTool(&calls, mcpremote.MCPToolResponse{Result: toolCallReply}), + }) + + resp, err := agentControl().ExecuteMCPTool(ctx, toolRequest) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.Result).To(Equal(toolCallReply)) + Expect(calls.Load()).To(Equal(int32(1))) + }) + + It("reaches an agent whose tunnel a PEER holds, by relaying through that peer", func() { + // The hop the queue group used to hide. This replica holds nothing, so + // the only way the request reaches the worker is the connection row + // naming the peer and the peer's relay splicing the stream onto the + // tunnel it holds. + var calls atomic.Int32 + startAgent("agent-remote", theirs, agentworker.Config{ + MCPTool: recordingTool(&calls, mcpremote.MCPToolResponse{Result: toolCallReply}), + }) + Expect(mine.Held()).To(BeEmpty(), "this spec is only about the relayed path") + + resp, err := agentControl().ExecuteMCPTool(ctx, toolRequest) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.Result).To(Equal(toolCallReply)) + Expect(calls.Load()).To(Equal(int32(1))) + }) + + It("answers discovery from the agent worker's own reply", func() { + startAgent("agent-disco", mine, agentworker.Config{ + MCPDiscovery: func(context.Context, json.RawMessage) (json.RawMessage, error) { + return json.Marshal(mcpremote.MCPDiscoveryResponse{ + Servers: []mcpremote.MCPServerInfo{{Name: "weather", Type: "remote", Tools: []string{"get_weather"}}}, + }) + }, + }) + + resp, err := agentControl().DiscoverMCPTools(ctx, mcpremote.MCPDiscoveryRequest{ModelName: "test-model"}) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.Servers).To(HaveLen(1)) + Expect(resp.Servers[0].Name).To(Equal("weather")) + }) + + It("returns a worker's own error answer and never offers the tool to a second worker", func() { + // The one thing the retry may not do. "This MCP server rejected your + // arguments" retried elsewhere becomes "the fleet is broken", and a + // tool with a side effect would run twice. + // + // Both workers answer the same refusal, so which one the random + // tie-break picks does not matter: what is asserted is that exactly ONE + // handler ran in total. + var first, second atomic.Int32 + refusal := mcpremote.MCPToolResponse{Error: "tool 'weather' is not configured"} + startAgent("agent-a", mine, agentworker.Config{MCPTool: recordingTool(&first, refusal)}) + startAgent("agent-b", mine, agentworker.Config{MCPTool: recordingTool(&second, refusal)}) + + _, err := agentControl().ExecuteMCPTool(ctx, toolRequest) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("tool 'weather' is not configured")) + Expect(first.Load()+second.Load()).To(Equal(int32(1)), + "the worker's own answer must not be offered to another worker") + }) + + It("retries against another agent when one refuses the stream, and the second answers", func() { + // Deterministic without touching the tie-break: the refusing worker is + // the one THIS replica holds, which the selector prefers, and the + // worker that answers is peer-held, which is where the retry has to + // fall back to. So the first pick always refuses, the second always + // relays, and both halves of the rule are exercised on the wire. + var refusals atomic.Int32 + var calls atomic.Int32 + startRefusingAgent("agent-dead", mine, &refusals) + startAgent("agent-alive", theirs, agentworker.Config{ + MCPTool: recordingTool(&calls, mcpremote.MCPToolResponse{Result: toolCallReply}), + }) + + resp, err := agentControl().ExecuteMCPTool(ctx, toolRequest) + Expect(err).ToNot(HaveOccurred()) + Expect(resp.Result).To(Equal(toolCallReply)) + Expect(refusals.Load()).To(Equal(int32(1)), "the first pick must have been the refusing worker") + Expect(calls.Load()).To(Equal(int32(1))) + }) + + It("gives up after three picks and returns the last refusal unchanged", func() { + // The bound, and the taxonomy. Returning the last error UNWRAPPED is + // what keeps cluster.IsWorkerAnswer able to see what the worker + // actually said; wrapping it in the unroutable umbrella here would make + // a fleet that is refusing look like a frontend with no route. + var refusals atomic.Int32 + for _, name := range []string{"agent-1", "agent-2", "agent-3", "agent-4"} { + startRefusingAgent(name, mine, &refusals) + } + + _, err := agentControl().ExecuteMCPTool(ctx, toolRequest) + Expect(err).To(MatchError(cluster.ErrStreamTargetUnavailable)) + Expect(cluster.IsWorkerAnswer(err)).To(BeTrue()) + Expect(errors.Is(err, nodes.ErrWorkerUnroutable)).To(BeFalse()) + Expect(errors.Is(err, nodes.ErrNoAgentWorker)).To(BeFalse(), + "three workers refused; that is not an empty fleet") + Expect(refusals.Load()).To(Equal(int32(3)), "exactly three picks, never a fourth") + }) + + It("reports an empty fleet as neither a route verdict nor a worker answer", func() { + registerAgent("agent-registered-but-offline") + + _, err := agentControl().ExecuteMCPTool(ctx, toolRequest) + Expect(err).To(MatchError(nodes.ErrNoAgentWorker)) + Expect(errors.Is(err, nodes.ErrWorkerUnroutable)).To(BeFalse()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("refuses on a nil client rather than panicking inside a request", func() { + // The interface a caller holds this through is satisfied by a typed nil + // pointer, which is not an untyped nil and passes every `!= nil` check + // a call site makes. + var missing *nodes.AgentControlClient + _, err := missing.ExecuteMCPTool(ctx, toolRequest) + Expect(err).To(MatchError(nodes.ErrNoAgentControl)) + _, err = missing.DiscoverMCPTools(ctx, mcpremote.MCPDiscoveryRequest{}) + Expect(err).To(MatchError(nodes.ErrNoAgentControl)) + }) +}) diff --git a/core/services/nodes/agent_selector.go b/core/services/nodes/agent_selector.go new file mode 100644 index 000000000..3ac924b9f --- /dev/null +++ b/core/services/nodes/agent_selector.go @@ -0,0 +1,171 @@ +// SPDX-License-Identifier: MIT + +package nodes + +import ( + "context" + "errors" + "fmt" + "math/rand/v2" + + "github.com/mudler/LocalAI/core/services/cluster" +) + +// ErrNoAgentWorker reports that no agent worker in this deployment currently +// holds a tunnel any live replica can reach. +// +// It is a deployment fact about this moment and NOT a verdict about any worker, +// which is why it is deliberately its own sentinel: it does not wrap +// ErrWorkerUnroutable, and cluster.IsWorkerAnswer does not accept it. Nothing +// was asked of any worker and nothing was learned about one, so no reap guard +// may act on it and nothing may be marked unhealthy because of it. +// +// The umbrella matters more than it looks. ErrWorkerUnroutable is what every +// path that DELETES a node_models row matches on, and an empty agent fleet is +// not evidence about a single node. +var ErrNoAgentWorker = errors.New("nodes: no agent worker holds a tunnel to this cluster") + +// AgentConnectionReader is the narrow port onto cluster.Registry.ConnectedAmong. +// +// An interface rather than the concrete registry so a spec can drive the +// selection by deciding what is connected, without also having to stage the +// claim rows and instance heartbeats that would produce that answer. The real +// read is pinned by its own specs, against the real database, in +// core/services/cluster. +type AgentConnectionReader interface { + ConnectedAmong(ctx context.Context, nodeIDs []string, owner string) (held []string, heldByOwner []string, err error) +} + +// AgentSelector picks an agent worker to send a control RPC to. +// +// It is what replaces a NATS queue group. The queue group was never a broker +// feature this design has to reproduce: what it did was SELECT one subscriber +// out of a set, and selection is a query. Asking it here is strictly better +// than letting a broker balance, for one concrete reason: this replica can +// prefer an agent whose tunnel it holds itself and skip the relay hop +// altogether, which hidden balancing could never do. +// +// It answers a ROUTING question and never an absence one. See ErrNoAgentWorker. +type AgentSelector struct { + registry *NodeRegistry + conns AgentConnectionReader + // selfInstanceID is this replica's id, the one the connection rows record + // as an owner. An empty one is not an error the selector can report: every + // call would simply take a relay hop with nothing saying so. It is refused + // where the selector is built instead. + selfInstanceID string +} + +// NewAgentSelector returns the selector for the agent workers registry knows +// about, reading presence through conns and preferring the tunnels +// selfInstanceID holds. +func NewAgentSelector(registry *NodeRegistry, conns AgentConnectionReader, selfInstanceID string) *AgentSelector { + return &AgentSelector{registry: registry, conns: conns, selfInstanceID: selfInstanceID} +} + +// PickConnected returns the id AND the node type of an agent node whose tunnel +// a live replica holds, preferring one THIS replica owns so the call skips the +// relay. +// +// The type is returned rather than looked up again because the caller that +// re-broadcasts a resulting progress stream needs it for every line, and the +// selector has ALREADY read the node rows to build the candidate list, so it +// costs nothing here. Resolving it through the registry per line would be a +// database read per progress line and, worse, a SECOND place in the tree that +// decides what a node's type is. +// +// It is always NodeTypeAgent today. It is returned as a value rather than +// asserted as a constant because the day a backend worker takes a dispatched +// claim, the caller that must not guess is this one. +func (s *AgentSelector) PickConnected(ctx context.Context) (nodeID, nodeType string, err error) { + return s.pickConnectedExcluding(ctx, nil) +} + +// pickConnectedExcluding is PickConnected with a set of ids already tried. +// +// The exclusion exists for the retry in AgentControlClient and is not exposed: +// a retry that could re-pick the worker that just failed to answer is not a +// retry, it is the same call again, and with a small fleet it would be the +// same call most of the time. +func (s *AgentSelector) pickConnectedExcluding(ctx context.Context, tried map[string]bool) (string, string, error) { + if s == nil || s.registry == nil || s.conns == nil { + // Not an error about any worker. A deployment with no selector has + // nothing to select from, which is the same answer as an empty fleet. + return "", "", fmt.Errorf("this deployment has no agent selector: %w", ErrNoAgentWorker) + } + agents, err := s.registry.selectableAgentNodes(ctx) + if err != nil { + // Deliberately plain. A database that would not answer says nothing + // about any worker, so this must carry neither ErrWorkerUnroutable nor + // anything cluster.IsWorkerAnswer accepts. + return "", "", fmt.Errorf("listing the agent workers of this deployment: %w", err) + } + typeOf := make(map[string]string, len(agents)) + ids := make([]string, 0, len(agents)) + for _, a := range agents { + if tried[a.ID] { + continue + } + typeOf[a.ID] = a.NodeType + ids = append(ids, a.ID) + } + held, heldByOwner, err := s.conns.ConnectedAmong(ctx, ids, s.selfInstanceID) + if err != nil { + return "", "", fmt.Errorf("reading which agent workers are connected: %w", err) + } + // heldByOwner first because that call takes no relay hop. Falling back to + // the rest rather than stopping there is what keeps a replica that holds no + // agent tunnel able to reach the fleet at all. + candidates := heldByOwner + if len(candidates) == 0 { + candidates = held + } + if len(candidates) == 0 { + return "", "", fmt.Errorf("of %d agent workers registered, none is connected: %w", len(ids), ErrNoAgentWorker) + } + // Random rather than round robin. A per-replica counter is per-replica + // state that says nothing about load, and with several replicas the + // counters agree on nothing anyway. + picked := candidates[rand.IntN(len(candidates))] + nodeType := typeOf[picked] + if nodeType == "" { + // A candidate the connection read named that the node read did not is + // this selector contradicting itself. Refused rather than returned with + // an empty type: an unknown node type is denied by every allow list + // downstream, and the only symptom of that is work that silently never + // happens. + return "", "", fmt.Errorf("agent worker %q was reported connected but has no node type: %w", picked, ErrNoAgentWorker) + } + return picked, nodeType, nil +} + +// selectableAgentNodes returns the agent nodes a control RPC may be sent to. +// +// Two statuses are excluded and no more. StatusPending is a node an admin has +// not approved, which cannot hold a tunnel anyway and must never be handed +// work; StatusDraining is one that has been asked to stop taking work, which is +// the operator's decision and not a routing fact. +// +// Every other status is left in ON PURPOSE, and this is where the invariant +// bites. StatusUnhealthy and StatusOffline are written by the health monitor +// and by shutdown, on their own clocks; whether a request can reach this worker +// RIGHT NOW is what ConnectedAmong answers, and it answers it from the +// connection rows. Filtering here on a health verdict would make the selection +// refuse a worker that is connected and answering because a probe was late, +// which is an absence read leaking into a routing one. +func (r *NodeRegistry) selectableAgentNodes(ctx context.Context) ([]BackendNode, error) { + var agents []BackendNode + if err := r.db.WithContext(ctx). + Where("node_type = ? AND status NOT IN ?", NodeTypeAgent, []string{StatusPending, StatusDraining}). + Order("id"). + Find(&agents).Error; err != nil { + return nil, fmt.Errorf("listing agent nodes: %w", err) + } + return agents, nil +} + +// The real reader is the cluster registry, asserted here so a drift between +// its signature and this port fails to COMPILE. Without it the only thing that +// would notice is the wiring in core/application, which opens a database and a +// bus and so has no unit spec at all. +var _ AgentConnectionReader = (*cluster.Registry)(nil) diff --git a/core/services/nodes/agent_selector_test.go b/core/services/nodes/agent_selector_test.go new file mode 100644 index 000000000..959762cfc --- /dev/null +++ b/core/services/nodes/agent_selector_test.go @@ -0,0 +1,271 @@ +// SPDX-License-Identifier: MIT + +package nodes + +import ( + "context" + "errors" + "runtime" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// stubConnections decides what is connected, so a selection spec can state the +// cluster state it is about in one line instead of staging claim rows and +// instance heartbeats to produce it. What the real read answers for a given +// cluster state is pinned against the real database in +// core/services/cluster/connected_among_test.go, and the relayed round trip in +// agent_control_test.go drives the real registry over a real transport. +type stubConnections struct { + held []string + heldByOwner []string + err error + + // owners records the owner id each call was made with. An empty one makes + // heldByOwner always empty in production, so every call takes a relay hop + // with nothing anywhere saying so. + owners []string + // asked records the candidate list of each call. + asked [][]string +} + +func (s *stubConnections) ConnectedAmong(_ context.Context, nodeIDs []string, owner string) ([]string, []string, error) { + s.owners = append(s.owners, owner) + s.asked = append(s.asked, append([]string(nil), nodeIDs...)) + if s.err != nil { + return nil, nil, s.err + } + // The real read answers only about ids it was given, so the stub must too: + // a stub that answered about a node the selector had excluded would hide an + // exclusion that stopped working. + offered := map[string]bool{} + for _, id := range nodeIDs { + offered[id] = true + } + var held, byOwner []string + for _, id := range s.held { + if offered[id] { + held = append(held, id) + } + } + for _, id := range s.heldByOwner { + if offered[id] { + byOwner = append(byOwner, id) + } + } + return held, byOwner, nil +} + +// The selection that replaces the queue group. Two ways it can be wrong, and +// both are the invariant rather than a preference: picking a node no live +// replica holds sends a request at a process that is gone, and refusing a node +// that is merely re-homing between replicas takes MCP down for a fleet that is +// fine. +var _ = Describe("AgentSelector", func() { + var ( + ctx context.Context + db *gorm.DB + registry *NodeRegistry + conns *stubConnections + ) + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + ctx = context.Background() + db = testutil.SetupTestDB() + var err error + registry, err = NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + conns = &stubConnections{} + }) + + register := func(name, nodeType string) string { + GinkgoHelper() + node := &BackendNode{Name: name, NodeType: nodeType, Address: name + ":50051"} + Expect(registry.Register(ctx, node, true)).To(Succeed()) + fetched, err := registry.GetByName(ctx, name) + Expect(err).ToNot(HaveOccurred()) + return fetched.ID + } + + // The repeat count every preference assertion uses. The tie-break is random + // by design, so a single call proves nothing: with two candidates a + // selector that ignored the preference entirely would pass one run in two. + const picks = 20 + + It("prefers an agent whose tunnel THIS replica holds, so the call skips the relay", func() { + mine := register("agent-mine", NodeTypeAgent) + theirs := register("agent-theirs", NodeTypeAgent) + conns.held = []string{mine, theirs} + conns.heldByOwner = []string{mine} + + sel := NewAgentSelector(registry, conns, "me") + for range picks { + id, nodeType, err := sel.PickConnected(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(id).To(Equal(mine)) + // The node type is what a later re-broadcast checks its allow list + // against, and an unknown type is DENIED there. An empty one would + // present only as progress that never reaches a browser. + Expect(nodeType).To(Equal(NodeTypeAgent)) + Expect(nodeType).ToNot(BeEmpty()) + } + Expect(conns.owners).To(HaveLen(picks)) + Expect(conns.owners[0]).To(Equal("me"), + "the connection read must be asked which tunnels THIS replica holds") + }) + + It("takes a peer-held agent when this replica holds none", func() { + theirs := register("agent-theirs", NodeTypeAgent) + conns.held = []string{theirs} + + sel := NewAgentSelector(registry, conns, "me") + for range picks { + id, nodeType, err := sel.PickConnected(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(id).To(Equal(theirs)) + Expect(nodeType).To(Equal(NodeTypeAgent)) + } + }) + + It("spreads its picks over the agents it owns rather than pinning one", func() { + // Random, not "the first one". A selector that returned candidates[0] + // would pass every assertion above and send the whole deployment's MCP + // traffic to one worker. + a := register("agent-a", NodeTypeAgent) + b := register("agent-b", NodeTypeAgent) + conns.held = []string{a, b} + conns.heldByOwner = []string{a, b} + + sel := NewAgentSelector(registry, conns, "me") + seen := map[string]int{} + // 40 draws of a fair two-way choice miss one side with probability + // 2^-39, which is far below the flake floor of anything else here. + for range 2 * picks { + id, _, err := sel.PickConnected(ctx) + Expect(err).ToNot(HaveOccurred()) + seen[id]++ + } + Expect(seen).To(HaveKey(a)) + Expect(seen).To(HaveKey(b)) + }) + + Describe("when nothing can be picked", func() { + assertNoWorker := func(id, nodeType string, err error) { + GinkgoHelper() + Expect(err).To(MatchError(ErrNoAgentWorker)) + // The taxonomy assertion, and the reason this sentinel is its own. + // ErrWorkerUnroutable is what every path that DELETES a node_models + // row matches on, and IsWorkerAnswer is what a reap guard acts on. + // An empty agent fleet is evidence about neither: nothing was asked + // of any worker. + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeFalse()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + // An id returned beside the error would be dialled by a caller that + // checked the id first, which is a caller mistake this makes + // impossible rather than one it documents. + Expect(id).To(BeEmpty()) + Expect(nodeType).To(BeEmpty()) + } + + It("refuses when registered agents exist but none is connected", func() { + register("agent-a", NodeTypeAgent) + sel := NewAgentSelector(registry, conns, "me") + assertNoWorker(sel.PickConnected(ctx)) + }) + + It("refuses when the deployment has no agent nodes at all", func() { + sel := NewAgentSelector(registry, conns, "me") + assertNoWorker(sel.PickConnected(ctx)) + }) + + It("refuses when it was built with no way to read connections", func() { + assertNoWorker(NewAgentSelector(registry, nil, "me").PickConnected(ctx)) + }) + }) + + It("never offers a BACKEND worker, which serves no MCP verb at all", func() { + backend := register("backend-1", NodeTypeBackend) + agent := register("agent-a", NodeTypeAgent) + // The connection read is told everything is connected; only the + // candidate list keeps the backend worker out, so a selector that + // listed every node would pick it and get a 404 it would read as a + // version skew. + conns.held = []string{backend, agent} + conns.heldByOwner = []string{backend, agent} + + sel := NewAgentSelector(registry, conns, "me") + for range picks { + id, _, err := sel.PickConnected(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(id).To(Equal(agent)) + } + Expect(conns.asked[0]).To(ConsistOf(agent)) + }) + + It("never offers an agent an admin has not approved", func() { + node := &BackendNode{Name: "agent-pending", NodeType: NodeTypeAgent, Address: "p:50051"} + Expect(registry.Register(ctx, node, false)).To(Succeed()) + pending, err := registry.GetByName(ctx, "agent-pending") + Expect(err).ToNot(HaveOccurred()) + Expect(pending.Status).To(Equal(StatusPending)) + conns.held = []string{pending.ID} + conns.heldByOwner = []string{pending.ID} + + sel := NewAgentSelector(registry, conns, "me") + _, _, err = sel.PickConnected(ctx) + Expect(err).To(MatchError(ErrNoAgentWorker)) + Expect(conns.asked[0]).To(BeEmpty()) + }) + + It("never offers an agent an operator has put into draining", func() { + id := register("agent-draining", NodeTypeAgent) + Expect(registry.MarkDraining(ctx, id)).To(Succeed()) + conns.held = []string{id} + conns.heldByOwner = []string{id} + + sel := NewAgentSelector(registry, conns, "me") + _, _, err := sel.PickConnected(ctx) + Expect(err).To(MatchError(ErrNoAgentWorker)) + }) + + It("still offers an agent the health monitor has marked unhealthy", func() { + // The invariant, at the one place in this task it can break subtly. + // Whether a request can reach this worker RIGHT NOW is the connection + // rows' answer; a health verdict is written on another clock, and + // letting it filter the candidates would refuse a worker that is + // connected and answering because a probe was late. + id := register("agent-a", NodeTypeAgent) + Expect(registry.MarkUnhealthy(ctx, id)).To(Succeed()) + conns.held = []string{id} + conns.heldByOwner = []string{id} + + sel := NewAgentSelector(registry, conns, "me") + picked, _, err := sel.PickConnected(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(picked).To(Equal(id)) + }) + + It("reports a connection read that failed as neither an answer nor a route verdict", func() { + register("agent-a", NodeTypeAgent) + conns.err = errors.New("the database would not answer") + + sel := NewAgentSelector(registry, conns, "me") + id, nodeType, err := sel.PickConnected(ctx) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("the database would not answer")) + // A database that would not answer says nothing about any worker. + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeFalse()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + Expect(errors.Is(err, ErrNoAgentWorker)).To(BeFalse()) + Expect(id).To(BeEmpty()) + Expect(nodeType).To(BeEmpty()) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index f2dffc88b..5d248c549 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -143,7 +143,7 @@ Both **backend** and **agent** nodes are issued one. Earlier releases minted a c An agent worker's tunnel carries only the `http` tag: it runs no backend processes, so it does not offer the `grpc` tag at all. Its control server binds `127.0.0.1` on a port chosen by the kernel and advertises it nowhere, so an agent worker still opens no inbound port. -**An agent worker still requires `--nats-url`.** The tunnel is an addition, not a replacement: agent jobs, MCP execution, MCP CI jobs and `nodes..backend.stop` all still travel on NATS. What the tunnel changes today is that the frontend *can* reach an agent worker directly, which is what later releases move those verbs onto. +**An agent worker still requires `--nats-url`.** The tunnel is not yet a replacement: agent jobs, MCP CI jobs and `nodes..backend.stop` still travel on NATS. MCP tool execution and MCP discovery no longer do - they are control RPCs on the tunnel, chosen by the frontend rather than by a queue group (see [MCP in Distributed Mode](#mcp-in-distributed-mode)) - and an agent worker's minted JWT no longer grants `mcp.tools.execute` or `mcp.discovery`. The tunnel lands on exactly one frontend replica, and that replica records itself as the owner of the worker's connection in the `node_connections` table. When the socket dies the claim is dropped, but the row stays behind with no owner and a `disconnected_at` stamp, so a worker that is re-dialling the load balancer can be told from one that has never connected. The row is deleted once that departure is older than ten liveness windows (five minutes). If the replica stalls long enough for its peers to reap it, it re-claims the tunnels it still holds on a live session as soon as it re-registers, skipping any whose socket has already gone. That re-claim needs the replica to have an advertised address: without one it never had an instance row to begin with, and its tunnels are usable only by the replica holding them. @@ -435,7 +435,7 @@ A frontend replica that dies mid-load does not wedge the model: the job row carr **This section is about agent workers and the frontend.** A serve-backend worker opens no NATS connection, so none of it applies to one; its own credential is the tunnel token it gets at registration, and its control plane is authenticated by `LOCALAI_REGISTRATION_TOKEN`. An agent worker now has both: a NATS credential for everything still on the bus, and a tunnel token plus the same `LOCALAI_REGISTRATION_TOKEN` bearer check in front of its control server. -By default, NATS connections are anonymous: any client that can reach port `4222` may publish the subjects still carried on it. Those are the agent-worker job subjects, MCP, and the frontend's own cross-replica events. `nodes..backend.install` and its nine siblings are **not** among them - they are HTTP routes on the worker's tunnel, see [The worker control plane](#the-worker-control-plane). Enable JWT auth to scope agent workers to their own subjects and give the frontend a dedicated service credential. +By default, NATS connections are anonymous: any client that can reach port `4222` may publish the subjects still carried on it. Those are the agent-worker job subjects and the frontend's own cross-replica events. `nodes..backend.install` and its nine siblings are **not** among them - they are HTTP routes on the worker's tunnel, see [The worker control plane](#the-worker-control-plane). Enable JWT auth to scope agent workers to their own subjects and give the frontend a dedicated service credential. | Flag | Env Var | Description | |------|---------|-------------| @@ -935,7 +935,7 @@ local-ai agent-worker \ Agent workers: - Execute agent chat messages dispatched via NATS - Run MCP CI jobs (with access to MCP servers via docker) -- Handle MCP tool discovery and execution requests from the frontend +- Handle MCP tool discovery and execution requests, which the frontend sends over the worker's own tunnel - Get auto-provisioned API keys during registration for calling the inference API In the docker-compose setup, the agent worker mounts the Docker socket so it can run MCP stdio servers (e.g., `docker run` commands): @@ -949,12 +949,31 @@ agent-worker-1: ## MCP in Distributed Mode -MCP servers configured in model configs work in distributed mode. The frontend routes MCP operations through NATS to agent workers: +MCP servers configured in model configs work in distributed mode. The frontend holds no MCP sessions of its own - creating one usually means running `docker`, which is what an agent worker is for - so it asks an agent worker instead: -- **MCP discovery** (`GET /v1/mcp/servers/:model`): routed to agent workers which create sessions and return server info -- **MCP tool execution** (during `/v1/chat/completions`): tool calls are routed to agent workers via NATS request-reply +- **MCP discovery** (`GET /v1/mcp/servers/:model`): the frontend picks a connected agent worker and asks it over that worker's tunnel; the worker creates the sessions and returns server info +- **MCP tool execution** (during `/v1/chat/completions`): the same, per tool call - **MCP CI jobs**: executed entirely on agent workers with access to docker for stdio-based MCP servers +### How a frontend picks an agent worker + +Discovery and tool execution used to be NATS request-reply onto a queue group, where the broker chose the worker and neither side could say which one had answered. They are now an ordinary control RPC plus a **selection**, because a queue group was only ever a way of choosing a subscriber, and choosing is a query: + +1. The frontend lists the approved agent nodes that are not draining. +2. It asks the `node_connections` table, in one statement joined against live replicas, which of those tunnels a **live** frontend replica currently holds. +3. It prefers one **this** replica holds, so the call skips the relay hop entirely, and otherwise takes any connected one at random. A broker's hidden balancing could not make that choice. +4. It issues the control RPC over that worker's tunnel, relayed through the owning replica when another one holds it. + +A worker that answers with an error - "no such tool", "that MCP server refused your arguments" - is the worker's own answer and is returned to you unchanged; it is never re-tried on a second worker, because that would run a tool twice. A call that never reached a worker is re-tried, against a different worker, at most three times. + +If no agent worker in the deployment currently holds a tunnel, the request fails with a message saying so. That is a statement about this moment, not about any particular worker: nothing is marked unhealthy and no model is evicted because of it. + +### MCP prompts and resources are not available in distributed mode + +`GET /v1/mcp/prompts/:model`, `POST /v1/mcp/prompts/:model/:prompt`, `GET /v1/mcp/resources/:model` and `POST /v1/mcp/resources/:model/read` are served **only** from MCP sessions held by the frontend process, and in distributed mode it holds none. There is no verb that carries prompts or resources to an agent worker. + +In distributed mode these four endpoints answer **501 Not Implemented** with the reason in the body. Earlier releases answered `200` with an empty list, which was indistinguishable from a model that genuinely has no prompts. This is a pre-existing gap rather than a consequence of moving MCP off the bus - tools and discovery had a carrier to an agent worker and these never did - and single-binary deployments are unaffected. + ## vLLM Multi-Node (Data-Parallel) A single vLLM model can span multiple GPU nodes via data parallelism: the head node serves the OpenAI API and runs the local DP ranks, follower nodes run vanilla `vllm serve --headless` and speak ZMQ directly to the head. LocalAI's role is starting the follower processes and surfacing them in the admin UI; the cross-rank tensor traffic is vLLM's own. @@ -1386,7 +1405,7 @@ Notes: - It is **not** the same as the worker being gone, and nothing acts on it as if it were. A model on an unroutable worker is not reaped, its rows are left alone, and the node is not demoted: doing any of those on a lost route is how a rolling frontend restart turns into a fleet-wide eviction. - Check the worker process is running and that it has an open tunnel (`opened a tunnelled stream to a worker` in the frontend log, and the worker's own dial/reconnect lines). A worker behind a load balancer that keeps reconnecting is usually an idle-timeout or WebSocket-upgrade problem at the proxy; see the tunnel section above. - **`no route` is not `gone`, and nothing in the frontend reads it as such.** A worker is declared **gone** by one mechanism only: no live frontend replica holds its tunnel *and* its departure is older than `--worker-reconnect-grace`. That is a fact recorded in the shared database, so every replica answers it identically. "No route" is one replica failing to reach a worker right now, and it is not evidence about the worker at all. -- Older releases decided absence from `nats: no responders available for request`, which was one frontend's observation that nobody answered *it* within a request budget. Two replicas asking at the same moment could disagree and demote each other's workers. That signal is gone from the scheduler; if you still see the message, it concerns only the subjects that remain on the bus (agent-worker jobs and MCP), never a serve-backend worker. +- Older releases decided absence from `nats: no responders available for request`, which was one frontend's observation that nobody answered *it* within a request budget. Two replicas asking at the same moment could disagree and demote each other's workers. That signal is gone from the scheduler; if you still see the message, it concerns only the subjects that remain on the bus (agent-worker jobs and MCP CI), never a serve-backend worker. **A worker fills its own disk over time:** - A request that carries a file (an image, an audio clip, a video) stages that file to the worker under `/../staging/ephemeral/`. The worker deletes these 6 hours after the request that needed them, and sweeps every 30 minutes plus once at startup, so a worker that crashed mid-request still reclaims the space. diff --git a/pkg/natsauth/permissions.go b/pkg/natsauth/permissions.go index 4c081c672..9145e0b3e 100644 --- a/pkg/natsauth/permissions.go +++ b/pkg/natsauth/permissions.go @@ -29,6 +29,12 @@ func WorkerPermissions(nodeID, nodeType string) (pubAllow, subAllow []string) { case "agent": // Agent workers consume queue workloads; they must not handle backend.install. // Keep this list in sync with the subscriptions in core/cli/agent_worker.go. + // + // MCP tool execution and discovery are NOT here any more: they are + // control RPCs on the tunnel the worker holds, chosen by the frontend + // rather than by a queue group. Removing them narrowed this list; it + // must never be narrowed to nothing, because NATS reads an EMPTY allow + // list as no restriction at all. subAllow = []string{ "agent.execute", "agent.*.cancel", @@ -37,9 +43,7 @@ func WorkerPermissions(nodeID, nodeType string) (pubAllow, subAllow []string) { "jobs.*.cancel", "jobs.*.progress", "jobs.*.result", - "jobs.mcp-ci.new", // MCP CI jobs dispatched to agent workers - "mcp.tools.execute", - "mcp.discovery", + "jobs.mcp-ci.new", // MCP CI jobs dispatched to agent workers prefix + ".backend.stop", // stop events drive MCP session cleanup "staging.*.progress", "_INBOX.>", diff --git a/pkg/natsauth/permissions_coverage_test.go b/pkg/natsauth/permissions_coverage_test.go index ab1b993c0..2c5d9e9b3 100644 --- a/pkg/natsauth/permissions_coverage_test.go +++ b/pkg/natsauth/permissions_coverage_test.go @@ -99,15 +99,33 @@ var _ = Describe("WorkerPermissions subject coverage", func() { Context("agent worker", func() { // node_type "agent"; subjects from core/cli/agent_worker.go. pub, sub := natsauth.WorkerPermissions(nodeID, "agent") - _ = pub subscribed := []string{ messaging.SubjectAgentExecute, // dispatcher (default --agent-subject) - messaging.SubjectMCPToolExecute, // QueueSubscribeReply - messaging.SubjectMCPDiscovery, // QueueSubscribeReply messaging.SubjectMCPCIJobsNew, // QueueSubscribe — jobs.mcp-ci.new messaging.SubjectNodeBackendStop(nodeID), // Subscribe — MCP session cleanup } + + // The half that catches a narrowing going too far. NATS reads an EMPTY + // allow list as NO restriction, so a branch trimmed to nothing does not + // lock an agent worker down, it opens the whole account to it. + It("keeps the agent worker's allow lists non-empty", func() { + Expect(sub).ToNot(BeEmpty(), + "an empty allow list is unrestricted in NATS, not restrictive") + Expect(pub).ToNot(BeEmpty(), + "an empty allow list is unrestricted in NATS, not restrictive") + }) + + // MCP execution and discovery are control RPCs on the worker's tunnel + // now, chosen by the frontend rather than by a queue group. An agent + // worker subscribes to neither subject, so a JWT that still granted + // them would be granting a subscription nothing serves. + for _, subject := range []string{"mcp.tools.execute", "mcp.discovery"} { + It("no longer grants an agent worker "+subject, func() { + Expect(anyAllows(sub, subject)).To(BeFalse(), + "agent JWT sub allow-list %v still covers %s", sub, subject) + }) + } for _, subject := range subscribed { It("allows subscribing to "+subject, func() { Expect(anyAllows(sub, subject)).To(BeTrue(), diff --git a/tests/e2e/distributed/mcp_nats_test.go b/tests/e2e/distributed/mcp_nats_test.go deleted file mode 100644 index e0e868e4d..000000000 --- a/tests/e2e/distributed/mcp_nats_test.go +++ /dev/null @@ -1,177 +0,0 @@ -package distributed_test - -import ( - "encoding/json" - "sync/atomic" - "time" - - "github.com/mudler/LocalAI/core/config" - mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" - mcpRemote "github.com/mudler/LocalAI/core/services/mcp" - "github.com/mudler/LocalAI/core/services/messaging" - "github.com/mudler/LocalAI/pkg/functions" - - . "github.com/onsi/ginkgo/v2" - . "github.com/onsi/gomega" -) - -var _ = Describe("MCP NATS Routing", Label("Distributed"), func() { - var ( - infra *TestInfra - ) - - BeforeEach(func() { - infra = SetupNATSOnly() - }) - - Context("MCP Tool Execution via NATS", func() { - It("should execute MCP tool call via NATS request-reply", func() { - // Mock worker: subscribe to tool execute requests - sub, err := infra.NC.QueueSubscribeReply(messaging.SubjectMCPToolExecute, messaging.QueueAgentWorkers, func(data []byte, reply func([]byte)) { - var req mcpRemote.MCPToolRequest - Expect(json.Unmarshal(data, &req)).To(Succeed()) - Expect(req.ModelName).To(Equal("test-model")) - Expect(req.ToolName).To(Equal("weather")) - Expect(req.Arguments).To(HaveKeyWithValue("city", "London")) - - resp, _ := json.Marshal(mcpRemote.MCPToolResponse{ - Result: "Weather in London: 15°C, cloudy", - }) - reply(resp) - }) - Expect(err).ToNot(HaveOccurred()) - defer sub.Unsubscribe() - - FlushNATS(infra.NC) - - // Frontend side: pass NATS client and call remote - result, err := mcpTools.ExecuteMCPToolCallRemote( - infra.Ctx, - infra.NC, - "test-model", - config.MCPGenericConfig[config.MCPRemoteServers]{}, - config.MCPGenericConfig[config.MCPSTDIOServers]{}, - "weather", - `{"city": "London"}`, - ) - Expect(err).ToNot(HaveOccurred()) - Expect(result).To(Equal("Weather in London: 15°C, cloudy")) - }) - - It("should propagate remote MCP tool errors", func() { - sub, err := infra.NC.QueueSubscribeReply(messaging.SubjectMCPToolExecute, messaging.QueueAgentWorkers, func(data []byte, reply func([]byte)) { - resp, _ := json.Marshal(mcpRemote.MCPToolResponse{ - Error: "tool 'unknown' not found", - }) - reply(resp) - }) - Expect(err).ToNot(HaveOccurred()) - defer sub.Unsubscribe() - - FlushNATS(infra.NC) - - _, err = mcpTools.ExecuteMCPToolCallRemote( - infra.Ctx, - infra.NC, - "test-model", - config.MCPGenericConfig[config.MCPRemoteServers]{}, - config.MCPGenericConfig[config.MCPSTDIOServers]{}, - "unknown", - "{}", - ) - Expect(err).To(HaveOccurred()) - Expect(err.Error()).To(ContainSubstring("tool 'unknown' not found")) - }) - }) - - Context("MCP Discovery via NATS", func() { - It("should discover MCP servers via NATS request-reply", func() { - sub, err := infra.NC.QueueSubscribeReply(messaging.SubjectMCPDiscovery, messaging.QueueAgentWorkers, func(data []byte, reply func([]byte)) { - var req mcpRemote.MCPDiscoveryRequest - Expect(json.Unmarshal(data, &req)).To(Succeed()) - Expect(req.ModelName).To(Equal("discovery-model")) - - resp, _ := json.Marshal(mcpRemote.MCPDiscoveryResponse{ - Servers: []mcpRemote.MCPServerInfo{ - {Name: "weather-server", Type: "remote", Tools: []string{"get_weather", "get_forecast"}}, - {Name: "db-server", Type: "stdio", Tools: []string{"query_db"}}, - }, - Tools: []mcpRemote.MCPToolDef{ - {ServerName: "weather-server", ToolName: "get_weather", Function: functions.Function{Name: "get_weather", Description: "Get weather"}}, - {ServerName: "weather-server", ToolName: "get_forecast", Function: functions.Function{Name: "get_forecast", Description: "Get forecast"}}, - {ServerName: "db-server", ToolName: "query_db", Function: functions.Function{Name: "query_db", Description: "Query database"}}, - }, - }) - reply(resp) - }) - Expect(err).ToNot(HaveOccurred()) - defer sub.Unsubscribe() - - FlushNATS(infra.NC) - - result, err := mcpTools.DiscoverMCPToolsRemote( - infra.Ctx, - infra.NC, - "discovery-model", - config.MCPGenericConfig[config.MCPRemoteServers]{}, - config.MCPGenericConfig[config.MCPSTDIOServers]{}, - ) - Expect(err).ToNot(HaveOccurred()) - Expect(result.Servers).To(HaveLen(2)) - Expect(result.Servers[0].Name).To(Equal("weather-server")) - Expect(result.Servers[0].Tools).To(ConsistOf("get_weather", "get_forecast")) - Expect(result.Tools).To(HaveLen(3)) - Expect(result.Tools[2].ToolName).To(Equal("query_db")) - }) - }) - - Context("QueueSubscribeReply", func() { - It("should support queue subscribe with request-reply round-trip", func() { - // Subscribe with queue group - sub, err := infra.NC.QueueSubscribeReply("test.echo", "echo-workers", func(data []byte, reply func([]byte)) { - // Echo back the request data with a prefix - reply(append([]byte("echo:"), data...)) - }) - Expect(err).ToNot(HaveOccurred()) - defer sub.Unsubscribe() - - FlushNATS(infra.NC) - - // Send request and wait for reply - replyData, err := infra.NC.Request("test.echo", []byte("hello"), 5*time.Second) - Expect(err).ToNot(HaveOccurred()) - Expect(string(replyData)).To(Equal("echo:hello")) - }) - - It("should load-balance requests across queue subscribers", func() { - var worker1Count, worker2Count atomic.Int32 - - sub1, _ := infra.NC.QueueSubscribeReply("test.lb", "lb-workers", func(data []byte, reply func([]byte)) { - worker1Count.Add(1) - reply([]byte("w1")) - }) - defer sub1.Unsubscribe() - - sub2, _ := infra.NC.QueueSubscribeReply("test.lb", "lb-workers", func(data []byte, reply func([]byte)) { - worker2Count.Add(1) - reply([]byte("w2")) - }) - defer sub2.Unsubscribe() - - FlushNATS(infra.NC) - - // Send multiple requests - for range 10 { - _, err := infra.NC.Request("test.lb", []byte("req"), 5*time.Second) - Expect(err).ToNot(HaveOccurred()) - } - - // Both workers should have handled some requests - total := worker1Count.Load() + worker2Count.Load() - Expect(total).To(Equal(int32(10))) - // NATS typically distributes evenly, but we just check both got work - Expect(worker1Count.Load()).To(BeNumerically(">", 0)) - Expect(worker2Count.Load()).To(BeNumerically(">", 0)) - }) - }) -}) diff --git a/tests/e2e/distributed/nats_jwt_test.go b/tests/e2e/distributed/nats_jwt_test.go index b6b234385..dd0c23253 100644 --- a/tests/e2e/distributed/nats_jwt_test.go +++ b/tests/e2e/distributed/nats_jwt_test.go @@ -108,13 +108,10 @@ var _ = Describe("NATS JWT Auth", Label("Distributed", "NatsJWT"), func() { Expect(err).ToNot(HaveOccurred()) DeferCleanup(nc.Close) - // Mirror core/cli/agent_worker.go exactly. - _, err = nc.QueueSubscribeReply(messaging.SubjectMCPToolExecute, messaging.QueueAgentWorkers, func([]byte, func([]byte)) {}) - Expect(err).ToNot(HaveOccurred(), "agent JWT must allow %s", messaging.SubjectMCPToolExecute) - - _, err = nc.QueueSubscribeReply(messaging.SubjectMCPDiscovery, messaging.QueueAgentWorkers, func([]byte, func([]byte)) {}) - Expect(err).ToNot(HaveOccurred(), "agent JWT must allow %s", messaging.SubjectMCPDiscovery) - + // Mirror core/cli/agent_worker.go exactly. MCP tool execution and + // discovery are absent because they are no longer bus subjects at all: + // the frontend selects an agent worker itself and reaches it with a + // control RPC over the tunnel that worker holds. _, err = nc.QueueSubscribe(messaging.SubjectMCPCIJobsNew, messaging.QueueWorkers, func([]byte) {}) Expect(err).ToNot(HaveOccurred(), "agent JWT must allow %s (MCP CI jobs)", messaging.SubjectMCPCIJobsNew) diff --git a/tests/e2e/distributed/nats_queue_reply_test.go b/tests/e2e/distributed/nats_queue_reply_test.go new file mode 100644 index 000000000..34c0b5c2e --- /dev/null +++ b/tests/e2e/distributed/nats_queue_reply_test.go @@ -0,0 +1,77 @@ +package distributed_test + +import ( + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// What is left here is the messaging layer's own request-reply behaviour. +// +// The MCP specs that used to live beside it went with their carrier: MCP tool +// execution and discovery were the only subjects that combined a queue group +// WITH a reply, and they are now a selection plus an ordinary control RPC over +// the agent worker's tunnel. Their round trip is exercised against a real +// tunnel, a real connection row and a real relay in +// core/services/nodes/agent_control_test.go. +var _ = Describe("NATS queue request-reply", Label("Distributed"), func() { + var ( + infra *TestInfra + ) + + BeforeEach(func() { + infra = SetupNATSOnly() + }) + + Context("QueueSubscribeReply", func() { + It("should support queue subscribe with request-reply round-trip", func() { + // Subscribe with queue group + sub, err := infra.NC.QueueSubscribeReply("test.echo", "echo-workers", func(data []byte, reply func([]byte)) { + // Echo back the request data with a prefix + reply(append([]byte("echo:"), data...)) + }) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = sub.Unsubscribe() }() + + FlushNATS(infra.NC) + + // Send request and wait for reply + replyData, err := infra.NC.Request("test.echo", []byte("hello"), 5*time.Second) + Expect(err).ToNot(HaveOccurred()) + Expect(string(replyData)).To(Equal("echo:hello")) + }) + + It("should load-balance requests across queue subscribers", func() { + var worker1Count, worker2Count atomic.Int32 + + sub1, _ := infra.NC.QueueSubscribeReply("test.lb", "lb-workers", func(data []byte, reply func([]byte)) { + worker1Count.Add(1) + reply([]byte("w1")) + }) + defer func() { _ = sub1.Unsubscribe() }() + + sub2, _ := infra.NC.QueueSubscribeReply("test.lb", "lb-workers", func(data []byte, reply func([]byte)) { + worker2Count.Add(1) + reply([]byte("w2")) + }) + defer func() { _ = sub2.Unsubscribe() }() + + FlushNATS(infra.NC) + + // Send multiple requests + for range 10 { + _, err := infra.NC.Request("test.lb", []byte("req"), 5*time.Second) + Expect(err).ToNot(HaveOccurred()) + } + + // Both workers should have handled some requests + total := worker1Count.Load() + worker2Count.Load() + Expect(total).To(Equal(int32(10))) + // NATS typically distributes evenly, but we just check both got work + Expect(worker1Count.Load()).To(BeNumerically(">", 0)) + Expect(worker2Count.Load()).To(BeNumerically(">", 0)) + }) + }) +})