Files
LocalAI/core/application/application.go
localai-org-maint-bot 9bfd71387b feat(stores): add Valkey Search vector store backend (#11196)
* feat: add Valkey Search vector store backend

Add a new built-in Go gRPC store backend 'valkey-store' that implements the
four Stores RPCs (Set/Get/Delete/Find) against the Valkey Search module (FT.*)
using the pure-Go github.com/valkey-io/valkey-go client. It is selected via the
existing per-request 'backend' field on /stores, so there is no proto or HTTP
API change, and it mirrors the in-memory local-store while adding persistence
across restarts and opt-in HNSW.

Each vector is a Valkey HASH keyed by hex(little-endian float32); the index is
created lazily on first Set (FLAT+COSINE by default), cosine similarity is
derived as 1-distance, and namespaces get a collision-resistant token. Includes
unit tests (valkey-go mock) and env-gated integration tests against
valkey/valkey-bundle, plus build/matrix/gallery wiring and docs.

Assisted-by: Kiro:claude-opus-4.8 golangci-lint
Signed-off-by: Daria Korenieva <daric2612@gmail.com>

* Address review feedback: recover persisted index dimension, harden Find

- Load now recovers the persisted vector DIM from FT.INFO (not just index
  existence), so a post-restart Set/Find validates against the real DIM
  instead of silently re-learning a wrong one and dropping mismatched
  vectors from the index. This also restores Find's dimension check after
  a restart.
- StoresFind treats a dropped/missing index as an empty store (empty
  result, no error) and clears the stale indexCreated flag, matching
  local-store's empty-store behaviour.
- StoresSet reuses checkDims for its per-key length check so the four RPCs
  share one dimension-guard implementation.
- Add unit tests for FT.INFO dimension recovery, loadIndexState, and the
  dropped-index Find path.

Assisted-by: Kiro:claude-opus-4.8
Signed-off-by: Daria Korenieva <daric2612@gmail.com>

* Address review feedback: TLS ServerName/CA, Find nil-check, config fail-fast

Addresses external review comments on the valkey-store backend:

- StoresFind now rejects a nil/empty query Key before dereferencing it,
  so a malformed gRPC request can no longer panic the backend.
- TLS: derive ServerName (SNI) from the VALKEY_ADDR host so certificate
  verification works for IP-addressed endpoints, and add VALKEY_TLS_CA_CERT
  (custom CA bundle) and VALKEY_TLS_SKIP_VERIFY (testing-only) knobs.
- Config integer parsing now fails fast on a malformed value (e.g.
  VALKEY_HNSW_M=1x6) instead of silently defaulting, matching the
  fail-fast behaviour of the index-algo/distance-metric validation.
- Add VALKEY_DB (SELECT n) support for logical-DB isolation.
- Cap the human-readable part of a namespace token at 64 chars so a very
  long model name cannot produce an unbounded key prefix / index name
  (the appended short hash keeps distinct namespaces collision-free).
- Document the KNN-query injection-safety invariant (fields are constants)
  and why StoresGet uses a single aggregate DoMulti deadline for reads.
- Unit tests for the Find nil/empty-key guard, fail-fast HNSW parsing,
  and VALKEY_DB parsing/validation; docs + .env updated for the new vars.

Assisted-by: Kiro:claude-opus-4.8 golangci-lint
Signed-off-by: Daria Korenieva <daric2612@gmail.com>

* Address review feedback: configure valkey-store via model config

richiejp asked that the valkey-store backend take its configuration from
a model config rather than process-wide VALKEY_* environment variables,
so multiple stores can each have their own Valkey config within one
LocalAI process. This removes every env access from the backend and
routes config through the model-config seam every other backend uses.

- config.go: loadConfig(opts *pb.ModelOptions) now parses the model
  config `options:` list (key:value strings, split on the first ':')
  instead of os.Getenv. Option keys mirror the old VALKEY_* names without
  the prefix (addr, index_algo, distance_metric, ...). Defaults, fail-fast
  validation and the mandatory client name are unchanged.
- store.go: Load threads opts into loadConfig; TLS comments/errors renamed
  off the VALKEY_* names.
- core/backend/stores.go: StoreBackend and NewVectorStore take a
  *config.ModelConfigLoader, resolve the per-store ModelConfig by store
  name, and pass its Options (and Backend when unset) to the backend via
  WithLoadGRPCLoadModelOpts. No config -> default backend + built-in
  defaults, preserving the zero-config experience.
- Endpoints/routes/application: thread the config loader to StoreBackend.
- Unit + integration tests: configure via options; the integration test
  passes addr through the model-config path (VALKEY_ADDR is now only the
  test harness locating the server).
- docs + .env: document the model-config options, drop the env var table.

Assisted-by: Kiro:claude-opus-4.8
Signed-off-by: Daria Korenieva <daric2612@gmail.com>

* Remove valkey-store informational comment from .env The backend is configured via model config, not env vars — the comment was unnecessary noise in .env. The configuration is already documented in docs/content/features/stores.md.

Signed-off-by: Daria Korenieva <daric2612@gmail.com>

* feat(valkey-store): gate Load on NamespacePrefix to refuse autoload probing Mirror local-store's pattern: reject model names without store.NamespacePrefix so the model loader's greedy autoload probe cannot bind an arbitrary model name to the vector store backend (the #9287 failure mode). Also adds unit tests for the gate covering: prefixed namespace, prefix alone, unprefixed model name, empty model, and nil opts.

Signed-off-by: Daria Korenieva <daric2612@gmail.com>

* feat(valkey-store): add username_env/password_env credential indirection Add support for resolving Valkey credentials from environment variables named in the model config, mirroring cloud-proxy's api_key_env pattern. This keeps secrets out of model YAML files and lets distinct store configs each reference their own credentials. Options: username_env / password_env name the env var holding the value. The direct username / password options still work and take precedence when both are set (backward compatible). Includes 5 unit tests and updated stores.md documentation.

Signed-off-by: Daria Korenieva <daric2612@gmail.com>

* fix: correct rebase artifacts in backend-matrix.yml and Makefile Fix two issues introduced by the conflict-resolution script during the rebase onto master: 1. .github/backend-matrix.yml: valkey-store entries were merged INTO the cloud-proxy entries (duplicate keys in same YAML map items) instead of being separate list items. This broke cloud-proxy Linux builds and the cloud-proxy darwin entry lost its build-type/lang. Fixed by making them standalone entries and restoring cloud-proxy exactly as on master. 2. Makefile: duplicated .NOTPARALLEL and docker-build-backends lines. Collapsed to single lines that are master's current content plus the valkey-store additions. Also adds the three optional pickups from #10801: - /valkey-store in .gitignore (the built binary) - valkey-store row in docs/content/reference/compatibility-table.md - valkey-store line in backend/README.md

Signed-off-by: Daria Korenieva <daric2612@gmail.com>

---------

Signed-off-by: Daria Korenieva <daric2612@gmail.com>
Co-authored-by: Daria Korenieva <daric2612@gmail.com>
2026-07-29 20:12:29 +02:00

666 lines
25 KiB
Go

package application
import (
"context"
"math/rand/v2"
"sync"
"sync/atomic"
"time"
corebackend "github.com/mudler/LocalAI/core/backend"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/http/auth"
mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
"github.com/mudler/LocalAI/core/services/agentpool"
"github.com/mudler/LocalAI/core/services/cloudproxy/mitm"
"github.com/mudler/LocalAI/core/services/facerecognition"
"github.com/mudler/LocalAI/core/services/galleryop"
"github.com/mudler/LocalAI/core/services/monitoring"
"github.com/mudler/LocalAI/core/services/nodes"
"github.com/mudler/LocalAI/core/services/routing/admission"
"github.com/mudler/LocalAI/core/services/routing/billing"
"github.com/mudler/LocalAI/core/services/routing/pii"
"github.com/mudler/LocalAI/core/services/routing/piidetector"
"github.com/mudler/LocalAI/core/services/routing/router"
"github.com/mudler/LocalAI/core/services/voiceprofile"
"github.com/mudler/LocalAI/core/services/voicerecognition"
"github.com/mudler/LocalAI/core/templates"
pkggrpc "github.com/mudler/LocalAI/pkg/grpc"
localaitools "github.com/mudler/LocalAI/pkg/mcp/localaitools"
localaiInproc "github.com/mudler/LocalAI/pkg/mcp/localaitools/inproc"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/LocalAI/pkg/signals"
"github.com/mudler/xlog"
"gorm.io/gorm"
)
// faceEmbeddingDim is the expected dimension for face embeddings.
// Set to 0 so the Registry accepts whatever dim the loaded recognizer
// produces — ArcFace R50 is 512-d, MBF is 512-d, SFace is 128-d, and
// the insightface backend can load any of them via LoadModel options.
// Locking this to a specific value would force a single recognizer
// family per deployment; we keep the door open instead.
const faceEmbeddingDim = 0
// voiceEmbeddingDim is the expected dimension for speaker embeddings.
// 0 so the Registry accepts whatever dim the loaded recognizer
// produces — ECAPA-TDNN is 192, WeSpeaker ResNet34 is 256, 3D-Speaker
// ERes2Net is 192, CAM++ is 512.
const voiceEmbeddingDim = 0
type Application struct {
backendLoader *config.ModelConfigLoader
modelLoader *model.ModelLoader
applicationConfig *config.ApplicationConfig
startupConfig *config.ApplicationConfig // Stores original config from env vars (before file loading)
templatesEvaluator *templates.Evaluator
galleryService *galleryop.GalleryService
agentJobService *agentpool.AgentJobService
agentPoolService atomic.Pointer[agentpool.AgentPoolService]
faceRegistry facerecognition.Registry
voiceRegistry voicerecognition.Registry
voiceProfileStore *voiceprofile.Store
authDB *gorm.DB
metricsService *monitoring.LocalAIMetricsService
statsRecorder *billing.Recorder
fallbackUser *auth.User
piiRedactor *pii.Redactor
piiEvents pii.EventStore
mitmCA atomic.Pointer[mitm.CA]
mitmServer atomic.Pointer[mitm.Server]
mitmMutex sync.Mutex // serializes Stop+Start; readers use atomic loads
// mitmHostConflicts records duplicate-host claims across model configs.
// Non-empty disables the MITM listener until resolved — the strict
// 1-to-1 host↔model invariant the dispatcher relies on. Read by
// /api/middleware/status so the admin UI can surface the cause.
mitmHostConflicts atomic.Pointer[map[string][]string]
routerDecisions router.DecisionStore
routerRegistry *router.Registry
admissionLimiter *admission.Limiter
watchdogMutex sync.Mutex
watchdogStop chan bool
p2pMutex sync.Mutex
p2pCtx context.Context
p2pCancel context.CancelFunc
agentJobMutex sync.Mutex
// Distributed mode services (nil when not in distributed mode)
distributed *DistributedServices
// Upgrade checker (background service for detecting backend upgrades)
upgradeChecker *UpgradeChecker
// LocalAI Assistant in-process MCP server. nil when DisableLocalAIAssistant
// is set; otherwise initialised in start() after galleryService.
localAIAssistant *mcpTools.LocalAIAssistantHolder
// startupComplete flips to true once New() has finished its whole startup
// sequence. It backs the /readyz probe.
//
// The expensive step is the model preload: since #10949 it materializes
// HuggingFace artifacts for managed backends, which is tens of GB for a
// large model (31 GB observed on a live cluster). Tracking it explicitly
// means readiness reports lifecycle state instead of "the handler was
// reachable", so a replica that is still starting can be kept out of a
// load balancer's rotation.
startupComplete atomic.Bool
shutdownOnce sync.Once
}
// Ready reports whether the application has finished starting up and can serve
// traffic. It backs the /readyz probe and is safe to call from any goroutine,
// including while startup is still running.
func (a *Application) Ready() bool { return a.startupComplete.Load() }
// markStartupComplete flips the application to ready. Called once, at the very
// end of New(), so every startup step — model preload included — has finished
// before the process advertises itself as able to serve.
func (a *Application) markStartupComplete() { a.startupComplete.Store(true) }
func newApplication(appConfig *config.ApplicationConfig) *Application {
ml := model.NewModelLoader(appConfig.SystemState)
// Apply the per-model load-failure cooldown (0 disables). Set here rather
// than in the watchdog block so it takes effect regardless of whether the
// watchdog/LRU limiter is enabled.
ml.SetLoadFailureCooldown(appConfig.ModelLoadFailureCooldown, 0)
// Close MCP sessions when a model is unloaded (watchdog eviction, manual shutdown, etc.)
ml.OnModelUnload(func(modelName string) {
mcpTools.CloseMCPSessions(modelName)
})
// Record a model_load backend trace for every real backend load, so the
// Traces UI shows which backend runtime served each model and how long
// the load took. Load failures are traced by the modality wrappers.
ml.SetLoadObserver(corebackend.ModelLoadTraceObserver(appConfig))
app := &Application{
backendLoader: config.NewModelConfigLoader(
appConfig.SystemState.Model.ModelsPath,
config.WithArtifactMaterializer(appConfig.ModelArtifactMaterializer),
config.WithPreloadDisplay(appConfig.ModelPreloadRenderMode, appConfig.DisableModelPreloadColor),
),
modelLoader: ml,
applicationConfig: appConfig,
templatesEvaluator: templates.NewEvaluator(appConfig.SystemState.Model.ModelsPath),
voiceProfileStore: voiceprofile.NewStore(appConfig.DataPath),
}
// Face-recognition registry backed by LocalAI's built-in vector store.
// The resolver closes over the ModelLoader so the Registry stays
// decoupled from loader plumbing; swapping in a postgres-backed
// implementation later is a single construction change here.
//
// `faceStoreName` is the default namespace passed to StoreBackend when
// the request doesn't override it. Face and voice MUST use distinct
// namespaces — the local-store gRPC surface rejects mixed dimensions
// inside one namespace ("Try to add key with length N when existing
// length is M"). ArcFace buffalo_l produces 512-dim embeddings while
// ECAPA-TDNN produces 192-dim; enrolling one after the other into a
// shared namespace is exactly how we hit that error.
const (
faceStoreName = "localai-face-biometrics"
voiceStoreName = "localai-voice-biometrics"
)
faceStoreResolver := func(_ context.Context, storeName string) (pkggrpc.Backend, error) {
return corebackend.StoreBackend(ml, appConfig, app.backendLoader, storeName, "")
}
app.faceRegistry = facerecognition.NewStoreRegistry(faceStoreResolver, faceStoreName, faceEmbeddingDim)
// Voice (speaker) recognition registry — same plumbing, separate
// namespace so embedding spaces stay isolated (a face vector and a
// speaker vector are not comparable and differ in dimensionality).
voiceStoreResolver := func(_ context.Context, storeName string) (pkggrpc.Backend, error) {
return corebackend.StoreBackend(ml, appConfig, app.backendLoader, storeName, "")
}
app.voiceRegistry = voicerecognition.NewStoreRegistry(voiceStoreResolver, voiceStoreName, voiceEmbeddingDim)
return app
}
func (a *Application) ModelConfigLoader() *config.ModelConfigLoader {
return a.backendLoader
}
func (a *Application) ModelLoader() *model.ModelLoader {
return a.modelLoader
}
func (a *Application) ApplicationConfig() *config.ApplicationConfig {
return a.applicationConfig
}
func (a *Application) TemplatesEvaluator() *templates.Evaluator {
return a.templatesEvaluator
}
func (a *Application) GalleryService() *galleryop.GalleryService {
return a.galleryService
}
func (a *Application) AgentJobService() *agentpool.AgentJobService {
return a.agentJobService
}
func (a *Application) UpgradeChecker() *UpgradeChecker {
return a.upgradeChecker
}
// LocalAIAssistant returns the in-process MCP holder used by the chat handler
// when an admin opts into the assistant modality. Returns nil when the feature
// is disabled at startup.
func (a *Application) LocalAIAssistant() *mcpTools.LocalAIAssistantHolder {
return a.localAIAssistant
}
// distributedDB returns the PostgreSQL database for distributed coordination,
// or nil in standalone mode.
func (a *Application) distributedDB() *gorm.DB {
if a.distributed != nil {
return a.authDB
}
return nil
}
func (a *Application) AgentPoolService() *agentpool.AgentPoolService {
return a.agentPoolService.Load()
}
// FaceRegistry returns the face-recognition registry used for 1:N
// identification. The current implementation is backed by the
// in-memory local-store backend; see core/services/facerecognition
// for the interface and the postgres TODO.
func (a *Application) FaceRegistry() facerecognition.Registry {
return a.faceRegistry
}
// VoiceRegistry returns the voice (speaker) recognition registry used
// for 1:N identification. Same in-memory local-store backing as
// FaceRegistry but a separate instance — voice embeddings live in
// their own vector space.
func (a *Application) VoiceRegistry() voicerecognition.Registry {
return a.voiceRegistry
}
// VoiceProfileStore returns the persistent library of reusable voice-cloning
// references. It is distinct from VoiceRegistry, which stores speaker
// recognition embeddings rather than synthesis reference audio.
func (a *Application) VoiceProfileStore() *voiceprofile.Store {
return a.voiceProfileStore
}
// AuthDB returns the auth database connection, or nil if auth is not enabled.
func (a *Application) AuthDB() *gorm.DB {
return a.authDB
}
// MetricsService returns the OTel + Prometheus metric service. nil when
// --disable-metrics is set or initialisation failed at startup.
//
// The service is created in startup.go before any counter is registered
// so that otel.SetMeterProvider runs early enough for the billing
// recorder's counters to bind to the Prom-backed provider rather than
// the no-op global. core/http/app.go reuses this instance instead of
// constructing its own — two providers would orphan one set of counters
// behind whichever provider lost the SetMeterProvider race.
func (a *Application) MetricsService() *monitoring.LocalAIMetricsService {
return a.metricsService
}
// StatsRecorder returns the billing recorder used by the usage
// middleware. It is non-nil whenever stats are not explicitly disabled
// — i.e., the no-auth single-user path still gets a working recorder
// (in-memory by default). Routes register UsageMiddleware against this
// recorder regardless of auth state.
func (a *Application) StatsRecorder() *billing.Recorder {
return a.statsRecorder
}
// FallbackUser is the synthetic "local" user that UsageMiddleware uses
// to attribute requests when no authenticated user is on the context
// (i.e., --auth is off). nil when auth is on, since real users are
// always available there.
func (a *Application) FallbackUser() *auth.User {
return a.fallbackUser
}
// PIIRedactor returns the regex-tier PII redactor or nil if PII
// filtering is disabled. The chat-route middleware uses this to apply
// redaction before dispatch.
func (a *Application) PIIRedactor() *pii.Redactor {
return a.piiRedactor
}
// PIIEvents returns the PII event store. Same nil-when-disabled
// semantics as PIIRedactor; admin REST and MCP read tools call List
// against it.
func (a *Application) PIIEvents() pii.EventStore {
return a.piiEvents
}
// PIINERResolver returns the resolver the chat PII middleware uses to
// turn a configured detector model name into a ready-to-use NERConfig:
// a token-classifier bound over the shared model loader (lazy — the
// model loads on first Detect) plus the detection policy read from that
// model's own pii_detection block. Unknown names resolve to (zero,
// false) so the middleware fails closed. Pass it via pii.WithNERResolver.
func (a *Application) PIINERResolver() pii.NERDetectorResolver {
return func(modelName string) (pii.NERConfig, bool) {
if modelName == "" {
return pii.NERConfig{}, false
}
cfg, ok := a.ModelConfigLoader().GetModelConfig(modelName)
if !ok {
return pii.NERConfig{}, false
}
// Pattern detectors match secrets with the restricted-regex tier
// in-process (no backend load). Build a pattern matcher instead of the
// gRPC token-classifier; on a compile error fail closed with an error
// detector so the request is blocked, not silently unscanned.
if cfg.IsPatternDetector() {
det, err := piidetector.NewPattern(cfg, a.ApplicationConfig())
if err != nil {
det = pii.NewErrNERDetector(err.Error())
}
return pii.NERConfigFromRaw(
det,
0, // patterns are deterministic — no confidence floor
cfg.PIIDetectionDefaultAction(),
patternEntityActions(cfg),
pii.SourcePattern,
), true
}
det := piidetector.New(a.ModelLoader(), cfg, a.ApplicationConfig())
return pii.NERConfigFromRaw(
det,
cfg.PIIDetectionMinScore(),
cfg.PIIDetectionDefaultAction(),
cfg.PIIDetectionEntityActions(),
pii.SourceNER,
), true
}
}
// patternEntityActions merges a pattern detector's per-pattern Action overrides
// into its entity_actions map. A pattern reports matches under its Name, so a
// per-pattern action is just an entity_actions[Name] entry; explicit
// entity_actions still win if both are set.
func patternEntityActions(cfg config.ModelConfig) map[string]string {
out := cfg.PIIDetectionEntityActions()
for _, p := range cfg.PIIDetection.Patterns {
if p.Action == "" || p.Name == "" {
continue
}
if out == nil {
out = map[string]string{}
}
if _, exists := out[p.Name]; !exists {
out[p.Name] = p.Action
}
}
return out
}
// ResolvePIIPolicy resolves the effective request-side PII policy for a
// consuming model, layering the instance-wide default detector
// (PIIDefaultDetectors, set via POST /api/settings) on top of the per-model
// config. It is the single decision point shared by the chat middleware (via
// WithPolicyResolver) and the MITM listener so both agree.
//
// - enabled: an explicit pii.enabled on the model always wins (true OR
// false). Otherwise PII is on when the backend defaults it on — today
// that means cloud-proxy models, which cross the network to a third party.
// - detectors: the model's own pii.detectors, or — when it lists none — the
// global PIIDefaultDetectors fallback. This is what makes cloud-proxy/MITM
// redaction work out of the box.
//
// appConfig is read live, so changes via the settings API take effect on the
// next request without a restart.
func (a *Application) ResolvePIIPolicy(cfg *config.ModelConfig) (enabled bool, detectors []string) {
if cfg == nil {
return false, nil
}
appCfg := a.ApplicationConfig()
// PIIIsEnabled already encodes "explicit pii.enabled wins, else backend
// default (cloud-proxy)" — the single source of that rule.
enabled = cfg.PIIIsEnabled()
if !enabled {
return false, nil
}
detectors = cfg.PIIDetectors()
if len(detectors) == 0 {
detectors = append([]string(nil), appCfg.PIIDefaultDetectors...)
}
return true, detectors // enabled is necessarily true past the !enabled guard
}
// PIIPolicyResolver adapts ResolvePIIPolicy to pii.PolicyResolver for
// pii.WithPolicyResolver. The middleware carries the resolved model config as
// `any` (the MODEL_CONFIG context value, a *config.ModelConfig); this asserts
// it back and applies the instance-wide defaults.
func (a *Application) PIIPolicyResolver() pii.PolicyResolver {
return func(modelCfg any) (bool, []string) {
cfg, ok := modelCfg.(*config.ModelConfig)
if !ok {
return false, nil
}
return a.ResolvePIIPolicy(cfg)
}
}
// MITMCA returns the cloudproxy MITM proxy's CA, or nil when the
// MITM listener is disabled.
func (a *Application) MITMCA() *mitm.CA { return a.mitmCA.Load() }
// MITMServer returns the running MITM proxy or nil.
func (a *Application) MITMServer() *mitm.Server { return a.mitmServer.Load() }
// MITMHostConflicts returns a snapshot of host→[]model-name pairs that
// are claimed by 2+ model configs. Empty when the 1-to-1 invariant
// holds. Non-empty disables the MITM listener — read by the admin
// status endpoint to explain why.
func (a *Application) MITMHostConflicts() map[string][]string {
p := a.mitmHostConflicts.Load()
if p == nil {
return nil
}
return *p
}
// MITMHostOwners returns the host→model-name map, useful for the
// admin status endpoint. The lookup is recomputed on each call to
// stay current with model-config edits without needing a
// MITMRestart.
func (a *Application) MITMHostOwners() map[string]string {
if a.backendLoader == nil {
return nil
}
return a.backendLoader.MITMHostOwners().Owners
}
// RouterDecisions returns the routing decision store. nil when stats
// are disabled (--disable-stats); the RouteModel middleware skips the
// log write in that case but still rewrites requests.
func (a *Application) RouterDecisions() router.DecisionStore {
return a.routerDecisions
}
// RouterClassifierRegistry returns the process-wide classifier cache.
// Shared between the OpenAI and Anthropic route middlewares so the
// admin stats endpoint sees every live classifier — and so a
// classifier built on the OpenAI route is reused on Anthropic.
func (a *Application) RouterClassifierRegistry() *router.Registry {
return a.routerRegistry
}
// AdmissionLimiter returns the per-model admission limiter. The
// admission middleware uses it to gate concurrent requests; the
// admin status surface reads InFlight/Capacity from it for live
// load visibility.
func (a *Application) AdmissionLimiter() *admission.Limiter {
return a.admissionLimiter
}
// StartupConfig returns the original startup configuration (from env vars, before file loading)
func (a *Application) StartupConfig() *config.ApplicationConfig {
return a.startupConfig
}
// Distributed returns the distributed services, or nil if not in distributed mode.
func (a *Application) Distributed() *DistributedServices {
return a.distributed
}
// IsDistributed returns true if the application is running in distributed mode.
func (a *Application) IsDistributed() bool {
return a.distributed != nil
}
// Shutdown stops backend gRPC processes and distributed services
// synchronously on the caller's stack. The context-cancel goroutine wired
// in New does the same work asynchronously, which races test-binary exit
// and CLI shutdown — orphaning spawned mock-backend / llama.cpp / etc.
// children to init. Callers that need a guarantee that cleanup has
// finished before they proceed (AfterSuite/AfterEach, signal handlers)
// must call this. Safe to call multiple times.
func (a *Application) Shutdown() error {
var err error
a.shutdownOnce.Do(func() {
a.distributed.Shutdown()
if a.modelLoader != nil {
err = a.modelLoader.StopAllGRPC()
}
if a.voiceProfileStore != nil {
if closeErr := a.voiceProfileStore.Close(); err == nil {
err = closeErr
}
}
})
return err
}
// waitForHealthyWorker blocks until at least one healthy backend worker is registered.
// This prevents the agent pool from failing during startup when workers haven't connected yet.
func (a *Application) waitForHealthyWorker() {
maxWait := a.applicationConfig.Distributed.WorkerWaitTimeoutOrDefault()
const basePoll = 2 * time.Second
xlog.Info("Waiting for at least one healthy backend worker before starting agent pool")
deadline := time.Now().Add(maxWait)
for time.Now().Before(deadline) {
registered, err := a.distributed.Registry.List(context.Background())
if err == nil {
for _, n := range registered {
if n.NodeType == nodes.NodeTypeBackend && n.Status == nodes.StatusHealthy {
xlog.Info("Healthy backend worker found", "node", n.Name)
return
}
}
}
// Add 0-1s jitter to prevent thundering-herd on the node registry
jitter := time.Duration(rand.Int64N(int64(time.Second)))
select {
case <-a.applicationConfig.Context.Done():
return
case <-time.After(basePoll + jitter):
}
}
xlog.Warn("No healthy backend worker found after waiting, proceeding anyway")
}
// InstanceID returns the unique identifier for this frontend instance.
func (a *Application) InstanceID() string {
return a.applicationConfig.Distributed.InstanceID
}
func (a *Application) start() error {
galleryService := galleryop.NewGalleryService(a.ApplicationConfig(), a.ModelLoader())
err := galleryService.Start(a.ApplicationConfig().Context, a.ModelConfigLoader(), a.ApplicationConfig().SystemState)
if err != nil {
return err
}
a.galleryService = galleryService
// LocalAI Assistant: in-process MCP server exposing admin tools. Initialised
// once at startup and reused across chat sessions that opt in via metadata.
if !a.applicationConfig.DisableLocalAIAssistant {
holder := mcpTools.NewLocalAIAssistantHolder()
assistantClient := localaiInproc.New(
a.applicationConfig,
a.applicationConfig.SystemState,
a.backendLoader,
a.modelLoader,
a.galleryService,
)
// Wire usage tracking so the assistant's get_usage_stats tool
// returns real data; nil values keep the tool returning a clear
// "unavailable" error if startup ran with --disable-stats.
assistantClient.StatsRecorder = a.statsRecorder
assistantClient.FallbackUser = a.fallbackUser
assistantClient.VoiceProfiles = a.voiceProfileStore
// PII filter — same nil-or-real wiring.
assistantClient.PIIRedactor = a.piiRedactor
assistantClient.PIIEvents = a.piiEvents
assistantClient.RouterDecisions = a.routerDecisions
if err := holder.Initialize(a.applicationConfig.Context, assistantClient, localaitools.Options{}); err != nil {
// Why log+continue instead of fail: the assistant is an optional
// feature; a failure here must not take down the whole server.
xlog.Warn("LocalAI Assistant initialisation failed; feature unavailable", "error", err)
} else {
a.localAIAssistant = holder
// Tear the in-memory transport pair down on SIGINT/SIGTERM so the
// goroutine ends cleanly. Mirrors how core/http/endpoints/mcp/tools.go
// closes its per-model MCP sessions on graceful termination.
signals.RegisterGracefulTerminationHandler(func() {
_ = holder.Close()
})
}
}
// Initialize agent job service (Start() is deferred to after distributed wiring)
agentJobService := agentpool.NewAgentJobService(
a.ApplicationConfig(),
a.ModelLoader(),
a.ModelConfigLoader(),
a.TemplatesEvaluator(),
)
a.agentJobService = agentJobService
return nil
}
// StartAgentPool initializes and starts the agent pool service (LocalAGI integration).
// This must be called after the HTTP server is listening, because backends like
// PostgreSQL need to call the embeddings API during collection initialization.
func (a *Application) StartAgentPool() {
if !a.applicationConfig.AgentPool.Enabled {
return
}
// Build options struct from available dependencies
opts := agentpool.AgentPoolOptions{
AuthDB: a.authDB,
}
if d := a.Distributed(); d != nil {
if d.DistStores != nil && d.DistStores.Skills != nil {
opts.SkillStore = d.DistStores.Skills
}
opts.NATSClient = d.Nats
opts.EventBridge = d.AgentBridge
opts.AgentStore = d.AgentStore
}
aps, err := agentpool.NewAgentPoolService(a.applicationConfig, opts)
if err != nil {
xlog.Error("Failed to create agent pool service", "error", err)
return
}
// Wire distributed mode components
if d := a.Distributed(); d != nil {
// Wait for at least one healthy backend worker before starting the agent pool.
// Collections initialization calls embeddings which require a worker.
if d.Registry != nil {
a.waitForHealthyWorker()
}
}
if err := aps.Start(a.applicationConfig.Context); err != nil {
xlog.Error("Failed to start agent pool", "error", err)
return
}
// Wire per-user scoped services so collections, skills, and jobs are isolated per user
usm := agentpool.NewUserServicesManager(
aps.UserStorage(),
a.applicationConfig,
a.modelLoader,
a.backendLoader,
a.templatesEvaluator,
)
// Wire distributed backends to per-user job services
if a.agentJobService != nil {
if d := a.agentJobService.Dispatcher(); d != nil {
usm.SetJobDispatcher(d)
}
if s := a.agentJobService.DBStore(); s != nil {
usm.SetJobDBStore(s)
}
}
// Keep per-user agent tasks consistent across replicas (nil in standalone).
if d := a.Distributed(); d != nil {
usm.SetJobSyncNATS(d.Nats)
}
aps.SetUserServicesManager(usm)
a.agentPoolService.Store(aps)
}