From 3b4858851ac9a7f8d23da75e8e7ca97f5fa3fec9 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 4 Sep 2026 05:18:02 +0000 Subject: [PATCH] feat(distributed): carry an agent cancel on the worker's own tunnel agent..cancel was the last family on a message bus, and the only reason an agent worker dialled one. Its subscriber is the worker running the execution, and a worker has no database, so the family could not move to the PostgreSQL fan-out carrier: a cancel published there would reach no worker while reporting that it had been sent. It is a control verb now. An agent worker mounts workerctl.PathAgentCancel on the loopback control plane behind its tunnel and applies the cancel to the same registry the executor registers a run on. The frontend issues it through nodes.AgentControlClient.CancelAgentRun. That call is a FAN-OUT and not a pick, because nothing records which worker holds a given execution: the claim row names the claiming replica, and it is deleted when the run ends. Every agent worker a live replica can reach is asked over its own tunnel, relayed by the peer mesh when a peer holds it, and each worker answers only for itself. The answers stay apart, which is why this family was held back. A cancel a worker made is nil. A cancel some worker could not be asked is ErrAgentCancelUndelivered, which is neither a refusal nor a missing run. A cancel every reachable worker declined to own is ErrAgentRunNotOnAnyWorker. A deployment with no agent worker is ErrNoAgentWorker. Neither new sentinel wraps ErrWorkerUnroutable and neither is a worker answer, so nothing is reaped, demoted or evicted because of a cancel. A worker in the ABSENT CONNECTION condition, one whose tunnel was lost inside the reconnect grace, counts as undelivered. It is not retried in the call and not queued: a retry would spend a budget the caller did not choose, and a queue would need durable state whose only consumer is a run whose control stream went with the tunnel. A worker whose departure has outlived the grace is the one routing fact a caller may act on and is excluded, or a single retired agent node would make every cancel undelivered for ever. The fan-out reads a different node set from the pick. A draining worker takes no new work but is still finishing what it holds, so it is offered the cancel; a pending one is refused by the tunnel route on every dial and is not. With that, nothing in LocalAI connects to NATS. The agent worker's dial, its credential ladder and its refresh loop are gone, and so is the frontend's cancel carrier. LOCALAI_NATS_URL is accepted and ignored everywhere, and distributed mode no longer requires it. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/application/agent_control_wiring.go | 7 +- core/application/agent_control_wiring_test.go | 36 ++++ core/application/dispatch_loop_wiring.go | 9 +- core/application/distributed.go | 121 +++++------- core/application/distributed_test.go | 15 +- core/application/fanout_wiring.go | 33 ++-- core/application/fanout_wiring_test.go | 83 ++++---- core/cli/agent_worker.go | 115 ++++------- core/cli/agent_worker_test.go | 97 ++++++--- core/cli/run.go | 4 +- core/config/distributed_config.go | 12 +- core/services/agents/dispatcher.go | 32 ++- core/services/agents/events.go | 186 +++++++++++------- core/services/agents/events_pg_test.go | 141 ++++++++----- core/services/agents/events_test.go | 2 +- core/services/agentworker/handlers.go | 16 +- core/services/agentworker/server_test.go | 13 +- core/services/messaging/client.go | 43 ++-- core/services/messaging/subjects.go | 33 ++++ core/services/nodes/agent_control.go | 106 ++++++++++ core/services/nodes/agent_control_test.go | 174 +++++++++++++++- core/services/nodes/agent_selector.go | 133 ++++++++++++- core/services/nodes/agent_selector_test.go | 158 +++++++++++++-- core/services/nodes/interfaces.go | 11 +- core/services/workerctl/paths.go | 7 +- docs/content/features/distributed-mode.md | 103 +++++----- docs/content/reference/cli-reference.md | 2 +- .../e2e/distributed/agent_distributed_test.go | 142 ++++++++++--- .../distributed/agent_native_executor_test.go | 76 +++---- tests/e2e/distributed/foundation_test.go | 50 +---- tests/e2e/distributed/sse_routes_test.go | 4 +- 31 files changed, 1369 insertions(+), 595 deletions(-) diff --git a/core/application/agent_control_wiring.go b/core/application/agent_control_wiring.go index f991ee875..af5e2113e 100644 --- a/core/application/agent_control_wiring.go +++ b/core/application/agent_control_wiring.go @@ -26,6 +26,11 @@ import ( // 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. +// +// This client is also the deployment's agent CANCELLER, which is why the agent +// event bridge takes it: a cancel is a control RPC on the tunnels the workers +// hold, and the reconnect grace it is built with is what decides whether a +// worker that is not connected makes a cancel undelivered or is simply gone. func newAgentControl(cfg config.DistributedConfig, registry *nodes.NodeRegistry, conns nodes.AgentConnectionReader, control *nodes.ControlClient) (*nodes.AgentControlClient, error) { if cfg.InstanceID == "" { @@ -37,5 +42,5 @@ func newAgentControl(cfg config.DistributedConfig, registry *nodes.NodeRegistry, 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 + return nodes.NewAgentControlClient(nodes.NewAgentSelector(registry, conns, cfg.InstanceID, cfg.WorkerReconnectGrace), control), nil } diff --git a/core/application/agent_control_wiring_test.go b/core/application/agent_control_wiring_test.go index 3fb340229..f6a98354a 100644 --- a/core/application/agent_control_wiring_test.go +++ b/core/application/agent_control_wiring_test.go @@ -5,12 +5,15 @@ package application import ( "context" "runtime" + "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/cluster" mcpremote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/testutil" ) @@ -24,6 +27,12 @@ import ( type recordingConnections struct { owners []string seen chan string + // graces records the reconnect window each presence read was made with. It + // is the other argument with no other symptom: with a zero one every worker + // that lost its tunnel a moment ago reads as GONE, and a cancel addressed to + // it is reported as a run no worker is running rather than as one this + // deployment could not deliver. + graces []time.Duration } func (r *recordingConnections) ConnectedAmong(_ context.Context, _ []string, owner string) ([]string, []string, error) { @@ -37,6 +46,11 @@ func (r *recordingConnections) ConnectedAmong(_ context.Context, _ []string, own return nil, nil, nil } +func (r *recordingConnections) Presence(_ context.Context, _ string, grace time.Duration) (cluster.Presence, error) { + r.graces = append(r.graces, grace) + return cluster.PresenceReconnecting, nil +} + // newRecordingConnections returns a reader whose channel is ready BEFORE any // loop can be started against it. Creating it lazily from the spec goroutine // would race the loop's own goroutine reading it. @@ -89,6 +103,28 @@ var _ = Describe("building the frontend's agent control client", func() { Expect(conns.owners).To(ConsistOf("replica-7")) }) + It("makes the reconnect grace the one a cancel measures a lost tunnel against", func() { + // The other silent argument. With a zero grace every registered agent + // worker that is not connected right now reads as GONE, so a cancel + // addressed to one that is merely re-homing is reported as "no worker + // is running that execution" rather than as one that was not delivered. + // The only symptom is a wrong answer to a caller. + Expect(registry.Register(ctx, &nodes.BackendNode{ + Name: "agent-grace", NodeType: nodes.NodeTypeAgent, Address: "a:50051", + }, true)).To(Succeed()) + + conns := newRecordingConnections() + client, err := newAgentControl( + config.DistributedConfig{InstanceID: "replica-7", WorkerReconnectGrace: 7 * time.Minute}, + registry, conns, nodes.NewControlClient(nil, "token")) + Expect(err).ToNot(HaveOccurred()) + + // Driven through a real cancel, so what is asserted is the value that + // reached the presence read rather than one stored on a field. + _ = client.CancelAgentRun(ctx, messaging.AgentCancelRequest{AgentName: "a", MessageID: "m"}) + Expect(conns.graces).To(ConsistOf(7 * time.Minute)) + }) + It("refuses to build with no instance id", func() { _, err := newAgentControl(config.DistributedConfig{}, registry, newRecordingConnections(), nodes.NewControlClient(nil, "token")) diff --git a/core/application/dispatch_loop_wiring.go b/core/application/dispatch_loop_wiring.go index b34fefc39..1dbf05e12 100644 --- a/core/application/dispatch_loop_wiring.go +++ b/core/application/dispatch_loop_wiring.go @@ -48,6 +48,13 @@ import ( // The SELECTOR is built here rather than borrowed from newAgentControl, and // deliberately: nodes.AgentSelector holds no per-caller state, and sharing one // would couple the dispatch loop's lifetime to MCP's for nothing. +// +// The reconnect grace it is built with is INERT on this path and is passed +// correctly anyway. The selector reads it only in Reachable, which is what a +// fan-out verb (an agent cancel) asks; this loop asks PickConnected, which +// never needs it, because a worker that is not connected cannot be picked +// whatever the reason. Passing a value this loop cannot observe is still +// cheaper than a second constructor. func startJobDispatchLoop(ctx context.Context, cfg config.DistributedConfig, db *gorm.DB, store *jobs.JobStore, registry *nodes.NodeRegistry, conns nodes.AgentConnectionReader, control *nodes.ControlClient, broadcast *nodes.Rebroadcaster) (*jobs.DispatchLoop, error) { @@ -63,7 +70,7 @@ func startJobDispatchLoop(ctx context.Context, cfg config.DistributedConfig, db loop, err := jobs.NewDispatchLoop(jobs.DispatchConfig{ DB: db, Owner: cfg.InstanceID, - Selector: nodes.NewAgentSelector(registry, conns, cfg.InstanceID), + Selector: nodes.NewAgentSelector(registry, conns, cfg.InstanceID, cfg.WorkerReconnectGrace), Control: control, // The allow list lives in nodes and is keyed on the worker's node type; // nothing here decides what a worker may broadcast on. diff --git a/core/application/distributed.go b/core/application/distributed.go index ce57b2535..5502d8bc4 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -31,26 +31,6 @@ import ( // DistributedServices holds all services initialized for distributed mode. type DistributedServices struct { - // CancelCarrier is the NATS connection, and it is here for exactly one - // family: agent..cancel. - // - // Every other family this deployment fans out on moved to Bus below. That - // one could not, and the reason is structural rather than incidental: its - // only subscriber is the agent WORKER, which has no database and so cannot - // join a carrier that rides PostgreSQL. A cancel published on Bus would be - // published where no worker listens, succeed, and be reported as sent. - // - // It is named CancelCarrier and not Nats so that reaching for it to carry - // anything else has to be a decision. The field used to be Nats, and - // *messaging.Client satisfies messaging.Broadcaster, so any adopter that - // took it instead of Broadcast() compiled, started, published, and was - // delivered onto the carrier this deployment is being taken off, with no - // error anywhere. The client's queue and request/reply halves are deleted - // now, so the smaller mistakes are build failures; this name is what is - // left to make the remaining one visible. - // - // It goes when a cancel rides the worker's tunnel as a control verb. - CancelCarrier *messaging.Client Store storage.ObjectStore Registry *nodes.NodeRegistry Router *nodes.SmartRouter @@ -154,11 +134,9 @@ func (ds *DistributedServices) Shutdown() { if closer, ok := ds.Store.(io.Closer); ok { closer.Close() } - // AgentBridge has no Close method: the cancel subscription it holds is - // cleaned up when the carrier below is closed. - if ds.CancelCarrier != nil { - ds.CancelCarrier.Close() - } + // AgentBridge has no Close method and needs none: it holds no + // process-lifetime subscription of its own beyond the observable + // persister, whose carrier is closed below. // The broadcast carrier holds a PostgreSQL session pinned for the life // of the process, plus the goroutine parked on it. A replica that // leaves one behind on every restart runs the server out of @@ -201,30 +179,16 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Distributed instance", "id", cfg.Distributed.InstanceID) - // Connect to the cancel carrier. + // No message bus is dialled here, and there is none left to dial. The last + // family that needed one was agent..cancel, whose subscriber is an + // agent worker that has no database and so could not join the broadcast + // carrier below; it is now a control RPC on the tunnel that worker holds. + // A distributed deployment needs PostgreSQL and the frontends' own HTTP + // listener, and nothing else. // - // This is the ONE bus connection a frontend replica still opens, and it - // carries one family: agent..cancel, whose subscriber is an agent - // worker that has no database and cannot read the broadcast carrier opened - // below. Everything else a replica fans out travels on that one. - natsAuth := cfg.Distributed.NatsAuthConfig() - if natsAuth.RequireAuth && (natsAuth.ServiceUserJWT == "" || natsAuth.ServiceUserSeed == "") { - return nil, fmt.Errorf("LOCALAI_NATS_REQUIRE_AUTH requires LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED") - } - natsOpts := cfg.Distributed.NatsMessagingOptions("", "") - natsClient, err := messaging.New(cfg.Distributed.NatsURL, natsOpts...) - if err != nil { - return nil, fmt.Errorf("connecting to NATS: %w", err) - } - xlog.Info("Connected to the agent cancel carrier", "url", sanitize.URL(cfg.Distributed.NatsURL)) - - // Ensure the carrier is closed if any subsequent initialization step fails. + // success guards the carriers opened below, which must not be left pinned + // when a later initialization step fails. success := false - defer func() { - if !success { - natsClient.Close() - } - }() // Initialize object storage var store storage.ObjectStore @@ -439,28 +403,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Distributed agent store initialized") - // The job dispatcher and the agent event bridge, both on the broadcast - // carrier. See newFanoutBridges for why the two constructors are reached - // through one function that names *pgbus.Bus. - dispatcher, agentBridge, rebroadcast, err := newFanoutBridges(bus, natsClient, jobStore, agentStore, authDB, cfg.Distributed.InstanceID) - if err != nil { - return nil, err - } - - // Initialize Phase 4 stores (MCP, Gallery, FineTune, Skills) - distStores, err := distributed.InitStores(authDB) - if err != nil { - return nil, fmt.Errorf("initializing distributed stores: %w", err) - } - - // Initialize file manager with local cache - cacheDir := cfg.DataPath + "/cache" - fileMgr, err := storage.NewFileManager(store, cacheDir) - if err != nil { - return nil, fmt.Errorf("initializing file manager: %w", err) - } - xlog.Info("File manager initialized", "cacheDir", cacheDir) - // The frontend's control plane client. It reaches every worker over that // worker's own tunnel, on the same `http` stream tag the file stager below // uses, so a control RPC to a worker another replica holds is relayed the @@ -482,6 +424,33 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return nil, fmt.Errorf("wiring the agent control client: %w", err) } + // The job dispatcher and the agent event bridge, both on the broadcast + // carrier. See newFanoutBridges for why the two constructors are reached + // through one function that names *pgbus.Bus. + // + // The bridge takes agentControl and not a carrier: a cancel is the one + // family whose far end is an agent worker, and it now rides that worker's + // tunnel as a control RPC. It is built ABOVE for that reason, rather than + // with the rest of the control plane below. + dispatcher, agentBridge, rebroadcast, err := newFanoutBridges(bus, agentControl, jobStore, agentStore, authDB, cfg.Distributed.InstanceID) + if err != nil { + return nil, err + } + + // Initialize Phase 4 stores (MCP, Gallery, FineTune, Skills) + distStores, err := distributed.InitStores(authDB) + if err != nil { + return nil, fmt.Errorf("initializing distributed stores: %w", err) + } + + // Initialize file manager with local cache + cacheDir := cfg.DataPath + "/cache" + fileMgr, err := storage.NewFileManager(store, cacheDir) + if err != nil { + return nil, fmt.Errorf("initializing file manager: %w", err) + } + xlog.Info("File manager initialized", "cacheDir", cacheDir) + // The consumer side of the claim queue, built and started in one act: see // startJobDispatchLoop for why those are not two lines. jobDispatch, err := startJobDispatchLoop(cfg.Context, cfg.Distributed, authDB, jobStore, registry, clusterRegistry, controlClient, rebroadcast) @@ -688,7 +657,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade success = true ds := &DistributedServices{ - CancelCarrier: natsClient, Store: store, Registry: registry, Router: router, @@ -752,15 +720,14 @@ func requireBroadcastCarrier(ds *DistributedServices) error { // carrier, and it exists so that "this family travels on the broadcast carrier // and not on NATS" is decided once instead of at every adopter. // -// It was five sites before this: the fine-tune service, the quantization +// It was five field reads before this: the fine-tune service, the quantization // service, the agent-task setter (twice, on two startup paths), the per-user // services manager and the Open Responses store. Every one of them takes a -// messaging.Broadcaster, and *messaging.Client satisfies that interface too, so -// a site left holding ds.CancelCarrier compiles, starts, publishes and is -// delivered - onto a carrier only agent workers read, and only for cancels. -// Nothing would fail until NATS went away. Collapsing the choice to one -// function makes it a fact a spec can pin, which five scattered field reads -// were not. +// messaging.Broadcaster, and the struct used to carry a second field that +// satisfied it, so a site that reached for the wrong one compiled, started, +// published and was delivered onto a carrier almost nothing read. That second +// field is gone with the last family that needed a bus. Collapsing the choice +// to one function is what keeps it a fact a spec can pin. // // The return is the interface and not *pgbus.Bus on purpose: handing a nil // *pgbus.Bus to an adopter would produce a non-nil interface wrapping a nil diff --git a/core/application/distributed_test.go b/core/application/distributed_test.go index d3eb8167a..95449c4b0 100644 --- a/core/application/distributed_test.go +++ b/core/application/distributed_test.go @@ -123,7 +123,7 @@ var _ = Describe("shutting the distributed services down", func() { // NATS did. Collapsing the choice into one function is what makes it a fact // these specs can hold. var _ = Describe("handing the broadcast carrier to its adopters", func() { - It("returns the carrier the deployment opened and never the cancel carrier", func() { + It("returns the carrier the deployment opened", func() { db, dsn := testutil.SetupTestDBWithDSN() cfg := &config.ApplicationConfig{} cfg.Auth.DatabaseURL = dsn @@ -131,13 +131,12 @@ var _ = Describe("handing the broadcast carrier to its adopters", func() { Expect(err).ToNot(HaveOccurred()) DeferCleanup(bus.Close) - // The cancel carrier is present on the struct, exactly as it is in a - // real deployment: agent..cancel is the one family that could not - // move, because its only subscriber is an agent worker and a worker has - // no database. Identity, not "is a Broadcaster": both fields satisfy - // that interface, which is the whole reason a field read could pick the - // wrong one and stay green. - ds := &DistributedServices{CancelCarrier: &messaging.Client{}, Bus: bus} + // Identity and not "is a Broadcaster". There is no second carrier on + // this struct any more: the family that needed one, agent..cancel, + // rides the agent worker's own tunnel now. The identity assertion stays + // because what it pins is that adopters get THIS bus rather than + // anything else that satisfies the interface. + ds := &DistributedServices{Bus: bus} Expect(ds.Broadcast()).To(BeIdenticalTo(messaging.Broadcaster(bus))) }) diff --git a/core/application/fanout_wiring.go b/core/application/fanout_wiring.go index b14d94300..15ef31ca9 100644 --- a/core/application/fanout_wiring.go +++ b/core/application/fanout_wiring.go @@ -7,7 +7,6 @@ import ( "github.com/mudler/LocalAI/core/services/agents" "github.com/mudler/LocalAI/core/services/jobs" - "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/pgbus" "github.com/mudler/xlog" @@ -44,15 +43,16 @@ import ( // Written as one expression shared with the dispatcher and the bridge, that // mis-wiring stops being a line a spec has to guess at: there is no second // carrier in scope to point it at. -// cancelCarrier is NOT the same carrier and must not be folded into bus. Every -// family this function wires has both of its ends on a frontend replica and so -// moves to the broadcast carrier whole, with one exception: the only subscriber -// to agent..cancel is the agent WORKER, which has no database and cannot -// join the PostgreSQL carrier at all. A cancel published on bus would therefore -// reach no worker, and every cancel of a worker-run agent would be lost while -// reporting success. It stays on the carrier the worker reads until a cancel -// rides the worker's tunnel instead of a broadcast. -func newFanoutBridges(bus *pgbus.Bus, cancelCarrier messaging.Broadcaster, +// canceller is NOT a carrier and must never become one. Every family this +// function wires has both of its ends on a frontend replica and so travels on +// the broadcast carrier, with one exception: the process that holds a +// worker-run agent's cancel function is the agent WORKER, which has no database +// and cannot join the PostgreSQL carrier at all. A cancel published on bus +// would reach no worker, and every cancel of a worker-run agent would be lost +// while reporting success. So it does not travel on bus, or on any bus: it is a +// control RPC on the tunnel the worker already holds, issued by the agent +// control client this takes. +func newFanoutBridges(bus *pgbus.Bus, canceller agents.AgentWorkerCanceller, jobStore *jobs.JobStore, agentStore *agents.AgentStore, db *gorm.DB, instanceID string) (*jobs.Dispatcher, *agents.EventBridge, *nodes.Rebroadcaster, error) { // A nil check on the CONCRETE pointer, before it is widened. Once it is a @@ -63,12 +63,12 @@ func newFanoutBridges(bus *pgbus.Bus, cancelCarrier messaging.Broadcaster, return nil, nil, nil, fmt.Errorf("the job and agent fan-out bridges were built with no broadcast carrier: every job's progress and every agent's events would reach no SSE stream in the deployment") } - if cancelCarrier == nil { - return nil, nil, nil, fmt.Errorf("the agent event bridge was built with no carrier for agent cancels: every cancel of a worker-run agent would be published where no worker listens and reported as sent") + if canceller == nil { + return nil, nil, nil, fmt.Errorf("the agent event bridge was built with no way to cancel a worker-run agent: every cancel of one would be sent nowhere") } dispatcher := jobs.NewDispatcher(jobStore, bus, db, instanceID) - bridge := agents.NewEventBridge(bus, agentStore, instanceID).WithCancelCarrier(cancelCarrier) + bridge := agents.NewEventBridge(bus, agentStore, instanceID, canceller) // Warned rather than refused, and deliberately: the persister needs a store // and a deployment without one still serves live SSE correctly. What it @@ -82,3 +82,10 @@ func newFanoutBridges(bus *pgbus.Bus, cancelCarrier messaging.Broadcaster, return dispatcher, bridge, nodes.NewRebroadcaster(bus), nil } + +// The real canceller, asserted where both packages are already imported. +// +// agents may not import nodes, so the port there is satisfied structurally and +// a signature drift between the two would otherwise surface as a nil field in +// this file's argument list, which compiles. Here it is a build failure. +var _ agents.AgentWorkerCanceller = (*nodes.AgentControlClient)(nil) diff --git a/core/application/fanout_wiring_test.go b/core/application/fanout_wiring_test.go index 6c3a77c3a..8f13fb7ac 100644 --- a/core/application/fanout_wiring_test.go +++ b/core/application/fanout_wiring_test.go @@ -61,17 +61,18 @@ var _ = Describe("wiring the job and agent fan-out bridges", func() { Expect(err).ToNot(HaveOccurred()) }) - // The cancel carrier is a SEPARATE argument, and its absence is refused - // separately. Folding it into bus would put every agent cancel on a carrier - // no worker can read, and CancelExecution would go on returning nil. - It("refuses to build with no carrier for agent cancels", func() { + // The canceller is a SEPARATE argument, and its absence is refused + // separately. A bridge built without one has nowhere to send the cancel of + // an agent running on a worker, and the failure would present as + // CancelExecution reporting success on a cancel that reached nobody. + It("refuses to build with no way to cancel a worker-run agent", func() { _, _, _, err := newFanoutBridges(busA, nil, jobStore, agentStore, db, "replica-1") Expect(err).To(HaveOccurred()) - Expect(err.Error()).To(ContainSubstring("agent cancels")) + Expect(err.Error()).To(ContainSubstring("cancel a worker-run agent")) }) It("refuses to build with no carrier", func() { - _, _, _, err := newFanoutBridges(nil, testutil.NewFakeBus(), jobStore, agentStore, db, "replica-1") + _, _, _, err := newFanoutBridges(nil, &stubCanceller{}, jobStore, agentStore, db, "replica-1") Expect(err).To(HaveOccurred()) Expect(err.Error()).To(ContainSubstring("broadcast carrier")) }) @@ -80,7 +81,7 @@ var _ = Describe("wiring the job and agent fan-out bridges", func() { // it can only do if its wildcard subscription is on the carrier the peer // published to. It("subscribes the job dispatcher to results a peer replica broadcasts", func() { - dispatcher, _, _, err := newFanoutBridges(busA, testutil.NewFakeBus(), jobStore, agentStore, db, "replica-1") + dispatcher, _, _, err := newFanoutBridges(busA, &stubCanceller{}, jobStore, agentStore, db, "replica-1") Expect(err).ToNot(HaveOccurred()) Expect(dispatcher.Start(ctx)).To(Succeed()) DeferCleanup(dispatcher.Stop) @@ -101,44 +102,35 @@ var _ = Describe("wiring the job and agent fan-out bridges", func() { }, "20s").Should(Equal("completed")) }) - // S1b. WHICH carrier the cancel family rides, asserted by receipt on both - // carriers rather than by reading the argument list. + // S1b. A cancel does NOT travel on a carrier, asserted by where it goes and + // by where it does not. // - // The nil refusal above only says a carrier was passed. Passing bus for it - // - one token's difference at the one call site, and the natural edit for - // anyone finishing the migration - satisfies that refusal, compiles, and - // publishes successfully onto the PostgreSQL carrier, where the agent - // worker that has to act on the cancel is not and cannot be: it has no - // database. Every unit suite in agents and in application stays green and - // every cancel of a worker-run agent is lost while CancelExecution returns - // nil. + // The nil refusal above only says a canceller was passed. What it cannot + // say is that the bridge uses it instead of publishing onto the broadcast + // carrier, which is the edit anyone finishing this migration would reach + // for: it compiles, it publishes successfully onto PostgreSQL, and the + // agent worker that has to act on the cancel is not and cannot be there. + // Every unit suite stays green and every cancel of a worker-run agent is + // lost while CancelExecution returns nil. // - // So this asserts the cancel ARRIVES on the worker's carrier and, in the - // same spec, that it does NOT arrive on a peer replica's broadcast carrier. - // The negative half is the load-bearing one: the positive half alone passes - // for a bridge wired to both. - It("publishes agent cancels on the worker's carrier and not on the broadcast carrier", func() { - workerCarrier := testutil.NewFakeBus() - _, bridge, _, err := newFanoutBridges(busA, workerCarrier, jobStore, agentStore, db, "replica-1") + // So this asserts the cancel reaches the CANCELLER and, in the same spec, + // that nothing is published on a peer replica's broadcast carrier. The + // negative half is the load-bearing one: the positive half alone passes for + // a bridge that does both. + It("sends an agent cancel to the agent workers and publishes nothing on the broadcast carrier", func() { + canceller := &stubCanceller{} + _, bridge, _, err := newFanoutBridges(busA, canceller, jobStore, agentStore, db, "replica-1") Expect(err).ToNot(HaveOccurred()) onBroadcast := make(chan []byte, 4) _, err = busB.Subscribe(messaging.SubjectAgentCancelWildcard, func(data []byte) { onBroadcast <- data }) Expect(err).ToNot(HaveOccurred()) - onWorker := make(chan []byte, 4) - _, err = workerCarrier.Subscribe(messaging.SubjectAgentCancelWildcard, func(data []byte) { onWorker <- data }) - Expect(err).ToNot(HaveOccurred()) + Expect(bridge.CancelExecution(ctx, "a1", "u1", "msg-1")).To(Succeed()) - Expect(bridge.CancelExecution("a1", "u1", "msg-1")).To(Succeed()) - - // Expect and not Eventually: the worker's carrier here is the shared - // in-memory double, which delivers synchronously inside Publish, so by - // the time CancelExecution has returned the handler has already run. An - // Eventually would turn a mis-wiring into a timeout, which reports as - // slowness rather than as the wiring fact it is. - Expect(onWorker).To(Receive(), - "the agent cancel did not reach the worker's carrier; a worker has no database and cannot read the broadcast carrier, so this cancel reached nobody and was reported as sent") + Expect(canceller.requests).To(ConsistOf(messaging.AgentCancelRequest{ + AgentName: "a1", UserID: "u1", MessageID: "msg-1", + }), "the cancel did not reach the agent workers, so it reached nobody and was reported as sent") Consistently(onBroadcast, "2s").ShouldNot(Receive(), "the agent cancel was published on the broadcast carrier, where no agent worker is or can be subscribed") }) @@ -147,7 +139,7 @@ var _ = Describe("wiring the job and agent fan-out bridges", func() { // only do if it was started AND is on the same carrier AND its filter has // the right number of tokens. It("subscribes the agent observable persister to events a peer replica broadcasts", func() { - _, bridge, _, err := newFanoutBridges(busA, testutil.NewFakeBus(), jobStore, agentStore, db, "replica-1") + _, bridge, _, err := newFanoutBridges(busA, &stubCanceller{}, jobStore, agentStore, db, "replica-1") Expect(err).ToNot(HaveOccurred()) Expect(bridge).ToNot(BeNil()) @@ -180,7 +172,7 @@ var _ = Describe("wiring the job and agent fan-out bridges", func() { // value, which is true for a publish that went nowhere. DescribeTable("re-broadcasts a worker's line onto the carrier a peer replica reads", func(subject string, payload string) { - _, _, rebroadcast, err := newFanoutBridges(busA, testutil.NewFakeBus(), jobStore, agentStore, db, "replica-1") + _, _, rebroadcast, err := newFanoutBridges(busA, &stubCanceller{}, jobStore, agentStore, db, "replica-1") Expect(err).ToNot(HaveOccurred()) delivered := make(chan []byte, 4) @@ -196,7 +188,7 @@ var _ = Describe("wiring the job and agent fan-out bridges", func() { ) It("re-broadcasts an agent's events onto the carrier a peer replica reads", func() { - _, _, rebroadcast, err := newFanoutBridges(busA, testutil.NewFakeBus(), jobStore, agentStore, db, "replica-1") + _, _, rebroadcast, err := newFanoutBridges(busA, &stubCanceller{}, jobStore, agentStore, db, "replica-1") Expect(err).ToNot(HaveOccurred()) delivered := make(chan []byte, 4) @@ -209,3 +201,16 @@ var _ = Describe("wiring the job and agent fan-out bridges", func() { Eventually(delivered, "20s").Should(Receive(MatchJSON(`{"event_type":"json_message"}`))) }) }) + +// stubCanceller stands in for the frontend's agent control client, which +// reaches workers over their tunnels and is driven over a real one in +// core/services/nodes. What these specs need from it is only that the bridge +// asks it at all. +type stubCanceller struct { + requests []messaging.AgentCancelRequest +} + +func (s *stubCanceller) CancelAgentRun(_ context.Context, req messaging.AgentCancelRequest) error { + s.requests = append(s.requests, req) + return nil +} diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index 161e67b89..6df16fc21 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -20,7 +20,6 @@ import ( "github.com/mudler/LocalAI/core/services/jobs" mcpRemote "github.com/mudler/LocalAI/core/services/mcp" "github.com/mudler/LocalAI/core/services/messaging" - "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/cogito" "github.com/mudler/cogito/clients" "github.com/mudler/xlog" @@ -40,20 +39,23 @@ import ( // control verbs by RPC without the worker opening an inbound port. No verb the // frontend addresses to THIS worker travels on the bus any more. // -// --nats-url is still required, and for one thing only: agent..cancel. -// That family runs the other way, from a frontend replica to whichever worker -// holds the execution, and it could not move to the broadcast carrier because -// that carrier rides PostgreSQL and this process has no database. A worker that -// came up without a bus would register, serve, run agents and ignore every -// cancel, with nothing in the deployment reporting it. The flag goes when a -// cancel rides the tunnel as a control verb. +// It dials NO message bus. The last family that needed one was +// agent..cancel, which ran the other way, from a frontend replica to +// whichever worker held the execution; it is now a control verb on this +// worker's own tunnel (workerctl.PathAgentCancel), so a cancel reaches the +// worker running the agent without either side touching a broker. // // Usage: // -// localai agent-worker --nats-url nats://... --register-to http://localai:8080 +// localai agent-worker --register-to http://localai:8080 type AgentWorkerCMD struct { - // NATS (required) - NatsURL string `env:"LOCALAI_NATS_URL" required:"" help:"NATS server URL" group:"distributed"` + // NatsURL is accepted and ignored, exactly as the backend worker's is (see + // core/services/worker/config.go). An agent worker connects to no message + // bus: every verb a frontend addresses to it arrives on the tunnel it + // dials, and a cancel now arrives the same way. It stays here, without + // required, so an existing command line or unit file that still carries + // --nats-url starts rather than failing to parse. + NatsURL string `env:"LOCALAI_NATS_URL" help:"Ignored. An agent worker connects to no message bus; the frontend reaches it over its outbound tunnel. Accepted so an existing worker command line still starts." group:"distributed" hidden:""` // Registration (required) RegisterTo string `env:"LOCALAI_REGISTER_TO" required:"" help:"Frontend URL for registration" group:"registration"` @@ -88,7 +90,7 @@ func (cmd *AgentWorkerCMD) natsAuthRequired() bool { } func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { - xlog.Info("Starting agent worker", "nats", sanitize.URL(cmd.NatsURL), "register_to", cmd.RegisterTo) + xlog.Info("Starting agent worker", "register_to", cmd.RegisterTo) // Resolve API URL apiURL := cmp.Or(cmd.APIURL, strings.TrimRight(cmd.RegisterTo, "/")) @@ -117,9 +119,14 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { shutdownCtx, shutdownCancel := context.WithCancel(context.Background()) defer shutdownCancel() - // Acquire credentials via (re)registration. When the bus requires auth and no - // static fallback is configured, wait through admin approval until the - // frontend mints credentials rather than starting unauthenticated. + // Register, and obtain this node's identity and its tunnel credential. + // + // The manager is still the NATS credential manager and still gated on the + // NATS auth flags, and that is deliberate rather than left over: what the + // gate decides is whether registration WAITS THROUGH ADMIN APPROVAL instead + // of returning a pending response, and that behaviour is unchanged by this + // worker no longer dialling a bus. What it no longer does is dial one: the + // only value read off it below is TunnelToken. credMgr := workerregistry.NewNATSCredentialManager( func(ctx context.Context) (*workerregistry.RegisterResponse, error) { return regClient.RegisterFull(ctx, registrationBody) @@ -147,42 +154,6 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { go regClient.HeartbeatLoop(shutdownCtx, nodeID, heartbeatInterval, func() map[string]any { return map[string]any{} }) - // Resolve the cancel carrier's credentials with precedence: explicit env - // override, then frontend-minted (auto-refreshed before expiry), then - // service fallback. Each static source must supply JWT and seed together. - natsTLS := messaging.TLSFiles{CA: cmd.NatsTLSCA, Cert: cmd.NatsTLSCert, Key: cmd.NatsTLSKey} - var natsOpts []messaging.Option - switch { - case cmd.NatsJWT != "" || cmd.NatsUserSeed != "": - if (cmd.NatsJWT == "") != (cmd.NatsUserSeed == "") { - return fmt.Errorf("LOCALAI_NATS_JWT and LOCALAI_NATS_USER_SEED must be set together") - } - natsOpts = append(natsOpts, messaging.WithUserJWT(cmd.NatsJWT, cmd.NatsUserSeed)) - case credMgr.HasCredentials(): - natsOpts = append(natsOpts, messaging.WithUserJWTProvider(credMgr.Provider())) - go func() { - if err := credMgr.RefreshLoop(shutdownCtx); err != nil { - xlog.Error("NATS credential refresh permanently failed; shutting down agent worker", "error", err) - shutdownCancel() - } - }() - case cmd.NatsServiceJWT != "" || cmd.NatsServiceSeed != "": - if (cmd.NatsServiceJWT == "") != (cmd.NatsServiceSeed == "") { - return fmt.Errorf("LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED must be set together") - } - natsOpts = append(natsOpts, messaging.WithUserJWT(cmd.NatsServiceJWT, cmd.NatsServiceSeed)) - case cmd.natsAuthRequired(): - return fmt.Errorf("NATS JWT+seed required: enable frontend minting or set LOCALAI_NATS_* env vars") - } - if natsTLS.Enabled() { - natsOpts = append(natsOpts, messaging.WithTLS(natsTLS)) - } - natsClient, err := messaging.New(cmd.NatsURL, natsOpts...) - if err != nil { - return fmt.Errorf("connecting to NATS: %w", err) - } - defer natsClient.Close() - // The executor and the event bridge the control plane serves, built BEFORE // the tunnel because a verb mounted with a nil handler answers a 404, which // a frontend reads as a worker too old to serve it. @@ -190,7 +161,7 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { // No ConfigProvider and no SkillStore: config and skills arrive in the // request body, exactly as they arrived in the job payload before, because // an agent worker still has no database. - eventBridge := agents.NewEventBridge(natsClient, nil, "agent-worker-"+nodeID) + eventBridge := agents.NewWorkerEventBridge("agent-worker-" + nodeID) executor := agents.NewWorkerExecutor(eventBridge, nil, apiURL, cmd.APIToken) mcpCIJobTimeout, err := time.ParseDuration(cmd.MCPCIJobTimeout) @@ -212,9 +183,9 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { // 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 - // captured value would lock this worker out of its own tunnel at the first - // JWT refresh. + // because every registration the manager performs ROTATES it: the frontend + // stores only the hash of the newest one, so a captured value would lock + // this worker out of its own tunnel after any re-registration. // // It is started AFTER registration, which is what supplies both the node // identity the dial names and the credential it presents. @@ -234,39 +205,20 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { } }() - // The cancel listener is the ONE thing still on the bus here, and the only - // reason this process dialled one. A cancel is a broadcast to every replica - // and every worker, because the replica holding the run is not the one the - // cancel request lands on, and this worker cannot join the carrier the rest - // of the deployment fans out on: that carrier is the auth database, and an - // agent worker has no database access at all. - cancelSub, err := eventBridge.StartCancelListener() - if err != nil { - xlog.Warn("Failed to start cancel listener", "error", err) - } else { - defer func() { _ = cancelSub.Unsubscribe() }() - } - xlog.Info("Agent worker ready, serving agent execution and MCP CI runs on its tunnel", "node", nodeID) - // Wait for an OS signal or an internal fatal condition (e.g. NATS - // credentials became unrenewable), so the worker restarts and re-acquires - // rather than lingering unable to serve. + // Wait for an OS signal. There is no internal fatal condition left to wait + // on: the one that existed was a NATS credential this worker could no + // longer renew, and it renews none. sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) - var runErr error - select { - case <-sigCh: - case <-shutdownCtx.Done(): - runErr = fmt.Errorf("agent worker shutting down: NATS credentials unavailable") - xlog.Error("Internal shutdown requested", "error", runErr) - } + <-sigCh xlog.Info("Shutting down agent worker") shutdownCancel() // stop heartbeat loop immediately mcpTools.CloseAllMCPSessions() regClient.GracefulDeregister(nodeID) - return runErr + return nil } // The MCP verbs, written ONCE and served on two carriers. @@ -612,6 +564,11 @@ func agentWorkerControlHandlers(executor *agents.WorkerExecutor, apiURL, apiToke return agentworker.Config{ MCPTool: serveMCPToolRequest, MCPDiscovery: serveMCPDiscoveryRequest, + // The verb that removed this process's last reason to dial a bus. It + // reaches the SAME cancel registry the executor registers a run on, + // because it is the same bridge: a cancel arriving on the tunnel has to + // find an execution that is publishing onto a control stream. + AgentCancel: executor.Cancel, // Drops the MCP sessions cached for a backend that went away, on the // path a backend worker serves by killing the process instead. BackendStop: dropMCPSessionsForBackend, diff --git a/core/cli/agent_worker_test.go b/core/cli/agent_worker_test.go index 7d2617d6c..1e537999a 100644 --- a/core/cli/agent_worker_test.go +++ b/core/cli/agent_worker_test.go @@ -2,6 +2,7 @@ package cli import ( "bufio" + "bytes" "context" "encoding/json" "net/http" @@ -17,7 +18,6 @@ import ( "github.com/mudler/LocalAI/core/services/agents" mcpRemote "github.com/mudler/LocalAI/core/services/mcp" "github.com/mudler/LocalAI/core/services/messaging" - "github.com/mudler/LocalAI/core/services/testutil" "github.com/mudler/LocalAI/core/services/workerctl" ) @@ -83,14 +83,14 @@ var _ = Describe("The agent worker's backend stop", func() { // verb rather than as a wiring mistake. Nothing else in this repo would notice. var _ = Describe("The agent worker's control-plane wiring", func() { var base string + var bridge *agents.EventBridge BeforeEach(func() { mux := http.NewServeMux() // Built exactly as Run builds it, from an executor and the MCP CI // timeout, so a field this function forgets to set is a 404 here. - executor := agents.NewWorkerExecutor( - agents.NewEventBridge(testutil.NewFakeBus(), nil, "agent-worker-spec"), - nil, "http://127.0.0.1:1", "token") + bridge = agents.NewWorkerEventBridge("agent-worker-spec") + executor := agents.NewWorkerExecutor(bridge, nil, "http://127.0.0.1:1", "token") agentWorkerControlHandlers(executor, "http://127.0.0.1:1", "token", time.Second).Register(mux) srv := httptest.NewServer(mux) DeferCleanup(srv.Close) @@ -113,8 +113,50 @@ var _ = Describe("The agent worker's control-plane wiring", func() { // in the tree can tell the two apart and only this spec can. Entry("agent execute", workerctl.PathAgentExecute), Entry("mcp ci run", workerctl.PathMCPCIRun), + // The verb that removed this process's last reason to dial a bus. + // Unwired it answers a 404, which the frontend reads as a worker too + // old to serve it, and every cancel of an agent this worker is running + // is then reported as one that could not be delivered - for ever. + Entry("agent cancel", workerctl.PathAgentCancel), ) + // The cancel verb end to end through the mux, because the thing that must + // be true is that the path reaches the SAME cancel registry the executor + // registers a run on. Two bridges would compile, mount, answer 200 and + // cancel nothing. + It("cancels a run registered on the executor's own bridge, and says so", func() { + cancelled := make(chan struct{}) + bridge.RegisterCancel("msg-1", func() { close(cancelled) }) + + post := func(messageID string) messaging.AgentCancelReply { + GinkgoHelper() + body, err := json.Marshal(messaging.AgentCancelRequest{AgentName: "a1", UserID: "u1", MessageID: messageID}) + Expect(err).ToNot(HaveOccurred()) + resp, err := http.Post(base+workerctl.PathAgentCancel, "application/json", bytes.NewReader(body)) //nolint:gosec,noctx // httptest server, no redirects to follow + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = resp.Body.Close() }) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + var reply messaging.AgentCancelReply + Expect(json.NewDecoder(resp.Body).Decode(&reply)).To(Succeed()) + return reply + } + + Expect(post("msg-nobody-is-running").Cancelled).To(BeFalse(), + "a worker answers only for itself, and false is that answer rather than a failure") + Expect(post("msg-1").Cancelled).To(BeTrue()) + Eventually(cancelled, "20s").Should(BeClosed()) + }) + + It("fails to serve a cancel it cannot read, rather than answering that it found nothing", func() { + // The two are different facts. A 200 with cancelled false would be read + // as this worker's own answer about the run; a body it could not decode + // is not an answer about anything. + resp, err := http.Post(base+workerctl.PathAgentCancel, "application/json", strings.NewReader(`{"message_id":`)) //nolint:gosec,noctx // httptest server, no redirects to follow + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = resp.Body.Close() }) + Expect(resp.StatusCode).To(Equal(http.StatusInternalServerError)) + }) + DescribeTable("answers a dispatched verb as a stream, so progress and the terminal line share one body", func(path string) { resp, err := http.Post(base+path, "application/json", strings.NewReader(`{}`)) //nolint:gosec,noctx // httptest server, no redirects to follow @@ -147,28 +189,23 @@ var _ = Describe("The agent worker's control-plane wiring", func() { ) }) -// The last reason an agent worker dials a message bus, pinned so that removing -// it is a decision rather than an accident. +// The agent worker dials NO message bus, pinned so that a flag reappearing as +// required is a decision rather than an accident. // -// Every verb a frontend addresses to THIS worker now arrives on the tunnel it -// dials, and the two queue groups it used to join are gone. One family is left, -// in the other direction: agent..cancel. Its publisher is a frontend -// replica and its ONLY subscriber is this process, and a worker has no database -// and so cannot join the PostgreSQL carrier every other family moved to. A -// worker that came up without a bus would register, serve, run agents, and -// ignore every cancel, returning nothing to say so - the cancel would be -// published, would succeed, and would reach nobody. +// Every verb a frontend addresses to this worker arrives on the tunnel it +// dials, and that now includes the cancel: agent..cancel was the last +// family in the other direction, from a frontend replica to whichever worker +// held the execution, and it could not move to the broadcast carrier because +// that carrier rides PostgreSQL and this process has no database. It is a +// control verb on the worker's own tunnel instead. // -// So --nats-url stays required here, and it is required for this and for -// nothing else. When a cancel rides the tunnel as a control verb, this spec is -// what has to be deleted for the flag to become optional, and deleting it is -// then the visible half of that change. -var _ = Describe("The agent worker's remaining bus requirement", func() { +// --nats-url is still ACCEPTED, and ignored, so an existing command line or +// unit file starts unchanged. +var _ = Describe("The agent worker's bus requirement", func() { parse := func(args ...string) error { - // kong resolves env: tags from the process environment, and a - // LOCALAI_NATS_URL that is SET BUT EMPTY satisfies a required flag. - // Left in place, this spec would pass on a developer's shell and on - // nothing else. + // kong resolves env: tags from the process environment, so a + // LOCALAI_NATS_URL inherited from a developer's shell would make the + // first spec below pass for the wrong reason. if prior, had := os.LookupEnv("LOCALAI_NATS_URL"); had { Expect(os.Unsetenv("LOCALAI_NATS_URL")).To(Succeed()) DeferCleanup(func() { _ = os.Setenv("LOCALAI_NATS_URL", prior) }) @@ -182,16 +219,14 @@ var _ = Describe("The agent worker's remaining bus requirement", func() { return err } - It("refuses to start without a bus to hear cancels on", func() { - // Refused at parse time and not at first use. A worker that started - // and only failed to subscribe would already have registered itself as - // available to run agents nobody can cancel. - Expect(parse("--register-to", "http://frontend:8080")). - To(MatchError(ContainSubstring("--nats-url")), - "the agent worker started with no bus: every cancel of an agent it runs would be published to nobody and reported as sent") + It("starts with no bus named at all", func() { + Expect(parse("--register-to", "http://frontend:8080")).To(Succeed(), + "an agent worker connects to no message bus and must not demand the URL of one") }) - It("parses once the bus is named", func() { + It("still accepts a command line that names one", func() { + // Ignored, not rejected. An operator upgrading a fleet must not have to + // edit every unit file in the same change. Expect(parse("--register-to", "http://frontend:8080", "--nats-url", "nats://bus:4222")).To(Succeed()) }) }) diff --git a/core/cli/run.go b/core/cli/run.go index 7af51b846..1c773f5c1 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -162,9 +162,9 @@ type RunCMD struct { DefaultAPIKeyExpiry string `env:"LOCALAI_DEFAULT_API_KEY_EXPIRY" help:"Default expiry for API keys (e.g. 90d, 1y; empty = no expiry)" group:"auth"` // Distributed / Horizontal Scaling - Distributed bool `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL + NATS)" group:"distributed"` + Distributed bool `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL; no message bus)" group:"distributed"` InstanceID string `env:"LOCALAI_INSTANCE_ID" help:"Unique instance ID for distributed mode (auto-generated UUID if empty)" group:"distributed"` - NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"` + NatsURL string `env:"LOCALAI_NATS_URL" help:"Ignored. No component of a distributed deployment connects to a message bus; state and fan-out ride PostgreSQL and workers are reached over their own tunnels. Accepted so an existing command line still starts." group:"distributed"` DistributedAdvertiseAddr string `env:"LOCALAI_DISTRIBUTED_ADVERTISE_ADDR" help:"host:port other frontend replicas dial to reach this one (peer link). Empty = derived from the local address that routes to PostgreSQL, which only works when the database is on another host." group:"distributed"` StorageURL string `env:"LOCALAI_STORAGE_URL" help:"S3-compatible storage endpoint URL (e.g., http://minio:9000)" group:"distributed"` StorageBucket string `env:"LOCALAI_STORAGE_BUCKET" default:"localai" help:"S3 bucket name for object storage" group:"distributed"` diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index af2b375bf..65334b86b 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -22,7 +22,9 @@ type DistributedConfig struct { // answer is only usable when the database is remote, so a deployment with // a local or sidecar database has to set this. AdvertiseAddr string // LOCALAI_DISTRIBUTED_ADVERTISE_ADDR - NatsURL string // --nats-url / LOCALAI_NATS_URL + // NatsURL is accepted and ignored. No component of a distributed + // deployment dials a message bus any more. + NatsURL string // --nats-url / LOCALAI_NATS_URL StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) // RegistrationRequireAuth fails startup when distributed mode is enabled but @@ -156,9 +158,11 @@ func (c DistributedConfig) Validate() error { if !c.Enabled { return nil } - if c.NatsURL == "" { - return fmt.Errorf("distributed mode requires --nats-url / LOCALAI_NATS_URL") - } + // No message-bus URL is required, and none is dialled. The last family + // that needed one was agent..cancel, which now rides the agent + // worker's own tunnel as a control verb; a distributed deployment needs + // PostgreSQL and the frontends' own HTTP listener. The flag is still + // accepted so an existing command line starts unchanged. // S3 credentials must be paired if (c.StorageAccessKey != "" && c.StorageSecretKey == "") || (c.StorageAccessKey == "" && c.StorageSecretKey != "") { diff --git a/core/services/agents/dispatcher.go b/core/services/agents/dispatcher.go index 6f46a0346..82daac4be 100644 --- a/core/services/agents/dispatcher.go +++ b/core/services/agents/dispatcher.go @@ -278,6 +278,32 @@ func (d *WorkerExecutor) Execute(ctx context.Context, raw json.RawMessage, pub m return json.Marshal(map[string]string{"status": status, "error": errMsg}) } +// Cancel answers workerctl.PathAgentCancel on an agent worker. +// +// The reply is this worker's OWN ANSWER and travels as bytes on a 200: +// cancelled true when this process was running the named execution and its +// context has now been cancelled, false when it was not. False says nothing +// about any other worker and nothing about whether the run exists, and the +// frontend that fans the cancel out is the only thing that may assemble those +// answers into a verdict. +// +// A body this worker cannot read is the one thing here that is NOT an answer: +// it has learned nothing about any execution, so it is returned as an error, +// becomes a non-2xx over the tunnel, and is counted by the caller as a cancel +// it could not deliver rather than as one that found nothing. +func (d *WorkerExecutor) Cancel(_ context.Context, raw json.RawMessage) (json.RawMessage, error) { + var req messaging.AgentCancelRequest + if err := json.Unmarshal(raw, &req); err != nil { + return nil, fmt.Errorf("reading an agent cancel request: %w", err) + } + cancelled := d.eventBridge.CancelLocalExecution(req.MessageID) + if cancelled { + xlog.Info("Cancelled an agent execution on this worker after a control cancel", + "agent", req.AgentName, "user", req.UserID, "messageID", req.MessageID) + } + return json.Marshal(messaging.AgentCancelReply{Cancelled: cancelled}) +} + func (d *WorkerExecutor) handleJob(ctx context.Context, evt AgentChatEvent, bridge *EventBridge) (status, errMsg string) { xlog.Info("Processing agent chat job", "agent", evt.AgentName, "user", evt.UserID) @@ -302,9 +328,9 @@ func (d *WorkerExecutor) handleJob(ctx context.Context, evt AgentChatEvent, brid ctx, cancel := context.WithCancel(ctx) defer cancel() - // Register cancellation on the SHARED registry, which the bus-backed cancel - // listener also reads: a cancel arrives on the bus and has to reach an - // execution that is publishing onto a stream. + // Register cancellation on the SHARED registry, which the cancel control + // verb also reads: a cancel arrives on this worker's tunnel and has to + // reach an execution that is publishing onto a control stream. bridge.RegisterCancel(evt.MessageID, cancel) defer bridge.DeregisterCancel(evt.MessageID) diff --git a/core/services/agents/events.go b/core/services/agents/events.go index c6938d40f..9b78f7f55 100644 --- a/core/services/agents/events.go +++ b/core/services/agents/events.go @@ -30,11 +30,24 @@ type AgentEvent struct { Timestamp int64 `json:"timestamp"` // Unix milliseconds (set by PublishEvent) } -// AgentCancelEvent is the broadcast payload for cancelling agent execution. -type AgentCancelEvent struct { - AgentName string `json:"agent_name"` - UserID string `json:"user_id"` - MessageID string `json:"message_id,omitempty"` +// AgentWorkerCanceller carries a cancel to the agent workers of this +// deployment, over the tunnels they hold. +// +// A narrow port rather than the control client itself, because this package +// must not import core/services/nodes, and because the only thing a cancel +// needs from the frontend's control plane is this one verb. +// +// Its error vocabulary is the caller's whole answer and each value means +// something different: nil is a worker's own answer that it cancelled the run, +// nodes.ErrAgentCancelUndelivered is a cancel that may have reached nobody, and +// nodes.ErrAgentRunNotOnAnyWorker is every reachable worker answering that it +// is not running that execution. Nothing here may collapse them. +// +// The real implementation is *nodes.AgentControlClient, asserted where both +// packages are already imported (core/application), so a signature drift is a +// build failure rather than a nil field. +type AgentWorkerCanceller interface { + CancelAgentRun(ctx context.Context, req messaging.AgentCancelRequest) error } // EventBridge bridges agent events between the broadcast carrier and SSE @@ -47,24 +60,20 @@ type EventBridge struct { // to queue. bus messaging.Broadcaster - // cancelBus is where agent..cancel is published and heard, and it is - // a SEPARATE field from bus because in this deployment the two ends of that - // family cannot be on the same carrier yet. + // workers is how a cancel reaches an agent WORKER, and it is a separate + // field from bus because the two travel on different carriers on purpose. // - // Every other family here has both ends on a frontend replica, so both move - // to the broadcast carrier together. This one does not: its only subscriber - // is the agent WORKER (core/cli/agent_worker.go), the cancel has to reach - // the worker actually running the execution, and a worker has no database - // and so cannot join the PostgreSQL carrier at all. Publishing a cancel - // where no worker is listening would report a cancel that reached nobody as - // a cancel the execution declined, which is the one thing this whole - // programme may not do. + // Every other family here has both of its ends on a frontend replica, so + // both live on the broadcast carrier. This one does not: the process that + // holds a worker-run execution's cancel function is the agent worker, and a + // worker has no database and so cannot join the PostgreSQL carrier at all. + // It holds an outward-dialled tunnel instead, and a cancel is an ordinary + // control RPC on it, addressed to the workers a live replica can reach. // - // So it is named, and it is separate, and it stays on the carrier the - // worker reads until a cancel rides the worker's tunnel instead. Folding it - // back into bus is not a simplification: it is silent loss of every cancel - // for every worker-run agent. - cancelBus messaging.Broadcaster + // It is nil on an agent worker, which has nobody to forward a cancel to: + // there, a cancel ARRIVES as that control verb and is applied to the + // registry below. + workers AgentWorkerCanceller // pub is where the events this bridge produces GO, which is not always the // bus. On an agent worker running a dispatched claim it is the response // body of the control RPC the claiming replica is reading, so the events @@ -88,52 +97,69 @@ type EventBridge struct { } // NewEventBridge creates a new EventBridge on the deployment's fan-out carrier. +// A cancel it cannot apply itself is forwarded to the agent workers through +// workers. // -// Cancels go on that same carrier unless WithCancelCarrier says otherwise, -// which is right for the agent worker, where there is only one carrier to be -// on, and wrong for a frontend replica, which reads fan-out from PostgreSQL and -// has to reach workers that cannot. -func NewEventBridge(bus messaging.Broadcaster, store *AgentStore, instanceID string) *EventBridge { +// The canceller is a CONSTRUCTOR PARAMETER and not a builder call, and that is +// the whole reason it is spelled here. As a WithWorkerCanceller line it is one +// statement whose loss compiles, passes every suite in this package, and turns +// every cancel of a worker-run agent into a cancel that reached nobody and +// reported nothing: the local registry has no entry, and there is no longer +// anywhere for the cancel to go. +// +// A nil workers is legitimate on an agent worker, which forwards nothing. See +// NewWorkerEventBridge, which is how a worker builds one. +func NewEventBridge(bus messaging.Broadcaster, store *AgentStore, instanceID string, workers AgentWorkerCanceller) *EventBridge { return &EventBridge{ bus: bus, - cancelBus: bus, pub: bus, + workers: workers, store: store, instanceID: instanceID, cancelRegistry: &messaging.CancelRegistry{}, } } -// WithCancelCarrier puts agent..cancel on carrier instead of on the -// fan-out bus, and returns the receiver so it can be written as one expression -// with the constructor. +// NewWorkerEventBridge returns the bridge an AGENT WORKER runs on. // -// A frontend replica needs it and an agent worker does not. The cancel has to -// reach the worker running the execution; a worker has no database and cannot -// join the PostgreSQL carrier; so a frontend that published its cancels there -// would publish them where no worker listens, and a cancel that reached nobody -// is not a cancel that was refused. +// It joins no carrier, because there is none for it to join: every event it +// produces is written onto the response body of the control RPC that asked for +// the work (see WithPublisher), and every cancel it must act on arrives as a +// control verb rather than as a broadcast. What it keeps is the cancel +// registry, which is the one piece of state a worker's bridge exists for. // -// A nil carrier leaves the bridge on the fan-out bus rather than on nothing, -// because a bridge that publishes cancels nowhere is the failure this exists to -// prevent. -func (b *EventBridge) WithCancelCarrier(carrier messaging.Broadcaster) *EventBridge { - if b == nil || carrier == nil { - return b - } - b.cancelBus = carrier - return b +// The publisher it holds until a verb hands it a stream REFUSES rather than +// drops. A worker that publishes with no stream to write to has produced an +// event that reached nobody, and a no-op default would make that indetectable. +func NewWorkerEventBridge(instanceID string) *EventBridge { + return NewEventBridge(unroutedCarrier{}, nil, instanceID, nil) +} + +// unroutedCarrier is the carrier an agent worker's bridge holds: there is none. +// +// Both methods fail rather than silently succeeding, because both would +// otherwise be undetectable. A dropped publish is an agent event nobody sees; a +// subscription that never delivers is a listener that never fires. +type unroutedCarrier struct{} + +func (unroutedCarrier) Publish(subject string, _ any) error { + return fmt.Errorf("agents: an agent worker tried to publish %q with no control stream to write it to: a worker joins no carrier, so this event would have reached nobody", subject) +} + +func (unroutedCarrier) Subscribe(subject string, _ func([]byte)) (messaging.Subscription, error) { + return nil, fmt.Errorf("agents: an agent worker tried to subscribe to %q: a worker joins no carrier, so nothing would ever be delivered", subject) } // WithPublisher returns a view of this bridge whose events go to pub. // // Everything else is SHARED with the receiver, the cancel registry above all: a -// cancel arriving on the bus must reach an execution that is publishing onto a -// stream, and a bridge that copied the registry would register the cancel where -// nothing looks for it. +// cancel arriving as a control verb must reach an execution that is publishing +// onto a stream, and a bridge that copied the registry would register the +// cancel where nothing looks for it. // // A nil pub returns the receiver unchanged rather than a bridge that publishes -// nowhere, because a handler that was given no writer still has the bus. +// nowhere: on a frontend that leaves the events on the fan-out carrier, and on +// a worker it leaves them on the carrier that refuses loudly. func (b *EventBridge) WithPublisher(pub messaging.Publisher) *EventBridge { if b == nil || pub == nil { return b @@ -238,26 +264,59 @@ func (b *EventBridge) PublishStreamEvent(agentName, userID string, data map[stri }) } -// CancelExecution publishes a cancel event and also checks the local registry. -func (b *EventBridge) CancelExecution(agentName, userID, messageID string) error { - // Try local cancel first +// CancelExecution stops one agent execution wherever in the deployment it is +// running, and reports which of three different things happened. +// +// The three, because a caller that cannot tell them apart is the defect this +// family has been held back for: +// +// - nil means the execution was cancelled. Either it was running in THIS +// process, or an agent worker answered that it had cancelled it. +// - nodes.ErrAgentCancelUndelivered means the cancel may have reached nobody: +// an agent worker that might be running it could not be reached. It is not +// a refusal and it is not a missing task. +// - nodes.ErrAgentRunNotOnAnyWorker means every agent worker this deployment +// could reach answered that it is not running that execution. +// +// The local registry is tried FIRST and short-circuits. A message id names +// exactly one execution, so a local hit is this process's own answer about it, +// and there is nothing a worker could add. +func (b *EventBridge) CancelExecution(ctx context.Context, agentName, userID, messageID string) error { if b.cancelRegistry.Cancel(messageID) { xlog.Info("Cancelled agent execution locally", "agent", agentName, "user", userID, "messageID", messageID) + return nil } - // Broadcast so the replica that actually holds the execution can act on it. - // - // The error says whether the request was PUBLISHED and nothing more. This - // carrier is at-most-once with no replay, so a cancel that reached nobody - // and a cancel an execution declined are different facts that cannot be - // told apart from here, and neither may be reported as the other. - return b.cancelBus.Publish(messaging.SubjectAgentCancel(agentName), AgentCancelEvent{ + if b.workers == nil { + // Not a cancel that was refused and not a task that does not exist: + // this process has nowhere to send the cancel, which is a fact about + // its own wiring and says nothing about the run. + return fmt.Errorf("cancelling agent %q for user %q: this process holds no way to reach an agent worker, so the cancel of message %q was sent nowhere", + agentName, userID, messageID) + } + + return b.workers.CancelAgentRun(ctx, messaging.AgentCancelRequest{ AgentName: agentName, UserID: userID, MessageID: messageID, }) } +// CancelLocalExecution cancels an execution running in THIS process and reports +// whether it found one. +// +// It is what an agent worker's cancel control verb applies, and the bool is +// that worker's whole answer: true is "I cancelled it", false is "I am not +// running it". False is deliberately not an error, because it is not one: a +// deployment fans a cancel out to every worker it can reach and all but one of +// them are expected to say no. +func (b *EventBridge) CancelLocalExecution(messageID string) bool { + if b == nil || messageID == "" { + return false + } + return b.cancelRegistry.Cancel(messageID) +} + // RegisterCancel registers a cancel function for a running agent execution. func (b *EventBridge) RegisterCancel(key string, cancel context.CancelFunc) { b.cancelRegistry.Register(key, cancel) @@ -268,17 +327,6 @@ func (b *EventBridge) DeregisterCancel(key string) { b.cancelRegistry.Deregister(key) } -// StartCancelListener subscribes to the cancel broadcasts every replica sees. -func (b *EventBridge) StartCancelListener() (messaging.Subscription, error) { - return messaging.SubscribeJSON(b.cancelBus, messaging.SubjectAgentCancelWildcard, func(evt AgentCancelEvent) { - if evt.MessageID != "" { - if b.cancelRegistry.Cancel(evt.MessageID) { - xlog.Info("Cancelled an agent execution on this replica after a broadcast cancel", "agent", evt.AgentName, "user", evt.UserID, "messageID", evt.MessageID) - } - } - }) -} - // StartObservablePersister subscribes to every agent's events and persists the // observable_update ones to the database. This runs on the frontend, to capture // observables published by workers, which have no database access. diff --git a/core/services/agents/events_pg_test.go b/core/services/agents/events_pg_test.go index db8176126..b565dfcd4 100644 --- a/core/services/agents/events_pg_test.go +++ b/core/services/agents/events_pg_test.go @@ -5,6 +5,7 @@ package agents import ( "bytes" "context" + "errors" "net/http" "net/http/httptest" "sync" @@ -34,6 +35,7 @@ var _ = Describe("the agent event bridge on the broadcast carrier", func() { busA, busB *pgbus.Bus store *AgentStore bridge *EventBridge + workers *recordingCanceller ) BeforeEach(func() { @@ -53,7 +55,8 @@ var _ = Describe("the agent event bridge on the broadcast carrier", func() { var err error store, err = NewAgentStore(db) Expect(err).ToNot(HaveOccurred()) - bridge = NewEventBridge(busA, store, "replica-a") + workers = &recordingCanceller{} + bridge = NewEventBridge(busA, store, "replica-a", workers) }) // observable is the AgentEvent shape the persister acts on, minus the @@ -132,7 +135,7 @@ var _ = Describe("the agent event bridge on the broadcast carrier", func() { Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = sub.Unsubscribe() }) - peer := NewEventBridge(busB, store, "replica-b") + peer := NewEventBridge(busB, store, "replica-b", workers) Expect(peer.PublishMessage("a1", "u2", "agent", "for another user", "m-other")).To(Succeed()) Expect(peer.PublishMessage("a1", "u1", "agent", "for me", "m-mine")).To(Succeed()) @@ -177,7 +180,7 @@ var _ = Describe("the agent event bridge on the broadcast carrier", func() { // Delivery first, so this is a spec about a stream that WORKED and // then closed, rather than one that never started. - peer := NewEventBridge(busB, store, "replica-b") + peer := NewEventBridge(busB, store, "replica-b", workers) Expect(peer.PublishMessage("a1", "u1", "agent", "hello", "m-1")).To(Succeed()) Eventually(out.String, "20s").Should(ContainSubstring("hello")) @@ -188,63 +191,103 @@ var _ = Describe("the agent event bridge on the broadcast carrier", func() { }) }) - Describe("cancel broadcasts", func() { - It("reaches a peer replica's cancel listener", func() { - peer := NewEventBridge(busB, store, "replica-b") - sub, err := peer.StartCancelListener() - Expect(err).ToNot(HaveOccurred()) - DeferCleanup(func() { _ = sub.Unsubscribe() }) - + // A cancel no longer travels on any carrier, so what these pin is the split + // between the three answers a caller may be given. The CARRIAGE of a cancel + // to a worker is pinned where it happens, over a real tunnel and a real + // yamux session, in core/services/nodes/agent_control_test.go. + Describe("cancelling an execution", func() { + It("cancels a run held by THIS process without disturbing any worker", func() { cancelled := make(chan struct{}) - peer.RegisterCancel("msg-1", func() { close(cancelled) }) + bridge.RegisterCancel("msg-local", func() { close(cancelled) }) - Expect(bridge.CancelExecution("a1", "u1", "msg-1")).To(Succeed()) + Expect(bridge.CancelExecution(ctx, "a1", "u1", "msg-local")).To(Succeed()) Eventually(cancelled, "20s").Should(BeClosed()) + Expect(workers.calls).To(BeZero(), + "a message id names exactly one execution, and this process was running it") + }) + + It("forwards a run it does not hold to the deployment's agent workers", func() { + Expect(bridge.CancelExecution(ctx, "a1", "u1", "msg-remote")).To(Succeed()) + + Expect(workers.calls).To(Equal(1)) + Expect(workers.last).To(Equal(messaging.AgentCancelRequest{ + AgentName: "a1", UserID: "u1", MessageID: "msg-remote", + })) + }) + + // The three answers, kept apart. A caller told any one of these in + // place of another acts on something that did not happen. + It("returns the workers' answer unchanged, whichever of the three it is", func() { + workers.err = errUndeliveredForSpec + Expect(bridge.CancelExecution(ctx, "a1", "u1", "m")).To(MatchError(errUndeliveredForSpec)) + + workers.err = errNotOnAnyWorkerForSpec + Expect(bridge.CancelExecution(ctx, "a1", "u1", "m")).To(MatchError(errNotOnAnyWorkerForSpec)) + }) + + It("refuses rather than reporting success when it has nowhere to send a cancel", func() { + // A bridge built with no canceller. Reporting Succeed here would be + // a cancel that reached nobody presented as one that was made, + // which is the exact failure this family was held back for. + orphan := NewEventBridge(busA, store, "replica-a", nil) + err := orphan.CancelExecution(ctx, "a1", "u1", "m") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("sent nowhere")) + }) + }) + + Describe("the bridge an agent worker runs", func() { + It("applies a cancel handed to it locally and reports whether it found the run", func() { + w := NewWorkerEventBridge("agent-worker-1") + cancelled := make(chan struct{}) + w.RegisterCancel("msg-worker", func() { close(cancelled) }) + + Expect(w.CancelLocalExecution("msg-other")).To(BeFalse(), + "a worker may only ever answer for itself") + Expect(w.CancelLocalExecution("msg-worker")).To(BeTrue()) + Eventually(cancelled, "20s").Should(BeClosed()) }) - // The one family whose two ends are NOT on the same carrier, and the - // spec that says so out loud. - // - // Its only subscriber is the agent WORKER, which has no database and - // therefore cannot join the PostgreSQL carrier. A frontend that - // published its cancels onto the fan-out bus would publish them where - // no worker is listening, every cancel of a worker-run agent would be - // lost, and CancelExecution would return nil throughout: a cancel that - // reached nobody reported as a cancel that was sent, and one step later - // as a cancel the execution declined. - It("publishes a cancel where the worker listens and not onto the fan-out carrier", func() { - // The worker's carrier. Not a second pgbus: the whole point is that - // a worker cannot have one. - workerCarrier := testutil.NewFakeBus() + It("refuses to publish when no control stream has been handed to it", func() { + // A worker joins no carrier. A default that DROPPED the event would + // make an agent run whose events reached nobody look identical to + // one whose events were delivered. + w := NewWorkerEventBridge("agent-worker-1") + err := w.PublishMessage("a1", "u1", "agent", "hello", "m-1") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("reached nobody")) + }) - frontend := NewEventBridge(busA, store, "replica-a").WithCancelCarrier(workerCarrier) - - worker := NewEventBridge(workerCarrier, nil, "agent-worker-1") - workerSub, err := worker.StartCancelListener() - Expect(err).ToNot(HaveOccurred()) - DeferCleanup(func() { _ = workerSub.Unsubscribe() }) - - cancelled := make(chan struct{}) - worker.RegisterCancel("msg-worker", func() { close(cancelled) }) - - // A listener on the FAN-OUT carrier, which must be shown nothing. - // Without this the spec would pass with the cancel published on - // both, which is the shape that hides the loss rather than fixing - // it. - onTheBus := make(chan []byte, 4) - _, err = busB.Subscribe(messaging.SubjectAgentCancelWildcard, func(data []byte) { onTheBus <- data }) - Expect(err).ToNot(HaveOccurred()) - - Expect(frontend.CancelExecution("a1", "u1", "msg-worker")).To(Succeed()) - - Eventually(cancelled, "20s").Should(BeClosed()) - Consistently(onTheBus, "500ms", "50ms").ShouldNot(Receive(), - "a cancel on the fan-out carrier reaches no worker and would be lost") + It("refuses to subscribe, rather than returning a listener that never fires", func() { + w := NewWorkerEventBridge("agent-worker-1") + _, err := w.SubscribeEvents("a1", "u1", func(AgentEvent) {}) + Expect(err).To(HaveOccurred()) }) }) }) +// recordingCanceller stands in for the frontend's agent control client. The +// real one is driven over a real tunnel in core/services/nodes; what this pins +// is that the bridge asks it at all, with the right request, and hands its +// answer back unchanged. +type recordingCanceller struct { + calls int + last messaging.AgentCancelRequest + err error +} + +func (r *recordingCanceller) CancelAgentRun(_ context.Context, req messaging.AgentCancelRequest) error { + r.calls++ + r.last = req + return r.err +} + +var ( + errUndeliveredForSpec = errors.New("this cancel could not be delivered") + errNotOnAnyWorkerForSpec = errors.New("no agent worker is running that execution") +) + // syncBody is an http.ResponseWriter a spec may read WHILE the handler is still // writing. httptest.ResponseRecorder's buffer is not safe for that, and reading // it mid-stream is the only way to know the handler reached its wait rather than diff --git a/core/services/agents/events_test.go b/core/services/agents/events_test.go index 95802bf8b..bc8ac5b7e 100644 --- a/core/services/agents/events_test.go +++ b/core/services/agents/events_test.go @@ -42,7 +42,7 @@ var _ = Describe("EventBridge", func() { // safe-integer range. The timestamp must be in milliseconds. It("emits the timestamp in Unix milliseconds", func() { fake := &recordingBus{} - bridge := NewEventBridge(fake, nil, "instance-1") + bridge := NewEventBridge(fake, nil, "instance-1", nil) before := time.Now().UnixMilli() err := bridge.PublishMessage("agent", "user", "agent", "hello", "msg-1") diff --git a/core/services/agentworker/handlers.go b/core/services/agentworker/handlers.go index b933362df..fa44f9314 100644 --- a/core/services/agentworker/handlers.go +++ b/core/services/agentworker/handlers.go @@ -57,10 +57,6 @@ type StreamHandler func(ctx context.Context, raw json.RawMessage, pub messaging. // as "this worker does not serve that verb" (nodes.ErrWorkerControlUnsupported) // rather than as absence. It is why a later task can add two handlers without // any other file changing shape. -// -// workerctl.PathAgentCancel is named in the path table with no field here. It -// gets the same 404 for the same reason, and the path is fixed now so the two -// sides cannot disagree about it later. type Config struct { // MCPTool answers workerctl.PathMCPToolExecute. MCPTool UnaryHandler @@ -74,6 +70,15 @@ type Config struct { // pick a carrier. BackendStop func(ctx context.Context, req messaging.BackendStopRequest) error + // AgentCancel answers workerctl.PathAgentCancel. It is the LAST family a + // worker needed a message bus for: the cancel used to be a broadcast this + // process had to dial a bus to hear, and is now an ordinary control RPC on + // the tunnel it already holds. + // + // Unary and not streaming: a cancel has one answer and no progress. That + // answer is messaging.AgentCancelReply and it speaks only for this worker. + AgentCancel UnaryHandler + // AgentExecute answers workerctl.PathAgentExecute and MCPCIRun answers // workerctl.PathMCPCIRun. Both are nil in this task and both are set later, // which is where the queue subjects they replace become claim rows. They are @@ -99,6 +104,9 @@ func (c Config) Register(mux *http.ServeMux) { if c.BackendStop != nil { mux.HandleFunc(workerctl.PathBackendStop, serveBackendStop(c.BackendStop)) } + if c.AgentCancel != nil { + mux.HandleFunc(workerctl.PathAgentCancel, serveUnary(workerctl.PathAgentCancel, c.AgentCancel)) + } if c.AgentExecute != nil { mux.HandleFunc(workerctl.PathAgentExecute, serveStream(workerctl.PathAgentExecute, c.AgentExecute)) } diff --git a/core/services/agentworker/server_test.go b/core/services/agentworker/server_test.go index e4f4f0021..9e278710e 100644 --- a/core/services/agentworker/server_test.go +++ b/core/services/agentworker/server_test.go @@ -93,6 +93,7 @@ func fullConfig() agentworker.Config { return agentworker.Config{ MCPTool: echoHandler(`{"result":"tool-ran"}`), MCPDiscovery: echoHandler(`{"servers":[]}`), + AgentCancel: echoHandler(`{"cancelled":true}`), BackendStop: func(context.Context, messaging.BackendStopRequest) error { return nil }, } } @@ -108,6 +109,13 @@ var _ = Describe("The agent worker's control verbs", func() { disc := post(base, workerctl.PathMCPDiscovery, `{"model_name":"m"}`) Expect(disc.StatusCode).To(Equal(http.StatusOK)) Expect(bodyOf(disc)).To(Equal(`{"servers":[]}`)) + + // The cancel is unary and its answer is the worker's own: a 200 body + // saying whether THIS worker cancelled the run. Anything else the + // frontend reads as an answer it did not obtain. + cancel := post(base, workerctl.PathAgentCancel, `{"message_id":"m"}`) + Expect(cancel.StatusCode).To(Equal(http.StatusOK)) + Expect(bodyOf(cancel)).To(Equal(`{"cancelled":true}`)) }) It("answers backend.stop with the 204 a BACKEND worker answers on that same path", func() { @@ -161,7 +169,6 @@ var _ = Describe("The agent worker's control verbs", func() { }, Entry("agent execute, whose handler is nil until it is wired", workerctl.PathAgentExecute), Entry("mcp ci run, whose handler is nil until it is wired", workerctl.PathMCPCIRun), - Entry("agent cancel, which has no handler field at all yet", workerctl.PathAgentCancel), Entry("a backend worker's verb, which an agent worker never serves", workerctl.PathBackendInstall), Entry("a path no build has ever named", workerctl.Prefix+"nothing/here"), ) @@ -178,6 +185,7 @@ var _ = Describe("The agent worker's control verbs", func() { Entry("mcp tool execute", workerctl.PathMCPToolExecute, func(c *agentworker.Config) { c.MCPTool = nil }), Entry("mcp discovery", workerctl.PathMCPDiscovery, func(c *agentworker.Config) { c.MCPDiscovery = nil }), Entry("backend stop", workerctl.PathBackendStop, func(c *agentworker.Config) { c.BackendStop = nil }), + Entry("agent cancel", workerctl.PathAgentCancel, func(c *agentworker.Config) { c.AgentCancel = nil }), ) // A rule stated at every verb has to be pinned at every verb. The exit @@ -213,6 +221,9 @@ var _ = Describe("The agent worker's control verbs", func() { return errors.New("the session cache is wedged") } }), + Entry("agent cancel", workerctl.PathAgentCancel, func(c *agentworker.Config) { + c.AgentCancel = failingHandler(errors.New("the cancel registry is unreadable")) + }), Entry("agent execute, before it has published anything", workerctl.PathAgentExecute, func(c *agentworker.Config) { c.AgentExecute = func(context.Context, json.RawMessage, messaging.Publisher) (json.RawMessage, error) { return nil, errors.New("no agent to run") diff --git a/core/services/messaging/client.go b/core/services/messaging/client.go index 64c9b9968..47a98862e 100644 --- a/core/services/messaging/client.go +++ b/core/services/messaging/client.go @@ -18,22 +18,25 @@ import ( // subscription was rejected (e.g. by JWT permissions) before returning to the caller. const subscribeConfirmTimeout = 5 * time.Second -// Client is a NATS connection, and it is the carrier for exactly one family. +// Client is a NATS connection, and NO PRODUCTION PATH CONSTRUCTS ONE. // -// agent..cancel is the one fan-out family that did not move to the -// PostgreSQL carrier. Its only subscriber is the agent WORKER, which has no -// database and cannot join that carrier at all, so a cancel published there -// would reach no worker while reporting that it was sent. Both ends of that -// family still dial this client, which is why it, its connect options and its -// TLS plumbing are all still here. +// The last family that needed a bus was agent..cancel, whose subscriber is +// the agent WORKER: it has no database, so it could never join the PostgreSQL +// carrier the rest of the deployment fans out on. It does not need a bus either +// now, because it holds an outward tunnel and a cancel is a control verb on it +// (workerctl.PathAgentCancel). Nothing in core/ or pkg/ calls messaging.New. // -// Everything else it used to carry is gone, and so are the methods that carried -// it: queue subscriptions became a claim on the job store, and request/reply -// became a streaming control RPC on the tunnel each worker dials. Deleting the -// METHODS rather than only the call sites is what makes putting a family back -// on this carrier a build error, instead of a line that compiles, publishes -// successfully, and is delivered onto a carrier the deployment is being taken -// off. +// What survives here is this type, its connect options and its TLS plumbing, +// still exercised by the NATS JWT permission specs. Deleting them is a +// demolition of its own, together with the JWT minting at registration and +// pkg/natsauth's permission tables. +// +// The methods that carried everything else are already gone: queue +// subscriptions became a claim on the job store, and request/reply became a +// streaming control RPC on the tunnel each worker dials. Deleting the METHODS +// rather than only the call sites is what makes putting a family back on this +// carrier a build error, instead of a line that compiles, publishes +// successfully, and is delivered onto a carrier nothing reads. type Client struct { conn *nats.Conn mu sync.RWMutex @@ -242,12 +245,12 @@ func (c *Client) confirmSubscription(subject string, mk func(*nats.Conn) (*nats. // different fact from the server refusing it, and neither is evidence about any // node. // -// No production path calls it. The carrier's production users publish cancels -// and act on the delivery, not on the verdict; what needs the verdict is -// pkg/natsauth's permission grants, which are asserted against a real enforcing -// server and would otherwise be asserted against nothing, since an allow list -// that is EMPTY means unrestricted in NATS and a spec that only checks -// IsConnected cannot tell a granted publish from a denied one. +// No production path calls it, and this carrier no longer has production users +// at all. What needs the verdict is pkg/natsauth's permission grants, which are +// asserted against a real enforcing server and would otherwise be asserted +// against nothing, since an allow list that is EMPTY means unrestricted in NATS +// and a spec that only checks IsConnected cannot tell a granted publish from a +// denied one. func (c *Client) ConfirmRoundTrip(timeout time.Duration) error { c.mu.RLock() conn := c.conn diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index d7b35278d..b44bdb168 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -270,6 +270,39 @@ type BackendStopRequest struct { Force bool `json:"force,omitempty"` } +// AgentCancelRequest is the body of an agent cancel control request. +// +// It carries the same three fields the agent..cancel broadcast carried, +// because it says the same thing; what changed is the carrier. The cancel used +// to be published onto a bus every agent worker had to dial, and is now a +// control RPC on the tunnel the worker already holds, which is why an agent +// worker needs no bus credentials. +// +// MessageID is what identifies the execution, and it identifies exactly one: +// an agent run registers its cancel function under it on the process that is +// running it, and nowhere else. +type AgentCancelRequest struct { + AgentName string `json:"agent_name"` + UserID string `json:"user_id"` + MessageID string `json:"message_id,omitempty"` +} + +// AgentCancelReply is an agent worker's OWN ANSWER to a cancel. +// +// Cancelled false is that answer too, and it means one thing only: this worker +// is not running that execution. It is deliberately not spelled as an error, +// and no caller may read it as "the run does not exist" on its own, because a +// worker can only speak for itself. +// +// There is no Error field, and its absence is stated rather than left to be +// inferred. A cancel this worker could not read, or could not serve, is a +// non-2xx like every other verb's failure to serve, which the frontend reads +// as an answer it did not obtain; there is no third thing a worker can learn +// by looking in its own cancel registry. +type AgentCancelReply struct { + Cancelled bool `json:"cancelled"` +} + type ModelStopRequest struct { ModelName string `json:"model_name"` ProcessKey string `json:"process_key"` diff --git a/core/services/nodes/agent_control.go b/core/services/nodes/agent_control.go index 846a5ec96..89c7f7a20 100644 --- a/core/services/nodes/agent_control.go +++ b/core/services/nodes/agent_control.go @@ -8,7 +8,9 @@ import ( "fmt" mcpremote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/workerctl" + "github.com/mudler/xlog" ) // maxAgentPicks bounds how many agent workers one verb is offered to before it @@ -29,6 +31,28 @@ const maxAgentPicks = 3 // and says nothing about any worker. var ErrNoAgentControl = errors.New("nodes: this deployment has no agent control client") +// ErrAgentCancelUndelivered reports that a cancel may have reached nobody: at +// least one agent worker that could be running the execution was not asked, or +// did not answer. +// +// It is the answer this whole family was held back for. A cancel that could not +// be delivered is NOT a cancel that was refused and NOT a run that does not +// exist, and a caller told any of those three in place of another acts on +// something that did not happen. It says nothing about any particular worker +// either, which is why it is its own sentinel and carries neither +// ErrWorkerUnroutable nor anything cluster.IsWorkerAnswer accepts: no node may +// be reaped, demoted or evicted because a cancel went undelivered. +var ErrAgentCancelUndelivered = errors.New("nodes: an agent cancel could not be delivered to every agent worker that might be running it") + +// ErrAgentRunNotOnAnyWorker reports that every agent worker this deployment +// could reach answered that it is not running the named execution. +// +// It is assembled ONLY from workers' own answers, and only when every worker +// was reached; the moment one was not, ErrAgentCancelUndelivered is the answer +// instead. It still does not say the run does not exist: it says no agent +// worker is running it, which is the largest claim the evidence supports. +var ErrAgentRunNotOnAnyWorker = errors.New("nodes: no agent worker of this deployment is running that execution") + // AgentControlClient issues the frontend's control RPCs to whichever agent // worker the selector picks. // @@ -64,6 +88,88 @@ func (a *AgentControlClient) DiscoverMCPTools(ctx context.Context, req mcpremote func(r *mcpremote.MCPDiscoveryResponse) string { return r.Error }) } +// CancelAgentRun asks the agent workers of this deployment to stop one +// execution, and reports which of three different things happened. +// +// A FAN-OUT and not a pick, which is what makes it the one agent verb that does +// not go through agentVerb. A message id names exactly one execution, running +// on exactly one worker, and no row in this deployment records which: the claim +// that dispatched it names the claiming REPLICA, not the worker, and it is +// deleted when the run ends. So the cancel is offered to every worker a live +// replica can reach, exactly as the broadcast it replaces was, and each worker +// answers only for itself. +// +// The three answers, and why none may stand in for another: +// +// - nil. A worker answered that it cancelled the run. That is a worker's own +// answer and it is conclusive, even if another worker could not be reached: +// the execution has been cancelled, and there is only one of it. +// - ErrAgentCancelUndelivered. Some worker that might have been running it +// was not reached. Nothing was learned, and the caller may not report the +// run as missing or the cancel as declined. +// - ErrAgentRunNotOnAnyWorker. Every worker was reached and every one of them +// answered that it is not running that execution. +// +// A worker that is RECONNECTING is counted undelivered, and this is the +// decision rather than an omission: it is not retried here and it is not +// queued. Retrying would hold the caller for the length of the reconnect grace +// with no bound it chose, and queueing would need durable state whose only +// consumer is a run whose control stream died with the tunnel. Reporting it, +// once, as a cancel that may not have arrived is the only thing this frontend +// actually knows, and it leaves the retry where the budget lives: with the +// caller. +func (a *AgentControlClient) CancelAgentRun(ctx context.Context, req messaging.AgentCancelRequest) error { + if a == nil || a.sel == nil || a.cc == nil { + return fmt.Errorf("control rpc %s: %w", workerctl.PathAgentCancel, ErrNoAgentControl) + } + reach, err := a.sel.Reachable(ctx) + if err != nil { + return err + } + + if len(reach.Connected) == 0 && len(reach.Absent) == 0 { + // Nothing was asked of anyone and nothing was learned. This must not + // become ErrAgentRunNotOnAnyWorker, which is assembled from workers' + // own answers: a deployment with no agent worker has produced no + // answers at all, and reporting one would tell a caller that a run it + // can still see is not running anywhere. + return fmt.Errorf("cancelling message %q of agent %q: %w", req.MessageID, req.AgentName, ErrNoAgentWorker) + } + + cancelled := false + // Seeded with the workers nobody could ask at all: a tunnel lost inside the + // reconnect grace, or a presence this replica could not read. They are part + // of the fleet this cancel did not finish asking, and dropping them here is + // what would turn an unfinished fan-out into "no worker is running it". + undelivered := append([]string(nil), reach.Absent...) + for _, nodeID := range reach.Connected { + var reply messaging.AgentCancelReply + if err := a.cc.Call(ctx, nodeID, workerctl.PathAgentCancel, req, &reply); err != nil { + // An unreachable peer, a refused stream, a tunnel that died between + // the connection read and the dial, or a worker too old to serve + // the verb. None of them is an answer about this execution. + xlog.Warn("An agent worker could not be asked to cancel an execution", + "nodeID", nodeID, "agent", req.AgentName, "messageID", req.MessageID, "error", err) + undelivered = append(undelivered, nodeID) + continue + } + if reply.Cancelled { + cancelled = true + } + } + + switch { + case cancelled: + return nil + case len(undelivered) > 0: + return fmt.Errorf("cancelling message %q of agent %q: %d of %d agent workers could not be asked (%v): %w", + req.MessageID, req.AgentName, len(undelivered), len(reach.Connected)+len(reach.Absent), undelivered, ErrAgentCancelUndelivered) + default: + return fmt.Errorf("cancelling message %q of agent %q: all %d reachable agent workers answered that they are not running it: %w", + req.MessageID, req.AgentName, len(reach.Connected), ErrAgentRunNotOnAnyWorker) + } +} + // 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. // diff --git a/core/services/nodes/agent_control_test.go b/core/services/nodes/agent_control_test.go index aaeaa81f1..773460786 100644 --- a/core/services/nodes/agent_control_test.go +++ b/core/services/nodes/agent_control_test.go @@ -11,6 +11,7 @@ import ( "net/http/httptest" "runtime" "sync/atomic" + "time" "github.com/gorilla/websocket" "github.com/libp2p/go-yamux/v5" @@ -21,6 +22,7 @@ import ( "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/messaging" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/testutil" "github.com/mudler/LocalAI/core/services/worker" @@ -39,6 +41,12 @@ import ( const agentControlToken = "agent-control-token" +// agentReconnectGrace is the window these specs read a lost tunnel as +// reconnecting rather than gone. Long enough that a worker registered in this +// spec and never connected is never GONE, which is what makes the undelivered +// answer reachable here. +const agentReconnectGrace = time.Hour + // fakeFrontend is the far side of a worker's tunnel: the real WebSocket // upgrade and the real yamux server handshake. type fakeFrontend struct { @@ -241,7 +249,7 @@ var _ = Describe("AgentControlClient", func() { agentControl := func() *nodes.AgentControlClient { return nodes.NewAgentControlClient( - nodes.NewAgentSelector(registry, clusterReg, selfInstance), control) + nodes.NewAgentSelector(registry, clusterReg, selfInstance, agentReconnectGrace), control) } toolRequest := mcpremote.MCPToolRequest{ @@ -362,6 +370,170 @@ var _ = Describe("AgentControlClient", func() { Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) }) + // The cancel, over the same real transport as everything above it. + // + // A cancel is a FAN-OUT and not a pick: no row in this deployment records + // which worker holds a given execution, so the cancel is offered to every + // worker a live replica can reach and each answers only for itself. What + // these pin is that the three answers stay apart, because a caller told any + // one of them in place of another acts on something that did not happen. + Describe("cancelling one agent run", func() { + // cancelHandler answers the cancel verb the way a worker does: true + // when it holds the named run, false when it does not. + cancelHandler := func(mine string, seen *atomic.Int32) agentworker.UnaryHandler { + return func(_ context.Context, raw json.RawMessage) (json.RawMessage, error) { + seen.Add(1) + var req messaging.AgentCancelRequest + Expect(json.Unmarshal(raw, &req)).To(Succeed()) + return json.Marshal(messaging.AgentCancelReply{Cancelled: req.MessageID == mine}) + } + } + + cancelOf := func(messageID string) messaging.AgentCancelRequest { + return messaging.AgentCancelRequest{AgentName: "a1", UserID: "u1", MessageID: messageID} + } + + It("reaches the worker holding the run over the tunnel THIS replica holds", func() { + var asked atomic.Int32 + startAgent("agent-holder", mine, agentworker.Config{ + AgentCancel: cancelHandler("msg-1", &asked), + }) + + Expect(agentControl().CancelAgentRun(ctx, cancelOf("msg-1"))).To(Succeed()) + Expect(asked.Load()).To(Equal(int32(1))) + }) + + It("reaches a worker whose tunnel a PEER holds, by relaying through that peer", func() { + // The hop the broadcast used to hide. This replica holds nothing, + // so the only way the cancel arrives is the connection row naming + // the peer and the peer's relay splicing the stream onto the tunnel + // it holds. + var asked atomic.Int32 + startAgent("agent-remote", theirs, agentworker.Config{ + AgentCancel: cancelHandler("msg-1", &asked), + }) + Expect(mine.Held()).To(BeEmpty(), "this spec is only about the relayed path") + + Expect(agentControl().CancelAgentRun(ctx, cancelOf("msg-1"))).To(Succeed()) + Expect(asked.Load()).To(Equal(int32(1))) + }) + + It("asks EVERY reachable worker, because nothing records which one holds the run", func() { + // A pick would ask one and, four times in five, report a run that + // is running as one no worker is running. + var first, second atomic.Int32 + startAgent("agent-a", mine, agentworker.Config{AgentCancel: cancelHandler("nothing-here", &first)}) + startAgent("agent-b", theirs, agentworker.Config{AgentCancel: cancelHandler("msg-1", &second)}) + + Expect(agentControl().CancelAgentRun(ctx, cancelOf("msg-1"))).To(Succeed()) + Expect(first.Load()).To(Equal(int32(1))) + Expect(second.Load()).To(Equal(int32(1))) + }) + + It("reports a run no reachable worker holds as exactly that, and never as undelivered", func() { + var asked atomic.Int32 + startAgent("agent-a", mine, agentworker.Config{AgentCancel: cancelHandler("some-other-run", &asked)}) + + err := agentControl().CancelAgentRun(ctx, cancelOf("msg-1")) + Expect(err).To(MatchError(nodes.ErrAgentRunNotOnAnyWorker)) + Expect(err).ToNot(MatchError(nodes.ErrAgentCancelUndelivered)) + Expect(asked.Load()).To(Equal(int32(1))) + }) + + It("reports a worker it could not reach as UNDELIVERED and never as a run that does not exist", func() { + // The refusing worker's tunnel is real and its control server will + // not open, which is what a worker with a dead control plane looks + // like from here. Nothing was learned about the run, so the answer + // may not be "no worker is running it". + var refusals atomic.Int32 + startRefusingAgent("agent-dead", mine, &refusals) + + err := agentControl().CancelAgentRun(ctx, cancelOf("msg-1")) + Expect(err).To(MatchError(nodes.ErrAgentCancelUndelivered)) + Expect(err).ToNot(MatchError(nodes.ErrAgentRunNotOnAnyWorker)) + Expect(refusals.Load()).To(Equal(int32(1))) + }) + + It("counts a RECONNECTING worker as undelivered, which is the decision this task made", func() { + // A registered, approved agent worker with no live tunnel and a + // departure inside the grace. Nobody may act on that condition, so + // the cancel is reported as one that may not have arrived rather + // than retried here, queued, or folded into "no worker holds it". + var asked atomic.Int32 + startAgent("agent-alive", mine, agentworker.Config{AgentCancel: cancelHandler("some-other-run", &asked)}) + registerAgent("agent-reconnecting") + + err := agentControl().CancelAgentRun(ctx, cancelOf("msg-1")) + Expect(err).To(MatchError(nodes.ErrAgentCancelUndelivered)) + Expect(err).ToNot(MatchError(nodes.ErrAgentRunNotOnAnyWorker)) + Expect(asked.Load()).To(Equal(int32(1)), "the worker that WAS reachable must still have been asked") + }) + + It("reports a worker's cancellation even when another worker could not be reached", func() { + // There is one execution and one worker cancelled it. Reporting + // undelivered here would tell the caller nothing happened when + // something did. + var refusals, asked atomic.Int32 + startRefusingAgent("agent-dead", mine, &refusals) + startAgent("agent-holder", theirs, agentworker.Config{AgentCancel: cancelHandler("msg-1", &asked)}) + + Expect(agentControl().CancelAgentRun(ctx, cancelOf("msg-1"))).To(Succeed()) + Expect(refusals.Load()).To(Equal(int32(1))) + Expect(asked.Load()).To(Equal(int32(1))) + }) + + It("reports an empty fleet as neither a route verdict nor a worker answer", func() { + err := agentControl().CancelAgentRun(ctx, cancelOf("msg-1")) + 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() { + var missing *nodes.AgentControlClient + Expect(missing.CancelAgentRun(ctx, cancelOf("msg-1"))).To(MatchError(nodes.ErrNoAgentControl)) + }) + + // The reap guard's two predicates, asserted on the errors a REAL + // fan-out over a real tunnel actually returns. + // + // A cancel is offered to many workers and its answer is assembled from + // all of them, so no answer it produces speaks about any one node. The + // guard that reaps, demotes and evicts decides on exactly these two + // predicates, and a cancel that satisfied either would let one worker + // whose control server happened to be down take a node out of the + // deployment. + // + // The edit that breaks this is small and reads like an improvement: + // wrapping the per-worker Call failure with %w instead of naming the + // nodes with %v. cluster.ErrStreamTargetUnavailable then travels out of + // the fan-out, every spec above still passes, and an undelivered cancel + // has become a routing verdict about a node. + DescribeTable("produces no answer a reap guard may act on, whichever of the three it is", + func(fleet func(), sentinel error) { + fleet() + err := agentControl().CancelAgentRun(ctx, cancelOf("msg-1")) + Expect(err).To(MatchError(sentinel)) + Expect(errors.Is(err, nodes.ErrWorkerUnroutable)).To(BeFalse(), + "a cancel answer was readable as a routing verdict about a worker, which the scheduler may act on") + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse(), + "a cancel answer was readable as a worker's own answer, which a reap guard may act on") + }, + Entry("a worker that could not be asked", func() { + var refusals atomic.Int32 + startRefusingAgent("agent-dead", mine, &refusals) + }, nodes.ErrAgentCancelUndelivered), + Entry("a worker that is reconnecting", func() { + registerAgent("agent-reconnecting") + }, nodes.ErrAgentCancelUndelivered), + Entry("every reachable worker answering that it does not hold the run", func() { + var asked atomic.Int32 + startAgent("agent-a", mine, agentworker.Config{AgentCancel: cancelHandler("some-other-run", &asked)}) + }, nodes.ErrAgentRunNotOnAnyWorker), + Entry("a deployment with no agent worker at all", func() {}, nodes.ErrNoAgentWorker), + ) + }) + 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 diff --git a/core/services/nodes/agent_selector.go b/core/services/nodes/agent_selector.go index 3ac924b9f..ecfca8534 100644 --- a/core/services/nodes/agent_selector.go +++ b/core/services/nodes/agent_selector.go @@ -3,12 +3,16 @@ package nodes import ( + "cmp" "context" "errors" "fmt" "math/rand/v2" + "time" + "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/xlog" ) // ErrNoAgentWorker reports that no agent worker in this deployment currently @@ -34,6 +38,16 @@ var ErrNoAgentWorker = errors.New("nodes: no agent worker holds a tunnel to this // core/services/cluster. type AgentConnectionReader interface { ConnectedAmong(ctx context.Context, nodeIDs []string, owner string) (held []string, heldByOwner []string, err error) + + // Presence is what separates a worker that is GONE from one that is + // RECONNECTING, and a fan-out verb needs that split where a pick does not. + // + // A pick only ever asks who can be reached now. A cancel has to say what + // happened to the workers it could NOT reach, and those are two different + // answers: a departure older than the grace is a routing fact the caller + // may disregard, and a departure inside it is an absent connection nobody + // may act on, including by reporting that the run was not found. + Presence(ctx context.Context, nodeID string, grace time.Duration) (cluster.Presence, error) } // AgentSelector picks an agent worker to send a control RPC to. @@ -54,13 +68,50 @@ type AgentSelector struct { // call would simply take a relay hop with nothing saying so. It is refused // where the selector is built instead. selfInstanceID string + // grace is how long a lost tunnel is read as reconnecting rather than gone. + // It is the operator's trade (DistributedConfig.WorkerReconnectGrace) and + // is only consulted by Reachable; a pick never needs it, because a worker + // that is not connected cannot be picked whatever the reason. + grace time.Duration } // 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} +// grace is the window a departure must outlive before Reachable calls a worker +// gone; a non-positive one takes config.DefaultWorkerReconnectGrace rather than +// zero, because a zero grace would call every worker that lost its tunnel a +// millisecond ago GONE and quietly drop the cancels addressed to it. +func NewAgentSelector(registry *NodeRegistry, conns AgentConnectionReader, selfInstanceID string, grace time.Duration) *AgentSelector { + return &AgentSelector{ + registry: registry, + conns: conns, + selfInstanceID: selfInstanceID, + grace: cmp.Or(grace, config.DefaultWorkerReconnectGrace), + } +} + +// AgentReach is what this deployment can say about its agent workers right now, +// and it is deliberately two lists rather than one plus a count. +// +// Connected are the workers a LIVE replica holds a tunnel to. A verb may be +// issued to each of them, relayed by the peer mesh when the holder is not this +// replica. +// +// Absent are the registered agent workers that no live replica holds AND whose +// departure has not outlived the grace, plus those whose presence could not be +// read at all. Nothing may be asked of them and, far more importantly, nothing +// may be CONCLUDED from their silence: a fan-out that ignored them would report +// "no worker is running that execution" about a fleet it had not finished +// asking. +// +// Workers whose departure IS older than the grace appear in neither list. That +// is the one routing fact this system lets a caller act on, and leaving them in +// Absent would make every deployment that has ever retired an agent worker +// report every cancel as undelivered for ever. +type AgentReach struct { + Connected []string + Absent []string } // PickConnected returns the id AND the node type of an agent node whose tunnel @@ -139,6 +190,84 @@ func (s *AgentSelector) pickConnectedExcluding(ctx context.Context, tried map[st return picked, nodeType, nil } +// Reachable splits this deployment's agent workers into the ones a verb can be +// issued to and the ones whose silence proves nothing. +// +// It exists for a cancel, which is a fan-out and not a pick: the frontend does +// not know which worker holds a given execution, so it asks every worker it can +// reach and assembles their answers. What it must never do is assemble an +// answer out of a fleet it only partly asked, which is why the workers it could +// not reach come back named rather than dropped. +func (s *AgentSelector) Reachable(ctx context.Context) (AgentReach, error) { + if s == nil || s.registry == nil || s.conns == nil { + return AgentReach{}, fmt.Errorf("this deployment has no agent selector: %w", ErrNoAgentWorker) + } + agents, err := s.registry.cancellableAgentNodes(ctx) + if err != nil { + return AgentReach{}, fmt.Errorf("listing the agent workers of this deployment: %w", err) + } + ids := make([]string, 0, len(agents)) + for _, a := range agents { + ids = append(ids, a.ID) + } + held, _, err := s.conns.ConnectedAmong(ctx, ids, s.selfInstanceID) + if err != nil { + return AgentReach{}, fmt.Errorf("reading which agent workers are connected: %w", err) + } + connected := make(map[string]bool, len(held)) + for _, id := range held { + connected[id] = true + } + reach := AgentReach{Connected: held} + for _, id := range ids { + if connected[id] { + continue + } + p, err := s.conns.Presence(ctx, id, s.grace) + if err != nil { + // A read that FAILED is not an answer about this worker, so it + // cannot license the caller to conclude anything. Counted absent, + // which is the value nobody may act on, rather than gone, which is + // the one value that would let a cancel be reported as a run that + // does not exist. + xlog.Warn("Could not read an agent worker's presence; counting it as one this cancel could not reach", + "nodeID", id, "error", err) + reach.Absent = append(reach.Absent, id) + continue + } + if p == cluster.PresenceGone { + continue + } + reach.Absent = append(reach.Absent, id) + } + return reach, nil +} + +// cancellableAgentNodes returns the agent nodes a fan-out verb may be OFFERED +// to, which is not the same set as the ones work may be PLACED on. +// +// It differs from selectableAgentNodes in exactly one status, and the +// difference is the whole reason it exists. A DRAINING worker may take no new +// work, which is the operator's decision, but it is still finishing the runs it +// already holds. A cancel is not new work: it is a request about a run that is +// executing right now, and excluding a draining worker would report the cancel +// of a live execution as a run that no worker is running. +// +// StatusPending stays excluded. An unapproved node is refused by the tunnel +// route on every dial, so it can hold no execution to cancel; counting it would +// make every cancel in a deployment with one unapproved node report undelivered +// for ever. +func (r *NodeRegistry) cancellableAgentNodes(ctx context.Context) ([]BackendNode, error) { + var agents []BackendNode + if err := r.db.WithContext(ctx). + Where("node_type = ? AND status <> ?", NodeTypeAgent, StatusPending). + Order("id"). + Find(&agents).Error; err != nil { + return nil, fmt.Errorf("listing agent nodes for a fan-out verb: %w", err) + } + return agents, 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 diff --git a/core/services/nodes/agent_selector_test.go b/core/services/nodes/agent_selector_test.go index 959762cfc..b1c77323f 100644 --- a/core/services/nodes/agent_selector_test.go +++ b/core/services/nodes/agent_selector_test.go @@ -6,11 +6,13 @@ import ( "context" "errors" "runtime" + "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "gorm.io/gorm" + "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/testutil" ) @@ -32,6 +34,31 @@ type stubConnections struct { owners []string // asked records the candidate list of each call. asked [][]string + + // presence is what Presence answers per node, and presenceErr makes it + // fail. graces records the window each Presence call was made with: a + // selector that passed zero would call every worker that lost its tunnel a + // moment ago GONE, and silently drop the cancels addressed to it. + presence map[string]cluster.Presence + presenceErr error + graces []time.Duration +} + +// testGrace is the reconnect window these specs pass. A real value rather than +// zero, because zero is what NewAgentSelector substitutes a default for, and a +// spec that passed it could not tell the substitution from the value. +const testGrace = 42 * time.Second + +// presenceOf answers what the deployment can say about a node's tunnel. The +// zero value is cluster.PresenceUnknown, which is the value nobody may act on, +// so a spec that forgets to state a node's presence cannot accidentally get the +// one answer (PresenceGone) that licenses a caller to conclude something. +func (s *stubConnections) Presence(_ context.Context, nodeID string, grace time.Duration) (cluster.Presence, error) { + s.graces = append(s.graces, grace) + if s.presenceErr != nil { + return cluster.PresenceUnknown, s.presenceErr + } + return s.presence[nodeID], nil } func (s *stubConnections) ConnectedAmong(_ context.Context, nodeIDs []string, owner string) ([]string, []string, error) { @@ -106,7 +133,7 @@ var _ = Describe("AgentSelector", func() { conns.held = []string{mine, theirs} conns.heldByOwner = []string{mine} - sel := NewAgentSelector(registry, conns, "me") + sel := NewAgentSelector(registry, conns, "me", testGrace) for range picks { id, nodeType, err := sel.PickConnected(ctx) Expect(err).ToNot(HaveOccurred()) @@ -126,7 +153,7 @@ var _ = Describe("AgentSelector", func() { theirs := register("agent-theirs", NodeTypeAgent) conns.held = []string{theirs} - sel := NewAgentSelector(registry, conns, "me") + sel := NewAgentSelector(registry, conns, "me", testGrace) for range picks { id, nodeType, err := sel.PickConnected(ctx) Expect(err).ToNot(HaveOccurred()) @@ -144,7 +171,7 @@ var _ = Describe("AgentSelector", func() { conns.held = []string{a, b} conns.heldByOwner = []string{a, b} - sel := NewAgentSelector(registry, conns, "me") + sel := NewAgentSelector(registry, conns, "me", testGrace) 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. @@ -177,17 +204,17 @@ var _ = Describe("AgentSelector", func() { It("refuses when registered agents exist but none is connected", func() { register("agent-a", NodeTypeAgent) - sel := NewAgentSelector(registry, conns, "me") + sel := NewAgentSelector(registry, conns, "me", testGrace) assertNoWorker(sel.PickConnected(ctx)) }) It("refuses when the deployment has no agent nodes at all", func() { - sel := NewAgentSelector(registry, conns, "me") + sel := NewAgentSelector(registry, conns, "me", testGrace) assertNoWorker(sel.PickConnected(ctx)) }) It("refuses when it was built with no way to read connections", func() { - assertNoWorker(NewAgentSelector(registry, nil, "me").PickConnected(ctx)) + assertNoWorker(NewAgentSelector(registry, nil, "me", testGrace).PickConnected(ctx)) }) }) @@ -201,7 +228,7 @@ var _ = Describe("AgentSelector", func() { conns.held = []string{backend, agent} conns.heldByOwner = []string{backend, agent} - sel := NewAgentSelector(registry, conns, "me") + sel := NewAgentSelector(registry, conns, "me", testGrace) for range picks { id, _, err := sel.PickConnected(ctx) Expect(err).ToNot(HaveOccurred()) @@ -219,7 +246,7 @@ var _ = Describe("AgentSelector", func() { conns.held = []string{pending.ID} conns.heldByOwner = []string{pending.ID} - sel := NewAgentSelector(registry, conns, "me") + sel := NewAgentSelector(registry, conns, "me", testGrace) _, _, err = sel.PickConnected(ctx) Expect(err).To(MatchError(ErrNoAgentWorker)) Expect(conns.asked[0]).To(BeEmpty()) @@ -231,7 +258,7 @@ var _ = Describe("AgentSelector", func() { conns.held = []string{id} conns.heldByOwner = []string{id} - sel := NewAgentSelector(registry, conns, "me") + sel := NewAgentSelector(registry, conns, "me", testGrace) _, _, err := sel.PickConnected(ctx) Expect(err).To(MatchError(ErrNoAgentWorker)) }) @@ -247,17 +274,126 @@ var _ = Describe("AgentSelector", func() { conns.held = []string{id} conns.heldByOwner = []string{id} - sel := NewAgentSelector(registry, conns, "me") + sel := NewAgentSelector(registry, conns, "me", testGrace) picked, _, err := sel.PickConnected(ctx) Expect(err).ToNot(HaveOccurred()) Expect(picked).To(Equal(id)) }) + // Reachable is what a FAN-OUT verb asks, and it differs from a pick in + // exactly one way that matters: it has to say what happened to the workers + // it could not reach. A cancel assembled from a fleet it only partly asked + // reports "no worker is running that execution" about a run that is. + Describe("Reachable", func() { + It("separates a worker that is GONE from one that is merely reconnecting", func() { + // The whole invariant in one spec. A departure older than the grace + // is the one routing fact a caller may act on, so it disappears; a + // departure inside it is an absent connection nobody may act on, so + // it must be named. + live := register("agent-live", NodeTypeAgent) + reconnecting := register("agent-reconnecting", NodeTypeAgent) + gone := register("agent-gone", NodeTypeAgent) + conns.held = []string{live} + conns.presence = map[string]cluster.Presence{ + reconnecting: cluster.PresenceReconnecting, + gone: cluster.PresenceGone, + } + + reach, err := NewAgentSelector(registry, conns, "me", testGrace).Reachable(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(reach.Connected).To(ConsistOf(live)) + Expect(reach.Absent).To(ConsistOf(reconnecting)) + }) + + It("counts a worker that has never dialled as absent, not as gone", func() { + // PresenceUnknown is a worker that has never dialled or whose + // departure aged out of retention, and the registry cannot say + // which. Reading it as gone would license a caller to conclude + // something about a worker it has learned nothing about. + never := register("agent-never", NodeTypeAgent) + + reach, err := NewAgentSelector(registry, conns, "me", testGrace).Reachable(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(reach.Connected).To(BeEmpty()) + Expect(reach.Absent).To(ConsistOf(never)) + }) + + It("counts a presence it could not read as absent, because a failed read is not an answer", func() { + id := register("agent-a", NodeTypeAgent) + conns.presenceErr = errors.New("the database would not answer") + + reach, err := NewAgentSelector(registry, conns, "me", testGrace).Reachable(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(reach.Absent).To(ConsistOf(id)) + }) + + It("measures a departure against the grace it was built with", func() { + // A zero window would call every worker that lost its tunnel a + // moment ago GONE, and the cancels addressed to it would be + // reported as a run nobody is running. + register("agent-a", NodeTypeAgent) + + _, err := NewAgentSelector(registry, conns, "me", testGrace).Reachable(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(conns.graces).To(ConsistOf(testGrace)) + }) + + It("substitutes the deployment default for a grace of zero", func() { + register("agent-a", NodeTypeAgent) + + _, err := NewAgentSelector(registry, conns, "me", 0).Reachable(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(conns.graces).To(ConsistOf(config.DefaultWorkerReconnectGrace)) + }) + + It("excludes a PENDING node, which can hold no execution to cancel", func() { + // An unapproved node is refused by the tunnel route on every dial. + // Counting it would make every cancel in a deployment with one + // unapproved agent node report undelivered for ever. + pending := &BackendNode{Name: "agent-pending", NodeType: NodeTypeAgent, Address: "p:50051"} + Expect(registry.Register(ctx, pending, false)).To(Succeed()) + + reach, err := NewAgentSelector(registry, conns, "me", testGrace).Reachable(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(reach.Connected).To(BeEmpty()) + Expect(reach.Absent).To(BeEmpty()) + }) + + It("still offers a DRAINING worker, which may take no new work but is finishing what it holds", func() { + // The one status where a fan-out and a pick must disagree. A + // draining worker is still running the executions it took, so + // leaving it out would report the cancel of a live run as a run + // that no worker is running. + id := register("agent-draining", NodeTypeAgent) + Expect(registry.MarkDraining(ctx, id)).To(Succeed()) + conns.held = []string{id} + + reach, err := NewAgentSelector(registry, conns, "me", testGrace).Reachable(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(reach.Connected).To(ConsistOf(id)) + + // And the pick still refuses it, which is what the two sets are for. + _, _, err = NewAgentSelector(registry, conns, "me", testGrace).PickConnected(ctx) + Expect(err).To(MatchError(ErrNoAgentWorker)) + }) + + 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") + + _, err := NewAgentSelector(registry, conns, "me", testGrace).Reachable(ctx) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeFalse()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + Expect(errors.Is(err, ErrNoAgentWorker)).To(BeFalse()) + }) + }) + 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") + sel := NewAgentSelector(registry, conns, "me", testGrace) id, nodeType, err := sel.PickConnected(ctx) Expect(err).To(HaveOccurred()) Expect(err.Error()).To(ContainSubstring("the database would not answer")) diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index eed5c0cfd..452979596 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -102,12 +102,11 @@ type NodeHealthStore interface { // ModelLocator is used by RemoteUnloaderAdapter for model discovery. // -// Get is here for one reason: backend.stop is the only control verb whose -// carrier depends on what KIND of worker it is addressed to, because an agent -// node's stop is still published on the bus rather than issued over the tunnel -// it now holds. The callers that come through NodeCommandSender carry a node id -// and nothing else, so the type is read here rather than threaded through every -// one of them. +// Get is here for the node reads that need the whole row rather than a routing +// answer. It is no longer needed to pick a CARRIER for backend.stop: an agent +// worker and a backend worker serve that verb on the same control path over +// their own tunnels, and the frontend states the fact without asking what kind +// of worker it reached (see RemoteUnloaderAdapter.stopBackend). type ModelLocator interface { FindNodesWithModel(ctx context.Context, modelName string) ([]BackendNode, error) RemoveNodeModel(ctx context.Context, nodeID, modelName string, replicaIndex int) error diff --git a/core/services/workerctl/paths.go b/core/services/workerctl/paths.go index fbc49ae12..8803e2478 100644 --- a/core/services/workerctl/paths.go +++ b/core/services/workerctl/paths.go @@ -88,10 +88,9 @@ func BackendPaths() []string { // cached for that backend. The frontend issues the same RPC to either and does // not branch on the node's type to pick a carrier. // -// PathAgentCancel is named here with nothing mounting it yet. The frontend -// therefore gets the catch-all's 404, which it already reads as "this worker -// does not serve that verb" rather than as absence, and the path is fixed now -// so the two sides cannot disagree about it later. +// PathAgentCancel is served by an agent worker and by nothing else. It is what +// replaced the agent..cancel broadcast, and moving that family here is +// what removed the last reason an agent worker dialled a message bus. func AgentPaths() []string { return []string{ PathMCPToolExecute, diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 7369c5a74..5c260f503 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -5,7 +5,7 @@ weight = 71 url = "/features/distributed-mode/" +++ -Distributed mode enables horizontal scaling of LocalAI across multiple machines using **PostgreSQL** for state, node registry and cross-replica fan-out. A **NATS** server is still needed for one thing: delivering an agent cancel to the agent worker running the execution. Unlike the [P2P/federation approach]({{% relref "features/distributed_inferencing" %}}), distributed mode is designed for production deployments and Kubernetes environments where you need centralized management, health monitoring, and deterministic routing. +Distributed mode enables horizontal scaling of LocalAI across multiple machines using **PostgreSQL** for state, node registry and cross-replica fan-out. No message bus is needed: a deployment runs PostgreSQL and the frontends' own HTTP listener, and every worker is reached over the tunnel it dials outward. Unlike the [P2P/federation approach]({{% relref "features/distributed_inferencing" %}}), distributed mode is designed for production deployments and Kubernetes environments where you need centralized management, health monitoring, and deterministic routing. {{% notice note %}} Distributed mode requires authentication enabled with a **PostgreSQL** database - SQLite is not supported. This is because the node registry, job store, and other distributed state are stored in PostgreSQL tables. @@ -41,7 +41,7 @@ Each model gets its own gRPC backend process, so a single worker can serve multi - Each frontend replica holds **one extra PostgreSQL session** beyond its connection pool, pinned for the life of the process, and creates a `bus_messages` table. Both belong to the broadcast carrier that is replacing NATS for cross-replica fan-out; it already carries the four `state.*.delta` families (see [Cross-replica in-memory state](#cross-replica-in-memory-state)). Size `max_connections` for one additional session per frontend replica. - That session reports an `application_name` of `localai_pgbus_`, so `SELECT count(*) FROM pg_stat_activity WHERE application_name LIKE 'localai_pgbus_%'` counts the replicas currently listening. If the carrier loses its session it redials and re-registers on its own; a broadcast published while it was down is not replayed, which is why nothing that must survive a gap is carried by a broadcast alone. - `bus_messages` holds only broadcasts too large for a PostgreSQL notification, and every replica retires rows older than ten minutes. The table is a spill buffer, not a log: it is not a place to read past events from. -- **NATS** server - used for ONE subject family, `agent..cancel`, which is the only thing a frontend still sends over a bus and the only thing an agent worker still subscribes to. Everything else a frontend broadcasts travels on the PostgreSQL the deployment already runs, and every verb a frontend addresses to a worker is an HTTP route on that worker's tunnel. **Serve-backend workers do not connect to it at all**: every verb they take, and file staging with it, is an HTTP route on the worker's tunnel. Set no `LOCALAI_NATS_URL` on a `local-ai worker`. The frontend and any `local-ai agent-worker` still need one. +- **No message bus.** Nothing in a distributed deployment connects to NATS any more. Everything a frontend broadcasts travels on the PostgreSQL the deployment already runs; every verb a frontend addresses to a worker, including an agent cancel, is an HTTP route on that worker's own tunnel. `LOCALAI_NATS_URL` is accepted and ignored everywhere - on the frontend, on `local-ai worker` and on `local-ai agent-worker` - so an existing command line still starts. - All services must be on the same network (or reachable via configured URLs) ## Quick Start with Docker Compose @@ -52,10 +52,10 @@ The easiest way to try distributed mode locally is with the provided Docker Comp docker compose -f docker-compose.distributed.yaml up ``` -This starts PostgreSQL, NATS, a LocalAI frontend, and one worker node. When you send an inference request, the SmartRouter automatically installs the needed backend on the worker and loads the model. See the file for details on adding GPU support, shared volumes, and additional workers. +This starts PostgreSQL, a LocalAI frontend, and one worker node. The compose file still stands a NATS container up; nothing connects to it and you may delete that service. When you send an inference request, the SmartRouter automatically installs the needed backend on the worker and loads the model. See the file for details on adding GPU support, shared volumes, and additional workers. {{% notice tip %}} -Use `docker-compose.distributed.yaml` for quick local testing. For production, deploy PostgreSQL and NATS as managed services and run frontends/workers on separate hosts. +Use `docker-compose.distributed.yaml` for quick local testing. For production, deploy PostgreSQL as a managed service and run frontends/workers on separate hosts. There is no message bus to deploy. {{% /notice %}} ## Frontend Configuration @@ -66,11 +66,11 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These |------|---------|---------|-------------| | `--distributed` | `LOCALAI_DISTRIBUTED` | `false` | Enable distributed mode | | `--instance-id` | `LOCALAI_INSTANCE_ID` | auto UUID | Unique instance ID for this frontend | -| `--nats-url` | `LOCALAI_NATS_URL` | *(required)* | NATS server URL (e.g., `nats://localhost:4222`) | +| `--nats-url` | `LOCALAI_NATS_URL` | *(ignored)* | **Accepted and ignored.** A frontend opens no message-bus connection. Kept so an existing command line still starts. | | `--distributed-advertise-addr` | `LOCALAI_DISTRIBUTED_ADVERTISE_ADDR` | *(derived)* | `host:port` the **other frontend replicas** dial to reach this one. See [Replica peer links](#replica-peer-links). | | `--registration-token` | `LOCALAI_REGISTRATION_TOKEN` | *(empty)* | Token that workers must provide to register | | `--registration-require-auth` | `LOCALAI_REGISTRATION_REQUIRE_AUTH` | `false` | Fail startup when distributed mode is enabled but the registration token is empty (node endpoints and worker file-transfer would otherwise be unauthenticated) | -| `--distributed-require-auth` | `LOCALAI_DISTRIBUTED_REQUIRE_AUTH` | `false` | **Umbrella switch.** Implies both `--nats-require-auth` and `--registration-require-auth` - one knob to lock down the NATS bus *and* the registration/file-transfer layer. Set this in production instead of the two granular flags. | +| `--distributed-require-auth` | `LOCALAI_DISTRIBUTED_REQUIRE_AUTH` | `false` | **Umbrella switch.** Implies `--registration-require-auth`, which is what guards registration, the worker control planes and file transfer. It also implies the inert `--nats-require-auth`. Set this in production instead of the granular flags. | | `--auto-approve-nodes` | `LOCALAI_AUTO_APPROVE_NODES` | `false` | Auto-approve new worker nodes (skip admin approval) | | `--distributed-shared-models` | `LOCALAI_DISTRIBUTED_SHARED_MODELS` | `false` | Assert that every node mounts the **same** models directory at the **same** path (a shared volume). When `true`, the router skips file staging entirely and workers load models directly from the shared path instead of re-downloading them. See [Shared models directory](#shared-models-directory). | | `--distributed-disk-headroom-check` | `LOCALAI_DISTRIBUTED_DISK_HEADROOM_CHECK` | `true` | Reject worker nodes that lack free space to store the model, at scheduling time rather than partway through staging. When `false`, node selection ignores free disk; the check still runs and warns when it would have rejected every node. Also toggleable at runtime via the `distributed_disk_headroom_check` setting. See [Disk headroom](#disk-headroom). | @@ -185,10 +185,9 @@ The same carrier moves the traffic whose subscriber is an open HTTP response rat | Job result | `jobs..result` | The result persister on every replica | | Job cancel | `jobs..cancel` | Every replica, so the one holding the run can stop it | | Agent events | `agent..events.` | `GET /api/agents/{name}/sse/distributed` on any replica, and the observable persister | -| Agent cancel | `agent..cancel` | Every agent worker, and so still on NATS: see below | | Open Responses cancel | `responses..cancel` | The replica holding the generation | -One family on that list is not on this carrier. An `agent..cancel` has to reach the agent WORKER running the execution, and an agent worker has no database, so it cannot listen on PostgreSQL at all; that cancel is published on NATS, where the worker is listening, and it stays there until a cancel rides the worker's tunnel like every other verb the frontend addresses to a worker. Everything else in the table is on the PostgreSQL carrier. +An agent cancel is deliberately NOT on that list. It has to reach the agent WORKER running the execution, and an agent worker has no database, so it can never listen on PostgreSQL; it is a control verb on that worker's own tunnel instead. See [Cancelling an agent run](#cancelling-an-agent-run). This is what lets a user watch a job or an agent on one frontend while the work runs against another. **No broadcast on this list is the only path to anything durable.** A job's terminal state is written to its row by the replica that claimed the work, before that claim is released, so a dropped result costs an open stream its promptness and never costs the job its answer: a stream that is still open re-reads the row and closes on it. A cancel is a request and not a verdict: if it reaches nobody it has not been refused, and nothing in the API reports it as such. @@ -262,7 +261,7 @@ A worker that presents a credential belonging to no node, or names a node ID the **A node that has not registered since upgrading cannot tunnel.** Its row has no tunnel credential and the column cannot be back-filled, because the plaintext only ever existed in the response that minted it. Such a node is refused with `401` until it registers again, which a worker restart does. The frontend does *not* fall back to the registration token for these nodes. -Unlike the agent worker's API key and its NATS credential, a tunnel credential **is** issued to a node still awaiting approval. It is inert until then: the tunnel route re-reads the node's status on every dial and refuses a pending one. Withholding it would instead strand workers that register exactly once, since approval on its own prompts no re-registration. +Unlike the agent worker's API key, a tunnel credential **is** issued to a node still awaiting approval. It is inert until then: the tunnel route re-reads the node's status on every dial and refuses a pending one. Withholding it would instead strand workers that register exactly once, since approval on its own prompts no re-registration. A tunnel credential does not replace `LOCALAI_REGISTRATION_TOKEN`. Without one, node registration itself is unauthenticated, so anyone who can reach the frontend can register a worker and be issued a tunnel credential for it. How far that gets them depends on auto-approve: with auto-approve on the node is healthy at once and the credential works immediately; with it off the node is pending and the credential is inert until an admin approves, so approval is the real gate. LocalAI warns about the missing token at startup. @@ -270,7 +269,21 @@ 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`, and one thing only still reaches it on the bus.** Every verb the frontend addresses to a specific agent worker is now a control RPC on the tunnel that worker holds: MCP tool execution, MCP discovery, the backend stop that flushes cached MCP sessions, agent execution, and MCP CI runs. The progress and result lines the worker asks the frontend to re-publish on its behalf travel back on that same response body, and the frontend re-publishes them on the PostgreSQL carrier, not on NATS. One subject in the other direction is still NATS and is the reason `--nats-url` is still required: `agent..cancel`, which the worker subscribes to so a cancel can reach the execution it is running. +**An agent worker no longer needs `--nats-url`, and connects to no message bus at all.** Every verb the frontend addresses to a specific agent worker is a control RPC on the tunnel that worker holds: MCP tool execution, MCP discovery, the backend stop that flushes cached MCP sessions, agent execution, MCP CI runs, and now the cancel. The progress and result lines the worker asks the frontend to re-publish on its behalf travel back on that same response body, and the frontend re-publishes them on the PostgreSQL carrier. + +#### Cancelling an agent run + +A cancel names one execution by its message id, and nothing in the deployment records which worker holds it: the claim that dispatched the run names the claiming *replica*, and it is deleted when the run ends. So the frontend offers the cancel to **every agent worker a live replica can reach**, over each worker's own tunnel (`POST /v1/control/agent/cancel`), and each worker answers only for itself. + +A caller gets one of three answers, and they are deliberately different facts: + +| Outcome | What it means | +|---------|---------------| +| success | A worker answered that it cancelled the run, or the run was held by the replica the request landed on. | +| *could not be delivered* | At least one agent worker that might have been running it was not reached: its tunnel was lost inside the [reconnect grace](#a-lost-tunnel-is-a-departure-not-an-absence), its control plane refused the stream, or a peer holding it was unreachable. Nothing was learned. It is **not** a refusal and **not** a missing run. | +| *no agent worker is running that execution* | Every agent worker was reached and every one of them answered that it does not hold the run. | + +A worker that is **reconnecting** always produces the second answer. The cancel is not retried inside the request and not queued: retrying would hold the caller for the length of the reconnect grace, and the retry belongs with whoever owns the budget. Re-issue the cancel once the worker is connected again. There is no `nodes..*` subject left, and an agent worker's minted JWT no longer grants `mcp.tools.execute`, `mcp.discovery`, `nodes..backend.stop`, `agent.execute` or `jobs.mcp-ci.new`. The `--agent-subject` and `--agent-queue` flags (`LOCALAI_AGENT_SUBJECT`, `LOCALAI_AGENT_QUEUE`) are gone: there is no subject for an agent worker to subscribe to and no queue group to be one of. @@ -465,7 +478,7 @@ Registering against an upgraded frontend **clears** a node's `address` and `http A worker on this release opens **no inbound listener on a routable interface**. Its backend gRPC processes and its HTTP file-transfer server all bind loopback, and the frontend reaches both through the tunnel. Concretely: -- **No inbound firewall rule, published port, Service or Ingress is needed for a worker.** A serve-backend worker needs outbound access to the frontend URL (`LOCALAI_REGISTER_TO`), and nothing else - not even to NATS. An agent worker binds only loopback too, and needs outbound access to both `LOCALAI_REGISTER_TO` (registration, heartbeats and now its tunnel) and `LOCALAI_NATS_URL`. +- **No inbound firewall rule, published port, Service or Ingress is needed for a worker.** A serve-backend worker needs outbound access to the frontend URL (`LOCALAI_REGISTER_TO`), and nothing else - not even to NATS. An agent worker binds only loopback too, and needs outbound access to `LOCALAI_REGISTER_TO` and nothing else either: registration, heartbeats, its tunnel and every verb the frontend addresses to it all go there. - **`LOCALAI_ADVERTISE_ADDR` and `LOCALAI_ADVERTISE_HTTP_ADDR` are gone.** There is nothing to advertise. Both are ignored if still set; remove them. - **`LOCALAI_ADDR` and `LOCALAI_SERVE_ADDR` are read for their port only.** The port is the base of the backend port range, and `port-1` is the HTTP file-transfer port. The host half names an interface nothing binds. - The node's `address` and `http_address` fields in `GET /api/nodes` are empty, and are cleared for nodes that reported them before the upgrade. @@ -581,50 +594,25 @@ The chat UI renders this state inline and retries automatically once the model r A frontend replica that dies mid-load does not wedge the model: the job row carries a heartbeat and another replica reclaims a job whose heartbeat has stopped. The heartbeat is time-based, not byte-based, because a checkpoint load legitimately transfers zero bytes for many minutes. {{% /notice %}} -### NATS JWT authentication (recommended for production) +### NATS credentials (inert) -**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, which covers the one subject family still on the bus, and a tunnel token plus the same `LOCALAI_REGISTRATION_TOKEN` bearer check in front of its control server. +**No LocalAI component connects to NATS.** The frontend's cross-replica fan-out is on PostgreSQL, a serve-backend worker takes every verb on its own tunnel, and an agent worker now does too, including the cancel that was the last family on a bus. -By default, NATS connections are anonymous: any client that can reach port `4222` may publish the one subject family still carried on it, `agent..cancel`. Anyone who can reach an unauthenticated bus can therefore cancel any running agent. Nothing else is on it: the agent-worker job subjects became rows in a claim table, the frontend's cross-replica events travel on PostgreSQL, and `nodes..backend.install` and its nine siblings 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. +Every `LOCALAI_NATS_*` setting is therefore accepted and inert, so an existing command line, unit file or Helm values file starts unchanged: -| Flag | Env Var | Description | -|------|---------|-------------| -| `--nats-account-seed` | `LOCALAI_NATS_ACCOUNT_SEED` | Account signing seed (`SU...`). The frontend mints a per-node user JWT at registration (`nats_jwt` in the register response). | -| `--nats-service-jwt` | `LOCALAI_NATS_SERVICE_JWT` | User JWT for the frontend (and optional fallback for agent workers). The frontend publishes one subject family with it: `agent..cancel`. | -| `--nats-service-seed` | `LOCALAI_NATS_SERVICE_SEED` | User signing seed (`SU...`) paired with the service JWT. | -| `--nats-worker-jwt-ttl` | `LOCALAI_NATS_WORKER_JWT_TTL` | Lifetime of minted worker JWTs (default `24h`). | -| `--nats-require-auth` | `LOCALAI_NATS_REQUIRE_AUTH` | Fail startup if JWT credentials are missing when distributed mode is enabled. | +| Flag | Env Var | Status | +|------|---------|--------| +| `--nats-url` | `LOCALAI_NATS_URL` | Accepted and ignored on the frontend, `local-ai worker` and `local-ai agent-worker`. | +| `--nats-account-seed` | `LOCALAI_NATS_ACCOUNT_SEED` | The frontend still mints a per-node user JWT at registration (`nats_jwt` in the register response). Nothing consumes it. | +| `--nats-service-jwt` / `--nats-service-seed` | `LOCALAI_NATS_SERVICE_JWT` / `LOCALAI_NATS_SERVICE_SEED` | Accepted, unused: the frontend opens no bus connection to present them on. | +| `--nats-worker-jwt-ttl` | `LOCALAI_NATS_WORKER_JWT_TTL` | Lifetime of the minted-but-unused worker JWTs. | +| `--nats-require-auth` | `LOCALAI_NATS_REQUIRE_AUTH` | On an agent worker this still makes registration **wait through admin approval** rather than starting against a pending node. It no longer gates any bus connection. | +| `--nats-tls-ca` / `--nats-tls-cert` / `--nats-tls-key` | `LOCALAI_NATS_TLS_*` | Accepted, unused. | -### NATS TLS / mTLS (optional) - -Use `tls://` in `--nats-url` / `LOCALAI_NATS_URL` for encrypted transport. When the server uses a private CA or requires client certificates, set: - -| Flag | Env Var | Description | -|------|---------|-------------| -| `--nats-tls-ca` | `LOCALAI_NATS_TLS_CA` | PEM file to verify the NATS server (private CA) | -| `--nats-tls-cert` | `LOCALAI_NATS_TLS_CERT` | Client certificate for NATS mTLS | -| `--nats-tls-key` | `LOCALAI_NATS_TLS_KEY` | Client private key (required with `--nats-tls-cert`) | - -The same env vars apply to backend workers and `local-ai agent-worker`. If the server cert is already trusted by the OS, `tls://` alone is enough. - -**Worker register response** (when minting is enabled and the node is approved): - -```json -{ - "id": "…", - "nats_jwt": "eyJ…", - "nats_user_seed": "SU…" -} -``` - -Agent workers connect with that JWT and seed automatically (shown once; store securely). Override with `LOCALAI_NATS_JWT` / `LOCALAI_NATS_USER_SEED` if needed. Set `LOCALAI_NATS_REQUIRE_AUTH=true` on an agent worker when the bus requires credentials; `local-ai worker` has no such flag, because it opens no connection to require credentials for. A JWT is still minted for a serve-backend node at registration and is simply unused; it grants nothing but that connection's own reply inbox. - -When `LOCALAI_NATS_REQUIRE_AUTH=true` and no static credentials are provided, an agent worker that registers while still **pending admin approval** keeps re-registering (with backoff) until an admin approves it and the frontend mints its JWT - it does not start unauthenticated. This retry is **bounded**: if the node is never approved (or no credentials are minted) after a large number of attempts, the worker exits non-zero so the failure is visible (a crash-looping or failed worker) rather than hanging silently. Minted worker JWTs are also **refreshed automatically** before they expire (the worker re-registers at ~75% of the JWT lifetime), so long-running workers survive past `LOCALAI_NATS_WORKER_JWT_TTL`; the NATS connection picks up the new JWT on its next reconnect. If refresh fails persistently, the worker exits (to restart and re-acquire) rather than drifting toward an expired, unrenewable JWT. Statically configured (`LOCALAI_NATS_JWT`) and service (`LOCALAI_NATS_SERVICE_JWT`) credentials are used as-is and not refreshed. - -Generate operator/account material with [`scripts/nats-auth-setup.sh`](https://github.com/mudler/LocalAI/blob/master/scripts/nats-auth-setup.sh) (requires [nsc](https://docs.nats.io/running-a-nats-service/configuration/securing_nats/auth_intro/nsc)). Configure the NATS server with account resolver JWTs before enabling `LOCALAI_NATS_REQUIRE_AUTH`. +You may stop running a NATS server, and remove these settings at your convenience. {{% notice note %}} -`LOCALAI_AUTH` (HTTP users/sessions) and NATS JWTs are separate: end-user API keys do not connect to NATS. HTTP registration still uses `LOCALAI_REGISTRATION_TOKEN`. +`LOCALAI_AUTH` (HTTP users/sessions) is unrelated. HTTP registration still uses `LOCALAI_REGISTRATION_TOKEN`, and every worker control plane sits behind that same bearer check. {{% /notice %}} ### Optional: S3 Object Storage @@ -674,7 +662,7 @@ during installation as well as the committed snapshot. {{% /notice %}} {{% notice warning %}} -The worker HTTP file transfer server is authenticated by `LOCALAI_REGISTRATION_TOKEN`. If the token is **empty**, the server **fails open** - anyone who can reach the port gets read/write access to the worker's models/staging/data directories (a remote model-poisoning / exfiltration vector), **and to the `/v1/control/` routes that install, upgrade and delete backends and stop the node**. The worker logs a loud warning at startup in this case. Always set `LOCALAI_REGISTRATION_TOKEN` in distributed mode, and set `LOCALAI_DISTRIBUTED_REQUIRE_AUTH=true` (frontend **and** workers) to make a missing token a hard startup error rather than a silent fail-open. On the frontend and on agent workers it also makes missing NATS credentials fatal; on a serve-backend worker it means the registration token alone, since that worker uses no bus credential. +The worker HTTP file transfer server is authenticated by `LOCALAI_REGISTRATION_TOKEN`. If the token is **empty**, the server **fails open** - anyone who can reach the port gets read/write access to the worker's models/staging/data directories (a remote model-poisoning / exfiltration vector), **and to the `/v1/control/` routes that install, upgrade and delete backends and stop the node**. The worker logs a loud warning at startup in this case. Always set `LOCALAI_REGISTRATION_TOKEN` in distributed mode, and set `LOCALAI_DISTRIBUTED_REQUIRE_AUTH=true` (frontend **and** workers) to make a missing token a hard startup error rather than a silent fail-open. On an agent worker it additionally makes registration wait through admin approval instead of starting against a pending node. By default the server binds loopback, so "anyone who can reach the port" means a process on the worker host, and no firewall rule is required. Setting `LOCALAI_HTTP_ADDR` to a routable address opts back out of that and puts the fail-open case back on the network - if you do it, firewall the port. {{% /notice %}} @@ -1078,12 +1066,11 @@ Agent workers are dedicated processes for executing agent chats and MCP CI jobs. ```bash local-ai agent-worker \ --register-to http://frontend:8080 \ - --nats-url nats://nats:4222 \ --registration-token changeme ``` Agent workers: -- Execute agent chat messages dispatched via NATS +- Execute agent chat messages dispatched to it as streaming control verbs on its tunnel - Run MCP CI jobs (with access to MCP servers via docker) - 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 @@ -1300,7 +1287,7 @@ local-ai worker \ --registration-token changeme ``` -**Multiple frontend replicas:** Run multiple LocalAI frontends behind a load balancer. Since all state is in PostgreSQL and coordination is via NATS, frontends are fully stateless and interchangeable. +**Multiple frontend replicas:** Run multiple LocalAI frontends behind a load balancer. Since all state is in PostgreSQL and coordination is via PostgreSQL and the workers' own tunnels, frontends are fully stateless and interchangeable. ## Model Scheduling @@ -1503,12 +1490,12 @@ Notes: |---|---|---| | **Discovery** | Automatic via libp2p token | Self-registration to frontend URL | | **State storage** | In-memory / ledger | PostgreSQL | -| **Coordination** | Gossip protocol | The worker's own tunnel for serve-backend work; PostgreSQL `LISTEN`/`NOTIFY` for cross-replica frontend events; NATS for agent workers | +| **Coordination** | Gossip protocol | Each worker's own tunnel for every verb addressed to it, agent workers included; PostgreSQL `LISTEN`/`NOTIFY` for cross-replica frontend events | | **Node management** | Automatic | REST API + WebUI | | **Health monitoring** | Peer heartbeats | Centralized HealthMonitor | | **Backend management** | Manual per node | Dynamic via the worker's `backend.install` control route | | **Best for** | Ad-hoc clusters, community sharing | Production, Kubernetes, managed infrastructure | -| **Setup complexity** | Minimal (share a token) | Requires PostgreSQL on the frontend, plus NATS if you run agent workers. Serve-backend workers need neither: only an outbound route to the frontend URL. | +| **Setup complexity** | Minimal (share a token) | Requires PostgreSQL on the frontend, and nothing else. Workers of either kind need only an outbound route to the frontend URL. | ## Troubleshooting @@ -1518,9 +1505,7 @@ Notes: - Ensure auth is enabled on the frontend (`LOCALAI_AUTH=true`) **NATS connection errors:** -- These concern the **frontend** and **agent workers** only. A `local-ai worker` opens no NATS connection; if one is failing to join, look at its tunnel and its `--register-to` instead. -- Confirm NATS is running and reachable (`nats-server --signal ldm` or check port 4222) -- Check that `--nats-url` uses the correct hostname/IP from that component's network perspective +- Nothing in LocalAI connects to NATS any more, on any component. If a release you are running still logs one, it predates the tunnel migration; on this release, look at the failing component's tunnel and its `--register-to` instead. **PostgreSQL connection errors:** - Verify the connection URL format: `postgresql://user:password@host:5432/dbname?sslmode=disable` @@ -1555,7 +1540,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 CI), 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, and no component opens a bus connection to produce it. **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/docs/content/reference/cli-reference.md b/docs/content/reference/cli-reference.md index ce68705f8..15b456bd5 100644 --- a/docs/content/reference/cli-reference.md +++ b/docs/content/reference/cli-reference.md @@ -211,7 +211,7 @@ LocalAI supports several subcommands beyond `run`: - `local-ai mcp-server` - Run the LocalAI admin tool surface as a stdio MCP server (controls a remote LocalAI instance over HTTP) - `local-ai worker` - Start a worker for distributed mode (generic, backend-agnostic; needs only an outbound route to the frontend, no message bus) - `local-ai p2p-worker` - Run workers to distribute workload via p2p (llama.cpp-only) -- `local-ai agent-worker` - Start an agent worker for distributed mode (serves agent execution and MCP CI runs as control verbs on the tunnel it holds; still needs `--nats-url`, for cancellation and nothing else: an agent cancel has to reach the worker running the execution, and a worker has no database and so cannot read the PostgreSQL carrier every other broadcast travels on) +- `local-ai agent-worker` - Start an agent worker for distributed mode (serves agent execution, MCP CI runs and cancellation as control verbs on the tunnel it holds; it connects to no message bus, so `--nats-url` is accepted and ignored) - `local-ai util` - Utility commands - `local-ai explorer` - Run P2P explorer - `local-ai federated` - Run LocalAI in federated mode diff --git a/tests/e2e/distributed/agent_distributed_test.go b/tests/e2e/distributed/agent_distributed_test.go index dfb90ea8d..d044eb89a 100644 --- a/tests/e2e/distributed/agent_distributed_test.go +++ b/tests/e2e/distributed/agent_distributed_test.go @@ -2,9 +2,18 @@ package distributed_test import ( "context" - "sync/atomic" + "net" + "net/http" + "net/http/httptest" + "time" + + "github.com/gorilla/websocket" + "github.com/libp2p/go-yamux/v5" "github.com/mudler/LocalAI/core/services/agents" + "github.com/mudler/LocalAI/core/services/agentworker" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -111,8 +120,8 @@ var _ = Describe("Phase 3: Agent Conversations & SSE", Label("Distributed"), fun // Two carriers, because the whole point of this bridge is that the // user's SSE connection and the replica running the agent are not // the same process. - watcher := agents.NewEventBridge(infra.Bus(), store, "instance-1") - runner := agents.NewEventBridge(infra.Bus(), store, "instance-2") + watcher := agents.NewEventBridge(infra.Bus(), store, "instance-1", nil) + runner := agents.NewEventBridge(infra.Bus(), store, "instance-2", nil) received := make(chan agents.AgentEvent, 16) sub, err := watcher.SubscribeEvents("my-agent", "user1", func(evt agents.AgentEvent) { @@ -141,35 +150,66 @@ var _ = Describe("Phase 3: Agent Conversations & SSE", Label("Distributed"), fun // Conversation persistence removed — chat history is browser-only. - // Two FRONTEND replicas, which share a carrier. A cancel bound for an - // agent WORKER does not travel this way: a worker has no database, so - // it cannot listen on the PostgreSQL carrier and the frontend publishes - // its cancels where the worker is listening instead. That pairing is - // pinned in core/services/agents; this pins the replica-to-replica half. - It("cancels a running agent from another replica", func() { - bridge := agents.NewEventBridge(infra.Bus(), store, "instance-1") - canceller := agents.NewEventBridge(infra.Bus(), store, "instance-2") + // The whole cancel path, end to end and with nothing doubled: a real + // agent worker holding a real WebSocket + yamux tunnel, a real + // connection row deciding which replica owns it, the real selection + // over the real node rows, and the real control client on top. + // + // A cancel does not travel on a carrier any more. Its far end is the + // agent WORKER, which has no database and so cannot join the carrier + // the rest of the deployment fans out on, and it holds an outward + // tunnel instead. This is what that is. + It("cancels an agent run on a real worker over the tunnel it holds", func() { + const replica = "instance-1" - // Start cancel listener - cancelSub, err := bridge.StartCancelListener() + registry, err := nodes.NewNodeRegistry(db) Expect(err).ToNot(HaveOccurred()) - defer cancelSub.Unsubscribe() + clusterReg := cluster.NewRegistry(db) + Expect(clusterReg.Register(infra.Ctx, replica, "10.0.0.1:8080", "v1")).To(Succeed()) + tunnels := cluster.NewTunnelRegistry(clusterReg, replica) - // Register a cancellable context - _, cancel := context.WithCancel(infra.Ctx) - var cancelled atomic.Bool - wrappedCancel := context.CancelFunc(func() { - cancelled.Store(true) - cancel() + // The worker's own bridge, and the run registered on it. This is + // the state a dispatched agent execution leaves on a worker. + workerBridge := agents.NewWorkerEventBridge("agent-worker-e2e") + executor := agents.NewWorkerExecutor(workerBridge, nil, "http://127.0.0.1:1", "token") + cancelled := make(chan struct{}) + workerBridge.RegisterCancel("msg-e2e", func() { close(cancelled) }) + + node := &nodes.BackendNode{Name: "agent-e2e", NodeType: nodes.NodeTypeAgent, Address: "agent-e2e:50051"} + Expect(registry.Register(infra.Ctx, node, true)).To(Succeed()) + registered, err := registry.GetByName(infra.Ctx, "agent-e2e") + Expect(err).ToNot(HaveOccurred()) + + frontend := newTunnelFrontend() + rt, err := agentworker.Start(infra.Ctx, agentworker.Options{ + FrontendURL: frontend.URL(), + NodeID: registered.ID, + TunnelToken: func() string { return "tunnel-secret" }, + ControlToken: "control-token", + Handlers: agentworker.Config{AgentCancel: executor.Cancel}, }) - bridge.RegisterCancel("test-msg-id", wrappedCancel) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = rt.Close() }) + _, err = tunnels.Attach(infra.Ctx, registered.ID, frontend.Session()) + Expect(err).ToNot(HaveOccurred()) - // Issued on the replica that does NOT hold the execution, which is - // the case the broadcast exists for. Its error says only that the - // request was published, so what is asserted is the effect. - Expect(canceller.CancelExecution("my-agent", "user1", "test-msg-id")).To(Succeed()) + control := nodes.NewControlClient(nodes.WorkerNetDialerFor(func(nodeID string) func(context.Context, string, string) (net.Conn, error) { + return cluster.NewWorkerDialer(tunnels, nil).DialerFor(nodeID, cluster.StreamTagHTTP) + }), "control-token") + agentControl := nodes.NewAgentControlClient( + nodes.NewAgentSelector(registry, clusterReg, replica, time.Hour), control) - Eventually(func() bool { return cancelled.Load() }, "20s").Should(BeTrue()) + // Issued through the frontend's own bridge, which is what a cancel + // request landing on a replica reaches. + bridge := agents.NewEventBridge(infra.Bus(), store, replica, agentControl) + Expect(bridge.CancelExecution(infra.Ctx, "my-agent", "user1", "msg-e2e")).To(Succeed()) + Eventually(cancelled, "20s").Should(BeClosed()) + + // And the second of the three answers, from the same live fleet: + // a run no worker holds is NOT reported as cancelled. + err = bridge.CancelExecution(infra.Ctx, "my-agent", "user1", "msg-nobody-holds") + Expect(err).To(MatchError(nodes.ErrAgentRunNotOnAnyWorker)) + Expect(err).ToNot(MatchError(nodes.ErrAgentCancelUndelivered)) }) // Agent execution is now dispatched via AgentPoolService.dispatchChat(), @@ -206,3 +246,53 @@ var _ = Describe("Phase 3: Agent Conversations & SSE", Label("Distributed"), fun }) }) }) + +// tunnelFrontend is the far side of a worker's tunnel: the real WebSocket +// upgrade and the real yamux server handshake, with no LocalAI frontend behind +// it. It is what lets these specs put a REAL agent worker on a REAL tunnel +// without starting a whole server. +type tunnelFrontend struct { + srv *httptest.Server + sessions chan *yamux.Session +} + +func newTunnelFrontend() *tunnelFrontend { + GinkgoHelper() + f := &tunnelFrontend{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 +} + +func (f *tunnelFrontend) URL() string { return f.srv.URL } + +// 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 *tunnelFrontend) Session() *yamux.Session { + GinkgoHelper() + var sess *yamux.Session + Eventually(f.sessions, "20s").Should(Receive(&sess)) + DeferCleanup(func() { _ = sess.Close() }) + return sess +} diff --git a/tests/e2e/distributed/agent_native_executor_test.go b/tests/e2e/distributed/agent_native_executor_test.go index 1d147bcef..1df3b32ad 100644 --- a/tests/e2e/distributed/agent_native_executor_test.go +++ b/tests/e2e/distributed/agent_native_executor_test.go @@ -316,7 +316,7 @@ var _ = Describe("Native Agent Executor", Label("Distributed", "AgentNative"), f Context("WorkerExecutor", func() { It("should dispatch chat via NATS and receive response", func() { - bridge := agents.NewEventBridge(infra.NC, nil, "test-instance") + bridge := agents.NewEventBridge(infra.NC, nil, "test-instance", nil) configs := &mockConfigProvider{configs: map[string]*agents.AgentConfig{ "test-agent": { @@ -378,28 +378,24 @@ var _ = Describe("Native Agent Executor", Label("Distributed", "AgentNative"), f }) It("should handle cancellation via EventBridge", func() { - bridge := agents.NewEventBridge(infra.NC, nil, "cancel-test") + // The local half: a run held by THIS process is cancelled without + // anything leaving it. The cancel of a run held by an agent WORKER + // travels over that worker's tunnel and is driven end to end, on a + // real one, in agent_distributed_test.go. + bridge := agents.NewEventBridge(infra.NC, nil, "cancel-test", nil) var cancelled atomic.Bool bridge.RegisterCancel("test-msg-id", func() { cancelled.Store(true) }) - // Start cancel listener - cancelSub, err := bridge.StartCancelListener() - Expect(err).ToNot(HaveOccurred()) - defer cancelSub.Unsubscribe() - - FlushNATS(infra.NC) - - // Cancel the execution - Expect(bridge.CancelExecution("test-agent", "user1", "test-msg-id")).To(Succeed()) + Expect(bridge.CancelExecution(infra.Ctx, "test-agent", "user1", "test-msg-id")).To(Succeed()) Eventually(func() bool { return cancelled.Load() }, "5s").Should(BeTrue()) }) It("should execute agent chat from enriched payload without ConfigProvider", func() { - bridge := agents.NewEventBridge(infra.NC, nil, "enriched-test") + bridge := agents.NewEventBridge(infra.NC, nil, "enriched-test", nil) // Executor with NO ConfigProvider (simulating DB-free worker) executor := agents.NewWorkerExecutor(bridge, nil, "http://localhost:8080", "test-key") @@ -576,7 +572,7 @@ var _ = Describe("Native Agent Executor", Label("Distributed", "AgentNative"), f Context("Full Distributed Chat Flow", func() { It("should dispatch chat via NATS, execute, and publish response via EventBridge", func() { - bridge := agents.NewEventBridge(infra.NC, nil, "flow-test") + bridge := agents.NewEventBridge(infra.NC, nil, "flow-test", nil) // Store agent config in PostgreSQL cfg := agents.AgentConfig{ @@ -676,7 +672,7 @@ var _ = Describe("Native Agent Executor", Label("Distributed", "AgentNative"), f }) It("should dispatch background run via NATS with system role", func() { - bridge := agents.NewEventBridge(infra.NC, nil, "bg-test") + bridge := agents.NewEventBridge(infra.NC, nil, "bg-test", nil) cfg := agents.AgentConfig{ Name: "bg-agent", @@ -996,7 +992,7 @@ var _ = Describe("Native Agent Executor", Label("Distributed", "AgentNative"), f llmURL, llmShutdown := startAgentMockLLMServer("I found the answer using my skills.") defer llmShutdown() - bridge := agents.NewEventBridge(infra.NC, nil, "full-e2e-test") + bridge := agents.NewEventBridge(infra.NC, nil, "full-e2e-test", nil) // Subscribe to agent events var receivedEvents []agents.AgentEvent @@ -1041,32 +1037,38 @@ var _ = Describe("Native Agent Executor", Label("Distributed", "AgentNative"), f // it travels on the response body rather than only as a publish. Expect(string(reply)).To(ContainSubstring(`"status":"completed"`)) - // Wait for the full execution: processing + agent response + completed - Eventually(func() int { + // Waited for by CONTENT and not by count. A run with streaming on + // publishes an unbounded number of stream_event lines before its + // terminal status, so "at least three events have arrived" is + // reached long before the completed status is, and the assertions + // below then read a snapshot that cannot contain it yet. That is a + // race in the spec rather than in the executor, and it fails about + // one run in five. + seen := func() (agentMessage, completed bool) { eventMu.Lock() defer eventMu.Unlock() - return len(receivedEvents) - }, "15s").Should(BeNumerically(">=", 3)) - - eventMu.Lock() - defer eventMu.Unlock() - - var hasAgentMessage, hasCompleted bool - for _, evt := range receivedEvents { - if evt.EventType == "json_message" && evt.Sender == "agent" { - hasAgentMessage = true - Expect(evt.Content).To(ContainSubstring("found the answer")) - } - if evt.EventType == "json_message_status" && evt.Metadata != "" { - var meta map[string]string - json.Unmarshal([]byte(evt.Metadata), &meta) - if meta["status"] == "completed" { - hasCompleted = true + for _, evt := range receivedEvents { + if evt.EventType == "json_message" && evt.Sender == "agent" { + agentMessage = true + Expect(evt.Content).To(ContainSubstring("found the answer")) + } + if evt.EventType == "json_message_status" && evt.Metadata != "" { + var meta map[string]string + // A metadata blob this spec cannot read is not a + // completed status, and saying so beats failing the + // whole run on one malformed event. + if err := json.Unmarshal([]byte(evt.Metadata), &meta); err == nil && meta["status"] == "completed" { + completed = true + } } } + return agentMessage, completed } - Expect(hasAgentMessage).To(BeTrue(), "should receive agent response message via EventBridge") - Expect(hasCompleted).To(BeTrue(), "should receive completed status via EventBridge") + + Eventually(func() bool { _, completed := seen(); return completed }, "15s"). + Should(BeTrue(), "should receive completed status via EventBridge") + agentMessage, _ := seen() + Expect(agentMessage).To(BeTrue(), "should receive agent response message via EventBridge") }) It("should execute background agent run via NATS dispatcher with mock LLM", func() { @@ -1074,7 +1076,7 @@ var _ = Describe("Native Agent Executor", Label("Distributed", "AgentNative"), f llmURL, llmShutdown := startAgentMockLLMServer("All systems operational. No issues detected.") defer llmShutdown() - bridge := agents.NewEventBridge(infra.NC, nil, "bg-e2e-test") + bridge := agents.NewEventBridge(infra.NC, nil, "bg-e2e-test", nil) // Subscribe to agent events var receivedEvents []agents.AgentEvent diff --git a/tests/e2e/distributed/foundation_test.go b/tests/e2e/distributed/foundation_test.go index 62c2a4b83..c96bd71e5 100644 --- a/tests/e2e/distributed/foundation_test.go +++ b/tests/e2e/distributed/foundation_test.go @@ -7,8 +7,6 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/advisorylock" - "github.com/mudler/LocalAI/core/services/agents" - "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/storage" . "github.com/onsi/ginkgo/v2" @@ -83,16 +81,6 @@ var _ = Describe("Phase 0: Foundation", Label("Distributed"), func() { }) }) - // The one carrier a deployment still dials besides PostgreSQL, and the one - // family left on it. - // - // agent..cancel did not move to the broadcast carrier: its only - // subscriber is the agent WORKER, which has no database and cannot join - // that carrier at all. So this round trip is no longer "the messaging layer - // works" - it is the cancel path for every agent a worker runs, and if it - // stops working the symptom is a cancel that is published, succeeds, and - // reaches nobody. - // // Two Its that used to sit here went with the halves of the client they // exercised. "should support queue subscriptions for load balancing" pinned // that work reaches exactly one of N consumers; a queue group no longer @@ -103,38 +91,12 @@ var _ = Describe("Phase 0: Foundation", Label("Distributed"), func() { // asserted nothing of the kind; the property is now // core/services/pgbus/listener_test.go, which actually kills the session // with pg_terminate_backend and waits for delivery to resume. - Context("the cancel carrier", func() { - It("connects, publishes and subscribes, which is the agent cancel path", func() { - client, err := messaging.New(infra.NatsURL) - Expect(err).ToNot(HaveOccurred()) - defer client.Close() - - Expect(client.IsConnected()).To(BeTrue()) - - // The real subject and the real filter, not a placeholder pair. - // A worker subscribes to the wildcard and a frontend publishes to - // one agent's subject, so a round trip on "test.subject" would - // stay green through a filter that no longer matches what the - // builder mints - which is the failure this family actually has. - received := make(chan []byte, 1) - sub, err := client.Subscribe(messaging.SubjectAgentCancelWildcard, func(data []byte) { - received <- data - }) - Expect(err).ToNot(HaveOccurred()) - defer sub.Unsubscribe() - - // Small delay to ensure subscription is active - FlushNATS(client) - - err = client.Publish(messaging.SubjectAgentCancel("a1"), agents.AgentCancelEvent{ - AgentName: "a1", UserID: "u1", MessageID: "msg-1", - }) - Expect(err).ToNot(HaveOccurred()) - - Eventually(received, "5s").Should(Receive()) - }) - - }) + // + // The third went with the family it carried. A cancel used to be published + // on NATS because its subscriber was an agent worker that could not read + // the broadcast carrier; it is a control verb on that worker's own tunnel + // now, and the deployment dials no message bus at all. The path is driven + // end to end, over a real tunnel, in agent_distributed_test.go. Context("ObjectStore filesystem adapter", func() { var store *storage.FilesystemStore diff --git a/tests/e2e/distributed/sse_routes_test.go b/tests/e2e/distributed/sse_routes_test.go index 5f93b9673..2e7b64a9a 100644 --- a/tests/e2e/distributed/sse_routes_test.go +++ b/tests/e2e/distributed/sse_routes_test.go @@ -86,8 +86,8 @@ var _ = Describe("SSE Routes", Label("Distributed"), func() { agentStore, err := agents.NewAgentStore(db) Expect(err).ToNot(HaveOccurred()) - frontend0 := agents.NewEventBridge(infra.Bus(), agentStore, "frontend-0") - frontend1 := agents.NewEventBridge(infra.Bus(), agentStore, "frontend-1") + frontend0 := agents.NewEventBridge(infra.Bus(), agentStore, "frontend-0", nil) + frontend1 := agents.NewEventBridge(infra.Bus(), agentStore, "frontend-1", nil) received := make(chan agents.AgentEvent, 16) sub, err := frontend1.SubscribeEvents("test-agent", "user1", func(evt agents.AgentEvent) {