test(distributed): pin the unreadable-request rule at all three file verbs

The rule "a body this worker could not parse is a non-2xx, never the
worker's answer" is written at three exits in control_files.go and only
ensure was pinned. Turning stage's or listdir's decode exit into a
200-with-error left worker and nodes entirely green, and what that
converts is a frontend's malformed request into the worker's own verdict
about a file, which passes cluster.IsWorkerAnswer and reaches a reap
guard. The production code was already right; nothing held it there.

The e2e NATS JWT spec was asserting the opposite of the code and passing.
It published nodes.<id>.files.in and called it an allowed subject after
that grant was deleted, and it could not tell: a permission violation
does not close the connection, so FlushTimeout and IsConnected both stay
happy. It now reads LastError, the way its sibling always has, and
asserts the denial plus the one publish right a backend worker has left.

Also pinned, each mutation-verified alone: the CreateTemp branch (an
existing staging-tmp at 0500 reaches it without a seam), the walk's
context check (a caller that gave up must fail the listing, never be
answered with a short one), and the cache and data directory layout.

The data directory was derived twice, once in worker.go and once for the
listdir verb; worker.go now reads the same helper, so a move cannot leave
a verb listing files the file server does not serve. The per-verb RPC
ceiling moves from an argument at five call sites into fileRPCBudget,
so no site can name the wrong one, and the two values are asserted.

The body-cap table now holds both directions locally and with two
different claims: a body exactly at the cap proves the bound is a
ceiling and not an off-by-one, and an absolute megabyte proves the cap
stays above real gallery traffic. Only the second notices a cap shrunk
to 64 KiB.

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-02 21:54:04 +00:00
1 parent 7fc617c8ed
commit 4c3e0deb19
8 files changed
+205 -30

No files matched your search

