From 3a02f29edc6def68169a76635882020e49e7082e Mon Sep 17 00:00:00 2001 From: mudler-agent Date: Tue, 27 Oct 2026 10:27:45 +0100 Subject: [PATCH] fix(proxy): propagate model env and accept direct keys (#12940) * feat(proxy): accept direct upstream API keys Support direct credentials for models created in the UI. Keep env and file references as mutually exclusive alternatives. Use model env for frontend probes and omit credentials from structured config logs. Assisted-by: nib:gpt-6.1-sol Signed-off-by: Ettore Di Giacinto * feat(distributed): isolate model environments in worker processes Forward complete model environment snapshots through backend installation. Restart worker children when overrides change or are removed, and include keyed snapshot identities in frontend install coalescing. Cover wire compatibility, routing, real worker children, and option copying; document rollout requirements and control-transport security. Assisted-by: nib:gpt-6.1-sol Signed-off-by: Ettore Di Giacinto * fix(worker): serialize installs by model replica Different backend names can target the same process slot. Hold a slot lock through startup, readiness, and operation attachment so concurrent installs cannot reuse another configuration or an unready child. Assisted-by: nib:gpt-6.1-sol Signed-off-by: Ettore Di Giacinto * test(nodes): update endpoint stub for model env Keep the endpoint command stub aligned with the environment-aware install interface. Assert interface conformance so future changes fail at the stub declaration rather than at endpoint call sites. Assisted-by: nib:gpt-6.1-sol Signed-off-by: Ettore Di Giacinto --------- Signed-off-by: Ettore Di Giacinto Co-authored-by: Ettore Di Giacinto --- backend/backend.proto | 6 +- backend/go/cloud-proxy/credentials_test.go | 48 +++++ backend/go/cloud-proxy/proxy.go | 18 +- backend/go/localai-proxy/credentials_test.go | 48 +++++ backend/go/localai-proxy/proxy.go | 16 +- core/backend/options.go | 1 + core/backend/options_internal_test.go | 10 + core/config/meta/registry.go | 7 + core/config/meta/registry_test.go | 6 + core/config/model_config.go | 48 ++++- core/config/proxy_credentials_test.go | 44 +++++ .../localai/nodes_backends_list_test.go | 6 +- .../src/components/ConfigFieldRenderer.jsx | 3 +- .../tests/unit/configpassword.test.mjs | 36 ++++ .../failover/model_env_credentials_test.go | 34 ++++ core/services/failover/prober.go | 13 +- core/services/nodes/env_test.go | 60 ++++++ core/services/nodes/interfaces.go | 2 +- core/services/nodes/load_worker_ops_test.go | 6 +- core/services/nodes/router.go | 31 ++- core/services/nodes/router_test.go | 11 +- core/services/nodes/unloader.go | 20 +- core/services/worker/env_test.go | 187 ++++++++++++++++++ core/services/worker/install.go | 26 ++- core/services/worker/lifecycle.go | 5 + core/services/worker/supervisor.go | 37 +++- core/services/workerctl/backend.go | 38 +++- core/services/workerctl/env_test.go | 29 +++ docs/content/features/distributed-mode.md | 27 +++ docs/content/operations/cloud-proxy.md | 29 ++- pkg/model/env_isolation_test.go | 18 ++ pkg/model/initializers.go | 5 +- pkg/model/loader_options.go | 4 +- 33 files changed, 820 insertions(+), 59 deletions(-) create mode 100644 backend/go/cloud-proxy/credentials_test.go create mode 100644 backend/go/localai-proxy/credentials_test.go create mode 100644 core/config/proxy_credentials_test.go create mode 100644 core/http/react-ui/tests/unit/configpassword.test.mjs create mode 100644 core/services/failover/model_env_credentials_test.go create mode 100644 core/services/nodes/env_test.go create mode 100644 core/services/worker/env_test.go create mode 100644 core/services/workerctl/env_test.go create mode 100644 pkg/model/env_isolation_test.go diff --git a/backend/backend.proto b/backend/backend.proto index 2d3f0802b..79be236b2 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -544,9 +544,8 @@ message ModelOptions { // ProxyOptions configures the cloud-proxy backend. UpstreamURL and // Mode are always meaningful; Provider only matters in translate mode. -// The two api_key_* fields are mutually exclusive and resolved by the -// backend at LoadModel — core forwards the references rather than the -// plaintext key. +// The api_key, api_key_env and api_key_file fields are mutually exclusive +// and resolved by the backend at LoadModel. api_key is plaintext admin config. message ProxyOptions { string upstream_url = 1; string mode = 2; @@ -561,6 +560,7 @@ message ProxyOptions { // Cuts input cost on repeated/agentic calls (cache read = 0.1x). Only // meaningful for mode=translate + provider=anthropic; ignored otherwise. bool cache_prompt = 8; + string api_key = 9; } message Result { diff --git a/backend/go/cloud-proxy/credentials_test.go b/backend/go/cloud-proxy/credentials_test.go new file mode 100644 index 000000000..b2d0eb6e1 --- /dev/null +++ b/backend/go/cloud-proxy/credentials_test.go @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: MIT +package main + +import ( + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "os" + "path/filepath" +) + +var _ = Describe("Proxy credential sources", func() { + It("resolves direct, absent, env and file credentials", func() { + key, err := resolveAPIKey("direct-canary", "", "") + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("direct-canary")) + key, err = resolveAPIKey("", "", "") + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(BeEmpty()) + DeferCleanup(os.Unsetenv, "LOCALAI_TEST_PROXY_CREDENTIAL") + Expect(os.Setenv("LOCALAI_TEST_PROXY_CREDENTIAL", "env-canary")).To(Succeed()) + key, err = resolveAPIKey("", "LOCALAI_TEST_PROXY_CREDENTIAL", "") + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("env-canary")) + dir, err := os.MkdirTemp("", "proxy-credentials-") + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(os.RemoveAll, dir) + path := filepath.Join(dir, "key") + Expect(os.WriteFile(path, []byte("file-canary\n"), 0600)).To(Succeed()) + key, err = resolveAPIKey("", "", path) + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("file-canary")) + _, err = resolveAPIKey("", "", path+"-missing") + Expect(err).To(HaveOccurred()) + }) + It("rejects every conflicting source without exposing the key", func() { + for _, args := range [][3]string{{"secret-canary", "ENV", ""}, {"secret-canary", "", "/missing"}, {"", "ENV", "/missing"}, {"secret-canary", "ENV", "/missing"}} { + _, err := resolveAPIKey(args[0], args[1], args[2]) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).NotTo(ContainSubstring("secret-canary")) + } + }) + It("loads direct credentials from protobuf options", func() { + proxy := NewCloudProxy() + Expect(proxy.Load(&pb.ModelOptions{Model: "alias", Proxy: &pb.ProxyOptions{UpstreamUrl: "http://example.invalid", ApiKey: "direct-canary"}})).To(Succeed()) + Expect(proxy.cfg.Load().apiKey).To(Equal("direct-canary")) + }) +}) diff --git a/backend/go/cloud-proxy/proxy.go b/backend/go/cloud-proxy/proxy.go index f6e3a9f1b..9c87e6068 100644 --- a/backend/go/cloud-proxy/proxy.go +++ b/backend/go/cloud-proxy/proxy.go @@ -95,7 +95,7 @@ func (c *CloudProxy) Load(opts *pb.ModelOptions) error { return fmt.Errorf("cloud-proxy: unknown mode %q", mode) } - key, err := resolveAPIKey(po.GetApiKeyEnv(), po.GetApiKeyFile()) + key, err := resolveAPIKey(po.GetApiKey(), po.GetApiKeyEnv(), po.GetApiKeyFile()) if err != nil { return err } @@ -120,8 +120,20 @@ func (c *CloudProxy) Load(opts *pb.ModelOptions) error { // resolveAPIKey mirrors config.ProxyConfig.ResolveAPIKey. Duplicated // (a few lines) rather than importing core/config from a backend // binary — keeps backends independent of core's package layout. -// Mutual-exclusion is enforced upstream in core/config.Validate. -func resolveAPIKey(envName, filePath string) (string, error) { +// Also validate here for callers that bypass core config validation. +func resolveAPIKey(directKey, envName, filePath string) (string, error) { + sources := 0 + for _, value := range []string{directKey, envName, filePath} { + if value != "" { + sources++ + } + } + if sources > 1 { + return "", fmt.Errorf("cloud-proxy: api_key, api_key_env and api_key_file are mutually exclusive") + } + if directKey != "" { + return directKey, nil + } if envName != "" { v := os.Getenv(envName) if v == "" { diff --git a/backend/go/localai-proxy/credentials_test.go b/backend/go/localai-proxy/credentials_test.go new file mode 100644 index 000000000..d57edf22e --- /dev/null +++ b/backend/go/localai-proxy/credentials_test.go @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: MIT +package main + +import ( + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "os" + "path/filepath" +) + +var _ = Describe("Proxy credential sources", func() { + It("resolves direct, absent, env and file credentials", func() { + key, err := resolveAPIKey("direct-canary", "", "") + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("direct-canary")) + key, err = resolveAPIKey("", "", "") + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(BeEmpty()) + DeferCleanup(os.Unsetenv, "LOCALAI_TEST_PROXY_CREDENTIAL") + Expect(os.Setenv("LOCALAI_TEST_PROXY_CREDENTIAL", "env-canary")).To(Succeed()) + key, err = resolveAPIKey("", "LOCALAI_TEST_PROXY_CREDENTIAL", "") + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("env-canary")) + dir, err := os.MkdirTemp("", "proxy-credentials-") + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(os.RemoveAll, dir) + path := filepath.Join(dir, "key") + Expect(os.WriteFile(path, []byte("file-canary\n"), 0600)).To(Succeed()) + key, err = resolveAPIKey("", "", path) + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("file-canary")) + _, err = resolveAPIKey("", "", path+"-missing") + Expect(err).To(HaveOccurred()) + }) + It("rejects every conflicting source without exposing the key", func() { + for _, args := range [][3]string{{"secret-canary", "ENV", ""}, {"secret-canary", "", "/missing"}, {"", "ENV", "/missing"}, {"secret-canary", "ENV", "/missing"}} { + _, err := resolveAPIKey(args[0], args[1], args[2]) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).NotTo(ContainSubstring("secret-canary")) + } + }) + It("loads direct credentials from protobuf options", func() { + proxy := NewLocalAIProxy() + Expect(proxy.Load(&pb.ModelOptions{Model: "alias", Proxy: &pb.ProxyOptions{UpstreamUrl: "http://example.invalid", ApiKey: "direct-canary"}})).To(Succeed()) + Expect(proxy.cfg.Load().apiKey).To(Equal("direct-canary")) + }) +}) diff --git a/backend/go/localai-proxy/proxy.go b/backend/go/localai-proxy/proxy.go index a1d58bb02..c081e0888 100644 --- a/backend/go/localai-proxy/proxy.go +++ b/backend/go/localai-proxy/proxy.go @@ -92,7 +92,7 @@ func (p *LocalAIProxy) Load(opts *pb.ModelOptions) error { "mode", po.GetMode(), "provider", po.GetProvider()) } - key, err := resolveAPIKey(po.GetApiKeyEnv(), po.GetApiKeyFile()) + key, err := resolveAPIKey(po.GetApiKey(), po.GetApiKeyEnv(), po.GetApiKeyFile()) if err != nil { return err } @@ -151,7 +151,19 @@ func (p *LocalAIProxy) model(req string) string { // resolveAPIKey mirrors config.ProxyConfig.ResolveAPIKey (and cloud-proxy's // copy). Duplicated so the backend binary does not depend on core's layout. -func resolveAPIKey(envName, filePath string) (string, error) { +func resolveAPIKey(directKey, envName, filePath string) (string, error) { + sources := 0 + for _, value := range []string{directKey, envName, filePath} { + if value != "" { + sources++ + } + } + if sources > 1 { + return "", fmt.Errorf("localai-proxy: api_key, api_key_env and api_key_file are mutually exclusive") + } + if directKey != "" { + return directKey, nil + } if envName != "" { v := os.Getenv(envName) if v == "" { diff --git a/core/backend/options.go b/core/backend/options.go index b2b652ac8..00a567bde 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -547,6 +547,7 @@ func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions { UpstreamUrl: c.Proxy.UpstreamURL, Mode: c.Proxy.Mode, Provider: c.Proxy.Provider, + ApiKey: c.Proxy.APIKey, ApiKeyEnv: c.Proxy.APIKeyEnv, ApiKeyFile: c.Proxy.APIKeyFile, UpstreamModel: c.Proxy.UpstreamModel, diff --git a/core/backend/options_internal_test.go b/core/backend/options_internal_test.go index a83ccef69..2fed92e93 100644 --- a/core/backend/options_internal_test.go +++ b/core/backend/options_internal_test.go @@ -469,3 +469,13 @@ var _ = Describe("grpcModelOpts option forwarding", func() { Expect(grpcModelOpts(cfg, "/models").Options).To(Equal(cfg.Options)) }) }) + +var _ = Describe("Proxy credential options", func() { + It("forwards the direct credential and preserves wire field 9", func() { + threads := 1 + cfg := config.ModelConfig{Threads: &threads, Backend: "cloud-proxy", Proxy: config.ProxyConfig{APIKey: "direct-canary", UpstreamURL: "http://example.invalid"}} + opts := grpcModelOpts(cfg, "") + Expect(opts.Proxy.GetApiKey()).To(Equal("direct-canary")) + Expect(int(opts.Proxy.ProtoReflect().Descriptor().Fields().ByName("api_key").Number())).To(Equal(9)) + }) +}) diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index af4717ef6..901c2693a 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -1204,6 +1204,13 @@ func DefaultRegistry() map[string]FieldMetaOverride { Component: "input", Order: 210, }, + "proxy.api_key": { + Section: "proxy", + Label: "Proxy API Key", + Description: "Plaintext administrator credential stored in model YAML/JSON. Use only one of api_key, api_key_env or api_key_file. Masking is not encryption.", + Component: "password", + Order: 211, + }, "proxy.api_key_env": { Section: "proxy", Label: "Proxy API Key Env Var", diff --git a/core/config/meta/registry_test.go b/core/config/meta/registry_test.go index bd8c43bb8..227170e30 100644 --- a/core/config/meta/registry_test.go +++ b/core/config/meta/registry_test.go @@ -97,3 +97,9 @@ var _ = Describe("backend-only upscale options", func() { } }) }) + +var _ = Describe("Proxy credential metadata", func() { + It("masks the direct key with a password input", func() { + Expect(meta.DefaultRegistry()["proxy.api_key"].Component).To(Equal("password")) + }) +}) diff --git a/core/config/model_config.go b/core/config/model_config.go index dff45dfe4..25dfab7ee 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -281,14 +281,18 @@ type ProxyConfig struct { // wire format there is whatever the client sent. Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` + // APIKey is a plaintext administrator-managed upstream credential. + // Mutually exclusive with APIKeyEnv and APIKeyFile. + APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"` + // APIKeyEnv names the environment variable holding the upstream - // API key. Mutually exclusive with APIKeyFile. Both empty is + // API key. Mutually exclusive with APIKey and APIKeyFile. All empty is // allowed (no-auth upstreams). APIKeyEnv string `yaml:"api_key_env,omitempty" json:"api_key_env,omitempty"` // APIKeyFile is a path to a file whose contents are the upstream // API key. Trailing whitespace is trimmed. Mutually exclusive - // with APIKeyEnv. The integration point for K8s secret mounts, + // with APIKey and APIKeyEnv. The integration point for K8s secret mounts, // Vault agent files, and similar external-secret workflows. APIKeyFile string `yaml:"api_key_file,omitempty" json:"api_key_file,omitempty"` @@ -329,12 +333,40 @@ const ( ProxyProviderAnthropic = "anthropic" ) -// ResolveAPIKey returns the upstream key from api_key_env or api_key_file, or +// LogValue uses an allowlist so future credential fields cannot leak through config logs. +func (c ModelConfig) LogValue() slog.Value { + return slog.GroupValue(slog.String("name", c.Name), slog.String("backend", c.Backend)) +} + +// LogValue deliberately omits credentials and potentially credential-bearing URLs. +func (p ProxyConfig) LogValue() slog.Value { + return slog.GroupValue(slog.String("mode", p.Mode), slog.String("provider", p.Provider)) +} + +func (p ProxyConfig) validateCredentialSources() error { + sources := 0 + for _, value := range []string{p.APIKey, p.APIKeyEnv, p.APIKeyFile} { + if value != "" { + sources++ + } + } + if sources > 1 { + return fmt.Errorf("proxy: api_key, api_key_env and api_key_file are mutually exclusive") + } + return nil +} + +// ResolveAPIKey returns the upstream key from api_key, api_key_env or api_key_file, or // "" when neither is set. Mirrored (not imported, to keep backends independent // of core's package layout) by resolveAPIKey in backend/go/cloud-proxy/proxy.go -// — keep the two in sync, empty-value handling included. +// and backend/go/localai-proxy/proxy.go — keep all copies in sync. func (p ProxyConfig) ResolveAPIKey(envLookup func(string) string) (string, error) { + if err := p.validateCredentialSources(); err != nil { + return "", err + } switch { + case p.APIKey != "": + return p.APIKey, nil case p.APIKeyEnv != "": var v string if envLookup != nil { @@ -2058,12 +2090,8 @@ func (c *ModelConfig) Validate() (bool, error) { } } - // Cloud-proxy: at most one of api_key_env / api_key_file may be - // set. Both empty means no Authorization header (no-auth upstream - // or a development passthrough). The mode field accepts the empty - // string (defaults to passthrough), "passthrough", or "translate". - if c.Proxy.APIKeyEnv != "" && c.Proxy.APIKeyFile != "" { - return false, fmt.Errorf("proxy: api_key_env and api_key_file are mutually exclusive") + if err := c.Proxy.validateCredentialSources(); err != nil { + return false, err } switch c.Proxy.Mode { case "", ProxyModePassthrough, ProxyModeTranslate: diff --git a/core/config/proxy_credentials_test.go b/core/config/proxy_credentials_test.go new file mode 100644 index 000000000..6ae3fd579 --- /dev/null +++ b/core/config/proxy_credentials_test.go @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: MIT +package config + +import ( + "bytes" + "encoding/json" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" + "log/slog" +) + +var _ = Describe("Proxy credentials", func() { + It("preserves and resolves a direct admin credential", func() { + var p ProxyConfig + Expect(json.Unmarshal([]byte(`{"api_key":"fake-direct-key"}`), &p)).To(Succeed()) + key, err := p.ResolveAPIKey(nil) + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("fake-direct-key")) + encoded, err := json.Marshal(p) + Expect(err).NotTo(HaveOccurred()) + Expect(string(encoded)).To(ContainSubstring(`"api_key":"fake-direct-key"`)) + }) + It("rejects direct credentials combined with another source", func() { + var c ModelConfig + Expect(json.Unmarshal([]byte(`{"proxy":{"api_key":"fake-key","api_key_env":"OTHER"}}`), &c)).To(Succeed()) + _, err := c.Validate() + Expect(err).To(HaveOccurred()) + }) +}) + +var _ = Describe("Proxy credential logging", func() { + It("keeps secrets out of structured config logs without changing storage", func() { + var output bytes.Buffer + logger := slog.New(slog.NewJSONHandler(&output, nil)) + c := ModelConfig{Proxy: ProxyConfig{APIKey: "fake-direct-secret"}, Environment: map[string]string{"KEY": "fake-env-secret"}} + logger.Info("config", "config", c, "proxy", c.Proxy) + Expect(output.String()).NotTo(ContainSubstring("fake-direct-secret")) + Expect(output.String()).NotTo(ContainSubstring("fake-env-secret")) + data, err := yaml.Marshal(c) + Expect(err).NotTo(HaveOccurred()) + Expect(string(data)).To(ContainSubstring("fake-direct-secret")) + }) +}) diff --git a/core/http/endpoints/localai/nodes_backends_list_test.go b/core/http/endpoints/localai/nodes_backends_list_test.go index bc9fa918d..a0d3a9160 100644 --- a/core/http/endpoints/localai/nodes_backends_list_test.go +++ b/core/http/endpoints/localai/nodes_backends_list_test.go @@ -22,7 +22,9 @@ type stubNodeCommandSender struct { listBackendsCalled bool } -func (s *stubNodeCommandSender) InstallBackend(_, _, _, _, _, _, _ string, _ int, _ string, _ func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { +var _ nodes.NodeCommandSender = (*stubNodeCommandSender)(nil) + +func (s *stubNodeCommandSender) InstallBackend(_, _, _, _, _, _, _ string, _ int, _ string, _ func(workerctl.BackendInstallProgressEvent), _ ...map[string]string) (*workerctl.BackendInstallReply, error) { return &workerctl.BackendInstallReply{}, nil } @@ -46,7 +48,7 @@ func (s *stubNodeCommandSender) UnloadModelOnNode(_, _ string) error { return ni func (s *stubNodeCommandSender) PingNode(_ string) error { return nil } // The load operation verbs. This stub never reaches them. -func (s *stubNodeCommandSender) InstallBackendOp(_, _, _, _ string, _ int, _, _ string, _ time.Duration, _ func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { +func (s *stubNodeCommandSender) InstallBackendOp(_, _, _, _ string, _ int, _, _ string, _ time.Duration, _ func(workerctl.BackendInstallProgressEvent), _ ...map[string]string) (*workerctl.BackendInstallReply, error) { return &workerctl.BackendInstallReply{}, nil } diff --git a/core/http/react-ui/src/components/ConfigFieldRenderer.jsx b/core/http/react-ui/src/components/ConfigFieldRenderer.jsx index cc7699559..276d948e0 100644 --- a/core/http/react-ui/src/components/ConfigFieldRenderer.jsx +++ b/core/http/react-ui/src/components/ConfigFieldRenderer.jsx @@ -498,7 +498,8 @@ export default function ConfigFieldRenderer({ field, value, onChange, onRemove, // Default: text input return ( } description={description}> - handleChange(e.target.value)} placeholder={field.placeholder} /> diff --git a/core/http/react-ui/tests/unit/configpassword.test.mjs b/core/http/react-ui/tests/unit/configpassword.test.mjs new file mode 100644 index 000000000..c1f9c57f2 --- /dev/null +++ b/core/http/react-ui/tests/unit/configpassword.test.mjs @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: MIT +import test from 'node:test' +import assert from 'node:assert/strict' +import { readFileSync } from 'node:fs' +import * as React from 'react' +import { renderToStaticMarkup } from 'react-dom/server' +import { transformWithOxc } from 'vite' + +test('config password input preserves the real value on rerender and edits', async () => { + // Exercise the renderer's actual default-input JSX without loading unrelated + // editors (CodeMirror and model selectors require a browser). + const source = readFileSync(new URL('../../src/components/ConfigFieldRenderer.jsx', import.meta.url), 'utf8') + const input = source.slice(source.lastIndexOf('') + 2) + const { code } = await transformWithOxc(`return (${input})`, 'field.jsx', { jsx: { runtime: 'classic' } }) + const renderInput = new Function('React', 'component', 'value', 'handleChange', 'field', code) + const secret = 'fake-secret<&"' + const values = { 'proxy.api_key': secret, name: 'before' } + const render = () => renderInput(React, 'password', values['proxy.api_key'], value => { + values['proxy.api_key'] = value + }, {}) + + const first = render() + assert.equal(first.props.type, 'password') + assert.equal(first.props.autoComplete, 'new-password') + assert.equal(first.props.value, secret) + assert.match(renderToStaticMarkup(first), /type="password"/) + values.name = 'after' + assert.equal(render().props.value, secret) + assert.equal(JSON.parse(JSON.stringify(values))['proxy.api_key'], secret) + render().props.onChange({ target: { value: 'replacement-secret' } }) + assert.equal(render().props.value, 'replacement-secret') + render().props.onChange({ target: { value: '' } }) + assert.equal(render().props.value, '') + assert.equal(renderInput(React, 'password', undefined, () => {}, {}).props.value, '') + assert.equal(renderInput(React, 'input', 'plain', () => {}, {}).props.type, 'text') +}) diff --git a/core/services/failover/model_env_credentials_test.go b/core/services/failover/model_env_credentials_test.go new file mode 100644 index 000000000..d8a89f163 --- /dev/null +++ b/core/services/failover/model_env_credentials_test.go @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: MIT +package failover + +import ( + "github.com/mudler/LocalAI/core/config" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "net/http" +) + +var _ = Describe("Model env credentials", func() { + It("authorizes with a direct credential or frontend fallback", func() { + p := &DefaultProber{EnvLookup: func(string) string { return "frontend-key" }} + req, err := http.NewRequest("GET", "http://example.invalid", nil) + Expect(err).NotTo(HaveOccurred()) + c := config.ModelConfig{Proxy: config.ProxyConfig{APIKey: "direct-key"}} + Expect(p.authorize(req, c)).To(Succeed()) + Expect(req.Header.Get("Authorization")).To(Equal("Bearer direct-key")) + c.Proxy = config.ProxyConfig{APIKeyEnv: "KEY"} + Expect(p.authorize(req, c)).To(Succeed()) + Expect(req.Header.Get("Authorization")).To(Equal("Bearer frontend-key")) + }) + It("uses model env before the frontend env, including an explicit empty value", func() { + p := &DefaultProber{EnvLookup: func(string) string { return "frontend-key" }} + c := config.ModelConfig{Environment: map[string]string{"ANTHROPIC_API_KEY": "model-key"}} + c.Proxy.APIKeyEnv = "ANTHROPIC_API_KEY" + req, err := http.NewRequest("GET", "http://example.invalid", nil) + Expect(err).NotTo(HaveOccurred()) + Expect(p.authorize(req, c)).To(Succeed()) + Expect(req.Header.Get("Authorization")).To(Equal("Bearer model-key")) + c.Environment["ANTHROPIC_API_KEY"] = "" + Expect(p.authorize(req, c)).NotTo(Succeed()) + }) +}) diff --git a/core/services/failover/prober.go b/core/services/failover/prober.go index cb7a53183..e69eb3076 100644 --- a/core/services/failover/prober.go +++ b/core/services/failover/prober.go @@ -110,12 +110,21 @@ func PrepareTarget(cfg *config.ModelConfig) { } func (p *DefaultProber) authorize(req *http.Request, cfg config.ModelConfig) error { - if cfg.Proxy.APIKeyEnv != "" && p.EnvLookup == nil { + _, hasModelEnv := cfg.Environment[cfg.Proxy.APIKeyEnv] + if cfg.Proxy.APIKeyEnv != "" && !hasModelEnv && p.EnvLookup == nil { if _, warned := p.unresolvableEnvWarned.LoadOrStore(cfg.Proxy.APIKeyEnv, struct{}{}); !warned { xlog.Warn("failover: remote target has api_key_env set but no credential lookup is configured; liveness and requests will fail authorization", "target", cfg.Name, "api_key_env", cfg.Proxy.APIKeyEnv) } } - key, err := cfg.Proxy.ResolveAPIKey(p.EnvLookup) + key, err := cfg.Proxy.ResolveAPIKey(func(name string) string { + if value, ok := cfg.Environment[name]; ok { + return value + } + if p.EnvLookup != nil { + return p.EnvLookup(name) + } + return "" + }) if err != nil || key == "" { return err } diff --git a/core/services/nodes/env_test.go b/core/services/nodes/env_test.go new file mode 100644 index 000000000..f7c8a5292 --- /dev/null +++ b/core/services/nodes/env_test.go @@ -0,0 +1,60 @@ +// SPDX-License-Identifier: MIT +package nodes + +import ( + "context" + "github.com/mudler/LocalAI/core/services/workerctl" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "time" +) + +var _ = Describe("Router install environment", func() { + It("keeps changed snapshots out of an in-flight identical slot", func() { + started := make(chan struct{}, 2) + release := make(chan struct{}) + defer close(release) + unloader := &fakeUnloader{installReply: &workerctl.BackendInstallReply{Success: true, Address: "127.0.0.1:1234"}, installHook: func() { started <- struct{}{}; <-release }} + router := NewSmartRouter(&fakeModelRouter{}, SmartRouterOptions{Unloader: unloader}) + done := make(chan error, 2) + for _, value := range []string{"first", "second"} { + go func(value string) { + _, err := router.installBackendOnNode(context.Background(), &BackendNode{ID: "node"}, "backend", "model", 0, map[string]string{"TOKEN": value}) + done <- err + }(value) + } + Eventually(started).Should(Receive()) + Eventually(started).Should(Receive()) + // Both leaders entered before either could finish, so they cannot have shared + // one install reply (and therefore one child's credentials). + go func() { release <- struct{}{}; release <- struct{}{} }() + Eventually(done).Should(Receive(BeNil())) + Eventually(done).Should(Receive(BeNil())) + Expect(unloader.installCalls).To(HaveLen(2)) + Expect([]string{unloader.installCalls[0].env["TOKEN"], unloader.installCalls[1].env["TOKEN"]}).To(ConsistOf("first", "second")) + }) + It("coalesces identical snapshots without placing credentials in the identity", func() { + a := map[string]string{"TOKEN": "test-value", "EMPTY": ""} + b := map[string]string{"EMPTY": "", "TOKEN": "test-value"} + Expect(installEnvIdentity(a)).To(Equal(installEnvIdentity(b))) + Expect(installEnvIdentity(a)).NotTo(ContainSubstring("test-value")) + Expect(installEnvIdentity(nil)).NotTo(Equal(installEnvIdentity(map[string]string{}))) + }) + It("forwards complete and explicitly empty snapshots through the operation-aware fake", func() { + fake := &fakeOpWorker{fakeUnloader: &fakeUnloader{}} + for _, env := range []map[string]string{{"TOKEN": "test-value"}, {}} { + _, err := fake.InstallBackendOp("node", "backend", "model", "", 0, "", "operation", time.Minute, nil, env) + Expect(err).NotTo(HaveOccurred()) + Expect(fake.installs[len(fake.installs)-1].EnvVars).To(Equal(env)) + Expect(fake.installCalls[len(fake.installCalls)-1].env).To(Equal(env)) + } + }) + It("copies snapshots and preserves absence separately from removal", func() { + Expect(installEnvSnapshot(nil)).To(BeNil()) + Expect(installEnvSnapshot([]map[string]string{nil})).NotTo(BeNil()) + env := map[string]string{"TOKEN": "test-value"} + snapshot := installEnvSnapshot([]map[string]string{env}) + env["TOKEN"] = "changed" + Expect(snapshot["TOKEN"]).To(Equal("test-value")) + }) +}) diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index 52ca43215..a9467083c 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -88,7 +88,7 @@ type ReplicaUnloader interface { // LoadOperationInstaller starts a backend as a load operation the worker bounds. type LoadOperationInstaller interface { - InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) + InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) } // LoadOperationStopper is the only path that kills remote load work. It stops diff --git a/core/services/nodes/load_worker_ops_test.go b/core/services/nodes/load_worker_ops_test.go index 2f0a3ca16..fb77d3f20 100644 --- a/core/services/nodes/load_worker_ops_test.go +++ b/core/services/nodes/load_worker_ops_test.go @@ -50,11 +50,11 @@ func (w *fakeOpWorker) StopModelReplica(_ context.Context, _ string, replica Nod return workerctl.ModelStopReply{Matched: true, Terminated: true}, nil } -func (w *fakeOpWorker) InstallBackendOp(nodeID, backend, modelID, galleries string, replica int, opID, operationID string, deadline time.Duration, progress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { +func (w *fakeOpWorker) InstallBackendOp(nodeID, backend, modelID, galleries string, replica int, opID, operationID string, deadline time.Duration, progress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) { w.mu.Lock() - w.installs = append(w.installs, workerctl.BackendInstallRequest{ModelID: modelID, ReplicaIndex: int32(replica), OperationID: operationID, DeadlineMs: deadline.Milliseconds()}) + w.installs = append(w.installs, workerctl.BackendInstallRequest{EnvVars: installEnvSnapshot(env), ModelID: modelID, ReplicaIndex: int32(replica), OperationID: operationID, DeadlineMs: deadline.Milliseconds()}) w.mu.Unlock() - reply, err := w.fakeUnloader.InstallBackend(nodeID, backend, modelID, galleries, "", "", "", replica, opID, progress) + reply, err := w.fakeUnloader.InstallBackend(nodeID, backend, modelID, galleries, "", "", "", replica, opID, progress, env...) if reply != nil { copy := *reply copy.ProcessInstance, copy.ReportsOperations = "", false diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 2ea18bc19..0e35e424e 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -2,6 +2,9 @@ package nodes import ( "context" + "crypto/hmac" + "crypto/rand" + "crypto/sha256" "encoding/json" "errors" "fmt" @@ -1275,7 +1278,7 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID // Send backend.install — the worker installs the backend if needed and // starts the gRPC process bound to a port for this (model, replica) slot. - addr, installErr := r.installBackendOnNode(ctx, node, backendType, modelID, replicaIdx) + addr, installErr := r.installBackendOnNode(ctx, node, backendType, modelID, replicaIdx, modelOpts.GetEnvVars()) if installErr != nil { // Roll back the reservation explicitly so the column is accurate // before the next heartbeat. Best-effort. @@ -1388,13 +1391,14 @@ func (r *SmartRouter) estimateModelVRAM(ctx context.Context, opts *pb.ModelOptio // Routine load: the worker's fast-path "already running → return current // address" is correct here. Upgrades go through // DistributedBackendManager.UpgradeBackend on the backend.upgrade subject. -func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNode, backendType, modelID string, replicaIndex int) (string, error) { +func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNode, backendType, modelID string, replicaIndex int, env ...map[string]string) (string, error) { if r.unloader == nil { return "", fmt.Errorf("no NATS connection for backend installation") } + snapshot := installEnvSnapshot(env) reportLoadPhase(ctx, LoadJobStateInstalling, node, replicaIndex) - key := fmt.Sprintf("%s|%s|%s|%d", node.ID, backendType, modelID, replicaIndex) + key := fmt.Sprintf("%s|%s|%s|%d|%s", node.ID, backendType, modelID, replicaIndex, installEnvIdentity(snapshot)) // DoChan rather than Do so this wait honors ctx cancellation. InstallBackend // blocks for its full NATS deadline (15m by default) when a worker accepts // the request but never replies (e.g. it died mid-install). Without ctx @@ -1410,9 +1414,9 @@ func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNod var reply *workerctl.BackendInstallReply var err error if owned { - reply, err = r.unloader.InstallBackendOp(node.ID, backendType, modelID, r.galleriesJSON, replicaIndex, "", ref.Generation, r.loadOperationDeadline(), nil) + reply, err = r.unloader.InstallBackendOp(node.ID, backendType, modelID, r.galleriesJSON, replicaIndex, "", ref.Generation, r.loadOperationDeadline(), nil, snapshot) } else { - reply, err = r.unloader.InstallBackend(node.ID, backendType, modelID, r.galleriesJSON, "", "", "", replicaIndex, "", nil) + reply, err = r.unloader.InstallBackend(node.ID, backendType, modelID, r.galleriesJSON, "", "", "", replicaIndex, "", nil, snapshot) } if err != nil { return "", err @@ -2262,3 +2266,20 @@ func (r *SmartRouter) loadOperationDeadline() time.Duration { } return modelLoadAbsoluteMax } + +// Keyed identities let identical installs coalesce without retaining credentials +// in singleflight keys or exposing a dictionary-testable credential hash. +var installEnvKey = func() []byte { + key := make([]byte, 32) + if _, err := rand.Read(key); err != nil { + panic(err) + } + return key +}() + +func installEnvIdentity(env map[string]string) string { + body, _ := json.Marshal(env) + mac := hmac.New(sha256.New, installEnvKey) + mac.Write(body) + return fmt.Sprintf("%x", mac.Sum(nil)) +} diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index e5a972050..b73dc902d 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -573,6 +573,7 @@ type fakeUnloader struct { // every call so tests can verify both presence and shape (e.g. that backend // is non-empty). type installCall struct { + env map[string]string nodeID string backend string modelID string @@ -585,7 +586,7 @@ type upgradeCall struct { replica int } -func (f *fakeUnloader) InstallBackend(nodeID, backend, modelID, _, _, _, _ string, replica int, _ string, _ func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { +func (f *fakeUnloader) InstallBackend(nodeID, backend, modelID, _, _, _, _ string, replica int, _ string, _ func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) { // installHook intentionally runs OUTSIDE the mutex: the hook may block // on a channel and we don't want to serialize concurrent callers, // which would defeat the singleflight-overlap test. @@ -593,15 +594,15 @@ func (f *fakeUnloader) InstallBackend(nodeID, backend, modelID, _, _, _, _ strin f.installHook() } f.mu.Lock() - f.installCalls = append(f.installCalls, installCall{nodeID, backend, modelID, replica}) + f.installCalls = append(f.installCalls, installCall{nodeID: nodeID, backend: backend, modelID: modelID, replica: replica, env: installEnvSnapshot(env)}) f.mu.Unlock() return f.installReply, f.installErr } // The load operation verbs of the carrier seam. The default fake is a worker // that names operations and acknowledges every stop. -func (f *fakeUnloader) InstallBackendOp(nodeID, backend, modelID, galleries string, replica int, opID, _ string, _ time.Duration, progress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { - reply, err := f.InstallBackend(nodeID, backend, modelID, galleries, "", "", "", replica, opID, progress) +func (f *fakeUnloader) InstallBackendOp(nodeID, backend, modelID, galleries string, replica int, opID, _ string, _ time.Duration, progress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) { + reply, err := f.InstallBackend(nodeID, backend, modelID, galleries, "", "", "", replica, opID, progress, env...) if reply != nil { withOps := *reply withOps.ReportsOperations = true @@ -1406,7 +1407,7 @@ var _ = Describe("SmartRouter", func() { done := make(chan error, 5) for i := 0; i < 5; i++ { go func() { - _, err := router.installBackendOnNode(context.Background(), node, "llama-cpp", "my-model", 0) + _, err := router.installBackendOnNode(context.Background(), node, "llama-cpp", "my-model", 0, map[string]string{"TOKEN": "same-snapshot"}) done <- err }() } diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index 6682437a7..1780a6290 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -31,7 +31,7 @@ import ( // UpgradeBackend returns ErrNoRoute on an old worker that does not serve // backend.upgrade, and the caller falls back to the legacy install. type NodeCommandSender interface { - InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) + InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendUpgradeReply, error) DeleteBackend(nodeID, backendName string) (*workerctl.BackendDeleteReply, error) ListBackends(nodeID string) (*workerctl.BackendListReply, error) @@ -158,8 +158,9 @@ const operationControlTimeout = 5 * time.Second // InstallBackendOp is InstallBackend for a load operation: the request carries // the operation id and the longest the load may run, so the worker can bound it. -func (a *RemoteUnloaderAdapter) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { +func (a *RemoteUnloaderAdapter) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) { return a.installBackend(nodeID, workerctl.BackendInstallRequest{ + EnvVars: installEnvSnapshot(env), Backend: backendType, ModelID: modelID, BackendGalleries: galleriesJSON, @@ -257,8 +258,10 @@ func (a *RemoteUnloaderAdapter) InstallBackend( replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent), + env ...map[string]string, ) (*workerctl.BackendInstallReply, error) { return a.installBackend(nodeID, workerctl.BackendInstallRequest{ + EnvVars: installEnvSnapshot(env), Backend: backendType, ModelID: modelID, BackendGalleries: galleriesJSON, @@ -672,3 +675,16 @@ func (a *RemoteUnloaderAdapter) StopNode(nodeID string) error { subject := messaging.SubjectNodeStop(nodeID) return a.nats.Publish(subject, nil) } + +// installEnvSnapshot preserves nil for legacy callers and an explicit empty +// snapshot for new callers that remove their last model override. +func installEnvSnapshot(env []map[string]string) map[string]string { + if len(env) == 0 { + return nil + } + result := make(map[string]string, len(env[0])) + for k, v := range env[0] { + result[k] = v + } + return result +} diff --git a/core/services/worker/env_test.go b/core/services/worker/env_test.go new file mode 100644 index 000000000..9e83fef21 --- /dev/null +++ b/core/services/worker/env_test.go @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: MIT +package worker + +import ( + "context" + "fmt" + "net" + "os" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/mudler/LocalAI/core/services/workerctl" + + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "google.golang.org/grpc" +) + +// The real worker starts this test executable as a tiny gRPC backend. Only the +// child's explicit environment enables helper mode; the test runner is untouched. +func init() { + if os.Getenv("LOCALAI_ENV_TEST_CHILD") != "1" { + return + } + var addr string + for i, arg := range os.Args { + if arg == "--addr" && i+1 < len(os.Args) { + addr = os.Args[i+1] + } + } + listener, err := net.Listen("tcp", addr) + if err != nil { + os.Exit(2) + } + value, present := os.LookupEnv("LOCALAI_ENV_TEST_VALUE") + if err := os.WriteFile(os.Getenv("LOCALAI_ENV_TEST_OUTPUT"), []byte(fmt.Sprintf("%t:%s", present, value)), 0600); err != nil { + os.Exit(3) + } + // A filesystem gate lets the parent hold a real child in startup while + // concurrent install requests arrive, without mutating the runner's env. + if gate := os.Getenv("LOCALAI_ENV_TEST_GATE"); gate != "" { + for { + if _, err := os.Stat(gate); err == nil { + break + } + time.Sleep(10 * time.Millisecond) + } + } + server := grpc.NewServer() + pb.RegisterBackendServer(server, envTestBackend{}) + if server.Serve(listener) != nil { + os.Exit(4) + } + os.Exit(0) +} + +type envTestBackend struct{ pb.UnimplementedBackendServer } + +func (envTestBackend) Health(context.Context, *pb.HealthMessage) (*pb.Reply, error) { + return &pb.Reply{Message: []byte("OK")}, nil +} +func (envTestBackend) Free(context.Context, *pb.HealthMessage) (*pb.Result, error) { + return &pb.Result{Success: true}, nil +} + +var _ = Describe("Worker model environment", func() { + It("serializes concurrent serveInstall requests for the exact slot across backend names", func() { + root := GinkgoT().TempDir() + exe, err := os.Executable() + Expect(err).NotTo(HaveOccurred()) + for _, name := range []string{"first", "second", "sibling"} { + Expect(os.WriteFile(filepath.Join(root, name), []byte("#!/bin/sh\nexec "+strconv.Quote(exe)+" \"$@\"\n"), 0700)).To(Succeed()) + } + state, err := system.GetSystemState(system.WithBackendPath(root), system.WithBackendSystemPath(root), system.WithModelPath(root)) + Expect(err).NotTo(HaveOccurred()) + sup := &backendSupervisor{cfg: &Config{BackendsPath: root}, ml: model.NewModelLoader(state), processes: map[string]*backendProcess{}, minPort: 32000, maxPort: 42000, nextPort: 32000} + gate := filepath.Join(root, "ready") + firstDone, secondDone, siblingDone := make(chan workerctl.BackendInstallReply, 1), make(chan workerctl.BackendInstallReply, 1), make(chan workerctl.BackendInstallReply, 1) + DeferCleanup(func() { + _ = os.WriteFile(gate, nil, 0600) + // Drain all handlers before inspecting/stopping their children. + Eventually(firstDone, "40s").Should(Receive()) + Eventually(secondDone, "40s").Should(Receive()) + Eventually(siblingDone, "40s").Should(Receive()) + for key := range sup.processes { + Expect(sup.stopBackendExact(key, true)).To(Succeed()) + } + }) + request := func(backend string, replica int) workerctl.BackendInstallRequest { + return workerctl.BackendInstallRequest{ModelID: "model", Backend: backend, ReplicaIndex: int32(replica), OperationID: backend, + EnvVars: map[string]string{"LOCALAI_ENV_TEST_CHILD": "1", "LOCALAI_ENV_TEST_OUTPUT": filepath.Join(root, backend+"-env"), "LOCALAI_ENV_TEST_VALUE": backend}} + } + first := request("first", 0) + first.EnvVars["LOCALAI_ENV_TEST_GATE"] = gate + go func() { firstDone <- sup.serveInstall(context.Background(), first, nil) }() + Eventually(func() bool { _, err := os.Stat(filepath.Join(root, "first-env")); return err == nil }, "10s").Should(BeTrue()) + sup.mu.Lock() + original := sup.processes["model#0"] + sup.mu.Unlock() + go func() { secondDone <- sup.serveInstall(context.Background(), request("second", 0), nil) }() + go func() { siblingDone <- sup.serveInstall(context.Background(), request("sibling", 1), nil) }() + // A different replica must reach readiness while replica 0 is gated. + var sibling workerctl.BackendInstallReply + Eventually(siblingDone, "10s").Should(Receive(&sibling)) + siblingDone <- sibling + Expect(sibling.Success).To(BeTrue(), sibling.Error) + Consistently(func() bool { + sup.mu.Lock() + defer sup.mu.Unlock() + return sup.processes["model#0"] == original && original.proc.IsAlive() && original.operationID == "" + }, "500ms").Should(BeTrue(), "the second install must not replace or attach to the unready first child") + Consistently(secondDone, "200ms").ShouldNot(Receive()) + Expect(os.WriteFile(gate, nil, 0600)).To(Succeed()) + var a, b workerctl.BackendInstallReply + Eventually(firstDone, "10s").Should(Receive(&a)) + firstDone <- a + Eventually(secondDone, "10s").Should(Receive(&b)) + secondDone <- b + Expect(a.Success).To(BeTrue(), a.Error) + Expect(b.Success).To(BeTrue(), b.Error) + Expect(a.ProcessInstance).To(Equal(original.instance)) + Expect(b.ProcessInstance).NotTo(Equal(a.ProcessInstance)) + body, err := os.ReadFile(filepath.Join(root, "second-env")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(body)).To(Equal("true:second")) + sup.mu.Lock() + defer sup.mu.Unlock() + Expect(sup.processes["model#0"].backendName).To(Equal("second")) + Expect(sup.processes["model#0"].operationID).To(Equal("second")) + Expect(sup.operations["second"].instance).To(Equal(b.ProcessInstance)) + Expect(sup.processes["model#1"].instance).To(Equal(sibling.ProcessInstance)) + }) + + It("passes isolated snapshots to real children and restarts on replacement and removal", func() { + root := GinkgoT().TempDir() + exe, err := os.Executable() + Expect(err).NotTo(HaveOccurred()) + backend := filepath.Join(root, "backend") + Expect(os.WriteFile(backend, []byte("#!/bin/sh\nexec "+strconv.Quote(exe)+" \"$@\"\n"), 0700)).To(Succeed()) + state, err := system.GetSystemState(system.WithBackendPath(root), system.WithBackendSystemPath(root), system.WithModelPath(root)) + Expect(err).NotTo(HaveOccurred()) + sup := &backendSupervisor{cfg: &Config{}, ml: model.NewModelLoader(state), processes: map[string]*backendProcess{}, minPort: 32000, maxPort: 42000, nextPort: 32000} + DeferCleanup(func() { + for key := range sup.processes { + Expect(sup.stopBackendExact(key, true)).To(Succeed()) + } + }) + snapshot := func(slot, value string) map[string]string { + return map[string]string{"LOCALAI_ENV_TEST_CHILD": "1", "LOCALAI_ENV_TEST_OUTPUT": filepath.Join(root, slot), "LOCALAI_ENV_TEST_VALUE": value} + } + first := snapshot("a", "first") + _, err = sup.startBackend("a#0", "test", backend, first) + Expect(err).NotTo(HaveOccurred()) + original := sup.processes["a#0"] + first["LOCALAI_ENV_TEST_VALUE"] = "caller mutation" + Expect(original.envVars["LOCALAI_ENV_TEST_VALUE"]).To(Equal("first")) + _, err = sup.startBackend("a#0", "test", backend, snapshot("a", "first")) + Expect(err).NotTo(HaveOccurred()) + Expect(sup.processes["a#0"]).To(BeIdenticalTo(original)) + _, err = sup.startBackend("b#0", "test", backend, snapshot("b", "second")) + Expect(err).NotTo(HaveOccurred()) + body, err := os.ReadFile(filepath.Join(root, "b")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(body)).To(Equal("true:second")) + _, err = sup.startBackend("a#0", "test", backend, snapshot("a", "changed")) + Expect(err).NotTo(HaveOccurred()) + Expect(sup.processes["a#0"]).NotTo(BeIdenticalTo(original)) + removed := snapshot("a", "") + delete(removed, "LOCALAI_ENV_TEST_VALUE") + _, err = sup.startBackend("a#0", "test", backend, removed) + Expect(err).NotTo(HaveOccurred()) + body, err = os.ReadFile(filepath.Join(root, "a")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(body)).To(Equal("false:")) + body, err = os.ReadFile(filepath.Join(root, "b")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(body)).To(Equal("true:second")) + Expect(strings.Join(os.Environ(), "\n")).NotTo(ContainSubstring("LOCALAI_ENV_TEST_CHILD=1")) + Expect(sup.processMatchesEnv("b#0", map[string]string{})).To(BeFalse()) + Expect(sup.processMatchesEnv("b#0", nil)).To(BeTrue()) + }) +}) diff --git a/core/services/worker/install.go b/core/services/worker/install.go index 46a4693fb..a0d670cce 100644 --- a/core/services/worker/install.go +++ b/core/services/worker/install.go @@ -38,7 +38,8 @@ func buildProcessKey(modelID, backend string, replicaIndex int) string { // upgrade path; force=false is the routine load path. // // The caller is responsible for holding s.lockBackend(req.Backend) for -// the duration of the call so the gallery directory isn't raced. +// the duration of the call so the gallery directory isn't raced. Production +// serveInstall also holds lockInstallSlot through readiness and attachment. // // 1. If already running for this (model, replica) slot AND force is false, // return existing address (the fast path used by routine load events that @@ -73,7 +74,7 @@ func (s *backendSupervisor) installBackend(req workerctl.BackendInstallRequest, // running" again, and the cluster loops on a dead replica forever. if addr := s.getAddr(processKey); addr != "" { switch { - case !s.processMatchesBackend(processKey, req.Backend): + case !s.processMatchesBackend(processKey, req.Backend) || !s.processMatchesEnv(processKey, req.EnvVars): // The slot is held by a process that must not serve this // backend: either it was started from a DIFFERENT backend // (this model's previous backend was deleted, or superseded by @@ -169,7 +170,7 @@ func (s *backendSupervisor) installBackend(req workerctl.BackendInstallRequest, // Start the gRPC process on a new port (keyed by model, not just backend). // req.Backend is recorded on the process so a later backend.delete/stop can // find it by backend name. - return s.startBackend(processKey, req.Backend, backendPath) + return s.startBackend(processKey, req.Backend, backendPath, req.EnvVars) } // upgradeBackend stops every running process for `backend`, force-reinstalls @@ -276,3 +277,22 @@ func (s *backendSupervisor) lockBackend(name string) func() { m.Lock() return m.Unlock } + +// lockInstallSlot protects one complete serveInstall transaction, including +// readiness and operation attachment. It is acquired after lockBackend. Stops +// (including upgrade's multi-slot stops) must not acquire it: they can interrupt +// startup, and must not wait on an install holding an artifact lock they need. +func (s *backendSupervisor) lockInstallSlot(key string) func() { + s.mu.Lock() + if s.installSlotLocks == nil { + s.installSlotLocks = make(map[string]*sync.Mutex) + } + m := s.installSlotLocks[key] + if m == nil { + m = &sync.Mutex{} + s.installSlotLocks[key] = m + } + s.mu.Unlock() + m.Lock() + return m.Unlock +} diff --git a/core/services/worker/lifecycle.go b/core/services/worker/lifecycle.go index beec763d7..9c85cb071 100644 --- a/core/services/worker/lifecycle.go +++ b/core/services/worker/lifecycle.go @@ -127,6 +127,11 @@ func (s *backendSupervisor) serveInstall(_ context.Context, req workerctl.Backen xlog.Info("Received NATS backend.install event") release := s.lockBackend(req.Backend) defer release() + // A model/replica process key does not include the backend name. Keep + // competing installs on that exact slot serialized through readiness and + // operation attachment, not just the gallery artifact download. + releaseSlot := s.lockInstallSlot(buildProcessKey(req.ModelID, req.Backend, int(req.ReplicaIndex))) + defer releaseSlot() downloadCb, flush := s.downloadProgress(req.OpID, req.Backend, progress) defer flush() diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index 09c918d4b..e44d9c47d 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -27,9 +27,10 @@ import ( // backendProcess represents a single gRPC backend process. type backendProcess struct { - proc *process.Process - addr string // gRPC address (host:port) - port int + envVars map[string]string + proc *process.Process + addr string // gRPC address (host:port) + port int // serving marks a process that has answered a gRPC health check and is // therefore actually listening on addr. It is the opening bracket of the // lifecycle that stopping closes. @@ -160,6 +161,10 @@ type backendSupervisor struct { // download path doesn't race itself on the same directory. backendLocks map[string]*sync.Mutex + // installSlotLocks serializes installs for a process key, even when their + // backend names (and therefore artifact locks) differ. Worker-local only. + installSlotLocks map[string]*sync.Mutex + // operations are the loads the watchdog bounds, by operation id. Guarded by // mu. See operations.go. operations map[string]*loadOperation @@ -497,7 +502,16 @@ func (s *backendSupervisor) releasePort(port int) { // `backend` is the process key (modelID#replica); `backendName` is the gallery // backend it was started from, recorded so delete/stop can find this process // by backend name later. -func (s *backendSupervisor) startBackend(backend, backendName, backendPath string) (string, error) { +func (s *backendSupervisor) startBackend(backend, backendName, backendPath string, env ...map[string]string) (string, error) { + var snapshot map[string]string + if len(env) > 0 { + snapshot = maps.Clone(env[0]) + } + if !s.processMatchesEnv(backend, snapshot) { + if err := s.stopBackendExact(backend, false); err != nil { + return "", err + } + } // Mirrors the working directory pkg/model.startProcess hands the child. backendDir := filepath.Dir(backendPath) @@ -537,7 +551,7 @@ func (s *backendSupervisor) startBackend(backend, backendName, backendPath strin bindAddr := fmt.Sprintf("0.0.0.0:%d", port) clientAddr := fmt.Sprintf("127.0.0.1:%d", port) - proc, err := s.ml.StartProcess(backendPath, backend, bindAddr, nil) + proc, err := s.ml.StartProcess(backendPath, backend, bindAddr, snapshot) if err != nil { s.releasePortForKey(backend, port) s.mu.Unlock() @@ -557,6 +571,7 @@ func (s *backendSupervisor) startBackend(backend, backendName, backendPath strin } s.processes[backend] = &backendProcess{ + envVars: maps.Clone(snapshot), proc: proc, addr: clientAddr, port: port, @@ -1167,3 +1182,15 @@ func (s *backendSupervisor) LoadedBackendAddresses() []string { } return addrs } + +// An absent snapshot retains old-controller reuse semantics. An explicit +// empty snapshot must still restart a slot that carries removed overrides. +func (s *backendSupervisor) processMatchesEnv(key string, env map[string]string) bool { + if env == nil { + return true + } + s.mu.Lock() + defer s.mu.Unlock() + bp := s.processes[key] + return bp == nil || maps.Equal(bp.envVars, env) +} diff --git a/core/services/workerctl/backend.go b/core/services/workerctl/backend.go index 46214c1e6..b5d8cf7cf 100644 --- a/core/services/workerctl/backend.go +++ b/core/services/workerctl/backend.go @@ -1,10 +1,15 @@ package workerctl +import "encoding/json" + // BackendInstallRequest is the payload for a backend.install control request. type BackendInstallRequest struct { - Backend string `json:"backend"` - ModelID string `json:"model_id,omitempty"` - BackendGalleries string `json:"backend_galleries,omitempty"` + // EnvVars is a complete per-process snapshot. nil means a legacy caller; + // a non-nil empty map explicitly removes previous model overrides. + EnvVars map[string]string `json:"-"` + Backend string `json:"backend"` + ModelID string `json:"model_id,omitempty"` + BackendGalleries string `json:"backend_galleries,omitempty"` // URI is set for external installs (OCI image, URL, or path). When non-empty // the worker routes to InstallExternalBackend instead of the gallery lookup. URI string `json:"uri,omitempty"` @@ -188,3 +193,30 @@ type BackendDeleteReply struct { // conclusion against an older worker. ReportsStoppedProcesses bool `json:"reports_stopped_processes,omitempty"` } + +// MarshalJSON omits legacy snapshots without collapsing an explicit empty map. +func (r BackendInstallRequest) MarshalJSON() ([]byte, error) { + type wire BackendInstallRequest + var env *map[string]string + if r.EnvVars != nil { + env = &r.EnvVars + } + return json.Marshal(struct { + wire + EnvVars *map[string]string `json:"env_vars,omitempty"` + }{wire(r), env}) +} + +func (r *BackendInstallRequest) UnmarshalJSON(data []byte) error { + type wire BackendInstallRequest + var decoded struct { + wire + EnvVars map[string]string `json:"env_vars"` + } + if err := json.Unmarshal(data, &decoded); err != nil { + return err + } + *r = BackendInstallRequest(decoded.wire) + r.EnvVars = decoded.EnvVars + return nil +} diff --git a/core/services/workerctl/env_test.go b/core/services/workerctl/env_test.go new file mode 100644 index 000000000..07e1d79f9 --- /dev/null +++ b/core/services/workerctl/env_test.go @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: MIT +package workerctl + +import ( + "encoding/json" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Install environment wire compatibility", func() { + It("preserves absent legacy bytes and explicit empty snapshots", func() { + for _, body := range []string{`{"backend":"test"}`, `{"backend":"test","env_vars":{}}`, `{"backend":"test","env_vars":{"EMPTY":"","TOKEN":"test-value"}}`} { + var req BackendInstallRequest + Expect(json.Unmarshal([]byte(body), &req)).To(Succeed()) + encoded, err := json.Marshal(req) + Expect(err).NotTo(HaveOccurred()) + Expect(string(encoded)).To(Equal(body)) + } + }) + It("lets an old receiver ignore the additive snapshot", func() { + body, err := json.Marshal(BackendInstallRequest{Backend: "test", EnvVars: map[string]string{"TOKEN": "test-value"}}) + Expect(err).NotTo(HaveOccurred()) + var legacy struct { + Backend string `json:"backend"` + } + Expect(json.Unmarshal(body, &legacy)).To(Succeed()) + Expect(legacy.Backend).To(Equal("test")) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 3439094c5..23d22d040 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -1522,3 +1522,30 @@ Further enhancements, surfaced from a survey of SGLang, vLLM production-stack, R - **Prefill/decode disaggregation routing** ([#10068](https://github.com/mudler/LocalAI/issues/10068)): route prefill and decode to separate pools with KV transfer. - **Per-user fairness (VTC)** ([#10069](https://github.com/mudler/LocalAI/issues/10069)): balance per-user token usage against pod load. - **Minor tuning + MCP parity** ([#10070](https://github.com/mudler/LocalAI/issues/10070)): per-model TTL override, probabilistic LRU updates, and MCP scheduling-config tool parity. + +## Per-model worker environment + +A model's `env` map is sent with each backend installation as a complete +per-process snapshot. The worker applies it only to that model replica's child +process, on top of the worker's inherited environment; it never changes the +frontend or worker daemon environment. Models sharing a backend executable can +therefore use different credentials or runtime settings without sharing them. +An empty string is an explicit value, not an instruction to remove a variable. + +When a slot is installed again with changed overrides (including removal of the +last override), the worker restarts its child rather than reusing a process with +stale settings. Removing an override restores the worker's inherited value, if +one exists. Unload/reload a running model after changing its environment; this +is not live environment mutation of an already loaded child. + +This requires updated frontends **and** workers. Legacy install requests omit +the snapshot and retain legacy reuse behavior; older workers ignore the added +field and cannot provide per-model environment isolation. Upgrade workers before +relying on this feature. Environment snapshots may contain secrets and travel +over the control transport: restrict NATS access and enable its authentication +and TLS for untrusted networks. They are not secret-store references. + +Install coalescing is per frontend instance and includes a keyed identity of +the snapshot, so concurrent installs with different environments do not share an +install reply. The worker owns process-local snapshots; they are rebuilt from +model configuration on each install, not persisted as a new cluster state store. diff --git a/docs/content/operations/cloud-proxy.md b/docs/content/operations/cloud-proxy.md index 42f03bb98..6b85db3e3 100644 --- a/docs/content/operations/cloud-proxy.md +++ b/docs/content/operations/cloud-proxy.md @@ -63,9 +63,28 @@ against - and two modes: `proxy.provider` selects the auth scheme and (in translate mode) the wire format. Supported values: `openai`, `anthropic`. -If the upstream requires an API key, configure either an environment variable -(`api_key_env`) or a file (`api_key_file`). The key never appears in the config -file or the admin UI. If the upstream requires no API key, omit both fields. +If the upstream requires an API key, configure one of `api_key`, `api_key_env` +(an environment variable name), or `api_key_file` (a file path). If the upstream +requires no API key, omit all three fields. + +`api_key` is a plaintext administrator-managed credential, supported by both +`cloud-proxy` and `localai-proxy`: + +```yaml +proxy: + api_key: your-upstream-key +``` + +The model editor masks this input but preserves its actual value on unrelated +saves. Masking is not encryption or vault storage: model YAML/JSON, admin +configuration responses, exports, and backups can contain the plaintext key, +just as they can contain model `env` values. Restrict access and protect files +and backups. Use `api_key_env` or `api_key_file` for external secret management. + +Direct keys require updated core and proxy backend binaries. In distributed +installations, update workers and their proxy backends before frontends; older +binaries can ignore this additive field and send unauthenticated requests. +This does not add distributed model-environment propagation. ### OpenAI passthrough @@ -146,7 +165,7 @@ proxy: ### Upstreams without an API key For an OpenAI-compatible upstream that accepts requests without authentication, -omit both `api_key_env` and `api_key_file`: +omit `api_key`, `api_key_env`, and `api_key_file`: ```yaml name: internal-chat-proxy @@ -239,7 +258,7 @@ proxy: ``` The file is read at backend load time and trimmed of surrounding whitespace. -`api_key_env` and `api_key_file` are mutually exclusive. +`api_key`, `api_key_env`, and `api_key_file` are mutually exclusive. ## Combining with the intelligent router diff --git a/pkg/model/env_isolation_test.go b/pkg/model/env_isolation_test.go new file mode 100644 index 000000000..aa7486f6a --- /dev/null +++ b/pkg/model/env_isolation_test.go @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: MIT +package model + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Model env isolation", func() { + It("copies the caller map for each options instance", func() { + env := map[string]string{"KEY": "original"} + option := WithEnvVars(env) + first, second := NewOptions(option), NewOptions(option) + env["KEY"] = "caller mutation" + first.envVars["KEY"] = "first mutation" + Expect(second.envVars["KEY"]).To(Equal("original")) + }) +}) diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go index f7fe46980..c24556b06 100644 --- a/pkg/model/initializers.go +++ b/pkg/model/initializers.go @@ -75,7 +75,7 @@ const ( func (ml *ModelLoader) grpcModel(backend string, o *Options) func(string, string, string) (*Model, error) { return func(modelID, modelName, modelFile string) (*Model, error) { - xlog.Debug("Loading Model with gRPC", "modelID", modelID, "file", modelFile, "backend", backend, "options", *o) + xlog.Debug("Loading Model with gRPC", "modelID", modelID, "file", modelFile, "backend", backend) // Distributed mode: delegate to the model router if set. No load // event is emitted here: this branch runs per inference request and @@ -180,7 +180,6 @@ func (ml *ModelLoader) spawnGRPCModel(backend, uri string, o *Options, modelID, } xlog.Debug("Wait for the service to start up") - xlog.Debug("Options", "options", o.gRPCOptions) // Wait for the service to start up ready := false @@ -212,7 +211,7 @@ func (ml *ModelLoader) spawnGRPCModel(backend, uri string, o *Options, modelID, options.ModelFile = modelFile options.ModelPath = ml.ModelPath - xlog.Debug("GRPC: Loading model with options", "options", options) + xlog.Debug("GRPC: Loading model", "modelID", modelID, "backend", backend) res, err := client.GRPC(o.parallelRequests, ml.wd).LoadModel(o.context, options) if err != nil { diff --git a/pkg/model/loader_options.go b/pkg/model/loader_options.go index 8c8bfffad..f376701ce 100644 --- a/pkg/model/loader_options.go +++ b/pkg/model/loader_options.go @@ -2,6 +2,7 @@ package model import ( "context" + "maps" pb "github.com/mudler/LocalAI/pkg/grpc/proto" ) @@ -116,8 +117,9 @@ func WithModelSizeBytes(bytes int64) Option { } func WithEnvVars(envVars map[string]string) Option { + snapshot := maps.Clone(envVars) return func(o *Options) { - o.envVars = envVars + o.envVars = maps.Clone(snapshot) } }