workloadmanager: remove session when its Sandbox is deleted

cmd/workload-manager/main.go

@@ -37,6 +37,7 @@ import (
 
 	runtimev1alpha1 "github.com/volcano-sh/agentcube/pkg/apis/runtime/v1alpha1"
 	"github.com/volcano-sh/agentcube/pkg/mtls"
+	"github.com/volcano-sh/agentcube/pkg/store"
 	"github.com/volcano-sh/agentcube/pkg/workloadmanager"
 )
 
@@ -100,6 +101,7 @@ func main() {
 	sandboxReconciler := &workloadmanager.SandboxReconciler{
 		Client: mgr.GetClient(),
 		Scheme: mgr.GetScheme(),
+		Store:  store.Storage(),
 	}
 
 	codeInterpreterReconciler := &workloadmanager.CodeInterpreterReconciler{
@@ -188,6 +190,8 @@ func setupControllers(mgr ctrl.Manager, sandboxReconciler *workloadmanager.Sandb
 	// Setup Sandbox controller
 	if err := ctrl.NewControllerManagedBy(mgr).
 		For(&sandboxv1alpha1.Sandbox{}).
+		Watches(&sandboxv1alpha1.Sandbox{}, sandboxReconciler.SessionCleanupHandler()).
+		Watches(&extensionsv1alpha1.SandboxClaim{}, sandboxReconciler.SessionCleanupHandler()).
 		Complete(sandboxReconciler); err != nil {
 		return fmt.Errorf("unable to create sandbox controller: %w", err)
 	}

docs/design/agentcube-proposal.md

@@ -362,6 +362,8 @@ When creation fails, the Sandbox API Server automatically reclaims the underlyin
 
 The Runtime Controller watches `AgentRuntime`, `CodeInterpreter`, `Sandbox` resources. It allows the Sandbox APIServer to obtain the status of a Sandbox as soon as possible. When the SandboxController detects that a Sandbox resource has transitioned to the `Running` state, it immediately sends a notification to the Sandbox APIServer.
 
+The SandboxController also watches `Sandbox` and `SandboxClaim` deletions. When one is deleted by something other than the garbage collector or the session DELETE API (for example `kubectl delete`, or a CodeInterpreter deletion that cascades to its SandboxClaims), it removes the matching session record from the KV storage, so the Router returns `SESSION_NOT_FOUND` instead of proxying to a sandbox that no longer exists.
+
 The CodeInterpreter Controller is responsible for automatically managing SandboxTemplate and SandboxWarmPool resources based on WarmPoolSize configuration to enable fast sandbox allocation for code interpreter sessions:
 
 1. If WarmPoolSize > 0: creates/updates SandboxTemplate (defines pod spec of warm pool) and SandboxWarmPool (maintains pool of pre-warmed Pods).

pkg/workloadmanager/sandbox_controller.go

@@ -18,20 +18,37 @@ package workloadmanager
 
 import (
 	"context"
+	"errors"
 	"sync"
+	"time"
 
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/client-go/util/workqueue"
 	"k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/event"
+	"sigs.k8s.io/controller-runtime/pkg/handler"
+	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 
 	sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
+	extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
+
+	commontypes "github.com/volcano-sh/agentcube/pkg/common/types"
+	"github.com/volcano-sh/agentcube/pkg/store"
 )
 
+// sessionCleanupTimeout bounds the store calls made when a Sandbox or
+// SandboxClaim is deleted.
+const sessionCleanupTimeout = 5 * time.Second
+
 type SandboxReconciler struct {
 	client.Client
 	Scheme *runtime.Scheme
+	// Store is used to remove a session when its Sandbox or SandboxClaim is
+	// deleted. Cleanup is skipped when it is nil.
+	Store store.Store
 
 	// Map of watchers waiting for sandbox to reach Running state
 	watchers map[types.NamespacedName]chan SandboxStatusUpdate
@@ -100,3 +117,61 @@ func (r *SandboxReconciler) UnWatchSandbox(namespace, name string) {
 	defer r.mu.Unlock()
 	delete(r.watchers, key)
 }
+
+// SessionCleanupHandler returns an event handler that removes a session from
+// the store when its Sandbox or SandboxClaim is deleted. GC and the session
+// DELETE API already remove the entry themselves; this covers deletions made
+// by anything else, such as kubectl or a CodeInterpreter deletion cascading
+// to its SandboxClaims. Without it the Router keeps proxying to the old pod
+// instead of returning SESSION_NOT_FOUND.
+func (r *SandboxReconciler) SessionCleanupHandler() handler.EventHandler {
+	return handler.Funcs{
+		DeleteFunc: func(ctx context.Context, e event.DeleteEvent, _ workqueue.TypedRateLimitingInterface[reconcile.Request]) {
+			r.cleanupSession(ctx, e.Object)
+		},
+	}
+}
+
+func (r *SandboxReconciler) cleanupSession(ctx context.Context, obj client.Object) {
+	if r.Store == nil || obj == nil {
+		return
+	}
+
+	var kind string
+	switch obj.(type) {
+	case *sandboxv1alpha1.Sandbox:
+		kind = commontypes.SandboxKind
+	case *extensionsv1alpha1.SandboxClaim:
+		kind = commontypes.SandboxClaimsKind
+	default:
+		return
+	}
+
+	sessionID := obj.GetLabels()[SessionIdLabelKey]
+	if sessionID == "" {
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, sessionCleanupTimeout)
+	defer cancel()
+
+	info, err := r.Store.GetSandboxBySessionID(ctx, sessionID)
+	if err != nil {
+		if !errors.Is(err, store.ErrNotFound) {
+			klog.Warningf("Failed to get session %s for deleted %s %s/%s: %v", sessionID, kind, obj.GetNamespace(), obj.GetName(), err)
+		}
+		return
+	}
+	// Only remove the entry if it still points at the deleted object, so an
+	// object that carries the same label cannot remove a session stored
+	// under a different Sandbox or SandboxClaim.
+	if info == nil || info.Kind != kind || info.SandboxNamespace != obj.GetNamespace() || info.Name != obj.GetName() {
+		return
+	}
+
+	if err := r.Store.DeleteSandboxBySessionID(ctx, sessionID); err != nil {
+		klog.Warningf("Failed to delete session %s for deleted %s %s/%s: %v", sessionID, kind, obj.GetNamespace(), obj.GetName(), err)
+		return
+	}
+	klog.Infof("%s %s/%s was deleted, removed session %s from store", kind, obj.GetNamespace(), obj.GetName(), sessionID)
+}

pkg/workloadmanager/sandbox_controller_test.go

@@ -0,0 +1,167 @@
+/*
+Copyright The Volcano Authors.
+
+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 workloadmanager
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/event"
+
+	sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
+	extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
+
+	"github.com/volcano-sh/agentcube/pkg/common/types"
+	"github.com/volcano-sh/agentcube/pkg/store"
+)
+
+// sessionFakeStore keeps sessions in a map for session cleanup tests.
+type sessionFakeStore struct {
+	nopStore
+	sessions  map[string]*types.SandboxInfo
+	getErr    error
+	deleteErr error
+	deleted   []string
+}
+
+func (f *sessionFakeStore) GetSandboxBySessionID(_ context.Context, sessionID string) (*types.SandboxInfo, error) {
+	if f.getErr != nil {
+		return nil, f.getErr
+	}
+	info, ok := f.sessions[sessionID]
+	if !ok {
+		return nil, store.ErrNotFound
+	}
+	return info, nil
+}
+
+func (f *sessionFakeStore) DeleteSandboxBySessionID(_ context.Context, sessionID string) error {
+	if f.deleteErr != nil {
+		return f.deleteErr
+	}
+	delete(f.sessions, sessionID)
+	f.deleted = append(f.deleted, sessionID)
+	return nil
+}
+
+func objectMeta(namespace, name, sessionID string) metav1.ObjectMeta {
+	meta := metav1.ObjectMeta{Namespace: namespace, Name: name}
+	if sessionID != "" {
+		meta.Labels = map[string]string{SessionIdLabelKey: sessionID}
+	}
+	return meta
+}
+
+func TestSessionCleanupHandler_Delete(t *testing.T) {
+	directSession := &types.SandboxInfo{
+		Kind:             types.SandboxKind,
+		SessionID:        "session-direct",
+		SandboxNamespace: "default",
+		Name:             "sandbox-a",
+	}
+	claimSession := &types.SandboxInfo{
+		Kind:             types.SandboxClaimsKind,
+		SessionID:        "session-claim",
+		SandboxNamespace: "default",
+		Name:             "claim-a",
+	}
+
+	tests := []struct {
+		name        string
+		obj         client.Object
+		getErr      error
+		deleteErr   error
+		wantDeleted []string
+	}{
+		{
+			name:        "deleted Sandbox removes its session",
+			obj:         &sandboxv1alpha1.Sandbox{ObjectMeta: objectMeta("default", "sandbox-a", "session-direct")},
+			wantDeleted: []string{"session-direct"},
+		},
+		{
+			name:        "deleted SandboxClaim removes its session",
+			obj:         &extensionsv1alpha1.SandboxClaim{ObjectMeta: objectMeta("default", "claim-a", "session-claim")},
+			wantDeleted: []string{"session-claim"},
+		},
+		{
+			name: "name mismatch keeps the session",
+			obj:  &sandboxv1alpha1.Sandbox{ObjectMeta: objectMeta("default", "sandbox-other", "session-direct")},
+		},
+		{
+			name: "kind mismatch keeps the session",
+			obj:  &sandboxv1alpha1.Sandbox{ObjectMeta: objectMeta("default", "claim-a", "session-claim")},
+		},
+		{
+			name: "namespace mismatch keeps the session",
+			obj:  &sandboxv1alpha1.Sandbox{ObjectMeta: objectMeta("other", "sandbox-a", "session-direct")},
+		},
+		{
+			name: "object without session label is ignored",
+			obj:  &sandboxv1alpha1.Sandbox{ObjectMeta: objectMeta("default", "sandbox-a", "")},
+		},
+		{
+			name: "session already gone from store",
+			obj:  &sandboxv1alpha1.Sandbox{ObjectMeta: objectMeta("default", "sandbox-b", "session-missing")},
+		},
+		{
+			name:   "store get error keeps the session",
+			obj:    &sandboxv1alpha1.Sandbox{ObjectMeta: objectMeta("default", "sandbox-a", "session-direct")},
+			getErr: errors.New("connection refused"),
+		},
+		{
+			name:      "store delete error is not fatal",
+			obj:       &sandboxv1alpha1.Sandbox{ObjectMeta: objectMeta("default", "sandbox-a", "session-direct")},
+			deleteErr: errors.New("connection refused"),
+		},
+		{
+			name: "other object types are ignored",
+			obj:  &corev1.Pod{ObjectMeta: objectMeta("default", "sandbox-a", "session-direct")},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			fake := &sessionFakeStore{
+				sessions: map[string]*types.SandboxInfo{
+					directSession.SessionID: directSession,
+					claimSession.SessionID:  claimSession,
+				},
+				getErr:    tt.getErr,
+				deleteErr: tt.deleteErr,
+			}
+			r := &SandboxReconciler{Store: fake}
+
+			r.SessionCleanupHandler().Delete(context.Background(), event.DeleteEvent{Object: tt.obj}, nil)
+
+			assert.Equal(t, tt.wantDeleted, fake.deleted)
+		})
+	}
+}
+
+func TestSessionCleanupHandler_NilStore(t *testing.T) {
+	r := &SandboxReconciler{}
+	obj := &sandboxv1alpha1.Sandbox{ObjectMeta: objectMeta("default", "sandbox-a", "session-direct")}
+
+	assert.NotPanics(t, func() {
+		r.SessionCleanupHandler().Delete(context.Background(), event.DeleteEvent{Object: obj}, nil)
+	})
+}