harden(distributed): defensive companion-artifact option handling (follow-ups to #11075)

core/application/distributed.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/mudler/LocalAI/core/backend"
 	"github.com/mudler/LocalAI/core/config"
 	"github.com/mudler/LocalAI/core/services/agents"
 	"github.com/mudler/LocalAI/core/services/distributed"
@@ -377,6 +378,22 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
 			cfg.Distributed.BackendInstallTimeoutOrDefault(),
 			cfg.Distributed.ModelLoadTimeoutOrDefault(),
 		),
+		// Re-derive managed companion options from the model's current config when
+		// the reconciler replays a stored ModelOptions blob (which never runs
+		// grpcModelOpts). Without this, a replica scaled up from a blob captured
+		// before the companion resolved loads without the companion option and the
+		// backend fetches its own wrong default. nil-safe: an unknown model or a
+		// config with no companions yields no options.
+		CompanionOptionsFor: func(modelName string) []string {
+			if configLoader == nil {
+				return nil
+			}
+			modelCfg, ok := configLoader.GetModelConfig(modelName)
+			if !ok {
+				return nil
+			}
+			return backend.CompanionArtifactOptions(modelCfg)
+		},
 	})
 
 	// Wire staging-progress broadcasting so file-staging shows up on every

core/backend/companion_options_test.go

@@ -113,11 +113,34 @@ var _ = Describe("companion artifact backend options", func() {
 		Expect(opts.Options).To(Equal([]string{"attention_backend:sdpa"}))
 	})
 
-	It("skips a companion that has not been resolved yet", func() {
+	It("names the source repository when the companion is not resolved yet", func() {
+		// A companion that reaches load time WITHOUT a resolved snapshot must not
+		// vanish silently: emitting no option lets the backend fall back to its own
+		// hardcoded default, which is how a distributed longcat-video worker ended
+		// up trying to load the wrong base model and failing "base_model must point
+		// to a LongCat-Video checkpoint". Naming the DECLARED repository instead
+		// points the backend at the artifact the config actually asked for. The
+		// snapshot path (the staged, no-download fast path) is still preferred
+		// whenever the companion IS resolved.
 		cfg := configWithCompanion()
 		cfg.Artifacts[1].Resolved = nil
 		opts := grpcModelOpts(cfg, "/models")
-		_, found := optionValue(opts.Options, "base_model")
-		Expect(found).To(BeFalse())
+
+		value, found := optionValue(opts.Options, "base_model")
+		Expect(found).To(BeTrue())
+		Expect(value).To(Equal("meituan-longcat/LongCat-Video"))
+		// The fallback is a repo reference, never a models-relative snapshot path.
+		Expect(value).ToNot(ContainSubstring(".artifacts"))
+	})
+
+	It("prefers the resolved snapshot path over the source repository", func() {
+		opts := grpcModelOpts(configWithCompanion(), "/models")
+		value, found := optionValue(opts.Options, "base_model")
+		Expect(found).To(BeTrue())
+		expected, err := modelartifacts.RelativeSnapshotPath(companionKey)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(value).To(Equal(expected))
+		// The resolved fast path must never degrade to a bare repo id.
+		Expect(value).ToNot(Equal("meituan-longcat/LongCat-Video"))
 	})
 })

core/backend/options.go

