feat(distributed): acquire and auto-refresh worker NATS credentials

Workers fetched NATS credentials once at startup, which broke two cases
under JWT auth: a worker that registered while still pending admin
approval never received a minted JWT (it connected unauthenticated and
gave up), and a long-running worker's 24h JWT expired with no way to renew
it.

Introduce workerregistry.NATSCredentialManager, built on idempotent
re-registration (the frontend preserves the node row and mints a fresh JWT
each call):

- Acquire re-registers through admin approval until the node is approved
  and credentials are minted (or returns the first success when auth is
  not required, preserving anonymous-NATS behavior).
- RefreshLoop re-registers before the JWT expires (~75% of its lifetime),
  updating the credentials served to the connection.
- Both are bounded (default 100 attempts / consecutive failures) and
  return an error on exhaustion, so an unapprovable or unrenewable worker
  exits non-zero and surfaces the problem instead of hanging or drifting
  toward an expired credential.

The messaging client gains WithUserJWTProvider, fetching credentials on
each (re)connect so the connection transparently adopts a refreshed JWT
when the server expires the old one. RegisterFull exposes the approval
status and full response; Register delegates to it.

Both the backend worker and the agent worker are wired to this: explicit
env credentials are used as-is, minted credentials are acquired-with-wait
and refreshed, and a permanent refresh failure shuts the worker down so it
restarts and re-acquires.

Tests cover Acquire (wait-through-pending, bounded give-up, context
cancel), RefreshLoop (refresh-before-expiry, bounded failure, no-expiry
exit) and jwtExpiry decoding. Docs updated in distributed-mode.md.

Assisted-by: Claude:claude-opus-4-8 [Claude Code]
Signed-off-by: Richard Palethorpe <io@richiejp.com>

Author: richiejp

core/cli/agent_worker.go

@@ -90,15 +90,30 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 		registrationBody["token"] = cmd.RegistrationToken
 	}
 
-	nodeID, apiToken, regNatsJWT, regNatsSeed, err := regClient.RegisterWithRetry(context.Background(), registrationBody, 10)
+	// Context cancelled on shutdown — used by registration waits, heartbeat, and
+	// other background goroutines.
+	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
+	defer shutdownCancel()
+
+	// Acquire credentials via (re)registration. When the bus requires auth and no
+	// static fallback is configured, wait through admin approval until the
+	// frontend mints credentials rather than starting unauthenticated.
+	credMgr := workerregistry.NewNATSCredentialManager(
+		func(ctx context.Context) (*workerregistry.RegisterResponse, error) {
+			return regClient.RegisterFull(ctx, registrationBody)
+		},
+		cmd.NatsRequireAuth && cmd.NatsJWT == "" && cmd.NatsServiceJWT == "",
+	)
+	res, err := credMgr.Acquire(shutdownCtx)
 	if err != nil {
 		return fmt.Errorf("registration failed: %w", err)
 	}
+	nodeID := res.ID
 	xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cmd.RegisterTo)
 
 	// Use provisioned API token if none was set
 	if cmd.APIToken == "" {
-		cmd.APIToken = apiToken
+		cmd.APIToken = res.APIToken
 	}
 
 	// Start heartbeat
@@ -107,22 +122,35 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 		xlog.Warn("invalid heartbeat interval, using default 10s", "input", cmd.HeartbeatInterval, "error", err)
 	}
 	heartbeatInterval = cmp.Or(heartbeatInterval, 10*time.Second)
-	// Context cancelled on shutdown — used by heartbeat and other background goroutines
-	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
-	defer shutdownCancel()
 
 	go regClient.HeartbeatLoop(shutdownCtx, nodeID, heartbeatInterval, func() map[string]any { return map[string]any{} })
 
-	// Connect to NATS
-	natsJWT := cmp.Or(cmd.NatsJWT, regNatsJWT, cmd.NatsServiceJWT)
-	natsSeed := cmp.Or(cmd.NatsUserSeed, regNatsSeed, cmd.NatsServiceSeed)
-	if cmd.NatsRequireAuth && (natsJWT == "" || natsSeed == "") {
-		return fmt.Errorf("NATS JWT+seed required: enable frontend minting or set LOCALAI_NATS_* env vars")
-	}
+	// Resolve NATS credentials with precedence: explicit env override, then
+	// frontend-minted (auto-refreshed before expiry), then service fallback.
+	// Each static source must supply JWT and seed together.
 	natsTLS := messaging.TLSFiles{CA: cmd.NatsTLSCA, Cert: cmd.NatsTLSCert, Key: cmd.NatsTLSKey}
 	var natsOpts []messaging.Option
