refactor(distributed): move job and agent fan-out onto the PostgreSQL carrier

Five of the six families whose subscriber is an open HTTP response rather
than a process-lifetime cache now travel on pgbus: jobs.<id>.progress,
jobs.<id>.result, jobs.<id>.cancel, agent.<name>.events.<user> and
responses.<id>.cancel. Both ends of each move together, so there is no
state where a publisher is on one carrier and its subscriber on the other.

agent.<name>.cancel does NOT move, and the plan was wrong about why. Its
only subscriber in the tree is the agent worker, which has no database and
so cannot join the PostgreSQL carrier at all. Publishing that cancel on
pgbus would have lost every cancel of a worker-run agent while returning
nil, which reports a cancel that reached nobody as a cancel that was sent.
EventBridge now names its cancel carrier separately, a frontend replica
sets it to the carrier the worker reads, and it stays there until a cancel
rides the worker's tunnel like every other verb addressed to a worker.

The carrier drops at 256 rather than blocking, which is not safe on its own
for a result: a lost result has no successor message. It is not the only
path. The claiming replica persists the terminal line before it releases
the claim, and an open progress stream re-reads the job row once after
subscribing and then periodically, so a dropped terminal broadcast costs
promptness and never the answer.

Both per-request subscriptions close in a defer instead of on one return
path, and pgbus grows Subscribers() so the leak they would otherwise cause
can be asserted. It has no other symptom: only the first subscriber of a
channel issues a LISTEN, so a leaked filter just adds one closure per
notification for every stream the replica has ever served. Subscribe now
issues its LISTEN before it registers, which makes that count a readiness
signal rather than a figure to compare against itself.

Two rules that were stated at several sites and pinned at none are now one
each. The re-broadcaster is built beside the dispatcher and the bridge and
handed to the dispatch loop, so no line is left that can point it at a
carrier nobody subscribes to while every spec stays green. The set of
statuses a job never leaves is one exported set that the SSE bridge and the
store both read. The last hand-written subject filter in production code
became messaging.SubjectAgentEventsWildcard.

Assisted-by: Claude Opus 5 [claude-code]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto committed 2026-09-04 01:47:34 +00:00
1 parent 245010f2f6
commit 5a95bb3c0a
30 files changed
+1901 -277

No files matched your search