@@ -294,33 +294,85 @@ func EffectiveBatchSize(c config.ModelConfig) int {
 //
 // An option the author set explicitly always wins: pinning a companion to a
 // local checkout has to beat the managed snapshot.
+//
+// A companion that is declared but NOT resolved falls back to its source
+// repository id rather than being dropped: a dropped companion is invisible to
+// the backend, which then loads its own hardcoded default and fails far away
+// from the cause. The repo-id fallback trades the staging fast path (the weights
+// are fetched on the worker) for correctness, and logs a warning so the missing
+// controller-side resolution is diagnosable.
 func withCompanionArtifactOptions(options []string, artifacts []modelartifacts.Spec) []string {
 	configured := make(map[string]struct{}, len(options))
 	for _, option := range options {
 		if name, _, found := strings.Cut(option, ":"); found {
 			configured[name] = struct{}{}
 		}
 	}
-
 	// Copy before appending: opts.Options would otherwise share (and could
 	// reallocate away from) the config's own slice.
-	combined := slices.Clone(options)
+	return append(slices.Clone(options), ManagedCompanionOptions(artifacts, configured)...)
+}
+
+// ManagedCompanionOptions synthesizes the "<artifact name>:<value>" option for
+// every companion artifact whose name is not already in `configured`, so callers
+// can append them to a ModelOptions.Options slice. It is exported because the
+// distributed reconciler replays a stored proto blob that never runs
+// grpcModelOpts, and it must re-derive the same companion options from the
+// current config so a scaled-up replica gets them too.
+//
+// A resolved companion is surfaced as its staged, models-relative snapshot
+// directory (the no-download fast path a remote worker resolves under its own
+// ModelPath). A companion that reached load time WITHOUT a resolved snapshot
+// falls back to its source repository id rather than being dropped: dropping it
+// is invisible to the backend, which then loads its OWN hardcoded default and
+// fails far from the cause (a distributed longcat-video worker fetched the wrong
+// base model and failed "base_model must point to a LongCat-Video checkpoint").
+func ManagedCompanionOptions(artifacts []modelartifacts.Spec, configured map[string]struct{}) []string {
+	var out []string
 	for _, artifact := range artifacts {
-		if artifact.Target != modelartifacts.TargetCompanion || artifact.Resolved == nil {
+		if artifact.Target != modelartifacts.TargetCompanion {
 			continue
 		}
-		if _, exists := configured[artifact.Name]; exists {
-			xlog.Debug("keeping the configured companion option over the managed snapshot", "artifact", artifact.Name)
-			continue
+		if configured != nil {
+			if _, exists := configured[artifact.Name]; exists {
+				xlog.Debug("keeping the configured companion option over the managed snapshot", "artifact", artifact.Name)
+				continue
+			}
 		}
-		snapshot, err := modelartifacts.RelativeSnapshotPath(artifact.Resolved.CacheKey)
-		if err != nil {
-			xlog.Warn("skipping companion artifact with an unusable cache key", "artifact", artifact.Name, "error", err)
+
+		if artifact.Resolved != nil {
+			if snapshot, err := modelartifacts.RelativeSnapshotPath(artifact.Resolved.CacheKey); err == nil {
+				out = append(out, artifact.Name+":"+snapshot)
+				continue
+			} else {
+				xlog.Warn("companion artifact has an unusable cache key; falling back to its source repository", "artifact", artifact.Name, "error", err)
+			}
+		}
+
+		if repo := strings.TrimSpace(artifact.Source.Repo); repo != "" {
+			xlog.Warn("companion artifact is not resolved on the controller; the backend will fetch it by repository id (no staging fast path)",
+				"artifact", artifact.Name, "repo", repo)
+			out = append(out, artifact.Name+":"+repo)
 			continue
 		}
-		combined = append(combined, artifact.Name+":"+snapshot)
+		xlog.Warn("companion artifact is neither resolved nor has a source repository; the backend will get no option for it", "artifact", artifact.Name)
+	}
+	return out
+}
+
+// CompanionArtifactOptions returns the managed companion options a model's
+// current config would contribute, skipping any companion name the config
+// already pins in Options. It is the reconciler's entry point for re-deriving
+// companion options when it replays a stored ModelOptions blob (which predates,
+// and so cannot carry, a companion resolved after that blob was captured).
+func CompanionArtifactOptions(c config.ModelConfig) []string {
+	configured := make(map[string]struct{}, len(c.Options))
+	for _, option := range c.Options {
+		if name, _, found := strings.Cut(option, ":"); found {
+			configured[name] = struct{}{}
+		}
 	}
-	return combined
+	return ManagedCompanionOptions(c.Artifacts, configured)
 }
 
 func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {

core/services/nodes/file_staging_loadmodel_test.go

@@ -0,0 +1,58 @@
+package nodes
+
+import (
+	"context"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	grpc "github.com/mudler/LocalAI/pkg/grpc"
+	pb "github.com/mudler/LocalAI/pkg/grpc/proto"
+	ggrpc "google.golang.org/grpc"
+)
+
+// capturingLoadBackend records the ModelOptions handed to LoadModel, standing in
+// for the remote gRPC backend the FileStagingClient forwards to.
+type capturingLoadBackend struct {
+	grpc.Backend
+	got *pb.ModelOptions
+}
+
+func (b *capturingLoadBackend) LoadModel(_ context.Context, in *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) {
+	b.got = in
+	return &pb.Result{Success: true}, nil
+}
+
+// FileStagingClient is the concrete type of `client` at router.go's
+// client.LoadModel call site (buildClientForAddr wraps the gRPC backend with it
+// whenever a file stager is configured, i.e. every distributed load). This pins
+// the boundary the companion-option investigation kept returning to: the wrapper
+// must forward ModelOptions.Options (the channel a managed companion option like
+// base_model rides on) to the backend byte-for-byte. If a companion option is
+// present in what the controller sends here but absent at the Python backend,
+// this test proves the loss is NOT in this Go handoff.
+var _ = Describe("FileStagingClient LoadModel option pass-through", func() {
+	It("forwards ModelOptions.Options to the backend unchanged", func(ctx SpecContext) {
+		backend := &capturingLoadBackend{}
+		client := NewFileStagingClient(backend, &fakeFileStager{}, "worker-1")
+
+		sent := &pb.ModelOptions{
+			Model:     "longcat-video-avatar-1.5",
+			ModelPath: "/models/longcat-video-avatar-1.5",
+			Options: []string{
+				"attention_backend:sdpa",
+				"use_distill:true",
+				"base_model:.artifacts/huggingface/deadbeef/snapshot",
+			},
+		}
+
+		_, err := client.LoadModel(ctx, sent)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(backend.got).NotTo(BeNil())
+		Expect(backend.got.Options).To(Equal(sent.Options))
+		Expect(backend.got.Options).To(ContainElement("base_model:.artifacts/huggingface/deadbeef/snapshot"))
+		// The nested staged ModelPath the backend joins the relative option against
+		// must survive too.
+		Expect(backend.got.ModelPath).To(Equal("/models/longcat-video-avatar-1.5"))
+	})
+})

core/services/nodes/router.go

@@ -105,6 +105,12 @@ type SmartRouterOptions struct {
 	// arriving, so a peer trickling bytes forever cannot pin the advisory lock
 	// indefinitely. Zero selects modelLoadAbsoluteMax (24h).
 	ModelLoadAbsoluteMax time.Duration
+	// CompanionOptionsFor re-derives a model's managed companion options from its
+	// current config (backend.CompanionArtifactOptions). The reconciler replays a
+	// stored ModelOptions blob that never runs grpcModelOpts; without this a
+	// companion resolved after the blob was captured never reaches a scaled-up
+	// replica. nil disables re-derivation.
+	CompanionOptionsFor func(modelName string) []string
 }
 
 // modelLoadStagingMargin is the slack ModelLoadCeilingFor adds on top of the
@@ -187,6 +193,13 @@ type SmartRouter struct {
 	// hard countdown into a progress-extended hold (see load_deadline.go).
 	stagingStallWindow   time.Duration
 	modelLoadAbsoluteMax time.Duration
+	// companionOptionsFor re-derives a model's managed companion options from its
+	// CURRENT config. The reconciler's scale-up path replays a stored ModelOptions
+	// proto blob that never runs grpcModelOpts, so a companion resolved AFTER that
+	// blob was captured would otherwise never reach the replica. nil disables the
+	// re-derivation (single-node and tests that don't wire it). See
+	// SmartRouterOptions.CompanionOptionsFor.
+	companionOptionsFor func(modelName string) []string
 }
 
 // probeCacheTTL is how long a successful gRPC HealthCheck on a backend is
@@ -238,6 +251,7 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter
 		// ceiling, so nothing to normalize here.
 		stagingStallWindow:   opts.StagingStallWindow,
 		modelLoadAbsoluteMax: opts.ModelLoadAbsoluteMax,
+		companionOptionsFor:  opts.CompanionOptionsFor,
 	}
 }
 
@@ -448,6 +462,39 @@ func (r *SmartRouter) reapAbandonedLoad(node *BackendNode, trackingKey string, r
 	}
 }
 
+// applyCompanionOptions merges a model's freshly re-derived managed companion
+// options into opts.Options, adding only those the blob does not already carry.
+// Merge, not replace: an option the blob already has (including an author-pinned
+// one) wins, so this only ever ADDS a missing companion. No-op when no resolver
+// is wired (single-node) or opts is nil.
+func (r *SmartRouter) applyCompanionOptions(modelName string, opts *pb.ModelOptions) {
+	if r.companionOptionsFor == nil || opts == nil {
+		return
+	}
+	present := make(map[string]struct{}, len(opts.Options))
+	for _, opt := range opts.Options {
+		if name, _, found := strings.Cut(opt, ":"); found {
+			present[name] = struct{}{}
+		}
+	}
+	for _, opt := range r.companionOptionsFor(modelName) {
+		name, _, found := strings.Cut(opt, ":")
+		if !found {
+			continue
+		}
+		if _, exists := present[name]; exists {
+			continue
+		}
+		// Debug (not info): a correctly-captured blob already carries the option,
+		// so this only fires on the rarer stale-blob recovery — worth surfacing to
+		// an operator tracing a companion, but not on every replay.
+		xlog.Debug("reconciler injecting companion option missing from the stored blob",
+			"model", modelName, "option", opt)
+		opts.Options = append(opts.Options, opt)
+		present[name] = struct{}{}
+	}
+}
+
 // ScheduleAndLoadModel implements ModelScheduler for the reconciler.
 // It retrieves stored model options from an existing replica and performs the
 // full load sequence (stage files, LoadModel, SetNodeModel) on a new node.
@@ -470,6 +517,12 @@ func (r *SmartRouter) ScheduleAndLoadModel(ctx context.Context, modelName string
 		return nil, fmt.Errorf("unmarshalling stored model options for %s: %w", modelName, err)
 	}
 
+	// Re-derive managed companion options from the CURRENT config and merge in any
+	// the stored blob is missing. This path replays a proto blob captured at first
+	// load and never runs grpcModelOpts, so a companion resolved after the blob was
+	// captured would otherwise never reach the scaled-up replica.
+	r.applyCompanionOptions(modelName, &modelOpts)
+
 	// initialInFlight=0: reconciler is pre-loading, not serving a request.
 	// scheduleAndLoad picks both the node and the replica slot internally.
 	result, err := r.scheduleAndLoad(ctx, backendType, modelName, modelName, &modelOpts, false, 0)

core/services/nodes/router_companion_replay_test.go

@@ -0,0 +1,61 @@
+package nodes
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	pb "github.com/mudler/LocalAI/pkg/grpc/proto"
+)
+
+// The reconciler's scale-up path replays a stored ModelOptions proto blob and
+// never runs grpcModelOpts, so a managed companion option (e.g. longcat-video's
+// base_model) is present only if it was captured in that blob. A blob stored
+// before the companion resolved permanently lacks it, and the replica loads
+// without the companion — the backend then fetches its own wrong default and
+// fails "base_model must point to a LongCat-Video checkpoint". applyCompanionOptions
+// re-derives the companion from the CURRENT config and injects any that are
+// missing, so a remote replica gets the same option a fresh grpcModelOpts would.
+var _ = Describe("reconciler companion option re-derivation", func() {
+	It("injects a companion option the stored blob is missing", func() {
+		router := &SmartRouter{
+			companionOptionsFor: func(_ string) []string {
+				return []string{"base_model:.artifacts/huggingface/deadbeef/snapshot"}
+			},
+		}
+		// A blob captured before the companion resolved: no base_model.
+		opts := &pb.ModelOptions{Options: []string{"attention_backend:sdpa"}}
+
+		router.applyCompanionOptions("longcat-video-avatar-1.5", opts)
+
+		Expect(opts.Options).To(ContainElement("base_model:.artifacts/huggingface/deadbeef/snapshot"))
+		Expect(opts.Options).To(ContainElement("attention_backend:sdpa"))
+	})
+
+	It("never overrides a companion option the blob already carries", func() {
+		router := &SmartRouter{
+			companionOptionsFor: func(_ string) []string {
+				return []string{"base_model:.artifacts/huggingface/managed/snapshot"}
+			},
+		}
+		// An author-pinned (or already-captured) base_model must win.
+		opts := &pb.ModelOptions{Options: []string{"base_model:/opt/checkouts/LongCat-Video"}}
+
+		router.applyCompanionOptions("longcat-video-avatar-1.5", opts)
+
+		Expect(opts.Options).To(ContainElement("base_model:/opt/checkouts/LongCat-Video"))
+		base := 0
+		for _, o := range opts.Options {
+			if len(o) >= len("base_model:") && o[:len("base_model:")] == "base_model:" {
+				base++
+			}
+		}
+		Expect(base).To(Equal(1))
+	})
+
+	It("is a no-op when no resolver is wired (single-node)", func() {
+		router := &SmartRouter{}
+		opts := &pb.ModelOptions{Options: []string{"attention_backend:sdpa"}}
+		router.applyCompanionOptions("m", opts)
+		Expect(opts.Options).To(Equal([]string{"attention_backend:sdpa"}))
+	})
+})