feat: add session observability GET endpoints for sandboxes

Signed-off-by: Abhinav Singh <abhinavsingh717073@gmail.com>

Author: Abhinav-kodes

pkg/router/session_manager_test.go

@@ -86,6 +86,10 @@ func (f *fakeStoreClient) ListExpiredSandboxes(_ context.Context, _ time.Time, _
 	return nil, nil
 }
 
+func (f *fakeStoreClient) ListSandboxesByKind(_ context.Context, _ string) ([]*types.SandboxInfo, error) {
+	return nil, nil
+}
+
 func (f *fakeStoreClient) ListInactiveSandboxes(_ context.Context, _ time.Time, _ int64) ([]*types.SandboxInfo, error) {
 	return nil, nil
 }

pkg/store/interface.go

@@ -34,6 +34,8 @@ type Store interface {
 	UpdateSandbox(ctx context.Context, sandboxStore *types.SandboxInfo) error
 	// DeleteSandboxBySessionID delete sandbox by session ID
 	DeleteSandboxBySessionID(ctx context.Context, sessionID string) error
+	// ListSandboxesByKind returns all active sandboxes matching the given kind.
+	ListSandboxesByKind(ctx context.Context, kind string) ([]*types.SandboxInfo, error)
 	// ListExpiredSandboxes returns up to limit sandboxes with ExpiresAt before the given time
 	ListExpiredSandboxes(ctx context.Context, before time.Time, limit int64) ([]*types.SandboxInfo, error)
 	// ListInactiveSandboxes returns up to limit sandboxes with last-activity time before the given time

pkg/store/store_redis.go

@@ -143,6 +143,57 @@ func (rs *redisStore) GetSandboxBySessionID(ctx context.Context, sessionID strin
 	return &sandboxRedis, nil
 }
 
+// ListSandboxesByKind returns all active sandboxes matching the given kind.
+// Uses SCAN to prevent blocking the Redis instance on large datasets.
+func (rs *redisStore) ListSandboxesByKind(ctx context.Context, kind string) ([]*types.SandboxInfo, error) {
+	allSandboxes := make([]*types.SandboxInfo, 0) // Initialize as empty array, not nil
+	var cursor uint64
+	matchPattern := rs.sessionPrefix + "*"
+	seenKeys := make(map[string]bool)
+
+	for {
+		keys, nextCursor, err := rs.cli.Scan(ctx, cursor, matchPattern, 100).Result()
+		if err != nil {
+			return nil, fmt.Errorf("ListSandboxesByKind scan failed: %w", err)
+		}
+
+		if len(keys) > 0 {
+			sessionIDs := make([]string, 0, len(keys))
+			for _, key := range keys {
+				if key == rs.expiryIndexKey || key == rs.lastActivityIndexKey {
+					continue
+				}
+				if seenKeys[key] {
+					continue
+				}
+				seenKeys[key] = true
+				sessionIDs = append(sessionIDs, strings.TrimPrefix(key, rs.sessionPrefix))
+			}
+
+			if len(sessionIDs) > 0 {
+				// Batch load the fetched keys
+				sandboxes, err := rs.loadSandboxesBySessionIDs(ctx, sessionIDs)
+				if err != nil {
+					return nil, err
+				}
+
+				// Filter by requested kind
+				for _, sb := range sandboxes {
+					if sb != nil && sb.Kind == kind {
+						allSandboxes = append(allSandboxes, sb)
+					}
+				}
+			}
+		}
+
+		cursor = nextCursor
+		if cursor == 0 {
+			break // Scan complete
+		}
+	}
+	return allSandboxes, nil
+}
+
 func (rs *redisStore) StoreSandbox(ctx context.Context, sandboxRedis *types.SandboxInfo) error {
 	if sandboxRedis == nil {
 		return errors.New("StoreSandbox: sandbox is nil")

pkg/store/store_redis_test.go

@@ -67,8 +67,8 @@ func newTestRedisClient(t *testing.T) (*redisStore, *miniredis.Miniredis) {
 	rs := &redisStore{
 		cli:                  redisv9.NewClient(&redisv9.Options{Addr: mr.Addr()}),
 		sessionPrefix:        "session:",
-		expiryIndexKey:       "sandbox:expiry",
-		lastActivityIndexKey: "sandbox:last_activity",
+		expiryIndexKey:       "session:expiry",
+		lastActivityIndexKey: "session:last_activity",
 	}
 	return rs, mr
 }
@@ -344,3 +344,45 @@ func TestUpdateSandboxLastActivity(t *testing.T) {
 		t.Fatalf("unexpected lastActivity score after update: got %v, want %v", score, newLastActivity.Unix())
 	}
 }
+
+func TestRedisStore_ListSandboxesByKind(t *testing.T) {
+	ctx := context.Background()
+	c, _ := newTestRedisClient(t)
+
+	now := time.Now().UTC()
+	sb1 := newTestSandbox("sb-1", "sess-1", now.Add(1*time.Hour))
+	sb1.Kind = types.AgentRuntimeKind
+
+	sb2 := newTestSandbox("sb-2", "sess-2", now.Add(1*time.Hour))
+	sb2.Kind = types.CodeInterpreterKind
+
+	sb3 := newTestSandbox("sb-3", "sess-3", now.Add(1*time.Hour))
+	sb3.Kind = types.AgentRuntimeKind
+
+	assert.NoError(t, c.StoreSandbox(ctx, sb1))
+	assert.NoError(t, c.StoreSandbox(ctx, sb2))
+	assert.NoError(t, c.StoreSandbox(ctx, sb3))
+
+	// List AgentRuntimeKind
+	list, err := c.ListSandboxesByKind(ctx, types.AgentRuntimeKind)
+	assert.NoError(t, err)
+	assert.Len(t, list, 2)
+
+	ids := map[string]bool{}
+	for _, sb := range list {
+		ids[sb.SessionID] = true
+	}
+	assert.True(t, ids["sess-1"])
+	assert.True(t, ids["sess-3"])
+
+	// List CodeInterpreterKind
+	list2, err := c.ListSandboxesByKind(ctx, types.CodeInterpreterKind)
+	assert.NoError(t, err)
+	assert.Len(t, list2, 1)
+	assert.Equal(t, "sess-2", list2[0].SessionID)
+
+	// List unknown kind
+	list3, err := c.ListSandboxesByKind(ctx, "unknown-kind")
+	assert.NoError(t, err)
+	assert.Len(t, list3, 0)
+}

pkg/store/store_valkey.go

@@ -111,19 +111,26 @@ func (vs *valkeyStore) loadSandboxesBySessionIDs(ctx context.Context, sessionIDs
 		sessionIDKeys = append(sessionIDKeys, vs.sessionKey(sessionID))
 	}
 	// MGet should in same slot
-	stingSliceResults, err := vs.cli.Do(ctx, vs.cli.B().Mget().Key(sessionIDKeys...).Build()).AsStrSlice()
+	valkeyResults, err := vs.cli.Do(ctx, vs.cli.B().Mget().Key(sessionIDKeys...).Build()).ToArray()
 	if err != nil {
 		return nil, fmt.Errorf("loadSandboxesBySessionIDs: Valkey MGet sandboxes failed: %w", err)
 	}
 
-	if len(stingSliceResults) > len(sessionIDKeys) {
-		return nil, fmt.Errorf("unexpected MGet result size: %d, param size: %d", len(stingSliceResults), len(sessionIDKeys))
+	if len(valkeyResults) != len(sessionIDKeys) {
+		return nil, fmt.Errorf("unexpected MGet result size: %d, expected: %d", len(valkeyResults), len(sessionIDKeys))
 	}
 
-	sandboxResults := make([]*types.SandboxInfo, 0, len(stingSliceResults))
-	for i, sandboxObjString := range stingSliceResults {
+	sandboxResults := make([]*types.SandboxInfo, 0, len(valkeyResults))
+	for i, msg := range valkeyResults {
+		if msg.IsNil() {
+			// key does not exist, ignore
+			continue
+		}
+		sandboxObjString, err := msg.ToString()
+		if err != nil {
+			return nil, fmt.Errorf("parse sandbox string failed: %w, index: %v, sessionID: %v", err, i, sessionIDs[i])
+		}
 		if len(sandboxObjString) == 0 {
-			// sandboxObjString is empty while sessionKey not exist, ignore
 			continue
 		}
 		var sandboxRedis types.SandboxInfo
@@ -168,6 +175,57 @@ func (vs *valkeyStore) GetSandboxBySessionID(ctx context.Context, sessionID stri
 	return &sandboxRedis, nil
 }
 
+// ListSandboxesByKind returns all active sandboxes matching the given kind.
+// Uses SCAN to prevent blocking the Valkey instance on large datasets.
+func (vs *valkeyStore) ListSandboxesByKind(ctx context.Context, kind string) ([]*types.SandboxInfo, error) {
+	allSandboxes := make([]*types.SandboxInfo, 0) // Initialize as empty array, not nil
+	cursor := uint64(0)
+	matchPattern := vs.sessionPrefix + "*"
+	seenKeys := make(map[string]bool)
+
+	for {
+		scanRes, err := vs.cli.Do(ctx, vs.cli.B().Scan().Cursor(cursor).Match(matchPattern).Count(100).Build()).AsScanEntry()
+		if err != nil {
+			return nil, fmt.Errorf("ListSandboxesByKind scan failed: %w", err)
+		}
+
+		if len(scanRes.Elements) > 0 {
+			sessionIDs := make([]string, 0, len(scanRes.Elements))
+			for _, key := range scanRes.Elements {
+				if key == vs.expiryIndexKey || key == vs.lastActivityIndexKey {
+					continue
+				}
+				if seenKeys[key] {
+					continue
+				}
+				seenKeys[key] = true
+				sessionIDs = append(sessionIDs, strings.TrimPrefix(key, vs.sessionPrefix))
+			}
+
+			if len(sessionIDs) > 0 {
+				// Batch load the fetched keys
+				sandboxes, err := vs.loadSandboxesBySessionIDs(ctx, sessionIDs)
+				if err != nil {
+					return nil, err
+				}
+
+				// Filter by requested kind
+				for _, sb := range sandboxes {
+					if sb != nil && sb.Kind == kind {
+						allSandboxes = append(allSandboxes, sb)
+					}
+				}
+			}
+		}
+
+		cursor = scanRes.Cursor
+		if cursor == 0 {
+			break // Scan complete
+		}
+	}
+	return allSandboxes, nil
+}
+
 // StoreSandbox store sandbox into storage
 func (vs *valkeyStore) StoreSandbox(ctx context.Context, sandboxStore *types.SandboxInfo) error {
 	if sandboxStore == nil {

pkg/store/store_valkey_test.go

@@ -135,8 +135,8 @@ func newValkeyTestClient(t *testing.T) (*valkeyStore, *miniredis.Miniredis) {
 	rs := &valkeyStore{
 		cli:                  client,
 		sessionPrefix:        "session:",
-		expiryIndexKey:       "sandbox:expiry",
-		lastActivityIndexKey: "sandbox:last_activity",
+		expiryIndexKey:       "session:expiry",
+		lastActivityIndexKey: "session:last_activity",
 	}
 	return rs, mr
 }
@@ -360,3 +360,79 @@ func TestValkeyStore_UpdateSandboxLastActivity(t *testing.T) {
 	assert.Error(t, err)
 	assert.True(t, errors.Is(err, ErrNotFound))
 }
+
+func TestValkeyStore_ListSandboxesByKind(t *testing.T) {
+	ctx := context.Background()
+	c, _ := newValkeyTestClient(t)
+
+	now := time.Now().UTC()
+	sb1 := newTestSandbox("sb-1", "sess-1", now.Add(1*time.Hour))
+	sb1.Kind = types.AgentRuntimeKind
+
+	sb2 := newTestSandbox("sb-2", "sess-2", now.Add(1*time.Hour))
+	sb2.Kind = types.CodeInterpreterKind
+
+	sb3 := newTestSandbox("sb-3", "sess-3", now.Add(1*time.Hour))
+	sb3.Kind = types.AgentRuntimeKind
+
+	assert.NoError(t, c.StoreSandbox(ctx, sb1))
+	assert.NoError(t, c.StoreSandbox(ctx, sb2))
+	assert.NoError(t, c.StoreSandbox(ctx, sb3))
+
+	// List AgentRuntimeKind
+	list, err := c.ListSandboxesByKind(ctx, types.AgentRuntimeKind)
+	assert.NoError(t, err)
+	assert.Len(t, list, 2)
+
+	ids := map[string]bool{}
+	for _, sb := range list {
+		ids[sb.SessionID] = true
+	}
+	assert.True(t, ids["sess-1"])
+	assert.True(t, ids["sess-3"])
+
+	// List CodeInterpreterKind
+	list2, err := c.ListSandboxesByKind(ctx, types.CodeInterpreterKind)
+	assert.NoError(t, err)
+	assert.Len(t, list2, 1)
+	assert.Equal(t, "sess-2", list2[0].SessionID)
+
+	// List unknown kind
+	list3, err := c.ListSandboxesByKind(ctx, "unknown-kind")
+	assert.NoError(t, err)
+	assert.Len(t, list3, 0)
+}
+
+// TestValkeyStore_LoadSandboxesBySessionIDs_OrphanedZSetEntry verifies that
+// loadSandboxesBySessionIDs skips session IDs whose hash key has been evicted
+// from Valkey (orphaned sorted-set entry) instead of aborting the entire batch.
+func TestValkeyStore_LoadSandboxesBySessionIDs_OrphanedZSetEntry(t *testing.T) {
+	ctx := context.Background()
+	c, mr := newValkeyTestClient(t)
+
+	now := time.Now().UTC().Truncate(time.Second)
+
+	sb1 := newTestSandbox("sb-orphan", "sess-orphan", now.Add(-1*time.Hour))
+	sb2 := newTestSandbox("sb-alive", "sess-alive", now.Add(-2*time.Hour))
+
+	if err := c.StoreSandbox(ctx, sb1); err != nil {
+		t.Fatalf("StoreSandbox sb1 error: %v", err)
+	}
+	if err := c.StoreSandbox(ctx, sb2); err != nil {
+		t.Fatalf("StoreSandbox sb2 error: %v", err)
+	}
+
+	// Simulate Valkey evicting the hash key for sb1 while leaving its zset entry.
+	mr.Del(c.sessionKey("sess-orphan"))
+
+	result, err := c.loadSandboxesBySessionIDs(ctx, []string{"sess-orphan", "sess-alive"})
+	if err != nil {
+		t.Fatalf("expected no error with orphaned zset entry, got: %v", err)
+	}
+	if len(result) != 1 {
+		t.Fatalf("expected 1 sandbox (the non-evicted one), got %d", len(result))
+	}
+	if result[0].SandboxID != "sb-alive" {
+		t.Fatalf("expected sb-alive, got %s", result[0].SandboxID)
+	}
+}

pkg/workloadmanager/auth_test.go

@@ -31,7 +31,7 @@ import (
 )
 
 const (
-	testToken = "test-token"
+	testToken          = "test-token"
 	testServiceAccount = "system:serviceaccount:default:test-sa"
 )
 
@@ -85,7 +85,7 @@ func TestAuthMiddleware_InvalidHeaderFormat(t *testing.T) {
 			name:             "no Bearer prefix",
 			header:           "token123",
 			expectedBodyPart: "Invalid authorization header format",
- 		},
+		},
 		{
 			name:             "wrong prefix",
 			header:           "Basic token123",

pkg/workloadmanager/client_cache_test.go

@@ -28,7 +28,7 @@ import (
 )
 
 const (
-	jwtHeader = `{"alg":"HS256","typ":"JWT"}`
+	jwtHeader    = `{"alg":"HS256","typ":"JWT"}`
 	testCacheKey = "default:test-sa"
 )
 

pkg/workloadmanager/garbage_collection_test.go

@@ -46,6 +46,9 @@ func (nopStore) DeleteSandboxBySessionID(_ context.Context, _ string) error  { r
 func (nopStore) ListExpiredSandboxes(_ context.Context, _ time.Time, _ int64) ([]*types.SandboxInfo, error) {
 	return nil, nil
 }
+func (nopStore) ListSandboxesByKind(_ context.Context, _ string) ([]*types.SandboxInfo, error) {
+	return nil, nil
+}
 func (nopStore) ListInactiveSandboxes(_ context.Context, _ time.Time, _ int64) ([]*types.SandboxInfo, error) {
 	return nil, nil
 }