-	if natsJWT != "" && natsSeed != "" {
-		natsOpts = append(natsOpts, messaging.WithUserJWT(natsJWT, natsSeed))
+	switch {
+	case cmd.NatsJWT != "" || cmd.NatsUserSeed != "":
+		if (cmd.NatsJWT == "") != (cmd.NatsUserSeed == "") {
+			return fmt.Errorf("LOCALAI_NATS_JWT and LOCALAI_NATS_USER_SEED must be set together")
+		}
+		natsOpts = append(natsOpts, messaging.WithUserJWT(cmd.NatsJWT, cmd.NatsUserSeed))
+	case credMgr.HasCredentials():
+		natsOpts = append(natsOpts, messaging.WithUserJWTProvider(credMgr.Provider()))
+		go func() {
+			if err := credMgr.RefreshLoop(shutdownCtx); err != nil {
+				xlog.Error("NATS credential refresh permanently failed; shutting down agent worker", "error", err)
+				shutdownCancel()
+			}
+		}()
+	case cmd.NatsServiceJWT != "" || cmd.NatsServiceSeed != "":
+		if (cmd.NatsServiceJWT == "") != (cmd.NatsServiceSeed == "") {
+			return fmt.Errorf("LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED must be set together")
+		}
+		natsOpts = append(natsOpts, messaging.WithUserJWT(cmd.NatsServiceJWT, cmd.NatsServiceSeed))
+	case cmd.NatsRequireAuth:
+		return fmt.Errorf("NATS JWT+seed required: enable frontend minting or set LOCALAI_NATS_* env vars")
 	}
 	if natsTLS.Enabled() {
 		natsOpts = append(natsOpts, messaging.WithTLS(natsTLS))
@@ -205,17 +233,25 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
 
 	xlog.Info("Agent worker ready, waiting for jobs", "subject", cmd.Subject, "queue", cmd.Queue)
 
-	// Wait for shutdown
+	// Wait for an OS signal or an internal fatal condition (e.g. NATS
+	// credentials became unrenewable), so the worker restarts and re-acquires
+	// rather than lingering unable to serve.
 	sigCh := make(chan os.Signal, 1)
 	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
-	<-sigCh
+	var runErr error
+	select {
+	case <-sigCh:
+	case <-shutdownCtx.Done():
+		runErr = fmt.Errorf("agent worker shutting down: NATS credentials unavailable")
+		xlog.Error("Internal shutdown requested", "error", runErr)
+	}
 
 	xlog.Info("Shutting down agent worker")
 	shutdownCancel() // stop heartbeat loop immediately
 	dispatcher.Stop()
 	mcpTools.CloseAllMCPSessions()
 	regClient.GracefulDeregister(nodeID)
-	return nil
+	return runErr
 }
 
 // handleMCPToolRequest handles a NATS request-reply for MCP tool execution.

core/cli/workerregistry/client.go

@@ -58,40 +58,53 @@ func (c *RegistrationClient) setAuth(req *http.Request) {
 
 // RegisterResponse is the JSON body returned by /api/node/register.
 type RegisterResponse struct {
-	ID       string `json:"id"`
+	ID           string `json:"id"`
+	Status       string `json:"status,omitempty"` // "pending" until an admin approves the node
 	APIToken     string `json:"api_token,omitempty"`
 	NatsJWT      string `json:"nats_jwt,omitempty"`
 	NatsUserSeed string `json:"nats_user_seed,omitempty"`
 }
 
-// Register sends a single registration request and returns the node ID and
-// optional credentials (API token for agent workers, NATS JWT when configured).
-func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) (nodeID, apiToken, natsJWT, natsSeed string, err error) {
+// RegisterFull sends a single registration request and returns the full
+// response (node ID, approval status, and optional API token / NATS creds).
+// Re-registration is idempotent: the frontend preserves the node row and mints
+// a fresh NATS JWT each call, so this doubles as the credential-refresh call.
+func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]any) (*RegisterResponse, error) {
 	jsonBody, _ := json.Marshal(body)
 	url := c.baseURL() + "/api/node/register"
 
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(jsonBody))
 	if err != nil {
-		return "", "", "", "", fmt.Errorf("creating request: %w", err)
+		return nil, fmt.Errorf("creating request: %w", err)
 	}
 	req.Header.Set("Content-Type", "application/json")
 	c.setAuth(req)
 
 	resp, err := c.httpClient().Do(req)
 	if err != nil {
-		return "", "", "", "", fmt.Errorf("posting to %s: %w", url, err)
+		return nil, fmt.Errorf("posting to %s: %w", url, err)
 	}
 	defer resp.Body.Close()
 
 	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		return "", "", "", "", fmt.Errorf("registration failed with status %d", resp.StatusCode)
