diff --git a/core/services/agents/stream_publisher.go b/core/services/agents/stream_publisher.go new file mode 100644 index 000000000..f3f2b6978 --- /dev/null +++ b/core/services/agents/stream_publisher.go @@ -0,0 +1,74 @@ +package agents + +import ( + "encoding/json" + "fmt" + "io" + "sync" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// StreamPublisher is a messaging.Publisher that writes progress lines onto an +// in-flight control response instead of onto a bus. +// +// It exists because a worker has no database and therefore cannot NOTIFY, and +// because it does not need to: every message a worker sends is produced inside +// a handler the frontend invoked, so there is always an open response to write +// on. +// +// A subject written here is a REQUEST and not a publish. Nothing on this side +// is authorised: the frontend reading the stream checks the subject against the +// allow list for the node's type and refuses what the worker has no business +// on. That refusal is silent to this end on purpose, since a worker learning +// which subjects a frontend will carry is the beginning of a worker that probes +// for them. +// +// Writes are serialized: cogito's stream callbacks run on the goroutine that is +// executing the agent, but the status callback and the tool-result callback can +// interleave with a flush, and two concurrent json.Encoder writes on one +// http.ResponseWriter interleave bytes and corrupt the NDJSON framing. +type StreamPublisher struct { + mu sync.Mutex + enc *json.Encoder + flush func() +} + +// Compile-time proof that this is the Publisher a caller can hand to anything +// that publishes. Asserted here rather than at the first adopter so a change to +// either side fails to build in the file that owns the type. +var _ messaging.Publisher = (*StreamPublisher)(nil) + +// NewStreamPublisher returns a StreamPublisher writing to w. flush may be nil, +// which is the right shape for a writer that has no buffering to push through. +func NewStreamPublisher(w io.Writer, flush func()) *StreamPublisher { + return &StreamPublisher{enc: json.NewEncoder(w), flush: flush} +} + +// Publish writes one Envelope carrying subject and the JSON encoding of data, +// then flushes, so a progress tick reaches the frontend while the handler is +// still running rather than when the response closes. +// +// data is encoded OUTSIDE the lock. Encoding is the expensive half and it +// touches nothing shared, so holding the lock across it would serialise every +// caller behind the slowest one's marshalling for no gain in framing. +func (p *StreamPublisher) Publish(subject string, data any) error { + raw, err := json.Marshal(data) + if err != nil { + return fmt.Errorf("encoding a stream message for %q: %w", subject, err) + } + + p.mu.Lock() + defer p.mu.Unlock() + if err := p.enc.Encode(workerctl.Envelope{Subject: subject, Progress: raw}); err != nil { + return fmt.Errorf("writing a stream message for %q: %w", subject, err) + } + // The flush is inside the lock and belongs there. A flush concurrent with + // an Encode pushes a partially written line at the reader, which is the + // same torn frame the lock exists to prevent. + if p.flush != nil { + p.flush() + } + return nil +} diff --git a/core/services/agents/stream_publisher_test.go b/core/services/agents/stream_publisher_test.go new file mode 100644 index 000000000..4932a45a4 --- /dev/null +++ b/core/services/agents/stream_publisher_test.go @@ -0,0 +1,166 @@ +package agents + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "runtime" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// chunkWriter is an http.ResponseWriter's concurrency contract and nothing +// else: it is NOT safe for concurrent use, and it splits every Write into +// single bytes with a scheduling point between them. +// +// Both halves are deliberate. A double that locked internally would make two +// unsynchronised Encode calls look atomic and could never fail the way a real +// response fails, which is the failure this spec exists to catch. Splitting the +// write is what turns "two goroutines might interleave" into "two goroutines +// do", so the concurrency spec reddens on the framing itself and not only under +// the race detector. +type chunkWriter struct { + buf []byte +} + +func (w *chunkWriter) Write(p []byte) (int, error) { + for _, b := range p { + w.buf = append(w.buf, b) + runtime.Gosched() + } + return len(p), nil +} + +// errWriter fails every write, for the one thing Publish promises about a +// response it cannot write on. +type errWriter struct{ err error } + +func (w errWriter) Write([]byte) (int, error) { return 0, w.err } + +// decodeLines reads every envelope out of a buffer, which is the assertion that +// fails on interleaving: a torn line does not decode, and a line that swallowed +// its neighbour does not produce two. +func decodeLines(raw []byte) ([]workerctl.Envelope, error) { + dec := json.NewDecoder(bytes.NewReader(raw)) + var out []workerctl.Envelope + for { + var env workerctl.Envelope + err := dec.Decode(&env) + if errors.Is(err, io.EOF) { + return out, nil + } + if err != nil { + return out, err + } + out = append(out, env) + } +} + +var _ = Describe("StreamPublisher", func() { + var ( + buf *bytes.Buffer + flushes int + pub *StreamPublisher + ) + + BeforeEach(func() { + buf = &bytes.Buffer{} + flushes = 0 + pub = NewStreamPublisher(buf, func() { flushes++ }) + }) + + It("writes exactly one NDJSON line whose subject and progress decode back to the inputs", func() { + Expect(pub.Publish("agent.a1.events.status", map[string]any{"state": "thinking"})).To(Succeed()) + + lines, err := decodeLines(buf.Bytes()) + Expect(err).ToNot(HaveOccurred()) + Expect(lines).To(HaveLen(1)) + Expect(lines[0].Subject).To(Equal("agent.a1.events.status")) + Expect(string(lines[0].Progress)).To(MatchJSON(`{"state":"thinking"}`)) + // A line the frontend would read as terminal ends the stream, so a + // progress publisher must never write one. + Expect(lines[0].Reply).To(BeNil()) + Expect(bytes.Count(buf.Bytes(), []byte("\n"))).To(Equal(1)) + }) + + It("flushes after every line, so a tick reaches the frontend while the handler still runs", func() { + Expect(pub.Publish("jobs.j1.progress", map[string]any{"percentage": 10})).To(Succeed()) + Expect(flushes).To(Equal(1)) + Expect(pub.Publish("jobs.j1.progress", map[string]any{"percentage": 20})).To(Succeed()) + Expect(flushes).To(Equal(2)) + }) + + It("writes a line at all when it was given no flush to call", func() { + plain := &bytes.Buffer{} + Expect(NewStreamPublisher(plain, nil).Publish("jobs.j1.result", map[string]any{"ok": true})).To(Succeed()) + lines, err := decodeLines(plain.Bytes()) + Expect(err).ToNot(HaveOccurred()) + Expect(lines).To(HaveLen(1)) + }) + + It("keeps every concurrent publish a complete, decodable line of its own", func() { + w := &chunkWriter{} + p := NewStreamPublisher(w, func() {}) + + // Four writers rather than two, and enough lines each that an + // unsynchronised run has to interleave rather than merely be allowed + // to. Measured: at two writers of twenty lines an unlocked publisher + // still produced a decodable buffer on some schedules, so the framing + // assertion below was a coin toss and only the race detector was + // reliably red. + subjects := []string{ + "agent.a1.events.status", + "agent.a2.events.status", + "jobs.j1.progress", + "jobs.j2.result", + } + const perGoroutine = 60 + var wg sync.WaitGroup + errs := make(chan error, len(subjects)*perGoroutine) + for _, subject := range subjects { + wg.Add(1) + go func(subject string) { + defer wg.Done() + for i := 0; i < perGoroutine; i++ { + if err := p.Publish(subject, map[string]any{"n": i, "subject": subject}); err != nil { + errs <- err + } + } + }(subject) + } + wg.Wait() + close(errs) + Expect(errs).ToNot(Receive()) + + lines, err := decodeLines(w.buf) + Expect(err).ToNot(HaveOccurred(), "a line was torn by an interleaved write") + Expect(lines).To(HaveLen(len(subjects) * perGoroutine)) + // Every line must still carry the pair it was written with. A decode + // that happened to succeed on spliced bytes would show up here as a + // subject that does not match its payload. + for _, line := range lines { + var payload struct { + Subject string `json:"subject"` + } + Expect(json.Unmarshal(line.Progress, &payload)).To(Succeed()) + Expect(payload.Subject).To(Equal(line.Subject)) + } + }) + + It("reports a value it cannot encode without writing anything", func() { + Expect(pub.Publish("jobs.j1.progress", make(chan int))).ToNot(Succeed()) + Expect(buf.Len()).To(BeZero()) + Expect(flushes).To(BeZero()) + }) + + It("reports a response it cannot write on, and does not flush a line it did not write", func() { + failing := NewStreamPublisher(errWriter{err: errors.New("connection reset")}, func() { flushes++ }) + Expect(failing.Publish("jobs.j1.progress", map[string]any{"percentage": 10})).ToNot(Succeed()) + Expect(flushes).To(BeZero()) + }) +}) diff --git a/core/services/nodes/control_client.go b/core/services/nodes/control_client.go index 6c943892b..fe0d7e5f0 100644 --- a/core/services/nodes/control_client.go +++ b/core/services/nodes/control_client.go @@ -11,10 +11,7 @@ import ( "sync" "time" - "github.com/mudler/xlog" - "github.com/mudler/LocalAI/core/services/cluster" - "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/workerctl" "github.com/mudler/LocalAI/pkg/httpclient" ) @@ -152,6 +149,14 @@ func (c *ControlClient) Call(ctx context.Context, nodeID, path string, req, repl // stream, invoking onProgress for each progress line and decoding the single // terminal reply line into reply. onProgress may be nil. // +// onProgress receives the line's SUBJECT alongside its bytes, so a caller can +// tell a private progress tick from a re-broadcast request without this client +// understanding either. An empty subject is a line for this caller alone, which +// is what every pre-existing progress line is. The client deliberately does no +// authorization and no decoding of its own: whether a named broadcast may be +// made is a policy question about the node, and what a private tick means is a +// question about the verb, and this type knows neither. +// // onProgress runs SYNCHRONOUSLY, on this goroutine. The NATS carrier ran each // progress callback on a goroutine of its own because a slow callback there // stalled the one reader thread every worker's events arrived on; here the only @@ -159,7 +164,7 @@ func (c *ControlClient) Call(ctx context.Context, nodeID, path string, req, repl // caller's business. Dropping the guard is also what makes the events arrive in // the order the worker sent them. func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, path string, - req, reply any, onProgress func(messaging.BackendInstallProgressEvent)) error { + req, reply any, onProgress func(subject string, raw json.RawMessage)) error { resp, err := c.do(ctx, nodeID, path, req) if err != nil { return err @@ -188,14 +193,7 @@ func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, path string, if env.Progress == nil || onProgress == nil { continue } - var ev messaging.BackendInstallProgressEvent - if err := json.Unmarshal(env.Progress, &ev); err != nil { - // Progress is transient by contract, so a line this frontend cannot - // read costs a tick and never the operation. - xlog.Debug("unreadable control progress line", "node", nodeID, "path", path, "error", err) - continue - } - onProgress(ev) + onProgress(env.Subject, env.Progress) } if reply == nil { return nil diff --git a/core/services/nodes/control_client_test.go b/core/services/nodes/control_client_test.go index d8abc1197..ea19c3eff 100644 --- a/core/services/nodes/control_client_test.go +++ b/core/services/nodes/control_client_test.go @@ -223,13 +223,58 @@ var _ = Describe("ControlClient", func() { _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":100}`)}) _ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)}) } - var seen []float64 + var seen []string + var subjects []string var reply messaging.BackendInstallReply err := client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall, messaging.BackendInstallRequest{Backend: "mock", OpID: "op-1"}, &reply, - func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) }) + func(subject string, raw json.RawMessage) { + subjects = append(subjects, subject) + seen = append(seen, string(raw)) + }) Expect(err).NotTo(HaveOccurred()) - Expect(seen).To(Equal([]float64{50, 100})) + Expect(seen).To(Equal([]string{`{"percentage":50}`, `{"percentage":100}`})) + // A line that named no broadcast reaches the caller with an empty + // subject, which is what every pre-existing progress line is. + Expect(subjects).To(Equal([]string{"", ""})) + Expect(reply.Success).To(BeTrue()) + }) + + It("carries the subject a worker named through to the caller, and acts on it itself not at all", func() { + // The client is not where the authorization decision lives, and it + // must not become where it lives by accident: it hands the subject + // over and keeps reading. Whether the broadcast is made is a + // question about the node's type that this type cannot answer. + handler = func(w http.ResponseWriter, _ *http.Request) { + enc := json.NewEncoder(w) + _ = enc.Encode(workerctl.Envelope{ + Subject: "agent.a1.events.status", + Progress: json.RawMessage(`{"state":"thinking"}`), + }) + _ = enc.Encode(workerctl.Envelope{ + // A subject no worker of any type is allowed. The client + // still carries it: refusing here would put the policy in + // two places. + Subject: "cache.invalidate.models", + Progress: json.RawMessage(`{"model":"m1"}`), + }) + _ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)}) + } + type line struct { + subject string + raw string + } + var seen []line + var reply messaging.BackendInstallReply + Expect(client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall, + struct{}{}, &reply, + func(subject string, raw json.RawMessage) { + seen = append(seen, line{subject: subject, raw: string(raw)}) + })).To(Succeed()) + Expect(seen).To(Equal([]line{ + {subject: "agent.a1.events.status", raw: `{"state":"thinking"}`}, + {subject: "cache.invalidate.models", raw: `{"model":"m1"}`}, + })) Expect(reply.Success).To(BeTrue()) }) @@ -239,11 +284,11 @@ var _ = Describe("ControlClient", func() { _ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)}) _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":10}`)}) } - var seen []float64 + var seen []string var reply messaging.BackendInstallReply Expect(client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall, struct{}{}, &reply, - func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) })).To(Succeed()) + func(_ string, raw json.RawMessage) { seen = append(seen, string(raw)) })).To(Succeed()) Expect(seen).To(BeEmpty()) Expect(reply.Success).To(BeTrue()) }) @@ -274,19 +319,26 @@ var _ = Describe("ControlClient", func() { Expect(errors.Is(err, io.ErrUnexpectedEOF)).To(BeTrue()) }) - It("keeps going past a progress line it cannot read, since progress is transient", func() { + It("hands over a progress line it could not have decoded, because it does not decode them", func() { + // The client used to unmarshal every progress line into an install + // event and drop the ones that would not parse. It no longer knows + // what a progress line means, so a line that is not an install + // event is carried like any other and the DECODE is the caller's. + // The rule that a line the frontend cannot read costs a tick and + // never the operation now lives in installProgressBridge, and + // unloader_test.go pins it there. handler = func(w http.ResponseWriter, _ *http.Request) { enc := json.NewEncoder(w) _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":"not a number"}`)}) _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":70}`)}) _ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)}) } - var seen []float64 + var seen []string var reply messaging.BackendInstallReply Expect(client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall, struct{}{}, &reply, - func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) })).To(Succeed()) - Expect(seen).To(Equal([]float64{70})) + func(_ string, raw json.RawMessage) { seen = append(seen, string(raw)) })).To(Succeed()) + Expect(seen).To(Equal([]string{`{"percentage":"not a number"}`, `{"percentage":70}`})) Expect(reply.Success).To(BeTrue()) }) diff --git a/core/services/nodes/control_worker_fake_test.go b/core/services/nodes/control_worker_fake_test.go index b76d59b02..f5459e351 100644 --- a/core/services/nodes/control_worker_fake_test.go +++ b/core/services/nodes/control_worker_fake_test.go @@ -43,6 +43,13 @@ type scriptedControlWorkers struct { matched map[string][]matchedControlReply progress map[string][]messaging.BackendInstallProgressEvent + // rawProgress holds progress lines a spec wrote out as envelopes, for the + // two things scriptProgress cannot express: a line NAMING a subject, and a + // line whose payload is not an install-progress event at all. Both are + // things a real worker on another build can put on the wire, so the double + // has to be able to put them on the wire too. + rawProgress map[string][]workerctl.Envelope + // unreachable and expired are keyed by NODE, not by verb, because they are // failures of the ROUTE and a route belongs to a node. They are what the // dialer answers with; see scriptUnroutable and scriptTimeout. @@ -74,6 +81,7 @@ func newScriptedControlWorkers() *scriptedControlWorkers { unsupported: map[string]bool{}, matched: map[string][]matchedControlReply{}, progress: map[string][]messaging.BackendInstallProgressEvent{}, + rawProgress: map[string][]workerctl.Envelope{}, unreachable: map[string]bool{}, expired: map[string]bool{}, hangs: map[string]bool{}, @@ -131,6 +139,7 @@ func (s *scriptedControlWorkers) serve(w http.ResponseWriter, r *http.Request) { reply := s.replies[key] matchers := s.matched[key] ticks := s.progress[key] + rawTicks := s.rawProgress[key] s.mu.Unlock() if hang { @@ -182,6 +191,9 @@ func (s *scriptedControlWorkers) serve(w http.ResponseWriter, r *http.Request) { } _ = enc.Encode(workerctl.Envelope{Progress: raw}) } + for _, env := range rawTicks { + _ = enc.Encode(env) + } _ = enc.Encode(workerctl.Envelope{Reply: reply}) } @@ -270,6 +282,14 @@ func (s *scriptedControlWorkers) scriptProgress(key string, events []messaging.B s.progress[key] = events } +// scriptRawProgress queues progress lines exactly as they go on the wire, after +// any scriptProgress lines for the same verb. +func (s *scriptedControlWorkers) scriptRawProgress(key string, lines []workerctl.Envelope) { + s.mu.Lock() + defer s.mu.Unlock() + s.rawProgress[key] = lines +} + // scriptHang makes one verb on one node accept the request and never answer, // so the call ends only when the caller's budget does. // diff --git a/core/services/nodes/rebroadcast.go b/core/services/nodes/rebroadcast.go new file mode 100644 index 000000000..10e7809c8 --- /dev/null +++ b/core/services/nodes/rebroadcast.go @@ -0,0 +1,118 @@ +package nodes + +import ( + "encoding/json" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/xlog" +) + +// workerBroadcastAllow lists, per node type, the subject filters a worker may +// ask a frontend to publish on its behalf. +// +// An absent node type and an empty list both DENY EVERYTHING. That is the +// opposite of the NATS allow list this replaces, where an empty list meant no +// restriction, and it is the reason deleting an entry here is safe where +// deleting a permissions branch there was a privilege escalation. Nothing may +// be added to this map without a spec that a worker of that type is refused the +// subjects NOT added. +var workerBroadcastAllow = map[string][]string{ + NodeTypeAgent: { + "jobs.*.progress", + "jobs.*.result", + "agent.*.events.*", + }, + // A backend worker asks for no broadcasts. Spelled as an empty list rather + // than omitted, so the reader sees the decision. + NodeTypeBackend: {}, +} + +// MayBroadcast reports whether a worker of nodeType may have subject published +// on its behalf. +func MayBroadcast(nodeType, subject string) bool { + return mayBroadcastIn(workerBroadcastAllow, nodeType, subject) +} + +// mayBroadcastIn is MayBroadcast against a caller-supplied table. +// +// It exists so a spec can hold the table itself constant while varying what is +// IN it, which is the only way to state "an empty list denies" as a property of +// this function rather than as a property of today's entries. Production has +// exactly one table and MayBroadcast is the only way to reach it. +func mayBroadcastIn(table map[string][]string, nodeType, subject string) bool { + // An empty subject is a line that asked for no broadcast at all, so there + // is nothing here to authorise. Refusing it is not defensive tidiness: it + // keeps a filter that happened to match the empty string from turning every + // ordinary private progress tick into a publish. + if subject == "" { + return false + } + // A missing node type indexes to nil and the loop then runs zero times, + // which is the same denial an empty list produces. One code path for both + // is deliberate: "we have never heard of this worker" and "this worker is + // allowed nothing" have the same safe answer, and giving them separate + // branches invites one of the two to acquire an exception. + for _, filter := range table[nodeType] { + if messaging.SubjectMatches(filter, subject) { + return true + } + } + return false +} + +// Rebroadcaster turns an allowed progress line into a broadcast. +// +// It holds a Broadcaster and not a MessagingClient because fan-out is all it +// needs and all it may have: a worker asking for a broadcast must not be able +// to reach a request/reply or a queue group through this path. +type Rebroadcaster struct { + bus messaging.Broadcaster +} + +// NewRebroadcaster returns a Rebroadcaster publishing on bus. +func NewRebroadcaster(bus messaging.Broadcaster) *Rebroadcaster { + return &Rebroadcaster{bus: bus} +} + +// The shape of Handle's answer is asserted here, not left to its callers. +// +// A refused or failed re-broadcast must never become an error, because the only +// errors that reach a scheduler from this direction are the worker's own answer +// and an unroutable peer, and a publish failure is neither. Changing the return +// to an error therefore has to fail to COMPILE in the file that states the +// rule, rather than fail a spec somewhere a reviewer might read as a test +// needing an update. +var _ func(string, string, json.RawMessage) bool = (*Rebroadcaster)(nil).Handle + +// Handle publishes raw on subject when the worker's type allows it, and returns +// false having published nothing when it does not. It never returns an error +// that a caller could mistake for the worker's answer: a refused or failed +// re-broadcast is logged and the RPC continues, because the RPC's outcome is +// the worker's verdict about the work and a publish failure says nothing about +// it. +// +// raw is published as it arrived. It is a json.RawMessage and every Broadcaster +// encodes what it is given with json.Marshal, which returns a RawMessage +// verbatim, so the subscriber reads the bytes the worker wrote rather than a +// re-encoding of this frontend's idea of them. +func (r *Rebroadcaster) Handle(nodeType, subject string, raw json.RawMessage) bool { + if !MayBroadcast(nodeType, subject) { + // Warn rather than Debug: this is a worker asking for something it has + // no business on, which is the shape of a compromised or mismatched + // worker and is worth seeing without turning logging up. + xlog.Warn("refusing a worker's re-broadcast request", + "nodeType", nodeType, "subject", subject) + return false + } + if r == nil || r.bus == nil { + xlog.Debug("no broadcaster to re-broadcast a worker's progress line on", + "nodeType", nodeType, "subject", subject) + return false + } + if err := r.bus.Publish(subject, raw); err != nil { + xlog.Warn("a worker's re-broadcast could not be published", + "nodeType", nodeType, "subject", subject, "error", err) + return false + } + return true +} diff --git a/core/services/nodes/rebroadcast_test.go b/core/services/nodes/rebroadcast_test.go new file mode 100644 index 000000000..6f01c8ac1 --- /dev/null +++ b/core/services/nodes/rebroadcast_test.go @@ -0,0 +1,153 @@ +package nodes + +import ( + "encoding/json" + "errors" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// failingBus is a Broadcaster whose Publish always fails, for the one thing +// Handle promises about a carrier that is unhappy: the failure stays here. +type failingBus struct { + testutil.FakeBus + err error +} + +func (b *failingBus) Publish(string, any) error { return b.err } + +var _ = Describe("worker re-broadcast authorization", func() { + Describe("MayBroadcast", func() { + DescribeTable("answers for the node type it was asked about", + func(nodeType, subject string, allowed bool) { + Expect(MayBroadcast(nodeType, subject)).To(Equal(allowed)) + }, + Entry("agent, job progress", NodeTypeAgent, "jobs.j1.progress", true), + Entry("agent, job result", NodeTypeAgent, "jobs.j1.result", true), + Entry("agent, agent events", NodeTypeAgent, "agent.a1.events.status", true), + // The subjects an agent worker is NOT allowed. Every entry added to + // the table above owes this list a line, which is the only thing + // keeping the allow list from growing into "everything an agent + // happened to publish once". + Entry("agent, cache invalidation", NodeTypeAgent, "cache.invalidate.models", false), + Entry("agent, node registration", NodeTypeAgent, "node.register", false), + Entry("agent, model load", NodeTypeAgent, "model.load.n1", false), + Entry("agent, a job subject one token too deep", NodeTypeAgent, "jobs.j1.progress.extra", false), + Entry("agent, a job subject one token too shallow", NodeTypeAgent, "jobs.progress", false), + // A backend worker asks for no broadcasts, including the ones an + // agent worker is allowed: the table is per node type and not a + // single global list with a type-shaped comment on it. + Entry("backend, job progress", NodeTypeBackend, "jobs.j1.progress", false), + Entry("backend, job result", NodeTypeBackend, "jobs.j1.result", false), + Entry("backend, agent events", NodeTypeBackend, "agent.a1.events.status", false), + // A node type nothing has heard of is denied rather than defaulted. + Entry("unknown type, job progress", "router", "jobs.j1.progress", false), + Entry("unknown type, agent events", "router", "agent.a1.events.status", false), + Entry("empty type, job progress", "", "jobs.j1.progress", false), + ) + + It("refuses a line that names no subject at all", func() { + // A progress line with no subject asked for no broadcast. It must + // not become one because some filter matched the empty string. + Expect(MayBroadcast(NodeTypeAgent, "")).To(BeFalse()) + Expect(MayBroadcast(NodeTypeBackend, "")).To(BeFalse()) + }) + + It("refuses the tail wildcard, which the shared matcher does not implement", func() { + // SubjectMatches is the one definition of matching in the tree and + // it makes a '>' filter match nothing. Asserting it here is what + // stops someone writing "agent.>" into the allow list and believing + // they widened it, when what they did was narrow it to nothing. + Expect(messaging.SubjectMatches("agent.>", "agent.a1.events.status")).To(BeFalse()) + }) + + DescribeTable("reads an EMPTY allow list as denying everything, which is the inversion of the NATS list this replaces", + // NATS read an empty allow list as NO RESTRICTION, which is why the + // permissions it replaced had to spell the backend worker's list as + // {"_INBOX.>"} rather than leave it empty. This table is the + // opposite, and this spec is what documents it: the assumption a + // reviewer carries over from the deleted code is exactly wrong. + func(subject string) { + emptied := map[string][]string{ + NodeTypeAgent: {}, + NodeTypeBackend: {}, + } + Expect(mayBroadcastIn(emptied, NodeTypeAgent, subject)).To(BeFalse()) + // The same subject against the real table, so the case cannot + // pass by naming something that was never allowed. + Expect(MayBroadcast(NodeTypeAgent, subject)).To(BeTrue()) + }, + Entry("job progress", "jobs.j1.progress"), + Entry("job result", "jobs.j1.result"), + Entry("agent events", "agent.a1.events.status"), + ) + + It("reads a MISSING node type the same way as an empty list", func() { + absent := map[string][]string{NodeTypeBackend: {}} + Expect(mayBroadcastIn(absent, NodeTypeAgent, "jobs.j1.progress")).To(BeFalse()) + }) + }) + + Describe("Rebroadcaster.Handle", func() { + var ( + bus *testutil.FakeBus + rb *Rebroadcaster + ) + + BeforeEach(func() { + bus = testutil.NewFakeBus() + rb = NewRebroadcaster(bus) + }) + + It("publishes an allowed subject and hands the subscriber the worker's own bytes", func() { + delivered := make(chan []byte, 1) + _, err := bus.Subscribe("agent.*.events.*", func(payload []byte) { delivered <- payload }) + Expect(err).ToNot(HaveOccurred()) + + Expect(rb.Handle(NodeTypeAgent, "agent.a1.events.status", + json.RawMessage(`{"state":"thinking"}`))).To(BeTrue()) + + Expect(bus.PublishCount("agent.a1.events.status")).To(Equal(1)) + // Verbatim, not re-encoded: a subscriber decodes the worker's DTO + // and a frontend that re-marshalled its own idea of the line would + // hand it a different shape. + Eventually(delivered).Should(Receive(MatchJSON(`{"state":"thinking"}`))) + }) + + It("refuses a subject the agent worker has no business on, and publishes nothing", func() { + Expect(rb.Handle(NodeTypeAgent, "cache.invalidate.models", + json.RawMessage(`{"model":"m1"}`))).To(BeFalse()) + Expect(bus.PublishCount("cache.invalidate.models")).To(Equal(0)) + }) + + It("refuses a BACKEND worker a subject an agent worker would be allowed", func() { + Expect(rb.Handle(NodeTypeBackend, "jobs.j1.progress", + json.RawMessage(`{"percentage":10}`))).To(BeFalse()) + Expect(bus.PublishCount("jobs.j1.progress")).To(Equal(0)) + }) + + It("refuses a node type nothing has heard of", func() { + Expect(rb.Handle("router", "jobs.j1.progress", + json.RawMessage(`{"percentage":10}`))).To(BeFalse()) + Expect(bus.PublishCount("jobs.j1.progress")).To(Equal(0)) + }) + + It("reports a carrier that could not publish as a refusal, not as anything the worker said", func() { + // Handle's return type is the whole point: there is no error here + // for a caller to confuse with the RPC's outcome. A publish that + // failed says nothing about whether the work succeeded. + rb = NewRebroadcaster(&failingBus{err: errors.New("carrier is unhappy")}) + Expect(rb.Handle(NodeTypeAgent, "jobs.j1.progress", + json.RawMessage(`{"percentage":10}`))).To(BeFalse()) + }) + + It("refuses rather than panicking when it holds no broadcaster", func() { + Expect(NewRebroadcaster(nil).Handle(NodeTypeAgent, "jobs.j1.progress", + json.RawMessage(`{"percentage":10}`))).To(BeFalse()) + }) + }) +}) diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index 6af2b8ce3..3002f9e55 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -2,6 +2,7 @@ package nodes import ( "context" + "encoding/json" "errors" "fmt" "time" @@ -185,6 +186,45 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, mo return unloadErr } +// installProgressBridge adapts an install or upgrade's progress sink to +// CallStreaming's line callback, which carries a subject the client does not +// interpret. +// +// It makes two decisions a caller of InstallBackend must not have to make. +// +// A line naming a SUBJECT is a re-broadcast request and not install progress, +// so it is dropped here rather than delivered as a tick. backend.install and +// backend.upgrade are a BACKEND worker's verbs, and MayBroadcast denies a +// backend worker every subject, so this path has nothing to publish and no +// broadcaster to publish it on. Delivering it to onProgress instead would put a +// worker's arbitrary JSON through a decode into an install-progress event and +// report whatever fell out as the state of a download. +// +// A line this frontend cannot decode costs a tick and never the operation, +// because progress is transient by contract while the reply is the worker's +// verdict. +// +// It returns nil for a nil sink so CallStreaming keeps its "no callback, no +// work" path, rather than a non-nil closure wrapping a nil function. +func installProgressBridge(nodeID, path string, onProgress func(messaging.BackendInstallProgressEvent)) func(string, json.RawMessage) { + if onProgress == nil { + return nil + } + return func(subject string, raw json.RawMessage) { + if subject != "" { + xlog.Warn("refusing a re-broadcast request on a backend control stream", + "node", nodeID, "path", path, "subject", subject) + return + } + var ev messaging.BackendInstallProgressEvent + if err := json.Unmarshal(raw, &ev); err != nil { + xlog.Debug("unreadable control progress line", "node", nodeID, "path", path, "error", err) + return + } + onProgress(ev) + } +} + // InstallBackend asks a worker node to install a backend and start its process. // Idempotent on the worker: if the (modelID, replica) process is already // running, the worker short-circuits and returns its address; if the binary @@ -223,7 +263,7 @@ func (a *RemoteUnloaderAdapter) InstallBackend( Alias: alias, ReplicaIndex: int32(replicaIndex), OpID: opID, - }, &reply, onProgress) + }, &reply, installProgressBridge(nodeID, workerctl.PathBackendInstall, onProgress)) if err != nil { if isRequestTimeout(err) { return nil, fmt.Errorf("%w (nodeID=%s backend=%s): %v", @@ -258,7 +298,7 @@ func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSO Alias: alias, ReplicaIndex: int32(replicaIndex), OpID: opID, - }, &reply, onProgress) + }, &reply, installProgressBridge(nodeID, workerctl.PathBackendUpgrade, onProgress)) if err != nil { if isRequestTimeout(err) { return nil, fmt.Errorf("%w (nodeID=%s backend=%s): %v", @@ -291,7 +331,7 @@ func (a *RemoteUnloaderAdapter) installWithForceFallback(nodeID, backendType, ga ReplicaIndex: int32(replicaIndex), Force: true, OpID: opID, - }, &reply, onProgress) + }, &reply, installProgressBridge(nodeID, workerctl.PathBackendInstall, onProgress)) if err != nil { if isRequestTimeout(err) { return nil, fmt.Errorf("%w (nodeID=%s backend=%s): %v", diff --git a/core/services/nodes/unloader_test.go b/core/services/nodes/unloader_test.go index 28ddaf37d..8e7c1a622 100644 --- a/core/services/nodes/unloader_test.go +++ b/core/services/nodes/unloader_test.go @@ -634,4 +634,104 @@ var _ = Describe("RemoteUnloaderAdapter install progress streaming", func() { Expect(err).ToNot(HaveOccurred()) Expect(reply.Success).To(BeTrue()) }) + + // The bridge from a control STREAM LINE to an install-progress event is one + // definition used at three call sites: InstallBackend, UpgradeBackend and + // installWithForceFallback. Each is exercised separately below, because a + // site that stopped calling it would leave the other two green. + DescribeTable("delivers a line that names no broadcast to the caller's install-progress sink, on every streaming verb", + func(path string, call func(*RemoteUnloaderAdapter, func(messaging.BackendInstallProgressEvent)) error) { + workers := newScriptedControlWorkers() + workers.scriptReply(controlKey("n1", path), messaging.BackendInstallReply{Success: true}) + workers.scriptReply(controlKey("n1", workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true}) + workers.scriptProgress(controlKey("n1", path), []messaging.BackendInstallProgressEvent{ + {OpID: "op-abc", NodeID: "n1", Backend: "vllm", Percentage: 25}, + }) + + adapter := NewRemoteUnloaderAdapter(nil, workers.controlClient(), time.Second, time.Second) + var received []messaging.BackendInstallProgressEvent + Expect(call(adapter, func(ev messaging.BackendInstallProgressEvent) { + received = append(received, ev) + })).To(Succeed()) + Expect(received).To(HaveLen(1)) + Expect(received[0].Percentage).To(Equal(float64(25))) + }, + Entry("backend.install", workerctl.PathBackendInstall, + func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error { + _, err := a.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "op-abc", cb) + return err + }), + Entry("backend.upgrade", workerctl.PathBackendUpgrade, + func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error { + _, err := a.UpgradeBackend("n1", "vllm", "[]", "", "", "", 0, "op-abc", cb) + return err + }), + Entry("the legacy force-install fallback", workerctl.PathBackendInstall, + func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error { + _, err := a.installWithForceFallback("n1", "vllm", "[]", "", "", "", 0, "op-abc", cb) + return err + }), + ) + + DescribeTable("refuses a line that NAMES a broadcast rather than delivering it as install progress, on every streaming verb", + func(path string, call func(*RemoteUnloaderAdapter, func(messaging.BackendInstallProgressEvent)) error) { + // backend.install and backend.upgrade are a BACKEND worker's verbs, + // and MayBroadcast denies a backend worker every subject, so a line + // naming one is not something this path publishes. What it must + // also not do is deliver it: the payload is a broadcast body and + // decoding it into an install-progress event reports a percentage + // nobody sent. + workers := newScriptedControlWorkers() + workers.scriptReply(controlKey("n1", path), messaging.BackendInstallReply{Success: true}) + workers.scriptReply(controlKey("n1", workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true}) + workers.scriptRawProgress(controlKey("n1", path), []workerctl.Envelope{ + {Subject: "jobs.j1.progress", Progress: json.RawMessage(`{"percentage":99}`)}, + {Subject: "agent.a1.events.status", Progress: json.RawMessage(`{"percentage":98}`)}, + }) + + adapter := NewRemoteUnloaderAdapter(nil, workers.controlClient(), time.Second, time.Second) + var received []messaging.BackendInstallProgressEvent + Expect(call(adapter, func(ev messaging.BackendInstallProgressEvent) { + received = append(received, ev) + })).To(Succeed()) + // The RPC still returned the worker's reply: a refused re-broadcast + // costs a line and never the operation. + Expect(received).To(BeEmpty()) + }, + Entry("backend.install", workerctl.PathBackendInstall, + func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error { + _, err := a.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "op-abc", cb) + return err + }), + Entry("backend.upgrade", workerctl.PathBackendUpgrade, + func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error { + _, err := a.UpgradeBackend("n1", "vllm", "[]", "", "", "", 0, "op-abc", cb) + return err + }), + Entry("the legacy force-install fallback", workerctl.PathBackendInstall, + func(a *RemoteUnloaderAdapter, cb func(messaging.BackendInstallProgressEvent)) error { + _, err := a.installWithForceFallback("n1", "vllm", "[]", "", "", "", 0, "op-abc", cb) + return err + }), + ) + + It("keeps going past a progress line it cannot read, since progress is transient", func() { + // The rule the control client used to carry. It moved to the bridge + // with the decode, and it has to be pinned where it now lives: a line + // this frontend cannot read costs a tick and never the operation. + workers := newScriptedControlWorkers() + workers.scriptReply(controlKey("n1", workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true}) + workers.scriptRawProgress(controlKey("n1", workerctl.PathBackendInstall), []workerctl.Envelope{ + {Progress: json.RawMessage(`{"percentage":"not a number"}`)}, + {Progress: json.RawMessage(`{"percentage":70}`)}, + }) + + adapter := NewRemoteUnloaderAdapter(nil, workers.controlClient(), time.Second, time.Second) + var seen []float64 + reply, err := adapter.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "op-abc", + func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) }) + Expect(err).ToNot(HaveOccurred()) + Expect(seen).To(Equal([]float64{70})) + Expect(reply.Success).To(BeTrue()) + }) }) diff --git a/core/services/worker/control_client_roundtrip_test.go b/core/services/worker/control_client_roundtrip_test.go index 52046f320..aae8d8498 100644 --- a/core/services/worker/control_client_roundtrip_test.go +++ b/core/services/worker/control_client_roundtrip_test.go @@ -2,6 +2,7 @@ package worker import ( "context" + "encoding/json" "errors" "net" "os" @@ -96,12 +97,23 @@ var _ = Describe("the frontend's control client against the real worker", func() } var seen []float64 + var subjects []string var reply messaging.BackendInstallReply err := client.CallStreaming(context.Background(), nodeID, workerctl.PathBackendInstall, messaging.BackendInstallRequest{Backend: "mock", OpID: "op-1"}, &reply, - func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) }) + func(subject string, raw json.RawMessage) { + var ev messaging.BackendInstallProgressEvent + Expect(json.Unmarshal(raw, &ev)).To(Succeed()) + subjects = append(subjects, subject) + seen = append(seen, ev.Percentage) + }) Expect(err).NotTo(HaveOccurred()) Expect(seen).To(Equal([]float64{50, 100})) + // A BACKEND worker's install progress names no broadcast, over a real + // tunnel and a real NDJSON body rather than in a double's imagination. + // It is the property MayBroadcast would refuse anyway, asserted here on + // the bytes the worker actually wrote. + Expect(subjects).To(Equal([]string{"", ""})) Expect(reply.Success).To(BeTrue()) Expect(reply.WorkerLocalAddress).To(Equal("127.0.0.1:41234")) }) diff --git a/core/services/workerctl/paths.go b/core/services/workerctl/paths.go index 7d31cd928..fbc49ae12 100644 --- a/core/services/workerctl/paths.go +++ b/core/services/workerctl/paths.go @@ -125,12 +125,12 @@ func AllPaths() []string { // Envelope is one line of a streaming control response. // -// Exactly one of the two is set. Zero or more Progress lines are followed by -// exactly ONE Reply line, and the Reply line is the last thing on the body. -// That ordering is the contract: it is what lets the frontend stop reading, and -// it is what replaces the subscribe-before-request dance the NATS carrier -// needed, since progress and reply now share one response and nothing can -// arrive before the caller is listening. +// Exactly one of Progress and Reply is set. Zero or more Progress lines are +// followed by exactly ONE Reply line, and the Reply line is the last thing on +// the body. That ordering is the contract: it is what lets the frontend stop +// reading, and it is what replaces the subscribe-before-request dance the NATS +// carrier needed, since progress and reply now share one response and nothing +// can arrive before the caller is listening. // // Progress carrying the reply's own bytes is also why the 8000-byte // notification cap that bounded the NATS progress subject has no analogue here: @@ -138,6 +138,20 @@ func AllPaths() []string { type Envelope struct { Progress json.RawMessage `json:"progress,omitempty"` Reply json.RawMessage `json:"reply,omitempty"` + + // Subject names the broadcast this progress line asks the frontend reading + // the stream to make on its behalf. Empty means the line is for this caller + // alone, which is what every pre-existing progress line is. + // + // It is a REQUEST and not an instruction: the frontend checks it against an + // allow list derived from the node's type before publishing anything. A + // worker naming a subject it has no business on is refused and logged, and + // the stream continues. + // + // It qualifies a Progress line and never a Reply line. A reply is the + // worker's verdict about the work, and there is no version of "publish my + // verdict for me" that this control plane has to carry. + Subject string `json:"subject,omitempty"` } // ContentTypeStream is the media type of a streaming control response. diff --git a/core/services/workerctl/paths_test.go b/core/services/workerctl/paths_test.go index d912354cc..0bfc31ea1 100644 --- a/core/services/workerctl/paths_test.go +++ b/core/services/workerctl/paths_test.go @@ -124,6 +124,37 @@ var _ = Describe("control plane paths on the wire", func() { Expect(string(b)).To(Equal(`{"progress":{"percentage":50}}`)) }) + It("spells the re-broadcast key on the wire as \"subject\"", func() { + // Hand-written literal, like the paths above and for the same reason: a + // worker and a frontend built from different commits read each other's + // lines, and a renamed key is a re-broadcast request that silently + // becomes a private progress tick. Deriving the expectation from the + // struct tag would pin nothing. + b, err := json.Marshal(workerctl.Envelope{ + Subject: "agent.a1.events.status", + Progress: json.RawMessage(`{"tick":1}`), + }) + Expect(err).NotTo(HaveOccurred()) + Expect(string(b)).To(Equal(`{"progress":{"tick":1},"subject":"agent.a1.events.status"}`)) + }) + + It("omits the subject key entirely when a progress line names no broadcast", func() { + // omitempty is the backward-compatibility half: every pre-existing + // progress line has no subject, and an older frontend reading a line + // that carried an empty subject key would be reading a field it does + // not know about on every tick. + b, err := json.Marshal(workerctl.Envelope{Progress: json.RawMessage(`{"tick":1}`)}) + Expect(err).NotTo(HaveOccurred()) + Expect(string(b)).NotTo(ContainSubstring("subject")) + }) + + It("reads a subject back off the wire, which is what the frontend decides on", func() { + var env workerctl.Envelope + Expect(json.Unmarshal([]byte(`{"progress":{"tick":1},"subject":"jobs.j1.progress"}`), &env)).To(Succeed()) + Expect(env.Subject).To(Equal("jobs.j1.progress")) + Expect(string(env.Progress)).To(Equal(`{"tick":1}`)) + }) + It("names the streaming media type", func() { Expect(workerctl.ContentTypeStream).To(Equal("application/x-ndjson")) })