+14 -4
View File
@@ -8,7 +8,6 @@ import (
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/services/jobs"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/nodes"
"gorm.io/gorm"
)
@@ -35,19 +34,30 @@ import (
// they reload. That is a whole feature lost to a nil field, with no error
// anywhere, so it is refused here.
//
// The re-broadcaster is taken already built, from newFanoutBridges, and is a
// *nodes.Rebroadcaster rather than the jobs.ProgressBroadcaster interface the
// loop stores it as. Both of those are deliberate. Taking it built leaves ONE
// expression in the tree that decides which carrier job and agent fan-out goes
// on, next to the dispatcher and the bridge that must read the same one, so
// there is no separate line here to point at a carrier nobody subscribes to:
// that mis-wiring publishes successfully, returns true, reddens no spec in any
// package, and shows up only as an SSE stream with no progress in it. Naming
// the concrete type is what makes the refusal below fire, too: widened to the
// interface, a nil re-broadcaster is a non-nil value holding a nil pointer.
//
// 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.
func startJobDispatchLoop(ctx context.Context, cfg config.DistributedConfig, db *gorm.DB, store *jobs.JobStore,
registry *nodes.NodeRegistry, conns nodes.AgentConnectionReader,
control *nodes.ControlClient, bus messaging.Broadcaster) (*jobs.DispatchLoop, error) {
control *nodes.ControlClient, broadcast *nodes.Rebroadcaster) (*jobs.DispatchLoop, error) {
if cfg.InstanceID == "" {
return nil, fmt.Errorf("the job dispatch loop was built with no instance id: its claims could not be told from ones a dead replica left")
}
if registry == nil || conns == nil {
return nil, fmt.Errorf("the job dispatch loop was built with no way to find a connected agent worker")
}
if bus == nil {
if broadcast == nil {
return nil, fmt.Errorf("the job dispatch loop was built with no broadcaster: every job would run with its progress and its result reaching no SSE stream in the deployment")
}
loop, err := jobs.NewDispatchLoop(jobs.DispatchConfig{
@@ -57,7 +67,7 @@ func startJobDispatchLoop(ctx context.Context, cfg config.DistributedConfig, db
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.
Broadcast: nodes.NewRebroadcaster(bus),
Broadcast: broadcast,
Store: store,
})
if err != nil {
+27 -3
View File
@@ -26,6 +26,7 @@ var _ = Describe("building the job dispatch loop", func() {
var registry *nodes.NodeRegistry
var conns *recordingConnections
var ctx context.Context
var broadcast *nodes.Rebroadcaster
BeforeEach(func() {
if runtime.GOOS == "darwin" {
@@ -36,6 +37,14 @@ var _ = Describe("building the job dispatch loop", func() {
registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB())
Expect(err).ToNot(HaveOccurred())
conns = newRecordingConnections()
// A double is enough HERE, and only here. Which carrier this
// re-broadcaster publishes on is not this function's decision any more:
// it is handed one already built by newFanoutBridges, and that is where
// the carrier is pinned, by receipt on a second connection. What is
// left for these to say is that the loop refuses to be built without
// one and starts when it is.
broadcast = nodes.NewRebroadcaster(testutil.NewFakeBus())
})
// The silent one. A loop with no broadcaster dispatches work perfectly
@@ -48,10 +57,25 @@ var _ = Describe("building the job dispatch loop", func() {
Expect(err.Error()).To(ContainSubstring("broadcaster"))
})
// The nil that the interface would have hidden. The loop stores its
// re-broadcaster as the jobs.ProgressBroadcaster interface, and widened to
// that here a nil *nodes.Rebroadcaster is a NON-nil value holding a nil
// pointer, so the refusal above would never fire for the way one is
// actually absent: newFanoutBridges returns a typed nil alongside its
// error. This drives that exact value, which is why the parameter is the
// concrete type.
It("refuses a typed-nil broadcaster, which an interface parameter would have accepted", func() {
var absent *nodes.Rebroadcaster
_, err := startJobDispatchLoop(ctx, config.DistributedConfig{InstanceID: "replica-7"},
testutil.SetupTestDB(), nil, registry, conns, nodes.NewControlClient(nil, "token"), absent)
Expect(err).To(HaveOccurred())
Expect(err.Error()).To(ContainSubstring("broadcaster"))
})
It("refuses to build with no instance id", func() {
_, err := startJobDispatchLoop(ctx, config.DistributedConfig{},
testutil.SetupTestDB(), nil, registry, conns, nodes.NewControlClient(nil, "token"),
testutil.NewFakeBus())
broadcast)
Expect(err).To(HaveOccurred())
Expect(err.Error()).To(ContainSubstring("instance id"))
})
@@ -59,7 +83,7 @@ var _ = Describe("building the job dispatch loop", func() {
It("refuses to build with nothing to read connections through", func() {
_, err := startJobDispatchLoop(ctx, config.DistributedConfig{InstanceID: "replica-7"},
testutil.SetupTestDB(), nil, registry, nil, nodes.NewControlClient(nil, "token"),
testutil.NewFakeBus())
broadcast)
Expect(err).To(HaveOccurred())
Expect(err.Error()).To(ContainSubstring("connected agent worker"))
})
@@ -81,7 +105,7 @@ var _ = Describe("building the job dispatch loop", func() {
Expect(err).ToNot(HaveOccurred())
loop, err := startJobDispatchLoop(ctx, config.DistributedConfig{InstanceID: "replica-7"},
db, nil, registry, conns, nodes.NewControlClient(nil, "token"), testutil.NewFakeBus())
db, nil, registry, conns, nodes.NewControlClient(nil, "token"), broadcast)
Expect(err).ToNot(HaveOccurred())
DeferCleanup(loop.Stop)
+7 -13
View File
@@ -406,9 +406,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
}
xlog.Info("Distributed job store initialized")
// Initialize job dispatcher
dispatcher := jobs.NewDispatcher(jobStore, natsClient, authDB, cfg.Distributed.InstanceID)
// Initialize agent store
agentStore, err := agents.NewAgentStore(authDB)
if err != nil {
@@ -416,15 +413,12 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
}
xlog.Info("Distributed agent store initialized")
// Initialize agent event bridge
agentBridge := agents.NewEventBridge(natsClient, agentStore, cfg.Distributed.InstanceID)
// Start observable persister — captures observable_update events from workers
// (which have no DB access) and persists them to PostgreSQL.
if err := agentBridge.StartObservablePersister(); err != nil {
xlog.Warn("Failed to start observable persister", "error", err)
} else {
xlog.Info("Observable persister started")
// 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)
@@ -464,7 +458,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
// 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, natsClient)
jobDispatch, err := startJobDispatchLoop(cfg.Context, cfg.Distributed, authDB, jobStore, registry, clusterRegistry, controlClient, rebroadcast)
if err != nil {
return nil, fmt.Errorf("wiring the job dispatch loop: %w", err)
}
+84
View File
@@ -0,0 +1,84 @@
// SPDX-License-Identifier: MIT
package application
import (
"fmt"
"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"
"gorm.io/gorm"
)
// newFanoutBridges builds the three surfaces whose cross-replica traffic is job
// and agent fan-out, on the deployment's ONE broadcast carrier.
//
// One function rather than two constructor calls in the start-up path, and the
// parameter type is the reason. jobs.NewDispatcher and agents.NewEventBridge
// both take a messaging.Broadcaster, which they must: neither may know which
// carrier a deployment runs, and their specs publish through a double. But that
// also means *messaging.Client satisfies them, so wiring either of them to the
// NATS client instead of the carrier COMPILES, passes every unit spec in both
// packages, and presents only as an SSE stream that stays empty while the work
// it is watching runs to completion on the other side of a carrier nobody is
// subscribed to. Naming *pgbus.Bus here is what makes that a build failure
// rather than a silent one, and it is why this exists as a function instead of
// as two lines and a comment asking the reader to be careful.
//
// The observable persister is STARTED here for the reason startJobDispatchLoop
// starts its loop: a bridge that is built and never subscribed captures nothing
// a worker publishes, and the only symptom is observables that quietly stop
// being written on a deployment that has workers.
//
// The re-broadcaster is built HERE and handed to the dispatch loop rather than
// built there from a carrier of its own, and that is the whole reason this
// function returns three things. It is the surface a reviewer skips: it is what
// turns an agent worker's progress line into a broadcast, so pointing it at a
// carrier nobody subscribes to leaves every unit spec in every package passing.
// The re-broadcaster publishes, the publish succeeds, Handle returns true, and
// the only symptom in the deployment is an SSE stream with no progress in it.
// 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.<name>.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,
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
// messaging.Broadcaster a nil *pgbus.Bus is a non-nil interface holding a
// nil pointer, so every guard downstream reads it as a carrier that is
// present and every publish through it panics on a request instead.
if bus == nil {
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")
}
dispatcher := jobs.NewDispatcher(jobStore, bus, db, instanceID)
bridge := agents.NewEventBridge(bus, agentStore, instanceID).WithCancelCarrier(cancelCarrier)
// Warned rather than refused, and deliberately: the persister needs a store
// and a deployment without one still serves live SSE correctly. What it
// loses is the durable copy of a worker's observables, which is degraded
// rather than broken.
if err := bridge.StartObservablePersister(); err != nil {
xlog.Warn("Failed to start observable persister", "error", err)
} else {
xlog.Info("Observable persister started")
}
return dispatcher, bridge, nodes.NewRebroadcaster(bus), nil
}
+169
View File
@@ -0,0 +1,169 @@
// SPDX-License-Identifier: MIT
package application
import (
"context"
"encoding/json"
"runtime"
"strings"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"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/LocalAI/core/services/testutil"
"gorm.io/gorm"
)
// The three fan-out surfaces, asserted from the OTHER replica's carrier.
//
// Everything here publishes on a second Bus and asserts on the effect the
// surface built on the first one had. One bus talking to itself would pass with
// the surfaces wired to any carrier at all, which is exactly the wiring defect
// these exist to catch: a dispatcher on one carrier and a publisher on another
// leaves every unit spec in both packages green and only an SSE stream empty.
var _ = Describe("wiring the job and agent fan-out bridges", func() {
var (
ctx context.Context
db *gorm.DB
busA, busB *pgbus.Bus
jobStore *jobs.JobStore
agentStore *agents.AgentStore
)
BeforeEach(func() {
if runtime.GOOS == "darwin" {
Skip("testcontainers requires Docker, not available on macOS CI")
}
ctx = context.Background()
var dsn string
db, dsn = testutil.SetupTestDBWithDSN()
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
newBus := func() *pgbus.Bus {
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
Expect(err).ToNot(HaveOccurred())
DeferCleanup(b.Close)
return b
}
busA, busB = newBus(), newBus()
var err error
jobStore, err = jobs.NewJobStore(db)
Expect(err).ToNot(HaveOccurred())
agentStore, err = agents.NewAgentStore(db)
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() {
_, _, _, err := newFanoutBridges(busA, nil, jobStore, agentStore, db, "replica-1")
Expect(err).To(HaveOccurred())
Expect(err.Error()).To(ContainSubstring("agent cancels"))
})
It("refuses to build with no carrier", func() {
_, _, _, err := newFanoutBridges(nil, testutil.NewFakeBus(), jobStore, agentStore, db, "replica-1")
Expect(err).To(HaveOccurred())
Expect(err.Error()).To(ContainSubstring("broadcast carrier"))
})
// S1. The dispatcher persists a terminal result broadcast by a peer, which
// 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")
Expect(err).ToNot(HaveOccurred())
Expect(dispatcher.Start(ctx)).To(Succeed())
DeferCleanup(dispatcher.Stop)
job := &jobs.JobRecord{TaskID: "t1", UserID: "u1", Status: "running", TriggeredBy: "manual"}
Expect(jobStore.CreateJob(job)).To(Succeed())
Expect(busB.Publish(messaging.SubjectJobResult(job.ID), jobs.JobResultEvent{
JobID: job.ID, Status: "completed", Result: "the answer",
})).To(Succeed())
Eventually(func() string {
stored, err := jobStore.GetJob(job.ID)
if err != nil {
return ""
}
return stored.Status
}, "20s").Should(Equal("completed"))
})
// S2. The observable persister writes what a peer broadcast, which it can
// 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")
Expect(err).ToNot(HaveOccurred())
Expect(bridge).ToNot(BeNil())
Expect(busB.Publish(messaging.SubjectAgentEvents("a1", "u1"), agents.AgentEvent{
AgentName: "a1",
UserID: "u1",
EventType: "observable_update",
EventSubType: "tool_result",
SourceInstance: "replica-2",
MessageID: "obs-1",
Metadata: `{"tool":"grep"}`,
})).To(Succeed())
Eventually(func() int {
records, err := agentStore.GetObservables(agents.AgentKey("u1", "a1"), 10)
if err != nil {
return 0
}
return len(records)
}, "20s").Should(Equal(1))
})
// S3, and it is the one this whole arrangement is for. The re-broadcaster
// is what turns an agent worker's progress line into a broadcast, and it is
// the surface a reviewer skips, because pointing it at a carrier nobody
// reads leaves every unit spec in every package green: it publishes, the
// publish succeeds, and Handle returns true.
//
// So these assert the RECEIPT on a peer's carrier and never Handle's return
// 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")
Expect(err).ToNot(HaveOccurred())
delivered := make(chan []byte, 4)
_, err = busB.Subscribe(subject, func(data []byte) { delivered <- data })
Expect(err).ToNot(HaveOccurred())
rebroadcast.Handle(nodes.NodeTypeAgent, strings.ReplaceAll(subject, "*", "j1"), json.RawMessage(payload))
Eventually(delivered, "20s").Should(Receive(MatchJSON(payload)))
},
Entry("a job's progress", messaging.SubjectJobProgressWildcard, `{"job_id":"j1","status":"running"}`),
Entry("a job's result", messaging.SubjectJobResultWildcard, `{"job_id":"j1","status":"completed"}`),
)
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")
Expect(err).ToNot(HaveOccurred())
delivered := make(chan []byte, 4)
_, err = busB.Subscribe(messaging.SubjectAgentEventsWildcard, func(data []byte) { delivered <- data })
Expect(err).ToNot(HaveOccurred())
rebroadcast.Handle(nodes.NodeTypeAgent, messaging.SubjectAgentEvents("a1", "u1"),
json.RawMessage(`{"event_type":"json_message"}`))
Eventually(delivered, "20s").Should(Receive(MatchJSON(`{"event_type":"json_message"}`)))
})
})
@@ -0,0 +1,114 @@
// SPDX-License-Identifier: MIT
package openresponses
import (
"context"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"gorm.io/gorm"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/pgbus"
"github.com/mudler/LocalAI/core/services/testutil"
)
// The same two-replica topology, on the carrier a deployment actually runs.
//
// The specs beside this one share ONE in-memory double, so a replica hears its
// peer through a function call. Here each replica holds its own LISTEN
// connection, which is the only arrangement in which a cancel can be published
// on a channel nobody listened to, arrive after the handler that would have
// applied it was closed, or be dropped for a subscriber that fell behind. A
// double cannot fail any of those ways.
var _ = Describe("ResponseStore cross-replica on the broadcast carrier", func() {
var (
ctx context.Context
db *gorm.DB
store *distributed.ResponseMetadataStore
replicaA *ResponseStore
replicaB *ResponseStore
)
BeforeEach(func() {
ctx = context.Background()
var dsn string
db, dsn = testutil.SetupTestDBWithDSN()
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
newBus := func() *pgbus.Bus {
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
Expect(err).ToNot(HaveOccurred())
DeferCleanup(b.Close)
return b
}
var err error
store, err = distributed.NewResponseMetadataStore(db)
Expect(err).ToNot(HaveOccurred())
replicaA = NewResponseStore(0)
replicaB = NewResponseStore(0)
Expect(replicaA.EnableDistributed(ctx, newBus(), "replica-a", store)).To(Succeed())
Expect(replicaB.EnableDistributed(ctx, newBus(), "replica-b", store)).To(Succeed())
})
AfterEach(func() {
Expect(replicaA.Close()).To(Succeed())
Expect(replicaB.Close()).To(Succeed())
})
It("reaches the CancelFunc held by the owning replica over two LISTEN connections", func() {
const id = "resp_cancel_pg"
cancelled := make(chan struct{})
replicaA.StoreBackground(id, &schema.OpenResponsesRequest{Model: "test-model"},
&schema.ORResponseResource{
ID: id, Object: "response", CreatedAt: time.Now().Unix(),
Status: schema.ORStatusInProgress, Model: "test-model",
}, func() { close(cancelled) }, false)
// Waited for, not assumed. On the real carrier the metadata delta is
// asynchronous, so a cancel issued before it lands answers "not found"
// even though the durable row exists: the SyncedMap reads its own
// memory and goes to the table only on re-hydrate. The in-memory double
// the sibling specs share delivers synchronously and hides that
// entirely, which is why this spec exists.
Eventually(func() error {
_, err := replicaB.Get(id)
return err
}, "20s").Should(Succeed())
// The cancel lands on the replica that does NOT hold the CancelFunc.
resp, err := replicaB.Cancel(id)
Expect(err).ToNot(HaveOccurred())
Expect(resp.Status).To(Equal(schema.ORStatusCancelled))
Eventually(cancelled, "20s").Should(BeClosed())
})
It("does not report a cancel that reached nobody as a cancel that was refused", func() {
// The owner is gone, so nothing applies the broadcast. This carrier is
// at-most-once with no replay, and there is no reply to wait for, so the
// caller must still get a prompt terminal answer rather than an error
// that reads as the generation having declined to stop.
const id = "resp_dead_owner_pg"
replicaA.StoreBackground(id, &schema.OpenResponsesRequest{Model: "test-model"},
&schema.ORResponseResource{
ID: id, Object: "response", CreatedAt: time.Now().Unix(),
Status: schema.ORStatusInProgress, Model: "test-model",
}, func() {}, false)
Eventually(func() error {
_, err := replicaB.Get(id)
return err
}, "20s").Should(Succeed())
Expect(replicaA.Close()).To(Succeed())
resp, err := replicaB.Cancel(id)
Expect(err).ToNot(HaveOccurred())
Expect(resp.Status).To(Equal(schema.ORStatusCancelled))
})
})
+31 -15
View File
@@ -1,6 +1,8 @@
package routes
import (
"context"
"github.com/labstack/echo/v4"
"github.com/mudler/LocalAI/core/application"
"github.com/mudler/LocalAI/core/config"
@@ -19,21 +21,8 @@ func RegisterOpenResponsesRoutes(app *echo.Echo,
// How the MCP endpoints reach an agent worker; nil outside distributed mode.
agentControl := mcpAgentControl(application)
if d := application.Distributed(); d != nil {
// Replicate response metadata across frontend replicas and subscribe to
// delegated cancels. Without this a GET, a previous_response_id lookup or
// a cancel that the load balancer sends to a replica other than the
// creator 404s, and the cancel never reaches the CancelFunc (#10993).
// Standalone deployments skip this entirely and stay process-local.
//
// The durable store is what a replica re-hydrates from after its
// subscription missed a delta; without it the same response_id answers
// 404 here and 200 on the peer that created it, forever.
var responseStore *distributed.ResponseMetadataStore
if d.DistStores != nil {
responseStore = d.DistStores.Responses
}
if err := openresponses.GetGlobalStore().EnableDistributed(
application.ApplicationConfig().Context, d.Broadcast(), application.InstanceID(), responseStore); err != nil {
if err := enableDistributedResponses(application.ApplicationConfig().Context, d,
openresponses.GetGlobalStore(), application.InstanceID()); err != nil {
xlog.Error("Failed to enable cross-replica Open Responses store", "error", err)
}
}
@@ -91,3 +80,30 @@ func setOpenResponsesRequestContext(re *middleware.RequestExtractor) echo.Middle
}
}
}
// enableDistributedResponses replicates response metadata across frontend
// replicas and subscribes to delegated cancels. Without it a GET, a
// previous_response_id lookup or a cancel that the load balancer sends to a
// replica other than the creator 404s, and the cancel never reaches the
// CancelFunc (#10993). Standalone deployments never reach here and stay
// process-local.
//
// The durable store is what a replica re-hydrates from after its subscription
// missed a delta; without it the same response_id answers 404 here and 200 on
// the peer that created it, forever.
//
// A named function rather than a block inside route registration, and that is
// the point of it. EnableDistributed takes a messaging.Broadcaster, as it must:
// its own specs publish through a double. So handing it the NATS client instead
// of the deployment's carrier COMPILES and reddens nothing anywhere, and the
// only symptom is a cancel that answers 404 on every replica but one. Registering
// routes needs a whole Application and therefore has no spec; this needs a
// DistributedServices and a store, and therefore has one.
func enableDistributedResponses(ctx context.Context, d *application.DistributedServices,
store *openresponses.ResponseStore, replicaID string) error {
var responseStore *distributed.ResponseMetadataStore
if d.DistStores != nil {
responseStore = d.DistStores.Responses
}
return store.EnableDistributed(ctx, d.Broadcast(), replicaID, responseStore)
}
@@ -0,0 +1,97 @@
// SPDX-License-Identifier: MIT
package routes
import (
"context"
"runtime"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/application"
"github.com/mudler/LocalAI/core/http/endpoints/openresponses"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/pgbus"
"github.com/mudler/LocalAI/core/services/testutil"
)
// Which carrier the Open Responses store is enabled on.
//
// EnableDistributed takes a messaging.Broadcaster, which it must: its own specs
// publish through a double, and it cannot be made to name a concrete carrier
// without dragging that dependency through the whole endpoint package. The
// consequence is that handing it the NATS client instead of the deployment's
// carrier compiles and reddens nothing, and the only symptom is a cancel that
// answers 404 on every replica but the creator. So it is pinned here, by
// watching what actually arrives on the carrier.
var _ = Describe("wiring the Open Responses store to a carrier", func() {
var (
ctx context.Context
busA, busB *pgbus.Bus
store *distributed.ResponseMetadataStore
)
BeforeEach(func() {
if runtime.GOOS == "darwin" {
Skip("testcontainers requires Docker, not available on macOS CI")
}
ctx = context.Background()
db, dsn := testutil.SetupTestDBWithDSN()
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
newBus := func() *pgbus.Bus {
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
Expect(err).ToNot(HaveOccurred())
DeferCleanup(b.Close)
return b
}
busA, busB = newBus(), newBus()
var err error
store, err = distributed.NewResponseMetadataStore(db)
Expect(err).ToNot(HaveOccurred())
})
It("enables it on the deployment's broadcast carrier and not on anything else it holds", func() {
// A DistributedServices holding BOTH, exactly as a running deployment
// does. That is what makes this an assertion about which one was
// chosen rather than about there being one at all.
d := &application.DistributedServices{
Bus: busA,
DistStores: &distributed.Stores{Responses: store},
}
responses := openresponses.NewResponseStore(0)
Expect(enableDistributedResponses(ctx, d, responses, "replica-a")).To(Succeed())
DeferCleanup(func() { _ = responses.Close() })
// A peer's carrier sees the metadata this replica mirrors, which it can
// only do if the store was enabled on the carrier and not on the other
// thing DistributedServices is holding.
mirrored := make(chan []byte, 8)
_, err := busB.Subscribe(messaging.SubjectSyncStateDelta("responses.metadata"), func(data []byte) {
mirrored <- data
})
Expect(err).ToNot(HaveOccurred())
const id = "resp_wiring"
responses.StoreBackground(id, &schema.OpenResponsesRequest{Model: "test-model"},
&schema.ORResponseResource{
ID: id, Object: "response", CreatedAt: time.Now().Unix(),
Status: schema.ORStatusInProgress, Model: "test-model",
}, func() {}, false)
Eventually(mirrored, "20s").Should(Receive(ContainSubstring(id)))
})
It("refuses to enable without the durable store a reconnecting replica re-hydrates from", func() {
d := &application.DistributedServices{Bus: busA}
responses := openresponses.NewResponseStore(0)
Expect(enableDistributedResponses(ctx, d, responses, "replica-a")).ToNot(Succeed())
})
})
+106 -34
View File
@@ -16,7 +16,7 @@ import (
"github.com/mudler/xlog"
)
// AgentEvent is the NATS message payload for agent SSE events.
// AgentEvent is the broadcast payload for agent SSE events.
type AgentEvent struct {
AgentName string `json:"agent_name"`
UserID string `json:"user_id"`
@@ -30,17 +30,41 @@ type AgentEvent struct {
Timestamp int64 `json:"timestamp"` // Unix milliseconds (set by PublishEvent)
}
// AgentCancelEvent is the NATS message payload for cancelling agent execution.
// 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"`
}
// EventBridge bridges agent events between NATS and SSE connections.
// It enables cross-instance SSE: user connects to Frontend 1, agent runs on Frontend 2.
// EventBridge bridges agent events between the broadcast carrier and SSE
// connections. It enables cross-instance SSE: a user connects to Frontend 1
// while the agent runs on Frontend 2.
type EventBridge struct {
nats messaging.MessagingClient
// bus is the fan-out carrier this bridge subscribes on. A Broadcaster and
// not a MessagingClient because fan-out is all a bridge needs: an agent's
// events go to whoever is watching, and there is nothing here to request or
// to queue.
bus messaging.Broadcaster
// cancelBus is where agent.<name>.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.
//
// 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.
//
// 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
// 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
@@ -58,21 +82,49 @@ type EventBridge struct {
// stream-backed one. A copied sync.Map would swallow every cancel.
cancelRegistry *messaging.CancelRegistry
// Background NATS subscriptions owned by this bridge
// The process-lifetime subscription this bridge owns. The per-request one
// SubscribeEvents opens belongs to the HTTP handler that opened it.
obsPersisterSub messaging.Subscription
}
// NewEventBridge creates a new EventBridge.
func NewEventBridge(nc messaging.MessagingClient, store *AgentStore, instanceID string) *EventBridge {
// NewEventBridge creates a new EventBridge on the deployment's fan-out carrier.
//
// 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 {
return &EventBridge{
nats: nc,
pub: nc,
bus: bus,
cancelBus: bus,
pub: bus,
store: store,
instanceID: instanceID,
cancelRegistry: &messaging.CancelRegistry{},
}
}
// WithCancelCarrier puts agent.<name>.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.
//
// 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.
//
// 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
}
// 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
@@ -91,7 +143,7 @@ func (b *EventBridge) WithPublisher(pub messaging.Publisher) *EventBridge {
return &view
}
// PublishEvent publishes an agent event to NATS for SSE bridging.
// PublishEvent broadcasts an agent event for SSE bridging.
//
// Timestamp is emitted in Unix milliseconds to match the local dispatcher's
// json_message events (see dispatcher.go) and the React UI, which feeds the
@@ -106,8 +158,8 @@ func (b *EventBridge) PublishEvent(agentName, userID string, evt AgentEvent) err
// PersistObservable publishes an observable_update SSE event for real-time UI
// updates and, if a database store is available, writes the record to the DB.
// When the store is nil (e.g. on agent workers), the NATS event is still
// published so the frontend can persist it via StartObservablePersister.
// When the store is nil (e.g. on agent workers), the event is still published so
// the frontend can persist it via StartObservablePersister.
func (b *EventBridge) PersistObservable(agentName, userID, eventType string, obs any) {
payload := dbutil.MarshalJSON(obs)
recordID := uuid.New().String()
@@ -123,7 +175,7 @@ func (b *EventBridge) PersistObservable(agentName, userID, eventType string, obs
})
}
// Always publish NATS event — enables real-time SSE and remote persistence.
// Always broadcast, which is what enables real-time SSE and remote persistence.
b.PublishEvent(agentName, userID, AgentEvent{
AgentName: agentName,
UserID: userID,
@@ -135,7 +187,7 @@ func (b *EventBridge) PersistObservable(agentName, userID, eventType string, obs
})
}
// PublishMessage publishes a chat message event via NATS for SSE bridging.
// PublishMessage broadcasts a chat message event for SSE bridging.
// Uses "json_message" event type to match the React UI's expected SSE format.
// Conversation history is managed client-side (browser localStorage), not server-side.
func (b *EventBridge) PublishMessage(agentName, userID, sender, content, messageID string) error {
@@ -172,10 +224,10 @@ func (b *EventBridge) PublishStatus(agentName, userID, status string) error {
// SubscribeEvents subscribes to agent events for a specific agent+user.
func (b *EventBridge) SubscribeEvents(agentName, userID string, handler func(AgentEvent)) (messaging.Subscription, error) {
subject := messaging.SubjectAgentEvents(agentName, userID)
return messaging.SubscribeJSON(b.nats, subject, handler)
return messaging.SubscribeJSON(b.bus, subject, handler)
}
// PublishStreamEvent publishes a stream event (reasoning, content, tool_call, done) via NATS.
// PublishStreamEvent broadcasts a stream event (reasoning, content, tool_call, done).
// These are forwarded as "stream_event" SSE events matching the React UI's expected format.
func (b *EventBridge) PublishStreamEvent(agentName, userID string, data map[string]any) error {
return b.PublishEvent(agentName, userID, AgentEvent{
@@ -193,8 +245,13 @@ func (b *EventBridge) CancelExecution(agentName, userID, messageID string) error
xlog.Info("Cancelled agent execution locally", "agent", agentName, "user", userID, "messageID", messageID)
}
// Also publish via NATS for other instances
return b.nats.Publish(messaging.SubjectAgentCancel(agentName), AgentCancelEvent{
// 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{
AgentName: agentName,
UserID: userID,
MessageID: messageID,
@@ -211,28 +268,32 @@ func (b *EventBridge) DeregisterCancel(key string) {
b.cancelRegistry.Deregister(key)
}
// StartCancelListener subscribes to NATS cancel events (broadcast to all instances).
// StartCancelListener subscribes to the cancel broadcasts every replica sees.
func (b *EventBridge) StartCancelListener() (messaging.Subscription, error) {
return messaging.SubscribeJSON(b.nats, messaging.SubjectAgentCancelWildcard, func(evt AgentCancelEvent) {
return messaging.SubscribeJSON(b.cancelBus, messaging.SubjectAgentCancelWildcard, func(evt AgentCancelEvent) {
if evt.MessageID != "" {
if b.cancelRegistry.Cancel(evt.MessageID) {
xlog.Info("Cancelled agent via NATS", "agent", evt.AgentName, "user", evt.UserID, "messageID", 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 all agent events via NATS and persists
// observable_update events to the database. This runs on the frontend to capture
// observables published by workers (which have no database access).
// The subscription is stored on the EventBridge and cleaned up when the NATS
// connection closes.
// 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.
//
// The subscription lives for the life of this bridge and is one per replica, not
// one per request.
func (b *EventBridge) StartObservablePersister() error {
if b.store == nil {
return fmt.Errorf("no store available for observable persistence")
}
// Subscribe to all agent events using wildcard: agent.*.events.*
sub, err := messaging.SubscribeJSON(b.nats, "agent.*.events.*", func(evt AgentEvent) {
// The filter is the constant next to the builder it has to match. It used
// to be a literal here, four tokens spelled by hand three files from
// SubjectAgentEvents, and a filter one token short of its subject matches
// nothing at all with no error anywhere.
sub, err := messaging.SubscribeJSON(b.bus, messaging.SubjectAgentEventsWildcard, func(evt AgentEvent) {
if evt.EventType != "observable_update" {
return
}
@@ -267,7 +328,7 @@ func (b *EventBridge) StartObservablePersister() error {
return nil
}
// HandleSSE bridges NATS agent events to SSE for a specific agent and user.
// HandleSSE bridges an agent's event broadcasts to SSE for one agent and user.
func (b *EventBridge) HandleSSE(c echo.Context, agentName, userID string) error {
if agentName == "" {
return c.JSON(http.StatusBadRequest, map[string]string{"error": "agent name required"})
@@ -275,8 +336,8 @@ func (b *EventBridge) HandleSSE(c echo.Context, agentName, userID string) error
return b.handleSSEInternal(c, agentName, userID)
}
// SSEHandler returns an Echo handler that bridges NATS agent events to SSE.
// This is the distributed version of the SSE endpoint.
// SSEHandler returns an Echo handler that bridges an agent's event broadcasts
// to SSE. This is the distributed version of the SSE endpoint.
func (b *EventBridge) SSEHandler() echo.HandlerFunc {
return func(c echo.Context) error {
agentName := c.Param("name")
@@ -348,9 +409,20 @@ func (b *EventBridge) handleSSEInternal(c echo.Context, agentName, userID string
writeSSE("json_error", `{"error":"failed to subscribe to agent events"}`)
return nil
}
// Deferred, not called on the way out. This subscription is opened and
// closed PER HTTP REQUEST, and on the PostgreSQL carrier only the first
// subscriber of a channel issues a LISTEN while every later one registers
// an in-process filter, so a return that skipped this would leave a replica
// running one extra closure per notification for every stream it has ever
// served, and nothing in the tree would fail.
defer func() {
closed.Store(true)
if uerr := sub.Unsubscribe(); uerr != nil {
xlog.Warn("Failed to close an agent event subscription", "agent", agentName, "user", userID, "error", uerr)
}
}()
// Wait for client disconnect
<-c.Request().Context().Done()
closed.Store(true)
sub.Unsubscribe()
return nil
}
+279
View File
@@ -0,0 +1,279 @@
// SPDX-License-Identifier: MIT
package agents
import (
"bytes"
"context"
"net/http"
"net/http/httptest"
"sync"
"github.com/labstack/echo/v4"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"gorm.io/gorm"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/pgbus"
"github.com/mudler/LocalAI/core/services/testutil"
)
// The event bridge against the carrier it runs on, across TWO connections.
//
// The wildcard rows below are the matcher's own rows re-asserted through the
// real carrier, and that is not duplication: a matcher that is right and a
// filter that is wrong look identical from inside the messaging package, and
// the symptom of the wrong filter is a persister that writes nothing at all
// with no error anywhere.
var _ = Describe("the agent event bridge on the broadcast carrier", func() {
var (
ctx context.Context
db *gorm.DB
busA, busB *pgbus.Bus
store *AgentStore
bridge *EventBridge
)
BeforeEach(func() {
ctx = context.Background()
var dsn string
db, dsn = testutil.SetupTestDBWithDSN()
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
newBus := func() *pgbus.Bus {
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
Expect(err).ToNot(HaveOccurred())
DeferCleanup(b.Close)
return b
}
busA, busB = newBus(), newBus()
var err error
store, err = NewAgentStore(db)
Expect(err).ToNot(HaveOccurred())
bridge = NewEventBridge(busA, store, "replica-a")
})
// observable is the AgentEvent shape the persister acts on, minus the
// subject it travels on, so a spec varies only the subject.
observable := func(agent, user, id string) AgentEvent {
return AgentEvent{
AgentName: agent,
UserID: user,
EventType: "observable_update",
EventSubType: "tool_result",
SourceInstance: "replica-b",
MessageID: id,
Metadata: `{"tool":"grep"}`,
}
}
countFor := func(user, agent string) func() int {
return func() int {
records, err := store.GetObservables(AgentKey(user, agent), 10)
if err != nil {
return -1
}
return len(records)
}
}
Describe("StartObservablePersister", func() {
BeforeEach(func() {
Expect(bridge.StartObservablePersister()).To(Succeed())
})
It("persists an observable a peer replica published on the built subject", func() {
Expect(busB.Publish(messaging.SubjectAgentEvents("a1", "u1"), observable("a1", "u1", "obs-1"))).To(Succeed())
Eventually(countFor("u1", "a1"), "20s").Should(Equal(1))
})
It("ignores a three-token subject that a shorter filter would have swallowed", func() {
// agent.a1.events has one token fewer than SubjectAgentEvents
// builds. A filter of that shape matches this and NOT the real
// subject, which is the exact inversion the token count prevents.
Expect(busB.Publish("agent.a1.events", observable("a1", "u1", "obs-short"))).To(Succeed())
// The sentinel, on the subject that certainly matches, orders the
// negative without a clock.
Expect(busB.Publish(messaging.SubjectAgentEvents("a1", "u1"), observable("a1", "u1", "obs-real"))).To(Succeed())
Eventually(countFor("u1", "a1"), "20s").Should(Equal(1))
records, err := store.GetObservables(AgentKey("u1", "a1"), 10)
Expect(err).ToNot(HaveOccurred())
Expect(records[0].ID).To(Equal("obs-real"))
})
It("ignores a five-token subject with the same leading tokens", func() {
Expect(busB.Publish("agent.a1.b.events.u1", observable("a1", "u1", "obs-long"))).To(Succeed())
Expect(busB.Publish(messaging.SubjectAgentEvents("a1", "u1"), observable("a1", "u1", "obs-real"))).To(Succeed())
Eventually(countFor("u1", "a1"), "20s").Should(Equal(1))
records, err := store.GetObservables(AgentKey("u1", "a1"), 10)
Expect(err).ToNot(HaveOccurred())
Expect(records[0].ID).To(Equal("obs-real"))
})
It("persists for every agent and every user, not only the first it saw", func() {
Expect(busB.Publish(messaging.SubjectAgentEvents("a1", "u1"), observable("a1", "u1", "obs-a1"))).To(Succeed())
Expect(busB.Publish(messaging.SubjectAgentEvents("a2", "u2"), observable("a2", "u2", "obs-a2"))).To(Succeed())
Eventually(countFor("u1", "a1"), "20s").Should(Equal(1))
Eventually(countFor("u2", "a2"), "20s").Should(Equal(1))
})
})
Describe("SubscribeEvents", func() {
It("receives the agent and user it asked for and not another's events", func() {
mine := make(chan AgentEvent, 8)
sub, err := bridge.SubscribeEvents("a1", "u1", func(evt AgentEvent) { mine <- evt })
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() { _ = sub.Unsubscribe() })
peer := NewEventBridge(busB, store, "replica-b")
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())
var first AgentEvent
Eventually(mine, "20s").Should(Receive(&first))
Expect(first.MessageID).To(Equal("m-mine"))
Consistently(mine, "500ms", "50ms").ShouldNot(Receive())
})
It("closes cleanly and leaves the register where it found it", func() {
before := busA.Subscribers()
sub, err := bridge.SubscribeEvents("a1", "u1", func(AgentEvent) {})
Expect(err).ToNot(HaveOccurred())
Expect(busA.Subscribers()).To(Equal(before + 1))
Expect(sub.Unsubscribe()).To(Succeed())
Expect(busA.Subscribers()).To(Equal(before))
})
})
Describe("the SSE handler's per-request subscription", func() {
// The second of the two subscriptions this deployment opens and closes
// PER HTTP REQUEST. A leaked filter has no symptom: the replica just
// runs one more closure per notification for every stream it has ever
// served, for as long as the process lives.
It("is closed however the handler returns", func() {
before := busA.Subscribers()
out := &syncBody{}
req := httptest.NewRequest(http.MethodGet, "/api/agents/a1/sse/distributed?user_id=u1", nil)
reqCtx, cancelReq := context.WithCancel(ctx)
c := echo.New().NewContext(req.WithContext(reqCtx), out)
done := make(chan struct{})
go func() {
defer GinkgoRecover()
defer close(done)
Expect(bridge.HandleSSE(c, "a1", "u1")).To(Succeed())
}()
Eventually(busA.Subscribers, "20s").Should(Equal(before + 1))
// 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")
Expect(peer.PublishMessage("a1", "u1", "agent", "hello", "m-1")).To(Succeed())
Eventually(out.String, "20s").Should(ContainSubstring("hello"))
cancelReq()
Eventually(done, "20s").Should(BeClosed())
Eventually(busA.Subscribers, "20s").Should(Equal(before),
"a stream that has returned must leave no handler behind")
})
})
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() })
cancelled := make(chan struct{})
peer.RegisterCancel("msg-1", func() { close(cancelled) })
Expect(bridge.CancelExecution("a1", "u1", "msg-1")).To(Succeed())
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()
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")
})
})
})
// 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
// assuming it did.
type syncBody struct {
mu sync.Mutex
buf bytes.Buffer
header http.Header
}
func (w *syncBody) Header() http.Header {
if w.header == nil {
w.header = http.Header{}
}
return w.header
}
func (w *syncBody) Write(p []byte) (int, error) {
w.mu.Lock()
defer w.mu.Unlock()
return w.buf.Write(p)
}
func (w *syncBody) WriteHeader(int) {}
func (w *syncBody) Flush() {}
func (w *syncBody) String() string {
w.mu.Lock()
defer w.mu.Unlock()
return w.buf.String()
}
+8 -24
View File
@@ -9,41 +9,25 @@ import (
. "github.com/onsi/gomega"
)
// fakeMessagingClient implements messaging.MessagingClient and captures the
// last published payload so tests can assert on it.
type fakeMessagingClient struct {
// recordingBus is a messaging.Broadcaster that captures the last published
// payload. A Broadcaster and not a MessagingClient because fan-out is the whole
// of what an EventBridge may do with a carrier.
type recordingBus struct {
lastSubject string
lastData any
}
func (f *fakeMessagingClient) Publish(subject string, data any) error {
func (f *recordingBus) Publish(subject string, data any) error {
f.lastSubject = subject
f.lastData = data
return nil
}
func (f *fakeMessagingClient) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
func (f *recordingBus) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
return &fakeSub{}, nil
}
func (f *fakeMessagingClient) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) {
return &fakeSub{}, nil
}
func (f *fakeMessagingClient) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) {
return &fakeSub{}, nil
}
func (f *fakeMessagingClient) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) {
return &fakeSub{}, nil
}
func (f *fakeMessagingClient) Request(string, []byte, time.Duration) ([]byte, error) {
return nil, nil
}
func (f *fakeMessagingClient) IsConnected() bool { return true }
func (f *fakeMessagingClient) Close() {}
var _ messaging.Broadcaster = (*recordingBus)(nil)
type fakeSub struct{}
@@ -57,7 +41,7 @@ var _ = Describe("EventBridge", func() {
// React UI both expect Unix milliseconds. Nanoseconds also overflow JS's
// safe-integer range. The timestamp must be in milliseconds.
It("emits the timestamp in Unix milliseconds", func() {
fake := &fakeMessagingClient{}
fake := &recordingBus{}
bridge := NewEventBridge(fake, nil, "instance-1")
before := time.Now().UnixMilli()
+77 -20
View File
@@ -67,7 +67,7 @@ type CancelEvent struct {
// no concurrency limiter: there is nothing here to limit.
type Dispatcher struct {
store *JobStore
nats messaging.MessagingClient
bus messaging.Broadcaster
db *gorm.DB
instanceID string
configLoader ModelConfigLoader // optional: to enrich job events with model config
@@ -75,21 +75,44 @@ type Dispatcher struct {
// Cancel registry (notetaker pattern)
cancelRegistry messaging.CancelRegistry
// NATS subscriptions
// The broadcast subscriptions this dispatcher owns for the life of the
// process. The per-request one SubscribeProgress opens is not here: it
// belongs to the HTTP handler that opened it and is closed with it.
cancelSub messaging.Subscription
resultSub messaging.Subscription
progressSub messaging.Subscription
// terminalRecheck is how often an open SSE stream re-reads the job row.
// Zero means DefaultTerminalRecheck. Set once, before Start, and read from
// HTTP handler goroutines afterwards.
terminalRecheck time.Duration
// Lifecycle
ctx context.Context
cancel context.CancelFunc
}
// SetTerminalRecheck sets how often an open job-progress SSE stream re-reads
// the job row while it waits.
//
// It exists so a spec can drive the recovery from a dropped terminal broadcast
// without waiting out DefaultTerminalRecheck. Call it before Start.
func (d *Dispatcher) SetTerminalRecheck(interval time.Duration) {
d.terminalRecheck = interval
}
// NewDispatcher creates a new distributed job Dispatcher.
func NewDispatcher(store *JobStore, nc messaging.MessagingClient, db *gorm.DB, instanceID string) *Dispatcher {
//
// The carrier is a messaging.Broadcaster because fan-out is all this type does
// with it and all it may have: everything else a job needs from another replica
// travels on the claim row or on the control RPC's response body, and a
// dispatcher that could reach a request/reply or a queue group through this
// field would be able to reintroduce the publish-to-nobody the claim queue
// replaced.
func NewDispatcher(store *JobStore, bus messaging.Broadcaster, db *gorm.DB, instanceID string) *Dispatcher {
return &Dispatcher{
store: store,
nats: nc,
bus: bus,
db: db,
instanceID: instanceID,
}
@@ -133,9 +156,15 @@ func (d *Dispatcher) Start(ctx context.Context) error {
var err error
// Subscribe to cancel events (broadcast to all — each instance checks its registry)
d.cancelSub, err = messaging.SubscribeJSON(d.nats, messaging.SubjectJobCancelWildcard, func(evt CancelEvent) {
// A cancel that does not arrive is not a cancel that was refused. This
// carrier is at-most-once with no replay, so a broadcast that is dropped or
// that lands while this replica is reconnecting reaches nobody, and the
// registry simply never hears about it. Nothing here may report that as the
// execution having declined to stop: the only thing this subscription can
// say is that a cancel DID arrive.
d.cancelSub, err = messaging.SubscribeJSON(d.bus, messaging.SubjectJobCancelWildcard, func(evt CancelEvent) {
if d.cancelRegistry.Cancel(evt.JobID) {
xlog.Info("Cancelled job via NATS", "jobID", evt.JobID)
xlog.Info("Cancelled a job on this replica after a broadcast cancel", "jobID", evt.JobID)
}
})
if err != nil {
@@ -144,15 +173,30 @@ func (d *Dispatcher) Start(ctx context.Context) error {
// Subscribe to job result events from workers (persist to DB)
if d.store != nil {
d.resultSub, err = messaging.SubscribeJSON(d.nats, messaging.SubjectJobResultWildcard, func(evt JobResultEvent) {
d.store.UpdateJobStatus(evt.JobID, evt.Status, evt.Result, evt.Error)
// The fan-out COPY of a terminal result, and never the only one.
//
// This carrier drops a broadcast rather than blocking when one
// subscriber falls behind, and a result has no successor message, so a
// subscriber that missed one would never hear about that job again.
// That is survivable here only because it is not the path the answer
// travels: the replica that claimed the work persists the terminal line
// through DispatchLoop.settleClaim BEFORE it releases the claim, so the
// job row already carries the answer when this broadcast is published.
// A dropped result therefore costs a live SSE stream its promptness,
// which jobs/sse.go recovers from by reading the row, and never costs
// the job its answer. Nothing may be moved onto this subscription that
// is not also written to a table first.
d.resultSub, err = messaging.SubscribeJSON(d.bus, messaging.SubjectJobResultWildcard, func(evt JobResultEvent) {
if err := d.store.UpdateJobStatus(evt.JobID, evt.Status, evt.Result, evt.Error); err != nil {
xlog.Error("Failed to persist a broadcast job result", "job_id", evt.JobID, "error", err)
}
})
if err != nil {
return fmt.Errorf("subscribing to result events: %w", err)
}
// Subscribe to trace events from workers (persist to DB)
d.progressSub, err = messaging.SubscribeJSON(d.nats, messaging.SubjectJobProgressWildcard, func(evt ProgressEvent) {
d.progressSub, err = messaging.SubscribeJSON(d.bus, messaging.SubjectJobProgressWildcard, func(evt ProgressEvent) {
if evt.TraceType != "" && evt.TraceContent != "" {
if err := d.store.AppendJobTrace(evt.JobID, evt.TraceType, evt.TraceContent); err != nil {
xlog.Error("Failed to append job trace", "job_id", evt.JobID, "trace_type", evt.TraceType, "error", err)
@@ -172,8 +216,8 @@ func (d *Dispatcher) Start(ctx context.Context) error {
return nil
}
// unsubscribeAll nil-checks, unsubscribes, and nils out each NATS subscription.
// Safe to call multiple times.
// unsubscribeAll nil-checks, unsubscribes, and nils out each broadcast
// subscription this dispatcher owns. Safe to call multiple times.
func (d *Dispatcher) unsubscribeAll() {
if d.cancelSub != nil {
d.cancelSub.Unsubscribe()
@@ -247,31 +291,44 @@ func (d *Dispatcher) Enqueue(jobID, taskID, userID string) error {
return nil
}
// Cancel publishes a cancel event to NATS (broadcast to all instances).
// Cancel broadcasts a cancel request to every replica.
//
// Its error says whether the request was PUBLISHED and nothing else. There is
// no reply, and there is deliberately no attempt to synthesise one: a cancel
// that reached no subscriber and a cancel an execution declined are different
// facts, and this carrier cannot tell them apart, so neither may be reported as
// the other.
func (d *Dispatcher) Cancel(jobID string) error {
return d.nats.Publish(messaging.SubjectJobCancel(jobID), CancelEvent{
return d.bus.Publish(messaging.SubjectJobCancel(jobID), CancelEvent{
JobID: jobID,
})
}
// PublishProgress publishes a progress event for SSE bridging.
// PublishProgress broadcasts a progress event for SSE bridging.
func (d *Dispatcher) PublishProgress(jobID, status, message string) error {
return d.nats.Publish(messaging.SubjectJobProgress(jobID), ProgressEvent{
return d.bus.Publish(messaging.SubjectJobProgress(jobID), ProgressEvent{
JobID: jobID,
Status: status,
Message: message,
})
}
// SubscribeProgress subscribes to progress events for a specific job (for SSE bridging).
// SubscribeProgress subscribes to progress events for ONE job, for SSE
// bridging.
//
// The subject is the exact one SubjectJobProgress builds and never the
// wildcard. On the wildcard this would be a stream showing every job in the
// deployment to every client watching any of them, which is a data-boundary
// rather than a display bug, and the subscription is opened and closed per HTTP
// request so the caller MUST close it.
func (d *Dispatcher) SubscribeProgress(jobID string, handler func(ProgressEvent)) (messaging.Subscription, error) {
return messaging.SubscribeJSON(d.nats, messaging.SubjectJobProgress(jobID), handler)
return messaging.SubscribeJSON(d.bus, messaging.SubjectJobProgress(jobID), handler)
}
// PublishTrace publishes a trace event for a running job via NATS.
// The frontend subscribes and persists traces to DB.
// PublishTrace broadcasts a trace event for a running job. The frontend
// subscribes and persists traces to the database.
func (d *Dispatcher) PublishTrace(jobID, traceType, traceContent string) error {
return d.nats.Publish(messaging.SubjectJobProgress(jobID), ProgressEvent{
return d.bus.Publish(messaging.SubjectJobProgress(jobID), ProgressEvent{
JobID: jobID,
TraceType: traceType,
TraceContent: traceContent,
+271
View File
@@ -0,0 +1,271 @@
// SPDX-License-Identifier: MIT
package jobs
import (
"bytes"
"context"
"net/http"
"net/http/httptest"
"sync"
"time"
"github.com/labstack/echo/v4"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"gorm.io/gorm"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/pgbus"
"github.com/mudler/LocalAI/core/services/testutil"
)
// The dispatcher against the carrier it runs on in production, across TWO
// connections.
//
// One carrier talking to itself is not the cross-replica case, and a double is
// less than that: it cannot drop, cannot reconnect and cannot fail the way a
// LISTEN connection fails. Everything below publishes on bus B and asserts on a
// dispatcher built on bus A.
var _ = Describe("the job dispatcher on the broadcast carrier", func() {
var (
ctx context.Context
db *gorm.DB
busA, busB *pgbus.Bus
store *JobStore
disp *Dispatcher
)
BeforeEach(func() {
ctx = context.Background()
var dsn string
db, dsn = testutil.SetupTestDBWithDSN()
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
newBus := func() *pgbus.Bus {
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
Expect(err).ToNot(HaveOccurred())
DeferCleanup(b.Close)
return b
}
busA, busB = newBus(), newBus()
var err error
store, err = NewJobStore(db)
Expect(err).ToNot(HaveOccurred())
disp = NewDispatcher(store, busA, db, "replica-a")
})
Describe("SubscribeProgress", func() {
// The exactness of the per-request subscription, which is a data
// boundary and not a display detail: on the wildcard, every SSE client
// watching any job would be shown every OTHER job's progress, including
// jobs belonging to other users.
It("receives the job it asked for and not another job's progress", func() {
mine := make(chan ProgressEvent, 8)
sub, err := disp.SubscribeProgress("job-mine", func(evt ProgressEvent) { mine <- evt })
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() { _ = sub.Unsubscribe() })
// The other job first, then mine. If the filter is a wildcard in
// disguise, the first thing received is the other job's event, and
// no clock is needed to say so.
Expect(busB.Publish(messaging.SubjectJobProgress("job-theirs"), ProgressEvent{
JobID: "job-theirs", Status: "running", Message: "not for me",
})).To(Succeed())
Expect(busB.Publish(messaging.SubjectJobProgress("job-mine"), ProgressEvent{
JobID: "job-mine", Status: "running", Message: "for me",
})).To(Succeed())
var first ProgressEvent
Eventually(mine, "20s").Should(Receive(&first))
Expect(first.JobID).To(Equal("job-mine"))
Expect(first.Message).To(Equal("for me"))
Consistently(mine, "500ms", "50ms").ShouldNot(Receive())
})
It("delivers what a PEER replica published, not only its own publishes", func() {
out := make(chan ProgressEvent, 8)
sub, err := disp.SubscribeProgress("job-peer", func(evt ProgressEvent) { out <- evt })
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() { _ = sub.Unsubscribe() })
peer := NewDispatcher(store, busB, db, "replica-b")
Expect(peer.PublishProgress("job-peer", "running", "from the other replica")).To(Succeed())
Eventually(out, "20s").Should(Receive(HaveField("Message", "from the other replica")))
})
})
Describe("the wildcard subscriptions Start opens", func() {
It("persists a trace a peer replica published", func() {
Expect(disp.Start(ctx)).To(Succeed())
DeferCleanup(disp.Stop)
job := &JobRecord{TaskID: "t", UserID: "u", Status: "running", TriggeredBy: "manual"}
Expect(store.CreateJob(job)).To(Succeed())
peer := NewDispatcher(store, busB, db, "replica-b")
Expect(peer.PublishTrace(job.ID, "reasoning", "thinking about it")).To(Succeed())
Eventually(func() int {
stored, err := store.GetJob(job.ID)
if err != nil || stored.TracesJSON == "" {
return 0
}
return len(stored.TracesJSON)
}, "20s").Should(BeNumerically(">", 0))
})
It("persists a terminal result a peer replica published", func() {
Expect(disp.Start(ctx)).To(Succeed())
DeferCleanup(disp.Stop)
job := &JobRecord{TaskID: "t", UserID: "u", Status: "running", TriggeredBy: "manual"}
Expect(store.CreateJob(job)).To(Succeed())
PublishJobResult(busB, job.ID, "completed", "the answer", "")
Eventually(func() string {
stored, err := store.GetJob(job.ID)
if err != nil {
return ""
}
return stored.Status
}, "20s").Should(Equal("completed"))
})
It("leaves the register where it found it once the dispatcher stops", func() {
// Three process-lifetime subscriptions, and no more than three: a
// dispatcher that leaked one per Start would accumulate silently
// across the restarts a reconfiguration does.
before := busA.Subscribers()
Expect(disp.Start(ctx)).To(Succeed())
Expect(busA.Subscribers()).To(Equal(before + 3))
disp.Stop()
Expect(busA.Subscribers()).To(Equal(before))
})
})
Describe("the SSE stream when the terminal broadcast never arrives", func() {
// The 256-deep per-subscriber queue DROPS rather than blocking, and a
// terminal event has no successor: a stream waiting only on the carrier
// would hang for ever and the user would read a job that finished as a
// job that produced nothing. Nothing terminal is published here at all,
// which is the strongest form of that loss.
It("ends the stream from the job row rather than waiting for an event that was dropped", func() {
disp.SetTerminalRecheck(50 * time.Millisecond)
job := &JobRecord{TaskID: "t", UserID: "u", Status: "running", TriggeredBy: "manual"}
Expect(store.CreateJob(job)).To(Succeed())
before := busA.Subscribers()
body, handlerDone, _ := startStream(ctx, disp, job.ID)
// This carrier has no replay, so a progress event published before
// the handler's channel is listened is simply gone. Subscribers()
// counts only live registrations, which is what makes it usable as
// the gate here rather than merely as a leak counter.
Eventually(busA.Subscribers, "20s").Should(Equal(before + 1))
// The stream is only proved to be PAST its post-subscribe read of
// the row once it has forwarded a live event. Without this the row
// could go terminal while the handler was still in that one-shot
// read, and the spec would pass with the periodic re-read deleted:
// it did, when this was written the other way round.
Expect(NewDispatcher(store, busB, db, "replica-b").
PublishProgress(job.ID, "running", "step 1")).To(Succeed())
Eventually(body, "20s").Should(ContainSubstring("step 1"))
// The terminal state reaches the TABLE and no broadcast is ever
// published for it. This is exactly what a dropped result looks
// like to a subscriber: nothing.
Expect(store.UpdateJobStatus(job.ID, "completed", "the answer", "")).To(Succeed())
Eventually(handlerDone, "20s").Should(BeClosed(),
"the stream must end on the row when its terminal broadcast never arrives")
Expect(body()).To(ContainSubstring("event: done"))
Expect(body()).To(ContainSubstring("completed"))
})
It("closes its per-request subscription however it returns", func() {
// Two subscriptions are opened and closed per HTTP REQUEST in this
// deployment, and a leaked filter has no symptom at all: the
// replica just runs one more closure per notification for every
// stream it has ever served.
disp.SetTerminalRecheck(50 * time.Millisecond)
job := &JobRecord{TaskID: "t", UserID: "u", Status: "running", TriggeredBy: "manual"}
Expect(store.CreateJob(job)).To(Succeed())
before := busA.Subscribers()
_, handlerDone, cancelReq := startStream(ctx, disp, job.ID)
Eventually(busA.Subscribers, "20s").Should(Equal(before + 1))
cancelReq()
Eventually(handlerDone, "20s").Should(BeClosed())
Eventually(busA.Subscribers, "20s").Should(Equal(before),
"a stream that has returned must leave no handler behind")
})
})
})
// 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 has reached its wait rather
// than assuming it has.
type syncBody struct {
mu sync.Mutex
buf bytes.Buffer
header http.Header
}
func (w *syncBody) Header() http.Header {
if w.header == nil {
w.header = http.Header{}
}
return w.header
}
func (w *syncBody) Write(p []byte) (int, error) {
w.mu.Lock()
defer w.mu.Unlock()
return w.buf.Write(p)
}
func (w *syncBody) WriteHeader(int) {}
func (w *syncBody) Flush() {}
func (w *syncBody) String() string {
w.mu.Lock()
defer w.mu.Unlock()
return w.buf.String()
}
// startStream runs the job-progress SSE handler for jobID on its own goroutine
// and returns a reader for what it has written so far, a channel closed when it
// returns, and the cancel that stands in for the client hanging up.
func startStream(ctx context.Context, d *Dispatcher, jobID string) (func() string, chan struct{}, context.CancelFunc) {
GinkgoHelper()
out := &syncBody{}
req := httptest.NewRequest(http.MethodGet, "/api/agent/jobs/"+jobID+"/progress", nil)
reqCtx, cancelReq := context.WithCancel(ctx)
c := echo.New().NewContext(req.WithContext(reqCtx), out)
c.SetParamNames("id")
c.SetParamValues(jobID)
done := make(chan struct{})
go func() {
defer GinkgoRecover()
defer close(done)
Expect(d.SSEHandler()(c)).To(Succeed())
}()
DeferCleanup(func() {
cancelReq()
Eventually(done, "20s").Should(BeClosed())
})
return out.String, done, cancelReq
}
+25 -27
View File
@@ -2,6 +2,7 @@ package jobs
import (
"encoding/json"
"sync"
"time"
. "github.com/onsi/ginkgo/v2"
@@ -19,38 +20,35 @@ type publishCall struct {
data any
}
// fakeMessagingClient implements messaging.MessagingClient and records published messages.
type fakeMessagingClient struct {
// recordingBus is a messaging.Broadcaster that records what was published.
//
// A Broadcaster and NOT a MessagingClient, which is the point of the type
// rather than tidiness: the dispatcher may only fan out, and a double that
// still offered Request or QueueSubscribe would let a spec exercise a surface
// the production type can no longer reach.
type recordingBus struct {
mu sync.Mutex
calls []publishCall
}
func (f *fakeMessagingClient) Publish(subject string, data any) error {
func (f *recordingBus) Publish(subject string, data any) error {
f.mu.Lock()
defer f.mu.Unlock()
f.calls = append(f.calls, publishCall{subject: subject, data: data})
return nil
}
func (f *fakeMessagingClient) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
func (f *recordingBus) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
return &fakeSub{}, nil
}
func (f *fakeMessagingClient) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) {
return &fakeSub{}, nil
func (f *recordingBus) published() []publishCall {
f.mu.Lock()
defer f.mu.Unlock()
return append([]publishCall(nil), f.calls...)
}
func (f *fakeMessagingClient) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) {
return &fakeSub{}, nil
}
func (f *fakeMessagingClient) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) {
return &fakeSub{}, nil
}
func (f *fakeMessagingClient) Request(string, []byte, time.Duration) ([]byte, error) {
return nil, nil
}
func (f *fakeMessagingClient) IsConnected() bool { return true }
func (f *fakeMessagingClient) Close() {}
var _ messaging.Broadcaster = (*recordingBus)(nil)
// fakeSub implements messaging.Subscription.
type fakeSub struct{}
@@ -221,7 +219,7 @@ var _ = Describe("Dispatcher", func() {
Describe("Enqueue claim kind", func() {
var (
store *JobStore
fake *fakeMessagingClient
fake *recordingBus
disp *Dispatcher
db *gorm.DB
)
@@ -231,7 +229,7 @@ var _ = Describe("Dispatcher", func() {
var err error
store, err = NewJobStore(db)
Expect(err).ToNot(HaveOccurred())
fake = &fakeMessagingClient{}
fake = &recordingBus{}
disp = NewDispatcher(store, fake, db, "test-instance")
})
@@ -265,7 +263,7 @@ var _ = Describe("Dispatcher", func() {
Expect(disp.Enqueue(job.ID, task.ID, "user-1")).To(Succeed())
Expect(onlyClaim(db).Kind).To(Equal(string(ClaimKindMCPCI)))
Expect(fake.calls).To(BeEmpty(), "enqueueing must not publish: a queue subject nobody joined swallows the job")
Expect(fake.published()).To(BeEmpty(), "enqueueing must not publish: a queue subject nobody joined swallows the job")
})
It("writes a plain task claim for a model without MCP servers", func() {
@@ -294,7 +292,7 @@ var _ = Describe("Dispatcher", func() {
Expect(disp.Enqueue(job.ID, task.ID, "user-1")).To(Succeed())
Expect(onlyClaim(db).Kind).To(Equal(string(ClaimKindTask)))
Expect(fake.calls).To(BeEmpty(), "enqueueing must not publish: a queue subject nobody joined swallows the job")
Expect(fake.published()).To(BeEmpty(), "enqueueing must not publish: a queue subject nobody joined swallows the job")
})
It("writes a plain task claim when the model config is not found", func() {
@@ -321,7 +319,7 @@ var _ = Describe("Dispatcher", func() {
Expect(disp.Enqueue(job.ID, task.ID, "user-1")).To(Succeed())
Expect(onlyClaim(db).Kind).To(Equal(string(ClaimKindTask)))
Expect(fake.calls).To(BeEmpty(), "enqueueing must not publish: a queue subject nobody joined swallows the job")
Expect(fake.published()).To(BeEmpty(), "enqueueing must not publish: a queue subject nobody joined swallows the job")
})
})
@@ -331,7 +329,7 @@ var _ = Describe("Dispatcher", func() {
Describe("Enqueue event enrichment", func() {
var (
store *JobStore
fake *fakeMessagingClient
fake *recordingBus
disp *Dispatcher
db *gorm.DB
)
@@ -341,7 +339,7 @@ var _ = Describe("Dispatcher", func() {
var err error
store, err = NewJobStore(db)
Expect(err).ToNot(HaveOccurred())
fake = &fakeMessagingClient{}
fake = &recordingBus{}
disp = NewDispatcher(store, fake, db, "test-instance")
})
+10 -2
View File
@@ -5,7 +5,15 @@ import (
"github.com/mudler/xlog"
)
// PublishJobResult publishes a terminal job result event and a progress event via NATS.
// PublishJobResult broadcasts a terminal job result and the matching progress
// event.
//
// Both are FAN-OUT copies. The answer itself reaches the frontend on the reply
// line of the control RPC the claiming replica is reading, and is persisted
// before that claim is released, so neither of these publishes failing can lose
// a job its result. Errors are logged rather than returned for exactly that
// reason: a caller that treated a publish failure as the work having failed
// would report a job that ran as a job that did not.
func PublishJobResult(pub messaging.Publisher, jobID, status, result, errMsg string) {
if err := pub.Publish(messaging.SubjectJobResult(jobID), JobResultEvent{
JobID: jobID,
@@ -24,7 +32,7 @@ func PublishJobResult(pub messaging.Publisher, jobID, status, result, errMsg str
}
}
// PublishJobProgress publishes a status-only update (no result) via NATS.
// PublishJobProgress broadcasts a status-only update, with no result attached.
func PublishJobProgress(pub messaging.Publisher, jobID, status, message string) {
if err := pub.Publish(messaging.SubjectJobProgress(jobID), ProgressEvent{
JobID: jobID,
+95 -24
View File
@@ -6,12 +6,26 @@ import (
"net/http"
"sync"
"sync/atomic"
"time"
"github.com/labstack/echo/v4"
"github.com/mudler/xlog"
)
// SSEBridge provides an HTTP handler that bridges NATS progress events to SSE.
// This follows the notetaker pattern: subscribe to NATS, forward to SSE client.
// DefaultTerminalRecheck is how often an open job-progress SSE stream re-reads
// the job row while it waits for the broadcast that would close it.
//
// It exists because the broadcast carrier is at-most-once and drops rather than
// blocks: a terminal progress event can be lost, and a lost terminal event has
// no successor, so a stream waiting only on the bus would hang until the client
// gave up and the user would read a finished job as one that produced nothing.
// The row is where the answer is written before it is ever broadcast, so the
// stream ends on the TABLE and merely ends FASTER on the broadcast.
const DefaultTerminalRecheck = 15 * time.Second
// SSEBridge provides an HTTP handler that bridges job progress broadcasts to
// SSE: subscribe to the carrier, forward to the SSE client, and close the
// stream when the job is done.
func (d *Dispatcher) SSEHandler() echo.HandlerFunc {
return func(c echo.Context) error {
jobID := c.Param("id")
@@ -49,18 +63,11 @@ func (d *Dispatcher) SSEHandler() echo.HandlerFunc {
}
// Send current job state first
job, err := d.store.GetJob(jobID)
if err == nil {
sendEvent("status", ProgressEvent{
JobID: jobID,
Status: job.Status,
})
// If already terminal, send done and close — no need to subscribe
if job.Status == "completed" || job.Status == "failed" || job.Status == "cancelled" {
sendEvent("done", ProgressEvent{
JobID: jobID,
Status: job.Status,
})
if status, terminal, ok := d.jobStatus(jobID); ok {
sendEvent("status", ProgressEvent{JobID: jobID, Status: status})
// Already finished, so there is nothing to subscribe to.
if terminal {
sendEvent("done", ProgressEvent{JobID: jobID, Status: status})
return nil
}
}
@@ -69,15 +76,18 @@ func (d *Dispatcher) SSEHandler() echo.HandlerFunc {
// handler can return promptly instead of waiting for client disconnect.
done := make(chan struct{})
closeOnce := sync.Once{}
finish := func(evt ProgressEvent) {
sendEvent("done", evt)
closeOnce.Do(func() { close(done) })
}
// Subscribe to progress events for this job
sub, err := d.SubscribeProgress(jobID, func(evt ProgressEvent) {
sendEvent("progress", evt)
// Close the stream on terminal states
if evt.Status == "completed" || evt.Status == "failed" || evt.Status == "cancelled" {
sendEvent("done", evt)
closeOnce.Do(func() { close(done) })
if IsTerminalJobStatus(evt.Status) {
finish(evt)
}
})
if err != nil {
@@ -85,13 +95,74 @@ func (d *Dispatcher) SSEHandler() echo.HandlerFunc {
sendEvent("error", map[string]string{"error": "failed to subscribe"})
return nil
}
// Wait for client disconnect or terminal state
select {
case <-c.Request().Context().Done():
case <-done:
// Deferred, not called on the way out. Every return below this line
// leaks a handler otherwise, and the leak has no symptom: on the
// PostgreSQL carrier only the first subscriber of a channel issues a
// LISTEN and the rest are in-process filters, so a replica that has
// served ten thousand of these streams simply runs ten thousand
// closures per notification and reports nothing.
defer func() {
closed.Store(true)
if uerr := sub.Unsubscribe(); uerr != nil {
xlog.Warn("Failed to close a job progress subscription", "job_id", jobID, "error", uerr)
}
}()
// The row, once, immediately after subscribing. A job that finished
// between the read above and this subscription published its terminal
// event into the window where nobody was listening, and this carrier
// has no replay, so without this read the stream waits for an event
// that has already been and gone.
if status, terminal, ok := d.jobStatus(jobID); ok && terminal {
finish(ProgressEvent{JobID: jobID, Status: status})
return nil
}
// And the row again, periodically, for the rest of the stream's life.
// The carrier drops a broadcast rather than blocking when a subscriber
// falls behind, and a dropped TERMINAL event is a stream that never
// ends. Reading the row is what keeps a lost broadcast from reading as
// a job that produced no result.
ticker := time.NewTicker(d.terminalRecheckInterval())
defer ticker.Stop()
for {
select {
case <-c.Request().Context().Done():
return nil
case <-done:
return nil
case <-ticker.C:
if status, terminal, ok := d.jobStatus(jobID); ok && terminal {
finish(ProgressEvent{JobID: jobID, Status: status})
return nil
}
}
}
closed.Store(true)
sub.Unsubscribe()
return nil
}
}
// jobStatus reads a job's status from the TABLE and reports whether it is
// terminal. The third return separates "the row says pending" from "there is no
// row to read", which the caller must not collapse: a job whose row cannot be
// read has not been shown to have produced nothing.
func (d *Dispatcher) jobStatus(jobID string) (status string, terminal, ok bool) {
if d.store == nil {
return "", false, false
}
job, err := d.store.GetJob(jobID)
if err != nil {
return "", false, false
}
return job.Status, IsTerminalJobStatus(job.Status), true
}
// terminalRecheckInterval is how often an open stream re-reads the row.
// Configurable so a spec can drive the recovery without waiting out the
// production interval, and never zero, which would spin.
func (d *Dispatcher) terminalRecheckInterval() time.Duration {
if d.terminalRecheck > 0 {
return d.terminalRecheck
}
return DefaultTerminalRecheck
}
+23 -2
View File
@@ -229,6 +229,27 @@ func (s *JobStore) DeleteJob(id string) error {
return s.db.Where("id = ?", id).Delete(&JobRecord{}).Error
}
// TerminalJobStatuses are the statuses a job never leaves.
//
// ONE spelling for the whole package, because the rule was stated at four call
// sites and pinned at none of them: twice in the SSE bridge, which decides when
// to close a stream, and twice here, which decides when to stamp completed_at
// and which rows are still writable. A set that drifts between those two is a
// stream that closes on a status the store still considers open, or a row that
// accepts a second terminal write. Adding a status now means adding it here,
// and every reader follows.
var TerminalJobStatuses = []string{"completed", "failed", "cancelled"}
// IsTerminalJobStatus reports whether a job in this status has finished.
func IsTerminalJobStatus(status string) bool {
for _, terminal := range TerminalJobStatuses {
if status == terminal {
return true
}
}
return false
}
// UpdateJobStatus updates just the status (and optionally result/error) of a job.
func (s *JobStore) UpdateJobStatus(id, status, result, errMsg string) error {
updates := map[string]any{
@@ -245,11 +266,11 @@ func (s *JobStore) UpdateJobStatus(id, status, result, errMsg string) error {
if status == "running" {
updates["started_at"] = &now
}
if status == "completed" || status == "failed" || status == "cancelled" {
if IsTerminalJobStatus(status) {
updates["completed_at"] = &now
}
return s.db.Model(&JobRecord{}).
Where("id = ? AND status NOT IN ?", id, []string{"completed", "failed", "cancelled"}).
Where("id = ? AND status NOT IN ?", id, TerminalJobStatuses).
Updates(updates).Error
}
+47
View File
@@ -402,3 +402,50 @@ var _ = Describe("JobStore", func() {
})
})
})
// The terminal-status set, asserted as ONE fact rather than as two agreeing
// lists.
//
// It was written out at four call sites and pinned at none: twice in the SSE
// bridge, which decides when to close a stream, and twice in the store, which
// decides when to stamp completed_at and which rows are still writable. Drift
// between those two is a stream that closes on a status the store still
// considers open, or a row that accepts a second terminal write. These drive
// both readers from the same exported set, so a status added to one of them
// cannot be missing from the other.
var _ = Describe("the statuses a job never leaves", func() {
var store *JobStore
BeforeEach(func() {
var err error
store, err = NewJobStore(testutil.SetupTestDB())
Expect(err).ToNot(HaveOccurred())
})
It("names at least one status", func() {
// An empty set would make every assertion below vacuous.
Expect(TerminalJobStatuses).ToNot(BeEmpty())
})
It("refuses a further write to a row in any of them", func() {
for _, status := range TerminalJobStatuses {
Expect(IsTerminalJobStatus(status)).To(BeTrue(), status)
job := &JobRecord{TaskID: "t", UserID: "u", Status: "running", TriggeredBy: "manual"}
Expect(store.CreateJob(job)).To(Succeed())
Expect(store.UpdateJobStatus(job.ID, status, "first", "")).To(Succeed())
Expect(store.UpdateJobStatus(job.ID, "running", "second", "")).To(Succeed())
stored, err := store.GetJob(job.ID)
Expect(err).ToNot(HaveOccurred())
Expect(stored.Status).To(Equal(status), "a settled job must not be reopened")
Expect(stored.CompletedAt).ToNot(BeNil(), status)
}
})
It("does not claim a status a job can still leave", func() {
for _, status := range []string{"pending", "running", "", "queued"} {
Expect(IsTerminalJobStatus(status)).To(BeFalse(), status)
}
})
})
+10
View File
@@ -43,6 +43,16 @@ func SubjectAgentEvents(agentName, userID string) string {
return subjectAgentEventsPrefix + sanitizeSubjectToken(agentName) + ".events." + sanitizeSubjectToken(userID)
}
// SubjectAgentEventsWildcard matches every agent's SSE events for every user.
//
// It is a constant here rather than the string literal it used to be inside
// agents.StartObservablePersister, because that literal was the only
// hand-written subject filter left in the tree, and a filter that is not next
// to the builder it must match is a filter that outlives it. SubjectAgentEvents
// makes four tokens; this makes four, and SubjectMatches matches on token count
// first, so a three-token filter here would deliver nothing at all.
const SubjectAgentEventsWildcard = "agent.*.events.*"
// SubjectJobProgress returns the NATS subject for job progress updates.
func SubjectJobProgress(jobID string) string {
return subjectJobProgressPrefix + sanitizeSubjectToken(jobID) + ".progress"
@@ -2,6 +2,7 @@ package messaging
import (
"encoding/json"
"strings"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
@@ -114,3 +115,35 @@ var _ = Describe("per-tenant SyncedMap subjects", func() {
Expect(ValidFilter(SubjectSyncStateTenantDelta("agent.tasks", ""))).To(MatchError(ErrUnsupportedFilter))
})
})
// The agent-events wildcard used to be a hand-written literal at the one place
// that subscribes to it, three files away from the builder it has to match.
// These pin the pairing itself, not the string: the filter is only useful if
// every subject SubjectAgentEvents can produce matches it, and if the near
// misses that a shorter filter would swallow do not.
var _ = Describe("the agent-events wildcard", func() {
It("matches what SubjectAgentEvents builds, for any agent and any user", func() {
for _, agent := range []string{"a", "my-agent", "agent.with.dots"} {
for _, user := range []string{"u", "", "user id"} {
subject := SubjectAgentEvents(agent, user)
Expect(SubjectMatches(SubjectAgentEventsWildcard, subject)).To(BeTrue(),
"filter %q must match %q", SubjectAgentEventsWildcard, subject)
}
}
})
It("has the same token count as the subjects it matches", func() {
// SubjectMatches compares token counts before anything else, so a
// filter one token short matches NOTHING and the persister goes silent
// with no error anywhere. Stated as a count so a builder that grows a
// token reddens here rather than in a suite that only checks delivery.
Expect(strings.Count(SubjectAgentEventsWildcard, ".")).
To(Equal(strings.Count(SubjectAgentEvents("a", "u"), ".")))
})
It("does not match a subject with the events token in another position", func() {
Expect(SubjectMatches(SubjectAgentEventsWildcard, "agent.a.events")).To(BeFalse())
Expect(SubjectMatches(SubjectAgentEventsWildcard, "agent.a.b.events.u")).To(BeFalse())
Expect(SubjectMatches(SubjectAgentEventsWildcard, "agent.a.cancel")).To(BeFalse())
})
})
+7 -3
View File
@@ -18,9 +18,13 @@ import (
// subjects NOT added.
var workerBroadcastAllow = map[string][]string{
NodeTypeAgent: {
"jobs.*.progress",
"jobs.*.result",
"agent.*.events.*",
// The constants and not the strings. A filter written out here is a
// filter that stops matching the day its builder grows a token, and
// SubjectMatches compares token counts first, so the drift presents as
// a worker being refused everything rather than as anything readable.
messaging.SubjectJobProgressWildcard,
messaging.SubjectJobResultWildcard,
messaging.SubjectAgentEventsWildcard,
},
// A backend worker asks for no broadcasts. Spelled as an empty list rather
// than omitted, so the reader sees the decision.
+75
View File
@@ -1,6 +1,7 @@
package nodes
import (
"context"
"encoding/json"
"errors"
@@ -8,6 +9,7 @@ import (
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/pgbus"
"github.com/mudler/LocalAI/core/services/testutil"
)
@@ -151,3 +153,76 @@ var _ = Describe("worker re-broadcast authorization", func() {
})
})
})
// The re-broadcast path across TWO carriers, which is the shape production has
// and the shape a single in-memory double cannot have.
//
// This is the wiring failure that leaves every unit spec in every package green:
// point the Rebroadcaster at one carrier while the dispatcher subscribes on
// another and the rebroadcaster publishes, the publish succeeds, Handle returns
// true, and the only symptom anywhere is an SSE stream with no progress in it.
// So these assert the RECEIPT and never the return value.
var _ = Describe("re-broadcasting onto the carrier the subscriber reads", func() {
var (
busA, busB *pgbus.Bus
rb *Rebroadcaster
)
BeforeEach(func() {
db, dsn := testutil.SetupTestDBWithDSN()
Expect(pgbus.Migrate(context.Background(), db)).To(Succeed())
newBus := func() *pgbus.Bus {
b, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db})
Expect(err).ToNot(HaveOccurred())
DeferCleanup(b.Close)
return b
}
busA, busB = newBus(), newBus()
rb = NewRebroadcaster(busA)
})
It("reaches a job-progress subscriber on the other replica's carrier", func() {
delivered := make(chan []byte, 4)
_, err := busB.Subscribe(messaging.SubjectJobProgressWildcard, func(payload []byte) { delivered <- payload })
Expect(err).ToNot(HaveOccurred())
rb.Handle(NodeTypeAgent, messaging.SubjectJobProgress("j1"), json.RawMessage(`{"job_id":"j1","status":"running"}`))
Eventually(delivered, "20s").Should(Receive(MatchJSON(`{"job_id":"j1","status":"running"}`)))
})
It("reaches a job-result subscriber on the other replica's carrier", func() {
delivered := make(chan []byte, 4)
_, err := busB.Subscribe(messaging.SubjectJobResultWildcard, func(payload []byte) { delivered <- payload })
Expect(err).ToNot(HaveOccurred())
rb.Handle(NodeTypeAgent, messaging.SubjectJobResult("j1"), json.RawMessage(`{"job_id":"j1","status":"completed"}`))
Eventually(delivered, "20s").Should(Receive(MatchJSON(`{"job_id":"j1","status":"completed"}`)))
})
It("reaches an agent-events subscriber on the other replica's carrier", func() {
delivered := make(chan []byte, 4)
_, err := busB.Subscribe(messaging.SubjectAgentEventsWildcard, func(payload []byte) { delivered <- payload })
Expect(err).ToNot(HaveOccurred())
rb.Handle(NodeTypeAgent, messaging.SubjectAgentEvents("a1", "u1"), json.RawMessage(`{"event_type":"json_message"}`))
Eventually(delivered, "20s").Should(Receive(MatchJSON(`{"event_type":"json_message"}`)))
})
It("delivers nothing at all for a subject the worker is refused", func() {
// The negative half without a clock: a refused subject followed by an
// allowed one, and the allowed one arriving first is the proof.
delivered := make(chan []byte, 4)
_, err := busB.Subscribe(messaging.SubjectJobProgressWildcard, func(payload []byte) { delivered <- payload })
Expect(err).ToNot(HaveOccurred())
rb.Handle(NodeTypeBackend, messaging.SubjectJobProgress("j1"), json.RawMessage(`{"job_id":"refused"}`))
rb.Handle(NodeTypeAgent, messaging.SubjectJobProgress("j2"), json.RawMessage(`{"job_id":"allowed"}`))
var first []byte
Eventually(delivered, "20s").Should(Receive(&first))
Expect(first).To(MatchJSON(`{"job_id":"allowed"}`))
})
})
+43 -22
View File
@@ -406,38 +406,59 @@ func (b *Bus) Subscribe(subject string, handler func([]byte)) (messaging.Subscri
done: make(chan struct{}),
}
first := len(b.subs[channel]) == 0
b.mu.Unlock()
// The LISTEN before the registration, not after it.
//
// Both orders satisfy "Subscribe has listened by the time it returns", so
// the CALLER cannot tell them apart. An observer can: with the registration
// first, Subscribers() counts a handler whose channel is not listened yet,
// so anything waiting on that count as a readiness signal proceeds into a
// window where a publish is silently lost, which is at-most-once delivery
// doing exactly what it says while looking like a broken subscription.
// Registering last makes the counter mean what its name says, and it
// removes the failed-LISTEN rollback: a registration never made needs no
// undoing.
if first {
b.barrier("issue", "LISTEN")
if err := b.command("LISTEN " + pgx.Identifier{channel}.Sanitize()); err != nil {
close(sub.stop)
return nil, err
}
}
b.mu.Lock()
if b.subs[channel] == nil {
b.subs[channel] = map[uint64]*subscription{}
}
b.subs[channel][sub.id] = sub
b.mu.Unlock()
go sub.run(handler)
if first {
b.barrier("issue", "LISTEN")
if err := b.command("LISTEN " + pgx.Identifier{channel}.Sanitize()); err != nil {
// Not Unsubscribe: the ordering lock is already held here, and the
// connection was never listening on this channel, so there is
// nothing to UNLISTEN.
b.forget(sub)
return nil, err
}
}
return sub, nil
}
// forget removes a registration without touching the channel's LISTEN state.
func (b *Bus) forget(sub *subscription) {
sub.once.Do(func() {
b.mu.Lock()
delete(b.subs[sub.channel], sub.id)
if len(b.subs[sub.channel]) == 0 {
delete(b.subs, sub.channel)
}
b.mu.Unlock()
close(sub.stop)
})
// Subscribers reports how many handlers are registered right now, across every
// channel. A counted handler is a LIVE one: its channel was listened before it
// was counted, so this number is usable as a readiness signal and not only as a
// figure to compare against itself.
//
// It exists to be asserted, because the leak it makes visible has no other
// symptom. Two subscriptions in this deployment are opened and closed PER HTTP
// REQUEST (a job's progress stream and an agent's event stream), and on this
// carrier only the FIRST subscriber of a channel issues a LISTEN: every later
// one just registers an in-process filter. An Unsubscribe that failed to remove
// its filter would therefore leave a replica that has served ten thousand SSE
// requests running ten thousand closures per notification, and nothing in the
// tree would fail. It would just get slower.
func (b *Bus) Subscribers() int {
b.mu.Lock()
defer b.mu.Unlock()
n := 0
for _, channel := range b.subs {
n += len(channel)
}
return n
}
// barrier is a test seam and does nothing in production.
+75
View File
@@ -191,6 +191,81 @@ var _ = Describe("the PostgreSQL broadcast carrier", func() {
})
})
Describe("subscriptions that come and go", func() {
// The accumulation assertion, and it is deliberately NOT a delivery
// one. Two subscriptions in this deployment are opened and closed per
// HTTP REQUEST, so a filter that outlives its Unsubscribe is a replica
// that gets steadily slower and never fails. Delivery specs cannot see
// it: the leaked handlers deliver correctly, there are just more of
// them every request.
It("registers nothing that outlives its Unsubscribe", func() {
delivered := make(chan []byte, 8)
survivor, err := sub.Subscribe("jobs.churn.progress", func(data []byte) { delivered <- data })
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() { _ = survivor.Unsubscribe() })
start := sub.Subscribers()
Expect(start).To(Equal(1))
for i := 0; i < 1000; i++ {
s, err := sub.Subscribe("jobs.churn.progress", func([]byte) {})
Expect(err).ToNot(HaveOccurred())
Expect(s.Unsubscribe()).To(Succeed())
}
Expect(sub.Subscribers()).To(Equal(start),
"a subscription that is opened and closed must leave the register where it found it")
// And the survivor is still exactly one subscriber, not zero and
// not a thousand. A counter rather than a boolean, because
// "something arrived" is true for every wrong answer here.
Expect(pub.Publish("jobs.churn.progress", map[string]string{"m": "after"})).To(Succeed())
Eventually(delivered).Should(Receive(MatchJSON(`{"m":"after"}`)))
Consistently(delivered, "500ms", "50ms").ShouldNot(Receive(),
"a leaked handler would deliver the same notification again")
})
It("keeps the channel listened while a subscriber remains", func() {
// The churn above cycles LISTEN and UNLISTEN a thousand times on a
// channel that still has a live subscriber. If the refcount and the
// UNLISTEN it implies were not decided as one step, the channel
// ends up unlistened with a live subscription on it, and it does
// not self-heal: the next Subscribe sees first == false and never
// re-LISTENs.
out, _ := received(sub, "jobs.abc.progress")
for i := 0; i < 50; i++ {
s, err := sub.Subscribe("jobs.other.progress", func([]byte) {})
Expect(err).ToNot(HaveOccurred())
Expect(s.Unsubscribe()).To(Succeed())
}
Expect(pub.Publish("jobs.abc.progress", map[string]string{"m": "still here"})).To(Succeed())
Eventually(out).Should(Receive(MatchJSON(`{"m":"still here"}`)))
})
})
Describe("what Subscribe has finished doing when it returns", func() {
// The per-request subscriptions are why this matters. An SSE client
// connects and then triggers the work; if Subscribe returns before the
// LISTEN is on the wire, the first events go nowhere and the stream
// looks hung with no error on either side.
//
// There is deliberately no Eventually around the PUBLISH. Retrying the
// publish is exactly the thing that would hide an asynchronous LISTEN.
It("has already issued the LISTEN, every time", func() {
for i := 0; i < 50; i++ {
out := make(chan []byte, 1)
s, err := sub.Subscribe("jobs.race.progress", func(data []byte) { out <- data })
Expect(err).ToNot(HaveOccurred())
Expect(pub.Publish("jobs.race.progress", map[string]int{"i": i})).To(Succeed())
Eventually(out, "10s").Should(Receive(), "iteration %d subscribed and then missed its own first message", i)
Expect(s.Unsubscribe()).To(Succeed())
}
})
})
Describe("a handler that does not return", func() {
It("does not stop delivery to the other subscribers", func() {
// A carrier whose handlers run on one goroutine is a carrier one
+20 -1
View File
@@ -134,12 +134,31 @@ Several features keep state in a frontend's process memory and surface it over t
| Agent tasks | `state.agent-tasks.delta` and `state.agent-tasks.<user_id>.delta` |
| Open Responses metadata | `state.responses-metadata.delta` |
### Job and agent streams across replicas
The same carrier moves the traffic whose subscriber is an open HTTP response rather than a cache: a job's progress stream, its result, its cancel, an agent's SSE events, an agent cancel and an Open Responses cancel.
| Family | Subject | Read by |
|--------|---------|---------|
| Job progress | `jobs.<job_id>.progress` | `GET /api/agent/jobs/{id}/progress` on any replica, and the trace persister on every replica |
| Job result | `jobs.<job_id>.result` | The result persister on every replica |
| Job cancel | `jobs.<job_id>.cancel` | Every replica, so the one holding the run can stop it |
| Agent events | `agent.<agent>.events.<user_id>` | `GET /api/agents/{name}/sse/distributed` on any replica, and the observable persister |
| Agent cancel | `agent.<agent>.cancel` | Every agent worker, and so still on NATS: see below |
| Open Responses cancel | `responses.<response_id>.cancel` | The replica holding the generation |
One family on that list is not on this carrier. An `agent.<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.
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.
Two carrier details are visible to an operator.
**The 8000-byte notification cap.** PostgreSQL refuses a `pg_notify` payload of 8000 bytes or more, and that limit is measured against the whole encoded notification, not just the value being replicated. A broadcast that does not fit is written to the `bus_messages` table and the notification carries the row id instead; the receiving replica reads the row and delivers the original bytes. This is an ordinary path and not an error: a fine-tune job carrying a long training message spills every time. Rows are retired ten minutes after they are written, by every replica, so `bus_messages` is a spill buffer and never a log of past events.
**A broadcast is at most once, and is never replayed.** A `NOTIFY` reaches the sessions that are listening when it is issued and nobody else. A replica whose session was down in that window never receives the change, and no error is reported anywhere. That is why every one of these maps is backed by a durable table: the broadcast says only that something changed, and the table says what it changed to. A replica re-reads its table when its listener reconnects, so a missed delta is a delay and never a value that reads as though it had never been set.
**A slow subscriber loses broadcasts rather than stalling the carrier.** Each subscription buffers 256 messages; past that, its broadcasts are dropped and logged at error level on the replica that took them. One blocked SSE writer must not be able to stop delivery for the whole deployment, which is what the alternative would mean. The same rule follows from it: what must survive a gap lives in a table.
### Open Responses across replicas
A response created by `POST /v1/responses` is held by the replica that served the request. A round-robin load balancer sends the follow-up poll, the `previous_response_id` chain and the cancel to any replica, so that metadata is replicated to every frontend and is also written to a `response_metadata` table in PostgreSQL.
@@ -210,7 +229,7 @@ Both **backend** and **agent** nodes are issued one. Earlier releases minted a c
An agent worker's tunnel carries only the `http` tag: it runs no backend processes, so it does not offer the `grpc` tag at all. Its control server binds `127.0.0.1` on a port chosen by the kernel and advertises it nowhere, so an agent worker still opens no inbound port.
**An agent worker still requires `--nats-url`, and nothing the frontend sends it travels 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. What the bus still carries for an agent worker is the other direction and the broadcasts: cancellation, and the progress and result lines a worker asks the frontend to re-publish on its behalf.
**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.<agent>.cancel`, which the worker subscribes to so a cancel can reach the execution it is running.
There is no `nodes.<id>.*` subject left, and an agent worker's minted JWT no longer grants `mcp.tools.execute`, `mcp.discovery`, `nodes.<id>.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.
+38 -25
View File
@@ -106,36 +106,49 @@ var _ = Describe("Phase 3: Agent Conversations & SSE", Label("Distributed"), fun
// Conversation history is managed client-side (browser localStorage).
// No server-side conversation storage tests needed.
Context("Agent SSE Events via NATS", func() {
It("should bridge agent SSE events via NATS", func() {
bridge := agents.NewEventBridge(infra.NC, store, "instance-1")
Context("Agent SSE events on the broadcast carrier", func() {
It("bridges an agent's SSE events from a peer replica's carrier", func() {
// 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")
var received []agents.AgentEvent
sub, err := bridge.SubscribeEvents("my-agent", "user1", func(evt agents.AgentEvent) {
received = append(received, evt)
received := make(chan agents.AgentEvent, 16)
sub, err := watcher.SubscribeEvents("my-agent", "user1", func(evt agents.AgentEvent) {
received <- evt
})
Expect(err).ToNot(HaveOccurred())
defer sub.Unsubscribe()
FlushNATS(infra.NC)
// Published on the OTHER replica, as an agent execution there would.
Expect(runner.PublishMessage("my-agent", "user1", "user", "What's the weather?", "msg-1")).To(Succeed())
Expect(runner.PublishStatus("my-agent", "user1", "processing")).To(Succeed())
Expect(runner.PublishMessage("my-agent", "user1", "agent", "The weather is sunny.", "msg-2")).To(Succeed())
Expect(runner.PublishStatus("my-agent", "user1", "completed")).To(Succeed())
// Publish events (simulating agent execution on another instance)
bridge.PublishMessage("my-agent", "user1", "user", "What's the weather?", "msg-1")
bridge.PublishStatus("my-agent", "user1", "processing")
bridge.PublishMessage("my-agent", "user1", "agent", "The weather is sunny.", "msg-2")
bridge.PublishStatus("my-agent", "user1", "completed")
Eventually(func() int { return len(received) }, "5s").Should(Equal(4))
Expect(received[0].EventType).To(Equal("json_message"))
Expect(received[0].Sender).To(Equal("user"))
Expect(received[1].EventType).To(Equal("json_message_status"))
Expect(received[2].Sender).To(Equal("agent"))
var evts []agents.AgentEvent
for i := 0; i < 4; i++ {
var evt agents.AgentEvent
Eventually(received, "20s").Should(Receive(&evt))
evts = append(evts, evt)
}
Expect(evts[0].EventType).To(Equal("json_message"))
Expect(evts[0].Sender).To(Equal("user"))
Expect(evts[1].EventType).To(Equal("json_message_status"))
Expect(evts[2].Sender).To(Equal("agent"))
})
// Conversation persistence removed — chat history is browser-only.
It("should cancel running agent via NATS", func() {
bridge := agents.NewEventBridge(infra.NC, store, "instance-1")
// 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")
// Start cancel listener
cancelSub, err := bridge.StartCancelListener()
@@ -151,12 +164,12 @@ var _ = Describe("Phase 3: Agent Conversations & SSE", Label("Distributed"), fun
})
bridge.RegisterCancel("test-msg-id", wrappedCancel)
FlushNATS(infra.NC)
// 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())
// Cancel via NATS
Expect(bridge.CancelExecution("my-agent", "user1", "test-msg-id")).To(Succeed())
Eventually(func() bool { return cancelled.Load() }, "5s").Should(BeTrue())
Eventually(func() bool { return cancelled.Load() }, "20s").Should(BeTrue())
})
// Agent execution is now dispatched via AgentPoolService.dispatchChat(),
+11 -6
View File
@@ -53,7 +53,7 @@ var _ = Describe("Job Dispatch", Label("Distributed"), func() {
// not be told from ones a dead replica left.
Expect(cluster.NewRegistry(db).Register(infra.Ctx, owner, "127.0.0.1:8080", "v1")).To(Succeed())
dispatcher := jobs.NewDispatcher(store, infra.NC, db, owner)
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, owner)
task := &jobs.TaskRecord{UserID: "u1", Name: "dispatch-task", Model: "m1", Prompt: "p1"}
store.CreateTask(task)
@@ -100,7 +100,7 @@ var _ = Describe("Job Dispatch", Label("Distributed"), func() {
const owner = "plain-instance"
Expect(cluster.NewRegistry(db).Register(infra.Ctx, owner, "127.0.0.1:8081", "v1")).To(Succeed())
dispatcher := jobs.NewDispatcher(store, infra.NC, db, owner)
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, owner)
task := &jobs.TaskRecord{UserID: "u1", Name: "plain-task", Model: "m1", Prompt: "p1"}
Expect(store.CreateTask(task)).To(Succeed())
job := &jobs.JobRecord{TaskID: task.ID, UserID: "u1", Status: "pending", TriggeredBy: "api"}
@@ -158,12 +158,18 @@ var _ = Describe("Job Dispatch", Label("Distributed"), func() {
})
})
Context("NATS job cancellation", func() {
Context("job cancellation", func() {
// Cancellation stays a BROADCAST and is not part of the claim queue: the
// replica holding a run is not the one an API cancel lands on, so the
// signal has to reach every replica and every worker.
//
// Asserted across TWO carriers, because one carrier hearing itself
// proves nothing about the replica that actually holds the execution.
// A cancel that does not arrive is not a cancel that was refused, so
// what is pinned here is arrival and never the publisher's error.
It("broadcasts a cancel for a job on the job's own cancel subject", func() {
dispatcher := jobs.NewDispatcher(store, infra.NC, db, "cancel-instance")
publisher, listener := infra.Bus(), infra.Bus()
dispatcher := jobs.NewDispatcher(store, publisher, db, "cancel-instance")
task := &jobs.TaskRecord{UserID: "u1", Name: "cancel-task", Model: "m1", Prompt: "p1"}
store.CreateTask(task)
@@ -171,7 +177,7 @@ var _ = Describe("Job Dispatch", Label("Distributed"), func() {
store.CreateJob(job)
seen := make(chan string, 1)
sub, err := infra.NC.Subscribe(messaging.SubjectJobCancelWildcard, func(data []byte) {
sub, err := listener.Subscribe(messaging.SubjectJobCancelWildcard, func(data []byte) {
var evt jobs.CancelEvent
if json.Unmarshal(data, &evt) == nil {
select {
@@ -182,7 +188,6 @@ var _ = Describe("Job Dispatch", Label("Distributed"), func() {
})
Expect(err).ToNot(HaveOccurred())
defer func() { _ = sub.Unsubscribe() }()
FlushNATS(infra.NC)
Expect(dispatcher.Cancel(job.ID)).To(Succeed())
Eventually(seen, "10s").Should(Receive(Equal(job.ID)))
+21 -21
View File
@@ -176,7 +176,7 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
const owner = "test-instance"
Expect(cluster.NewRegistry(db).Register(infra.Ctx, owner, "127.0.0.1:8090", "v1")).To(Succeed())
dispatcher := jobs.NewDispatcher(store, infra.NC, db, owner)
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, owner)
task := &jobs.TaskRecord{UserID: "u1", Name: "dispatch-test", Model: "m1", Prompt: "p1"}
store.CreateTask(task)
@@ -204,7 +204,7 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
const owner = "lossy-instance"
Expect(cluster.NewRegistry(db).Register(infra.Ctx, owner, "127.0.0.1:8091", "v1")).To(Succeed())
dispatcher := jobs.NewDispatcher(store, infra.NC, db, owner)
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, owner)
task := &jobs.TaskRecord{UserID: "u1", Name: "lossy-test", Model: "m1", Prompt: "p1"}
store.CreateTask(task)
job := &jobs.JobRecord{TaskID: task.ID, UserID: "u1", Status: "pending", TriggeredBy: "api"}
@@ -231,9 +231,13 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
})
It("broadcasts a cancel on the job's own cancel subject", func() {
dispatcher := jobs.NewDispatcher(store, infra.NC, db, "test-instance")
// Two carriers, because the replica that holds the execution is
// never the one an API cancel lands on. A cancel that does not
// arrive is not a cancel that was refused, so this pins ARRIVAL.
publisher, listener := infra.Bus(), infra.Bus()
dispatcher := jobs.NewDispatcher(store, publisher, db, "test-instance")
seen := make(chan string, 1)
sub, err := infra.NC.Subscribe(messaging.SubjectJobCancelWildcard, func(data []byte) {
sub, err := listener.Subscribe(messaging.SubjectJobCancelWildcard, func(data []byte) {
var evt jobs.CancelEvent
if json.Unmarshal(data, &evt) == nil {
select {
@@ -244,14 +248,13 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
})
Expect(err).ToNot(HaveOccurred())
defer func() { _ = sub.Unsubscribe() }()
FlushNATS(infra.NC)
Expect(dispatcher.Cancel("job-to-cancel")).To(Succeed())
Eventually(seen, "10s").Should(Receive(Equal("job-to-cancel")))
})
It("reports job progress on the job's own progress subject", func() {
dispatcher := jobs.NewDispatcher(store, infra.NC, db, "test-instance")
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, "test-instance")
var progressEvents []jobs.ProgressEvent
var mu sync.Mutex
@@ -262,7 +265,6 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
})
Expect(err).ToNot(HaveOccurred())
defer func() { _ = sub.Unsubscribe() }()
FlushNATS(infra.NC)
Expect(dispatcher.PublishProgress("progress-job", "running", "step 1")).To(Succeed())
Expect(dispatcher.PublishProgress("progress-job", "running", "step 2")).To(Succeed())
@@ -318,9 +320,9 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
})
})
Context("Progress Streaming (NATS → SSE bridge)", func() {
It("should bridge NATS progress events", func() {
dispatcher := jobs.NewDispatcher(store, infra.NC, db, "test-instance")
Context("Progress streaming to the SSE bridge", func() {
It("bridges progress events to a per-job subscription", func() {
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, "test-instance")
dCtx, dCancel := context.WithCancel(infra.Ctx)
defer dCancel()
@@ -335,8 +337,6 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
Expect(err).ToNot(HaveOccurred())
defer func() { _ = sub.Unsubscribe() }()
FlushNATS(infra.NC)
// Publish progress events
dispatcher.PublishProgress("job-123", "running", "processing")
dispatcher.PublishProgress("job-123", "running", "almost done")
@@ -348,7 +348,7 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
})
It("should filter SSE events by job ID", func() {
dispatcher := jobs.NewDispatcher(store, infra.NC, db, "test-instance")
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, "test-instance")
dCtx, dCancel := context.WithCancel(infra.Ctx)
defer dCancel()
@@ -362,8 +362,6 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
})
defer subA.Unsubscribe()
FlushNATS(infra.NC)
// Publish to both job-A and job-B
dispatcher.PublishProgress("job-A", "running", "A progress")
dispatcher.PublishProgress("job-B", "running", "B progress")
@@ -379,7 +377,7 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
Context("Enriched claim payload (DB-free worker)", func() {
It("stores the full Job and Task on the claim row, so the worker needs no database", func() {
dispatcher := jobs.NewDispatcher(store, infra.NC, db, "enrichment-test")
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, "enrichment-test")
task := &jobs.TaskRecord{UserID: "u1", Name: "enrich-task", Model: "m1", Prompt: "hello {{.name}}"}
store.CreateTask(task)
@@ -406,7 +404,8 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
// every replica persists them, because the SSE stream a user is watching
// may be open on a replica that claimed nothing.
It("persists a result a worker's re-broadcast carried, whichever replica reads it", func() {
dispatcher := jobs.NewDispatcher(store, infra.NC, db, "result-test")
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, "result-test")
peer := infra.Bus()
dCtx, dCancel := context.WithCancel(infra.Ctx)
defer dCancel()
Expect(dispatcher.Start(dCtx)).To(Succeed())
@@ -416,9 +415,11 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
store.CreateTask(task)
job := &jobs.JobRecord{TaskID: task.ID, UserID: "u1", Status: "running", TriggeredBy: "api"}
store.CreateJob(job)
FlushNATS(infra.NC)
jobs.PublishJobResult(infra.NC, job.ID, "completed", "job finished successfully", "")
// Published by a PEER replica's carrier: this is the fan-out copy of
// a terminal line the claiming replica already persisted, and the
// replica asserted on here claimed nothing.
jobs.PublishJobResult(peer, job.ID, "completed", "job finished successfully", "")
Eventually(func() string {
j, _ := store.GetJob(job.ID)
@@ -430,7 +431,7 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
})
It("appends a trace a worker's re-broadcast carried", func() {
dispatcher := jobs.NewDispatcher(store, infra.NC, db, "trace-test")
dispatcher := jobs.NewDispatcher(store, infra.Bus(), db, "trace-test")
dCtx, dCancel := context.WithCancel(infra.Ctx)
defer dCancel()
Expect(dispatcher.Start(dCtx)).To(Succeed())
@@ -440,7 +441,6 @@ var _ = Describe("Phase 2: Jobs & Tasks", Label("Distributed"), func() {
store.CreateTask(task)
job := &jobs.JobRecord{TaskID: task.ID, UserID: "u1", Status: "running", TriggeredBy: "api"}
store.CreateJob(job)
FlushNATS(infra.NC)
Expect(dispatcher.PublishTrace(job.ID, "reasoning", "thinking about the problem")).To(Succeed())
Expect(dispatcher.PublishTrace(job.ID, "tool_call", "calling search tool")).To(Succeed())
+58 -31
View File
@@ -32,63 +32,90 @@ var _ = Describe("SSE Routes", Label("Distributed"), func() {
})
Context("Job progress SSE endpoint", func() {
It("should register job progress SSE endpoint when dispatcher active", func() {
// Frontend 0 publishes, frontend 1 is where the SSE reader is attached.
// One carrier hearing itself would pass with the dispatcher wired to
// any carrier at all, which is the wiring defect that leaves every unit
// spec green and only the stream empty.
It("delivers a job's progress to a reader attached to another frontend", func() {
jobStore, err := jobs.NewJobStore(db)
Expect(err).ToNot(HaveOccurred())
dispatcher := jobs.NewDispatcher(jobStore, infra.NC, db, "sse-instance")
frontend0 := jobs.NewDispatcher(jobStore, infra.Bus(), db, "frontend-0")
frontend1 := jobs.NewDispatcher(jobStore, infra.Bus(), db, "frontend-1")
dCtx, dCancel := context.WithCancel(infra.Ctx)
defer dCancel()
Expect(dispatcher.Start(dCtx)).To(Succeed())
defer dispatcher.Stop()
Expect(frontend1.Start(dCtx)).To(Succeed())
defer frontend1.Stop()
// Subscribe to progress for a job — verifies the dispatcher can bridge
// NATS progress events that an SSE endpoint would consume
var events []jobs.ProgressEvent
sub, err := dispatcher.SubscribeProgress("job-sse-test", func(evt jobs.ProgressEvent) {
events = append(events, evt)
mine := make(chan jobs.ProgressEvent, 16)
sub, err := frontend1.SubscribeProgress("job-sse-test", func(evt jobs.ProgressEvent) {
mine <- evt
})
Expect(err).ToNot(HaveOccurred())
defer sub.Unsubscribe()
FlushNATS(infra.NC)
// A reader on a DIFFERENT job, which must be shown nothing. The
// per-request subscription is a data boundary: on the wildcard,
// every client watching any job sees every other job.
theirs := make(chan jobs.ProgressEvent, 16)
otherSub, err := frontend1.SubscribeProgress("job-someone-else", func(evt jobs.ProgressEvent) {
theirs <- evt
})
Expect(err).ToNot(HaveOccurred())
defer func() { _ = otherSub.Unsubscribe() }()
dispatcher.PublishProgress("job-sse-test", "running", "step 1")
dispatcher.PublishProgress("job-sse-test", "running", "step 2")
dispatcher.PublishProgress("job-sse-test", "completed", "done")
Expect(frontend0.PublishProgress("job-sse-test", "running", "step 1")).To(Succeed())
Expect(frontend0.PublishProgress("job-sse-test", "running", "step 2")).To(Succeed())
Expect(frontend0.PublishProgress("job-sse-test", "completed", "done")).To(Succeed())
Eventually(func() int { return len(events) }, "5s").Should(Equal(3))
Expect(events[0].Status).To(Equal("running"))
Expect(events[2].Status).To(Equal("completed"))
var seen []jobs.ProgressEvent
for i := 0; i < 3; i++ {
var evt jobs.ProgressEvent
Eventually(mine, "20s").Should(Receive(&evt))
seen = append(seen, evt)
}
Expect(seen[0].Status).To(Equal("running"))
Expect(seen[2].Status).To(Equal("completed"))
Expect(theirs).ToNot(Receive(), "a reader on another job id must be shown nothing")
})
})
Context("Agent SSE endpoint", func() {
It("should register agent SSE endpoint when event bridge active", func() {
It("delivers an agent's events to a reader attached to another frontend", func() {
agentStore, err := agents.NewAgentStore(db)
Expect(err).ToNot(HaveOccurred())
bridge := agents.NewEventBridge(infra.NC, agentStore, "sse-instance")
frontend0 := agents.NewEventBridge(infra.Bus(), agentStore, "frontend-0")
frontend1 := agents.NewEventBridge(infra.Bus(), agentStore, "frontend-1")
// Subscribe to agent events — verifies the bridge can deliver
// NATS events that an SSE endpoint would consume
var received []agents.AgentEvent
sub, err := bridge.SubscribeEvents("test-agent", "user1", func(evt agents.AgentEvent) {
received = append(received, evt)
received := make(chan agents.AgentEvent, 16)
sub, err := frontend1.SubscribeEvents("test-agent", "user1", func(evt agents.AgentEvent) {
received <- evt
})
Expect(err).ToNot(HaveOccurred())
defer sub.Unsubscribe()
FlushNATS(infra.NC)
other := make(chan agents.AgentEvent, 16)
otherSub, err := frontend1.SubscribeEvents("test-agent", "user2", func(evt agents.AgentEvent) {
other <- evt
})
Expect(err).ToNot(HaveOccurred())
defer func() { _ = otherSub.Unsubscribe() }()
bridge.PublishMessage("test-agent", "user1", "user", "Hello", "msg-1")
bridge.PublishStatus("test-agent", "user1", "processing")
bridge.PublishMessage("test-agent", "user1", "agent", "Hi!", "msg-2")
Expect(frontend0.PublishMessage("test-agent", "user1", "user", "Hello", "msg-1")).To(Succeed())
Expect(frontend0.PublishStatus("test-agent", "user1", "processing")).To(Succeed())
Expect(frontend0.PublishMessage("test-agent", "user1", "agent", "Hi!", "msg-2")).To(Succeed())
Eventually(func() int { return len(received) }, "5s").Should(Equal(3))
Expect(received[0].EventType).To(Equal("json_message"))
Expect(received[1].EventType).To(Equal("json_message_status"))
var seen []agents.AgentEvent
for i := 0; i < 3; i++ {
var evt agents.AgentEvent
Eventually(received, "20s").Should(Receive(&evt))
seen = append(seen, evt)
}
Expect(seen[0].EventType).To(Equal("json_message"))
Expect(seen[1].EventType).To(Equal("json_message_status"))
Expect(other).ToNot(Receive(), "another user's stream must be shown nothing")
})
})
@@ -97,7 +124,7 @@ var _ = Describe("SSE Routes", Label("Distributed"), func() {
appCfg := config.NewApplicationConfig()
Expect(appCfg.Distributed.Enabled).To(BeFalse())
// Without distributed mode, NATS-backed SSE routes are not registered.
// Without distributed mode, carrier-backed SSE routes are not registered.
// Agent SSE events use the in-process LocalAGI SSE manager instead.
// Job progress is tracked in-memory.
Expect(appCfg.Distributed.NatsURL).To(BeEmpty())
+26
View File
@@ -9,6 +9,7 @@ import (
"time"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/pgbus"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
@@ -236,3 +237,28 @@ func FlushNATS(nc *messaging.Client) {
GinkgoHelper()
Expect(nc.Conn().Flush()).To(Succeed())
}
// Bus opens a broadcast carrier on THIS spec's database.
//
// It is the carrier the job, agent and response families travel on, and it is
// what these specs must build their dispatchers and bridges with. Publishing on
// one carrier while the subscriber reads another is a defect with no error
// anywhere: the publish succeeds and the SSE stream is simply empty, so a spec
// that used the NATS client here would keep passing after production had gone
// silent.
//
// Every call returns a SEPARATE carrier on the same database, so a spec can
// build two and assert across them, which is the shape a deployment has.
func (i *TestInfra) Bus() *pgbus.Bus {
GinkgoHelper()
Expect(i.PGURL).ToNot(BeEmpty(), "Bus needs a database; use SetupInfra rather than SetupNATSOnly")
db, err := gorm.Open(postgres.Open(i.PGURL), &gorm.Config{Logger: gormlogger.Default.LogMode(gormlogger.Silent)})
Expect(err).ToNot(HaveOccurred())
Expect(pgbus.Migrate(i.Ctx, db)).To(Succeed())
bus, err := pgbus.New(i.Ctx, pgbus.Config{DSN: i.PGURL, DB: db})
Expect(err).ToNot(HaveOccurred())
DeferCleanup(bus.Close)
return bus
}