fix(llm): refresh ESS assertion credentials (#1143)

Signed-off-by: Mike Camp <mcamp@nvidia.com>

Author: mikeyrcamp

src/compute-plane-services/nvca/vendor/github.com/NVIDIA/nvcf/src/libraries/go/lib/pkg/icms-translate/translate/function/llm.go

@@ -1,10 +1,23 @@
 # worker-llm-credentials
 
-Sidecar that maintains a fresh NVCF worker credential token on disk for the
-LLM inference container. It connects to NVCF over gRPC using the worker token,
-runs a background refresher that periodically fetches a new token, and writes
-it atomically to the configured token path (owner-only `0600` permissions). The
-process runs until its context is cancelled.
+Sidecar that maintains fresh NVCF credentials on disk for explicit LLM
+functions. It connects to NVCF over gRPC using the worker token and always
+refreshes the worker credential used by the LLM router client. When
+`ESS_ASSERTION_TOKEN_PATH` is set, it also refreshes the ESS assertion used by
+the ESS Agent. The process runs until its context is cancelled.
+
+Each refresher fetches its credential immediately, schedules later refreshes
+from the returned expiration, and atomically replaces the configured file. The
+shared files use `0644` permissions because the consuming infrastructure
+sidecars run as different non-root users. Customer workload containers do not
+mount these credential volumes.
+
+## Configuration
+
+- `WORKER_TOKEN_PATH` sets the worker credential file. It defaults to
+  `/var/run/llm/worker-token`.
+- `ESS_ASSERTION_TOKEN_PATH` optionally sets the ESS assertion file. When it is
+  unset, the sidecar does not request or write an ESS assertion.
 
 ## Build
 

src/compute-plane-services/worker-llm-credentials/README.md

@@ -1,4 +1,4 @@
-load("@rules_go//go:def.bzl", "go_library")
+load("@rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "configs",
@@ -7,6 +7,12 @@ go_library(
     visibility = ["//visibility:public"],
 )
 
+go_test(
+    name = "configs_test",
+    srcs = ["configs_test.go"],
+    embed = [":configs"],
+)
+
 alias(
     name = "go_default_library",
     actual = ":configs",

src/compute-plane-services/worker-llm-credentials/configs/BUILD.bazel

@@ -18,14 +18,15 @@ limitations under the License.
 package configs
 
 type Config struct {
-	NvcfFqdnGrpc      string `mapstructure:"NVCF_FQDN_GRPC"`
-	NvcfWorkerToken   string `mapstructure:"NVCF_WORKER_TOKEN"`
-	FunctionId        string `mapstructure:"FUNCTION_ID"`
-	FunctionVersionId string `mapstructure:"FUNCTION_VERSION_ID"`
-	NcaId             string `mapstructure:"NCA_ID"`
-	InstanceId        string `mapstructure:"INSTANCE_ID"`
-	SharedConfigDir   string `mapstructure:"SHARED_CONFIG_DIR"`
-	WorkerTokenPath   string `mapstructure:"WORKER_TOKEN_PATH"`
+	NvcfFqdnGrpc          string `mapstructure:"NVCF_FQDN_GRPC"`
+	NvcfWorkerToken       string `mapstructure:"NVCF_WORKER_TOKEN"`
+	FunctionId            string `mapstructure:"FUNCTION_ID"`
+	FunctionVersionId     string `mapstructure:"FUNCTION_VERSION_ID"`
+	NcaId                 string `mapstructure:"NCA_ID"`
+	InstanceId            string `mapstructure:"INSTANCE_ID"`
+	SharedConfigDir       string `mapstructure:"SHARED_CONFIG_DIR"`
+	WorkerTokenPath       string `mapstructure:"WORKER_TOKEN_PATH"`
+	ESSAssertionTokenPath string `mapstructure:"ESS_ASSERTION_TOKEN_PATH"`
 }
 
 const DefaultSharedConfigDir = "/config/shared"

src/compute-plane-services/worker-llm-credentials/configs/configs.go

@@ -0,0 +1,33 @@
+/*
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package configs
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestConfigSupportsOptionalESSAssertionTokenPath(t *testing.T) {
+	field, ok := reflect.TypeOf(Config{}).FieldByName("ESSAssertionTokenPath")
+	if !ok {
+		t.Fatal("Config must expose an optional ESSAssertionTokenPath field")
+	}
+	if got := field.Tag.Get("mapstructure"); got != "ESS_ASSERTION_TOKEN_PATH" {
+		t.Fatalf("ESSAssertionTokenPath mapstructure tag = %q, want %q", got, "ESS_ASSERTION_TOKEN_PATH")
+	}
+}

src/compute-plane-services/worker-llm-credentials/configs/configs_test.go

@@ -80,6 +80,9 @@ func (w *Worker) Run(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
+	if w.config.ESSAssertionTokenPath != "" {
+		w.client.StartAssertionTokenRefresher(ctx, w.config.ESSAssertionTokenPath, true)
+	}
 
 	token.StartTokenRefresher(ctx, "llm worker token", true,
 		func(ctx context.Context) (token.Token, error) {

src/compute-plane-services/worker-llm-credentials/internal/worker/worker.go

@@ -23,6 +23,7 @@ import (
 	"net"
 	"os"
 	"path/filepath"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -39,6 +40,7 @@ const testWorkerToken = "test-worker-token"
 
 type mockNVCFServer struct {
 	pb.UnimplementedWorkerServer
+	requestSecretCredentials func(context.Context, *pb.SecretCredentialsRequest) (*pb.SecretCredentialsResponse, error)
 }
 
 func (s *mockNVCFServer) ConnectOnce(_ context.Context, _ *pb.WorkerConnect) (*pb.WorkerConnectOnceResponse, error) {
@@ -49,19 +51,65 @@ func (s *mockNVCFServer) ConnectOnce(_ context.Context, _ *pb.WorkerConnect) (*p
 	}, nil
 }
 
+func (s *mockNVCFServer) RequestSecretCredentials(
+	ctx context.Context,
+	req *pb.SecretCredentialsRequest,
+) (*pb.SecretCredentialsResponse, error) {
+	if s.requestSecretCredentials == nil {
+		return s.UnimplementedWorkerServer.RequestSecretCredentials(ctx, req)
+	}
+	return s.requestSecretCredentials(ctx, req)
+}
+
 func startMockNVCFServer(t *testing.T) string {
+	return startMockNVCFServerWithImplementation(t, &mockNVCFServer{})
+}
+
+func startMockNVCFServerWithImplementation(t *testing.T, implementation *mockNVCFServer) string {
 	t.Helper()
-	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	var listenConfig net.ListenConfig
+	lis, err := listenConfig.Listen(context.Background(), "tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatalf("failed to listen: %v", err)
 	}
 	srv := grpc.NewServer()
-	pb.RegisterWorkerServer(srv, &mockNVCFServer{})
-	go srv.Serve(lis)
-	t.Cleanup(srv.GracefulStop)
+	pb.RegisterWorkerServer(srv, implementation)
+	serveErr := make(chan error, 1)
+	go func() {
+		serveErr <- srv.Serve(lis)
+	}()
+	t.Cleanup(func() {
+		srv.GracefulStop()
+		if err := <-serveErr; err != nil {
+			t.Errorf("serve mock NVCF server: %v", err)
+		}
+	})
 	return fmt.Sprintf("http://%s", lis.Addr().String())
 }
 
+func waitForFileContent(path, want string, timeout time.Duration) bool {
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		content, err := os.ReadFile(path)
+		if err == nil && string(content) == want {
+			return true
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	content, err := os.ReadFile(path)
+	return err == nil && string(content) == want
+}
+
+func cancelAndWaitForRun(t *testing.T, cancel context.CancelFunc, runErr <-chan error) {
+	t.Helper()
+	cancel()
+	select {
+	case <-runErr:
+	case <-time.After(2 * time.Second):
+		t.Log("Run did not stop within the cleanup timeout")
+	}
+}
+
 func TestRun_WritesTokenToDisk(t *testing.T) {
 	addr := startMockNVCFServer(t)
 	tmpDir := t.TempDir()
@@ -110,3 +158,134 @@ func TestRun_WritesTokenToDisk(t *testing.T) {
 		t.Fatalf("expected token %q, got %q", testWorkerToken, string(content))
 	}
 }
+
+func TestRun_RefreshesESSAssertionTokenUntilCancelled(t *testing.T) {
+	tmpDir := t.TempDir()
+	assertionTokenPath := filepath.Join(tmpDir, "ess", "jwt.token")
+	if err := os.MkdirAll(filepath.Dir(assertionTokenPath), 0755); err != nil {
+		t.Fatalf("create assertion token directory: %v", err)
+	}
+	if err := os.WriteFile(assertionTokenPath, []byte("worker-init-token"), 0600); err != nil {
+		t.Fatalf("write initial assertion token: %v", err)
+	}
+
+	var requestCount atomic.Int32
+	addr := startMockNVCFServerWithImplementation(t, &mockNVCFServer{
+		requestSecretCredentials: func(
+			_ context.Context,
+			_ *pb.SecretCredentialsRequest,
+		) (*pb.SecretCredentialsResponse, error) {
+			call := requestCount.Add(1)
+			tokenValue := "refreshed-assertion-1"
+			expiration := time.Now().Add(300 * time.Millisecond)
+			if call > 1 {
+				tokenValue = "refreshed-assertion-2"
+				expiration = time.Now().Add(time.Hour)
+			}
+			return &pb.SecretCredentialsResponse{
+				SecretCredentialsToken: tokenValue,
+				Expiration:             timestamppb.New(expiration),
+			}, nil
+		},
+	})
+
+	cfg := configs.Config{
+		NvcfFqdnGrpc:          addr,
+		NvcfWorkerToken:       "initial-token",
+		FunctionId:            "test-function-id",
+		FunctionVersionId:     "test-function-version-id",
+		NcaId:                 "test-nca-id",
+		InstanceId:            "test-instance-id",
+		SharedConfigDir:       tmpDir,
+		WorkerTokenPath:       filepath.Join(tmpDir, "worker-token"),
+		ESSAssertionTokenPath: assertionTokenPath,
+	}
+
+	w, err := New(cfg)
+	if err != nil {
+		t.Fatalf("New: %v", err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	runErr := lo.Async(func() error { return w.Run(ctx) })
+
+	if !waitForFileContent(assertionTokenPath, "refreshed-assertion-1", 5*time.Second) {
+		cancelAndWaitForRun(t, cancel, runErr)
+		t.Fatal("existing assertion token was not replaced by the immediate refresh")
+	}
+	info, err := os.Stat(assertionTokenPath)
+	if err != nil {
+		cancelAndWaitForRun(t, cancel, runErr)
+		t.Fatalf("stat refreshed assertion token: %v", err)
+	}
+	if got := info.Mode().Perm(); got != 0644 {
+		cancelAndWaitForRun(t, cancel, runErr)
+		t.Fatalf("assertion token mode = %o, want 0644", got)
+	}
+	if !waitForFileContent(assertionTokenPath, "refreshed-assertion-2", 5*time.Second) {
+		cancelAndWaitForRun(t, cancel, runErr)
+		t.Fatal("assertion token was not rotated on the next refresh cycle")
+	}
+
+	cancel()
+	select {
+	case err := <-runErr:
+		if err != nil {
+			t.Fatalf("Run after cancellation: %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Run did not stop cleanly after cancellation")
+	}
+}
+
+func TestRun_WithoutESSAssertionTokenPathSkipsSecretCredentialRefresh(t *testing.T) {
+	var requestCount atomic.Int32
+	addr := startMockNVCFServerWithImplementation(t, &mockNVCFServer{
+		requestSecretCredentials: func(
+			_ context.Context,
+			_ *pb.SecretCredentialsRequest,
+		) (*pb.SecretCredentialsResponse, error) {
+			requestCount.Add(1)
+			return &pb.SecretCredentialsResponse{
+				SecretCredentialsToken: "unexpected-token",
+				Expiration:             timestamppb.New(time.Now().Add(time.Hour)),
+			}, nil
+		},
+	})
+
+	tmpDir := t.TempDir()
+	assertionTokenPath := filepath.Join(tmpDir, "ess", "jwt.token")
+	workerTokenPath := filepath.Join(tmpDir, "worker-token")
+	cfg := configs.Config{
+		NvcfFqdnGrpc:      addr,
+		NvcfWorkerToken:   "initial-token",
+		FunctionId:        "test-function-id",
+		FunctionVersionId: "test-function-version-id",
+		NcaId:             "test-nca-id",
+		InstanceId:        "test-instance-id",
+		SharedConfigDir:   tmpDir,
+		WorkerTokenPath:   workerTokenPath,
+	}
+
+	w, err := New(cfg)
+	if err != nil {
+		t.Fatalf("New: %v", err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	runErr := lo.Async(func() error { return w.Run(ctx) })
+	if !waitForFileContent(workerTokenPath, testWorkerToken, 5*time.Second) {
+		cancelAndWaitForRun(t, cancel, runErr)
+		t.Fatal("worker token was not written")
+	}
+	cancel()
+	if err := <-runErr; err != nil {
+		t.Fatalf("Run: %v", err)
+	}
+
+	if got := requestCount.Load(); got != 0 {
+		t.Fatalf("RequestSecretCredentials calls = %d, want 0", got)
+	}
+	if _, err := os.Stat(assertionTokenPath); !os.IsNotExist(err) {
+		t.Fatalf("assertion token file should not exist without ESS_ASSERTION_TOKEN_PATH, stat error: %v", err)
+	}
+}