+		return nil, fmt.Errorf("registration failed with status %d", resp.StatusCode)
 	}
 
 	var result RegisterResponse
 	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-		return "", "", "", "", fmt.Errorf("decoding response: %w", err)
+		return nil, fmt.Errorf("decoding response: %w", err)
+	}
+	return &result, nil
+}
+
+// Register sends a single registration request and returns the node ID and
+// optional credentials (API token for agent workers, NATS JWT when configured).
+func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) (nodeID, apiToken, natsJWT, natsSeed string, err error) {
+	res, err := c.RegisterFull(ctx, body)
+	if err != nil {
+		return "", "", "", "", err
 	}
-	return result.ID, result.APIToken, result.NatsJWT, result.NatsUserSeed, nil
+	return res.ID, res.APIToken, res.NatsJWT, res.NatsUserSeed, nil
 }
 
 // RegisterWithRetry retries registration with exponential backoff.

core/cli/workerregistry/credentials.go

@@ -0,0 +1,200 @@
+package workerregistry
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/mudler/LocalAI/pkg/natsauth"
+	"github.com/mudler/xlog"
+)
+
+// statusPending mirrors nodes.StatusPending. It is duplicated rather than
+// imported so the lightweight registration client does not pull in the nodes
+// package (and its gorm/DB dependencies).
+const statusPending = "pending"
+
+// defaultMaxAttempts bounds how many times Acquire registers (and how many
+// consecutive times RefreshLoop may fail) before giving up. It is high enough
+// to ride out a slow admin approval or a transient frontend outage, but finite
+// so an unauthorized/unapprovable worker exits and surfaces the problem (via a
+// non-zero exit and the resulting restart) rather than waiting forever.
+const defaultMaxAttempts = 100
+
+// RegisterFunc performs one idempotent registration round-trip.
+type RegisterFunc func(ctx context.Context) (*RegisterResponse, error)
+
+// NATSCredentialManager acquires NATS credentials at startup — waiting through
+// admin approval when required — and refreshes them before the minted JWT
+// expires, by re-registering (which mints a fresh JWT). The live NATS
+// connection adopts a refreshed JWT on its next reconnect via Provider. Safe
+// for concurrent use.
+//
+// It addresses two failure modes: a worker that needs credentials but registers
+// while still pending approval (it would otherwise give up and never connect),
+// and a long-running worker whose 24h JWT expires with no way to renew it.
+type NATSCredentialManager struct {
+	register     RegisterFunc
+	requireCreds bool // block until credentials are present (frontend minting in use)
+
+	// Tunables; defaults set by NewNATSCredentialManager, overridable in tests.
+	initialBackoff time.Duration
+	maxBackoff     time.Duration
+	maxAttempts    int     // bound on Acquire attempts / consecutive refresh failures (<=0 = unlimited)
+	refreshLead    float64 // refresh once this fraction of the JWT lifetime has elapsed
+	refreshRetry   time.Duration
+	expiryOf       func(jwt string) (time.Time, bool)
+
+	mu     sync.RWMutex
+	jwt    string
+	seed   string
+	nodeID string
+}
+
+// NewNATSCredentialManager builds a manager over register. When requireCreds is
+// true, Acquire blocks until the node is approved and credentials are minted.
+func NewNATSCredentialManager(register RegisterFunc, requireCreds bool) *NATSCredentialManager {
+	return &NATSCredentialManager{
+		register:       register,
+		requireCreds:   requireCreds,
+		initialBackoff: 2 * time.Second,
+		maxBackoff:     30 * time.Second,
+		maxAttempts:    defaultMaxAttempts,
+		refreshLead:    0.75,
+		refreshRetry:   30 * time.Second,
+		expiryOf:       jwtExpiry,
+	}
+}
+
+// jwtExpiry decodes the expiry of a minted user JWT. ok is false when the token
+// is empty/undecodable or carries no expiry (e.g. a non-expiring service JWT).
+func jwtExpiry(token string) (time.Time, bool) {
+	if token == "" {
+		return time.Time{}, false
+	}
+	uc, err := natsauth.DecodeUserClaims(token)
+	if err != nil || uc.Expires == 0 {
+		return time.Time{}, false
+	}
+	return time.Unix(uc.Expires, 0), true
+}
+
+func (m *NATSCredentialManager) store(res *RegisterResponse) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.nodeID = res.ID
+	if res.NatsJWT != "" && res.NatsUserSeed != "" {
+		m.jwt, m.seed = res.NatsJWT, res.NatsUserSeed
+	}
+}
+
+// Current returns the latest NATS credentials (both empty until acquired).
+func (m *NATSCredentialManager) Current() (jwt, seed string) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.jwt, m.seed
+}
+
+// NodeID returns the node ID from the most recent registration.
+func (m *NATSCredentialManager) NodeID() string {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.nodeID
+}
+
+// Provider returns a callback compatible with messaging.WithUserJWTProvider,
+// supplying the current credentials on each (re)connect.
+func (m *NATSCredentialManager) Provider() func() (string, string) {
+	return m.Current
+}
+
+// HasCredentials reports whether complete NATS credentials have been obtained.
+func (m *NATSCredentialManager) HasCredentials() bool {
+	jwt, seed := m.Current()
+	return jwt != "" && seed != ""
+}
+
+// Acquire registers and, when requireCreds is set, keeps re-registering with
+// exponential backoff until the node is approved (status != pending) and
+// credentials are minted. Without requireCreds it returns the first successful
+// response (the historical one-shot behavior, preserved for anonymous NATS).
+func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, error) {
+	backoff := m.initialBackoff
+	var lastReason error
+	for attempt := 1; m.maxAttempts <= 0 || attempt <= m.maxAttempts; attempt++ {
+		res, err := m.register(ctx)
+		switch {
+		case err != nil:
+			lastReason = err
+			xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err)
+		case !m.requireCreds:
+			m.store(res)
+			return res, nil
+		case res.Status == statusPending:
+			lastReason = fmt.Errorf("node %s still pending admin approval", res.ID)
+			xlog.Info("Node pending admin approval; waiting", "node", res.ID, "attempt", attempt, "next_retry", backoff)
+		case res.NatsJWT == "" || res.NatsUserSeed == "":
+			lastReason = fmt.Errorf("node %s approved but NATS credentials not minted", res.ID)
+			xlog.Info("Node approved but NATS credentials not yet minted; waiting", "node", res.ID, "attempt", attempt, "next_retry", backoff)
+		default:
+			m.store(res)
+			return res, nil
+		}
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(backoff):
+		}
+		backoff = min(backoff*2, m.maxBackoff)
+	}
+	return nil, fmt.Errorf("giving up acquiring NATS credentials after %d attempts: %w", m.maxAttempts, lastReason)
+}
+
+// RefreshLoop re-registers to mint a fresh JWT before the current one expires,
+// updating the credentials returned by Current/Provider so the NATS connection
+// adopts them on its next reconnect. It returns nil when ctx is cancelled or
+// when the current credential has no expiry (nothing to refresh), and a non-nil
+// error after maxAttempts consecutive refresh failures — letting the caller
+// exit the worker so it restarts and re-acquires (or surfaces the outage)
+// rather than silently drifting toward an expired, unrenewable JWT.
+func (m *NATSCredentialManager) RefreshLoop(ctx context.Context) error {
+	failures := 0
+	for {
+		jwt, _ := m.Current()
+		exp, ok := m.expiryOf(jwt)
+		if !ok {
+			xlog.Debug("NATS credential has no expiry; refresh loop exiting")
+			return nil
+		}
+		wait := max(time.Duration(float64(time.Until(exp))*m.refreshLead), 0)
+		select {
+		case <-ctx.Done():
+			return nil
+		case <-time.After(wait):
+		}
+
+		res, err := m.register(ctx)
+		if err == nil && res.NatsJWT != "" && res.NatsUserSeed != "" {
+			m.store(res)
+			failures = 0
+			xlog.Info("Refreshed NATS credentials", "node", res.ID)
+			continue
+		}
+		failures++
+		if err != nil {
+			xlog.Warn("NATS credential refresh failed; will retry", "attempt", failures, "error", err)
+		} else {
+			xlog.Warn("NATS credential refresh returned no credentials; will retry", "attempt", failures)
+		}
+		if m.maxAttempts > 0 && failures >= m.maxAttempts {
+			return fmt.Errorf("NATS credential refresh failed %d times in a row", failures)
+		}
+		// Back off before retrying so a persistent failure near expiry does not spin.
+		select {
+		case <-ctx.Done():
+			return nil
+		case <-time.After(m.refreshRetry):
+		}
+	}
+}