fix(workloadmanager): refresh user client when token changes

pkg/workloadmanager/client_cache.go

@@ -30,6 +30,7 @@ import (
 type clientCacheEntry struct {
 	key         string
 	client      *UserK8sClient
+	token       string
 	tokenExpiry time.Time // Token expiration time parsed from JWT
 	element     *list.Element
 }
@@ -96,15 +97,17 @@ func NewClientCache(maxSize int) *ClientCache {
 
 // Get retrieves a client from cache based on key (service account)
 // Returns the client if found and cached token is not expired, nil otherwise
-// Different tokens for the same service account can share the same client
-func (c *ClientCache) Get(key string) *UserK8sClient {
+func (c *ClientCache) Get(key, token string) *UserK8sClient {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
 	entry, exists := c.cache[key]
 	if !exists {
 		return nil
 	}
+	if entry.token != token {
+		return nil
+	}
 
 	// Check if cached entry's token is expired
 	if !entry.tokenExpiry.IsZero() && time.Now().After(entry.tokenExpiry) {
@@ -133,6 +136,7 @@ func (c *ClientCache) Set(key, token string, client *UserK8sClient) {
 	if entry, exists := c.cache[key]; exists {
 		// Update existing entry
 		entry.client = client
+		entry.token = token
 		entry.tokenExpiry = tokenExpiry
 		c.lruList.MoveToFront(entry.element)
 		return
@@ -147,6 +151,7 @@ func (c *ClientCache) Set(key, token string, client *UserK8sClient) {
 	entry := &clientCacheEntry{
 		key:         key,
 		client:      client,
+		token:       token,
 		tokenExpiry: tokenExpiry,
 	}
 	entry.element = c.lruList.PushFront(entry)

pkg/workloadmanager/client_cache_test.go

@@ -171,7 +171,7 @@ func TestClientCache_SetAndGet(t *testing.T) {
 	cache := NewClientCache(10)
 
 	// Test Get on non-existent key
-	client := cache.Get("nonexistent-key")
+	client := cache.Get("nonexistent-key", "token")
 	assert.Nil(t, client)
 
 	// Test Set and Get
@@ -186,7 +186,7 @@ func TestClientCache_SetAndGet(t *testing.T) {
 
 	cache.Set(key, token, newClient)
 
-	retrieved := cache.Get(key)
+	retrieved := cache.Get(key, token)
 	assert.NotNil(t, retrieved)
 	assert.Equal(t, newClient, retrieved)
 	assert.Equal(t, 1, cache.Size())
@@ -208,7 +208,7 @@ func TestClientCache_Get_ExpiredToken(t *testing.T) {
 	cache.Set(key, token, client)
 
 	// Should return nil because token is expired
-	retrieved := cache.Get(key)
+	retrieved := cache.Get(key, token)
 	assert.Nil(t, retrieved)
 	assert.Equal(t, 0, cache.Size(), "Expired entry should be removed")
 }
@@ -230,11 +230,30 @@ func TestClientCache_Get_TokenWithoutExpiry(t *testing.T) {
 	cache.Set(key, token, client)
 
 	// Should return client because tokenExpiry is zero (no expiry check)
-	retrieved := cache.Get(key)
+	retrieved := cache.Get(key, token)
 	assert.NotNil(t, retrieved)
 	assert.Equal(t, client, retrieved)
 }
 
+func TestClientCache_Get_TokenMismatch(t *testing.T) {
+	cache := NewClientCache(10)
+
+	key := testCacheKey
+	token := createTestJWT(time.Now().Add(1 * time.Hour).Unix())
+	scheme := runtime.NewScheme()
+	dynamicClient := dynamicfake.NewSimpleDynamicClient(scheme)
+	client := &UserK8sClient{
+		dynamicClient: dynamicClient,
+		namespace:     "default",
+	}
+
+	cache.Set(key, token, client)
+
+	retrieved := cache.Get(key, token+"-new")
+	assert.Nil(t, retrieved)
+	assert.Equal(t, 1, cache.Size())
+}
+
 func TestClientCache_UpdateExisting(t *testing.T) {
 	cache := NewClientCache(10)
 
@@ -259,7 +278,10 @@ func TestClientCache_UpdateExisting(t *testing.T) {
 	cache.Set(key, token2, client2)
 	assert.Equal(t, 1, cache.Size(), "Size should not increase when updating")
 
-	retrieved := cache.Get(key)
+	oldClient := cache.Get(key, token1)
+	assert.Nil(t, oldClient, "Old token should not match updated cache entry")
+
+	retrieved := cache.Get(key, token2)
 	assert.Equal(t, client2, retrieved, "Should return updated client")
 }
 
@@ -268,11 +290,11 @@ func TestClientCache_Eviction(t *testing.T) {
 
 	scheme := runtime.NewScheme()
 	dynamicClient := dynamicfake.NewSimpleDynamicClient(scheme)
+	token := createTestJWT(time.Now().Add(1 * time.Hour).Unix())
 
 	// Fill cache to max size
 	for i := 0; i < 3; i++ {
 		key := "default:sa" + string(rune('0'+i))
-		token := createTestJWT(time.Now().Add(1 * time.Hour).Unix())
 		client := &UserK8sClient{
 			dynamicClient: dynamicClient,
 			namespace:     "default",
@@ -294,9 +316,9 @@ func TestClientCache_Eviction(t *testing.T) {
 	assert.Equal(t, 3, cache.Size(), "Size should remain at max")
 
 	// First key should be evicted
-	assert.Nil(t, cache.Get("default:sa0"))
+	assert.Nil(t, cache.Get("default:sa0", token))
 	// New key should be present
-	assert.NotNil(t, cache.Get(newKey))
+	assert.NotNil(t, cache.Get(newKey, newToken))
 }
 
 func TestClientCache_LRUBehavior(t *testing.T) {
@@ -317,7 +339,7 @@ func TestClientCache_LRUBehavior(t *testing.T) {
 	}
 
 	// Access first entry (should move to front)
-	cache.Get("default:sa0")
+	cache.Get("default:sa0", token)
 
 	// Add new entry - should evict sa1 (least recently used)
 	newKey := "default:sa3"
@@ -328,13 +350,13 @@ func TestClientCache_LRUBehavior(t *testing.T) {
 	cache.Set(newKey, token, newClient)
 
 	// sa0 should still be present (was accessed)
-	assert.NotNil(t, cache.Get("default:sa0"))
+	assert.NotNil(t, cache.Get("default:sa0", token))
 	// sa1 should be evicted
-	assert.Nil(t, cache.Get("default:sa1"))
+	assert.Nil(t, cache.Get("default:sa1", token))
 	// sa2 should be present
-	assert.NotNil(t, cache.Get("default:sa2"))
+	assert.NotNil(t, cache.Get("default:sa2", token))
 	// sa3 should be present
-	assert.NotNil(t, cache.Get(newKey))
+	assert.NotNil(t, cache.Get(newKey, token))
 }
 
 func TestClientCache_Remove(t *testing.T) {
@@ -374,7 +396,7 @@ func TestClientCache_Remove(t *testing.T) {
 
 			cache.Remove(tt.key)
 			assert.Equal(t, 0, cache.Size())
-			assert.Nil(t, cache.Get(tt.key))
+			assert.Nil(t, cache.Get(tt.key, "token"))
 		})
 	}
 }

pkg/workloadmanager/k8s_client.go

@@ -197,7 +197,7 @@ func (c *K8sClient) GetOrCreateUserK8sClient(userToken, namespace, serviceAccoun
 	cacheKey := makeCacheKey(namespace, serviceAccountName)
 
 	// Try to get from cache
-	if cachedClient := c.clientCache.Get(cacheKey); cachedClient != nil {
+	if cachedClient := c.clientCache.Get(cacheKey, userToken); cachedClient != nil {
 		return cachedClient, nil
 	}
 

pkg/workloadmanager/k8s_client_test.go

@@ -25,6 +25,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	listersv1 "k8s.io/client-go/listers/core/v1"
+	"k8s.io/client-go/rest"
 )
 
 // Helper function to create a pod with owner reference
@@ -210,3 +211,31 @@ func TestGetSandboxPodIP_InvalidPodStatus(t *testing.T) {
 		})
 	}
 }
+
+func TestGetOrCreateUserK8sClient_CacheByToken(t *testing.T) {
+	client := &K8sClient{
+		baseConfig:  &rest.Config{Host: "https://example.com"},
+		clientCache: NewClientCache(10),
+	}
+
+	firstClient, err := client.GetOrCreateUserK8sClient("first-token", "default", "runner")
+	assert.NoError(t, err)
+
+	sameTokenClient, err := client.GetOrCreateUserK8sClient("first-token", "default", "runner")
+	assert.NoError(t, err)
+	if firstClient != sameTokenClient {
+		t.Errorf("expected same token to reuse cached client")
+	}
+
+	secondClient, err := client.GetOrCreateUserK8sClient("second-token", "default", "runner")
+	assert.NoError(t, err)
+	if firstClient == secondClient {
+		t.Errorf("expected changed token to create a new client")
+	}
+
+	otherServiceAccountClient, err := client.GetOrCreateUserK8sClient("first-token", "default", "worker")
+	assert.NoError(t, err)
+	if secondClient == otherServiceAccountClient {
+		t.Errorf("expected different service account to create a new client")
+	}
+}