+5 -2
View File
@@ -5,11 +5,14 @@ import "context"
// FileStager abstracts file transfer between frontend and backend nodes
// in distributed mode. Two implementations exist:
//
// 1. S3NATSFileStager (primary): Both sides have FileManager with same S3.
// Frontend uploads to S3, sends NATS request-reply to backend to download locally.
// 1. S3FileStager (primary): Both sides have FileManager with same S3.
// Frontend uploads to S3, then calls the worker's control plane over its
// tunnel to have it download locally.
//
// 2. HTTPFileStager (fallback): Frontend pushes/pulls files directly over
// HTTP to a small file transfer server on the backend node (no S3 needed).
//
// Both reach the worker through its tunnel and neither uses NATS.
type FileStager interface {
// EnsureRemote ensures a local file is available on the remote node.
// Returns the remote-local path.
+33 -14
View File
@@ -43,6 +43,29 @@ const (
fileMetadataRPCTimeout = 30 * time.Second
)
// fileRPCBudget is the ceiling one file-staging verb's RPC gets.
//
// The mapping lives here rather than at the call sites, and that is the same
// argument callWorker is written under: five sites each naming their own
// constant is five chances to name the wrong one, and a stage verb given the
// metadata ceiling would abandon a multi-gigabyte copy after thirty seconds
// while the worker went on making it.
//
// A path this function does not know gets the SHORTER ceiling. That is the safe
// direction: a verb wrongly given 30 seconds fails visibly and is retried,
// while one wrongly given ten minutes parks a caller on a verb that was never
// meant to be slow.
func fileRPCBudget(path string) time.Duration {
switch path {
case workerctl.PathFilesEnsure, workerctl.PathFilesStage:
return fileTransferRPCTimeout
case workerctl.PathFilesTemp, workerctl.PathFilesListDir:
return fileMetadataRPCTimeout
default:
return fileMetadataRPCTimeout
}
}
// Control request/reply message types. Their JSON is the shape the
// nodes.<id>.files.* subjects carried, so the worker's handler bodies did not
// have to change when the carrier did.
@@ -83,15 +106,16 @@ type fileListDirReply struct {
}
// callWorker issues one file-staging RPC under a deadline derived from the
// caller's context.
// caller's context and bounded by the verb's own ceiling.
//
// It exists so the derivation is written ONCE. Five call sites each repeating
// context.WithTimeout is five chances for one of them to start from a
// background context instead, and that one site would then keep commanding a
// worker after its caller had gone, with nothing else in the suite any redder
// for it.
func (s *S3FileStager) callWorker(ctx context.Context, nodeID, path string, budget time.Duration, req, reply any) error {
rpcCtx, cancel := context.WithTimeout(ctx, budget)
// for it. The budget comes from the path rather than from an argument for the
// same reason; see fileRPCBudget.
func (s *S3FileStager) callWorker(ctx context.Context, nodeID, path string, req, reply any) error {
rpcCtx, cancel := context.WithTimeout(ctx, fileRPCBudget(path))
defer cancel()
return s.control.Call(rpcCtx, nodeID, path, req, reply)
}
@@ -115,8 +139,7 @@ func (s *S3FileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key
}
var reply fileEnsureReply
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesEnsure, fileTransferRPCTimeout,
fileEnsureRequest{Key: key}, &reply); err != nil {
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesEnsure, fileEnsureRequest{Key: key}, &reply); err != nil {
return "", err
}
if reply.Error != "" {
@@ -145,8 +168,7 @@ func (s *S3FileStager) FetchRemoteByKey(ctx context.Context, nodeID, key, localD
func (s *S3FileStager) fetchRemoteWithKey(ctx context.Context, nodeID, remotePath, key, localDst string, cleanup bool) error {
var reply fileStageReply
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesStage, fileTransferRPCTimeout,
fileStageRequest{LocalPath: remotePath, Key: key}, &reply); err != nil {
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesStage, fileStageRequest{LocalPath: remotePath, Key: key}, &reply); err != nil {
return err
}
if reply.Error != "" {
@@ -177,8 +199,7 @@ func (s *S3FileStager) fetchRemoteWithKey(ctx context.Context, nodeID, remotePat
// AllocRemoteTemp asks the worker to allocate a temp file.
func (s *S3FileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (string, error) {
var reply fileTempReply
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesTemp, fileMetadataRPCTimeout,
fileTempRequest{}, &reply); err != nil {
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesTemp, fileTempRequest{}, &reply); err != nil {
return "", err
}
if reply.Error != "" {
@@ -197,8 +218,7 @@ func (s *S3FileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (stri
// listing would read to the caller as files the worker does not have.
func (s *S3FileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error) {
var reply fileListDirReply
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesListDir, fileMetadataRPCTimeout,
fileListDirRequest{KeyPrefix: keyPrefix}, &reply); err != nil {
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesListDir, fileListDirRequest{KeyPrefix: keyPrefix}, &reply); err != nil {
return nil, err
}
if reply.Error != "" {
@@ -211,8 +231,7 @@ func (s *S3FileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix stri
// StageRemoteToStore tells the worker to upload a local file to shared storage.
func (s *S3FileStager) StageRemoteToStore(ctx context.Context, nodeID, remotePath, key string) error {
var reply fileStageReply
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesStage, fileTransferRPCTimeout,
fileStageRequest{LocalPath: remotePath, Key: key}, &reply); err != nil {
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesStage, fileStageRequest{LocalPath: remotePath, Key: key}, &reply); err != nil {
return err
}
if reply.Error != "" {
@@ -5,6 +5,7 @@ import (
"errors"
"os"
"path/filepath"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
@@ -44,6 +45,26 @@ var _ = Describe("the S3 file stager's control RPCs", func() {
Expect(os.WriteFile(local, []byte("weights"), 0o600)).To(Succeed())
})
// The budgets are written out BY HAND and deliberately not derived from the
// constants under test. They are the timeouts the NATS request-reply calls
// carried, and keeping them is a compatibility fact rather than a taste:
// an operator whose staging of a 35 GB checkpoint fits inside ten minutes
// today must not find it cut to thirty seconds by the carrier changing.
DescribeTable("gives each verb the ceiling its NATS timeout carried",
func(path string, want time.Duration) { Expect(fileRPCBudget(path)).To(Equal(want)) },
Entry("ensure moves bytes", workerctl.PathFilesEnsure, 10*time.Minute),
Entry("stage moves bytes", workerctl.PathFilesStage, 10*time.Minute),
Entry("temp is metadata", workerctl.PathFilesTemp, 30*time.Second),
Entry("listdir is metadata", workerctl.PathFilesListDir, 30*time.Second),
)
It("gives a verb it does not know the shorter ceiling", func() {
// The safe direction: a verb wrongly given thirty seconds fails visibly
// and is retried, while one wrongly given ten minutes parks a caller on
// a verb that was never meant to be slow.
Expect(fileRPCBudget(workerctl.Prefix + "invented")).To(Equal(30 * time.Second))
})
DescribeTable("addresses the verb that verb's path names",
func(path string, reply any, call func(*S3FileStager) error) {
workers.scriptReply(controlKey(nodeID, path), reply)
+3 -1
View File
@@ -22,7 +22,9 @@ import (
// for distributed mode. Input files are staged on the backend node before the
// gRPC call. Output files are retrieved from the backend after the call.
//
// Uses the FileStager interface — agnostic to transport (S3+NATS or gRPC).
// Uses the FileStager interface — agnostic to transport (an object store, or
// direct HTTP to the worker), and in both cases reached over the worker's
// tunnel.
// The caller gets a grpc.Backend that behaves identically to a local one —
// no changes needed in core/backend/*.go.
//
+111 -5
View File
@@ -73,6 +73,17 @@ var _ = Describe("worker file-staging control routes", func() {
DeferCleanup(srv.Close)
})
// The on-disk layout, written out BY HAND rather than derived from the
// helpers under test. What these directories ARE is an operator-facing
// contract: a deployment mounts a volume per directory and sizes it, so a
// join that moved would put staged bytes on a volume nobody provisioned.
// Deriving the expectation from the helper would pin nothing.
DescribeTable("derives its directories from the models path, once each",
func(got, want string) { Expect(got).To(Equal(want)) },
Entry("cache", (&Config{ModelsPath: "/srv/localai/models"}).stagingCacheDir(), "/srv/localai/cache"),
Entry("data", (&Config{ModelsPath: "/srv/localai/models"}).stagingDataDir(), "/srv/localai/data"),
)
It("allocates a temp path and returns it", func() {
resp := post(workerctl.PathFilesTemp, struct{}{})
Expect(resp.StatusCode).To(Equal(http.StatusOK))
@@ -220,6 +231,25 @@ var _ = Describe("worker file-staging control routes", func() {
Expect(reply.LocalPath).To(BeEmpty())
})
It("reports a temp file it cannot create as a 200 with an error field", func() {
// The staging directory ALREADY EXISTS and is unwritable, so MkdirAll
// succeeds and CreateTemp is the branch that fails. It is a second
// statement of the same rule inside the same handler as the MkdirAll
// spec above, and a spec that only ever reaches MkdirAll leaves this
// one free to answer a non-2xx for the worker's own verdict.
Expect(os.MkdirAll(filepath.Join(cacheDir, "staging-tmp"), 0o500)).To(Succeed())
resp := post(workerctl.PathFilesTemp, struct{}{})
Expect(resp.StatusCode).To(Equal(http.StatusOK))
var reply struct {
LocalPath string `json:"local_path"`
Error string `json:"error"`
}
decode(resp, &reply)
Expect(reply.Error).NotTo(BeEmpty())
Expect(reply.LocalPath).To(BeEmpty())
})
It("refuses a key prefix that climbs out of the directories it serves", func() {
resp := post(workerctl.PathFilesListDir, map[string]string{"key_prefix": "../../../etc"})
Expect(resp.StatusCode).To(Equal(http.StatusOK))
@@ -282,11 +312,87 @@ var _ = Describe("worker file-staging control routes", func() {
Entry("listdir", workerctl.PathFilesListDir),
)
It("reports a request body it cannot read as a rejection, never as a file that is not there", func() {
resp, err := http.Post(srv.URL+workerctl.PathFilesEnsure, "application/json", strings.NewReader("{not json"))
Expect(err).NotTo(HaveOccurred())
DeferCleanup(func() { _ = resp.Body.Close() })
Expect(resp.StatusCode).To(Equal(http.StatusBadRequest))
// The other direction of the same bound, kept beside the refusals so the
// pair is local rather than spread across the suite.
//
// The two entries hold DIFFERENT halves and both are needed. The
// cap-relative body holds that the boundary is inclusive, so the bound is a
// ceiling and not an off-by-one; it moves with the cap and therefore says
// nothing about the cap's size. The absolute body holds that the cap is
// large enough for real traffic: BackendInstallRequest.BackendGalleries is
// a serialized gallery list of a few hundred kilobytes, so a cap tightened
// below a megabyte would start refusing ordinary requests, and only an
// entry written in absolute bytes notices that.
DescribeTable("serves a body that fits inside the control cap",
func(path string, size int) {
body := append([]byte(`{"key":"`), bytes.Repeat([]byte("a"), size-len(`{"key":""}`))...)
body = append(body, []byte(`"}`)...)
Expect(body).To(HaveLen(size))
resp, err := http.Post(srv.URL+path, "application/json", bytes.NewReader(body))
Expect(err).NotTo(HaveOccurred())
DeferCleanup(func() { _ = resp.Body.Close() })
Expect(resp.StatusCode).To(Equal(http.StatusOK))
},
Entry("ensure, exactly at the cap", workerctl.PathFilesEnsure, maxControlRequestBytes),
Entry("stage, exactly at the cap", workerctl.PathFilesStage, maxControlRequestBytes),
Entry("temp, exactly at the cap", workerctl.PathFilesTemp, maxControlRequestBytes),
Entry("listdir, exactly at the cap", workerctl.PathFilesListDir, maxControlRequestBytes),
Entry("ensure, a megabyte of real traffic", workerctl.PathFilesEnsure, 1<<20),
Entry("stage, a megabyte of real traffic", workerctl.PathFilesStage, 1<<20),
Entry("temp, a megabyte of real traffic", workerctl.PathFilesTemp, 1<<20),
Entry("listdir, a megabyte of real traffic", workerctl.PathFilesListDir, 1<<20),
)
// A body this worker could not PARSE is the frontend's request being wrong,
// not this worker's verdict about a file. The two live in opposite buckets:
// a non-2xx is mapped onto ErrWorkerUnroutable, which nothing may act on,
// while a 200 with the error field set passes through unwrapped so
// cluster.IsWorkerAnswer sees it and a reap guard MAY act on it. A handler
// that answered `{"error":"invalid request"}` for an unparseable body would
// hand a malformed request to a reap guard as evidence about a file.
//
// Pinned at every verb that decodes a body, because each one writes the
// exit out for itself. The base's NATS handlers shipped the wrong answer
// here; adopting the shared door fixed it, and this is what keeps it fixed.
DescribeTable("reports a request body it cannot read as a rejection, never as a file that is not there",
func(path string) {
resp, err := http.Post(srv.URL+path, "application/json", strings.NewReader("{not json"))
Expect(err).NotTo(HaveOccurred())
DeferCleanup(func() { _ = resp.Body.Close() })
Expect(resp.StatusCode).To(Equal(http.StatusBadRequest))
// The status alone is not the whole rule: a 400 whose body still
// decodes as a reply with an error field would be read correctly by
// today's client and wrongly by anything that reads the body first.
var reply struct {
Error string `json:"error"`
}
Expect(json.NewDecoder(resp.Body).Decode(&reply)).NotTo(Succeed())
},
Entry("ensure", workerctl.PathFilesEnsure),
Entry("stage", workerctl.PathFilesStage),
Entry("listdir", workerctl.PathFilesListDir),
// temp decodes no body, so it has no such exit to state.
)
It("fails a listing the caller abandoned rather than answering a short one", func() {
// A caller that gave up must not be answered with the files walked so
// far. A partial listing is the one shape the frontend cannot tell from
// a directory that really is that size, and it would read as files the
// worker does not have. So the walk returns the context error and the
// verb reports a FAILED listing.
dir := filepath.Join(modelsDir, "abandoned")
Expect(os.MkdirAll(dir, 0o750)).To(Succeed())
for i := range 32 {
Expect(os.WriteFile(filepath.Join(dir, fmt.Sprintf("f-%02d", i)), []byte("x"), 0o600)).To(Succeed())
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
files, err := listStagedFiles(ctx, dir)
Expect(err).To(MatchError(context.Canceled))
Expect(files).To(BeEmpty())
})
It("mounts every file verb, so none can be dropped from the set the frontend calls", func() {
+5 -1
View File
@@ -186,7 +186,11 @@ func Run(ctx *cliContext.Context, cfg *Config) error {
// the top of Run so the worker fails before registering.)
httpAddr := cfg.resolveHTTPAddr()
stagingDir := filepath.Join(cfg.ModelsPath, "..", "staging")
dataDir := filepath.Join(cfg.ModelsPath, "..", "data")
// Derived through the same helper the listdir verb resolves `data/` keys
// against, and not a second time here. Two independent joins that agreed
// today would each stay self consistent if one moved, and the symptom would
// be a verb that lists files the file server does not serve.
dataDir := cfg.stagingDataDir()
// The readiness gate is created here but only armed once NATS is up, below.
// Until then /readyz reports ready, which is correct: reaching this line
// means the worker has already registered with the frontend, so it is
+5 -3
View File
@@ -180,9 +180,11 @@ carries file staging and backend logs, behind the same `LOCALAI_REGISTRATION_TOK
bearer check. They replace the fourteen `nodes.<id>.*` NATS subjects a worker
used to subscribe to - the ten backend and model lifecycle verbs, plus the four
object-store staging verbs (`POST /v1/control/files/{ensure,stage,temp,listdir}`,
mounted only when the deployment configured an object store). The request and
reply bodies are unchanged, so nothing an operator inspects on the wire has a new
shape. Agent workers still take `nodes.<id>.backend.stop` over NATS.
mounted only when the deployment configured an object store). The request bodies
are unchanged and the reply fields keep their names and types, so nothing an
operator inspects on the wire has a new shape. The one difference is that a
worker now OMITS an empty reply field where the NATS handlers always emitted it,
which a client reading a missing field as the zero value cannot tell apart. Agent workers still take `nodes.<id>.backend.stop` over NATS.
`files/listdir` is the verb the change is most visible on. Its reply used to be
sized against what the bus would carry, which put a wide model directory close to
+22 -4
View File
@@ -17,14 +17,32 @@ var _ = Describe("NATS JWT Auth", Label("Distributed", "NatsJWT"), func() {
infra = SetupJWTInfra()
})
It("connects with a minted backend worker JWT and publishes on allowed subjects", func() {
// Backend workers may publish under nodes.<id>.files.> (see pkg/natsauth permissions).
subject := nodeSubjectPrefix(infra.NodeID) + ".files.in"
Expect(infra.NC.Publish(subject, map[string]string{"path": "/tmp/model"})).To(Succeed())
It("connects with a minted backend worker JWT and publishes on its one remaining allowed subject", func() {
// A backend worker's publish grant is `_INBOX.>` and nothing else now.
// Every verb a frontend gives it, file staging included, is an HTTP
// route on its tunnel, so the `nodes.<id>.files.>` grant went with the
// subjects. See pkg/natsauth.WorkerPermissions.
Expect(infra.NC.Publish("_INBOX.probe", map[string]string{"path": "/tmp/model"})).To(Succeed())
Expect(infra.NC.Conn().FlushTimeout(2 * time.Second)).To(Succeed())
Expect(infra.NC.Conn().LastError()).ToNot(HaveOccurred())
Expect(infra.NC.Conn().IsConnected()).To(BeTrue())
})
It("denies a backend worker the file-staging subjects it no longer serves", func() {
// This spec used to assert the OPPOSITE, and kept passing after the
// grant was deleted. A NATS permission violation does not close the
// connection, so a spec that checks only FlushTimeout and IsConnected
// cannot tell an allowed publish from a denied one; LastError is what
// actually reads the server's verdict, which is why the sibling below
// has always used it.
subject := nodeSubjectPrefix(infra.NodeID) + ".files.stage"
Expect(infra.NC.Publish(subject, map[string]string{"path": "/tmp/model"})).To(Succeed())
Eventually(func() error {
_ = infra.NC.Conn().FlushTimeout(500 * time.Millisecond)
return infra.NC.Conn().LastError()
}, "3s", "50ms").Should(HaveOccurred())
})
It("allows backend subscribe on the node prefix", func() {
wild := nodeSubjectPrefix(infra.NodeID) + ".>"
sub, err := infra.NC.Subscribe(wild, func(_ []byte) {})