Fix/agent runtime delete session

sdk-python/agentcube/agent_runtime.py

@@ -31,6 +31,8 @@ def __init__(
         session_id: Optional[str] = None,
         timeout: int = 120,
         connect_timeout: float = 5.0,
+        workload_manager_url: Optional[str] = None,
+        auth_token: Optional[str] = None,
     ):
         self.agent_name = agent_name
         self.namespace = namespace
@@ -48,7 +50,23 @@ def __init__(
             )
         self.router_url = router_url
 
+        # Initialize Control Plane client for session deletion
+        if workload_manager_url or os.getenv("AGENTCUBE_WORKLOAD_MANAGER_URL"):
+            self.workload_manager_url = (
+                workload_manager_url
+                or os.getenv("AGENTCUBE_WORKLOAD_MANAGER_URL")
+            )
+            from agentcube.clients.control_plane import ControlPlaneClient
+            self._control_plane = ControlPlaneClient(
+                workload_manager_url=self.workload_manager_url,
+                auth_token=auth_token or os.getenv("AGENTCUBE_AUTH_TOKEN"),
+            )
+        else:
+            self._control_plane = None
+
         self.session_id: Optional[str] = session_id
+        self._owned_session = session_id is None
+
         self.dp_client = AgentRuntimeDataPlaneClient(
             router_url=self.router_url,
             namespace=self.namespace,
@@ -70,16 +88,17 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
+        self.stop()
 
-    def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None) -> Any:
+    def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None, path: str = "") -> Any:
         if not self.session_id:
             raise ValueError("AgentRuntime session_id is not initialized")
 
         resp = self.dp_client.invoke(
             session_id=self.session_id,
             payload=payload,
             timeout=timeout,
+            path=path,
         )
         resp.raise_for_status()
 
@@ -91,3 +110,20 @@ def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None) -> An
     def close(self) -> None:
         if self.dp_client:
             self.dp_client.close()
+        if self._control_plane:
+            self._control_plane.close()
+
+    def stop(self) -> None:
+        """Close local connection and delete server-side session if owned."""
+        try:
+            self.close()
+        except Exception as e:
+            self.logger.warning(f"Error closing local connection: {e}")
+
+        if self._owned_session and self.session_id and self._control_plane:
+            try:
+                self._control_plane.delete_agent_runtime_session(self.session_id)
+                self.logger.info(f"Deleted AgentRuntime session: {self.session_id}")
+                self.session_id = None
+            except Exception as e:
+                self.logger.warning(f"Error deleting AgentRuntime session: {e}")

sdk-python/agentcube/clients/control_plane.py

@@ -12,14 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import os
+from typing import TYPE_CHECKING, Dict, Any, Optional
+
 import requests
-from typing import Dict, Any, Optional
 
 from agentcube.utils.log import get_logger
 from agentcube.utils.utils import read_token_from_file
 from agentcube.utils.http import create_session
 
+if TYPE_CHECKING:
+    from agentcube.auth import AuthProvider
+
 class ControlPlaneClient:
     """Client for AgentCube Control Plane (WorkloadManager).
     Handles creation and deletion of Code Interpreter sessions.
@@ -33,6 +39,7 @@ def __init__(
         connect_timeout: float = 5.0,
         pool_connections: int = 10,
         pool_maxsize: int = 10,
+        auth: Optional["AuthProvider"] = None,
     ):
         """Initialize the Control Plane client.
 
@@ -43,6 +50,7 @@ def __init__(
             connect_timeout: Connection timeout in seconds (default: 5).
             pool_connections: Number of connection pools to cache (default: 10).
             pool_maxsize: Maximum connections per pool (default: 10).
+            auth: Optional AuthProvider instance (takes priority over auth_token).
         """
         # Prioritize argument -> env var
         self.base_url = workload_manager_url or os.getenv("WORKLOAD_MANAGER_URL")
@@ -52,9 +60,21 @@ def __init__(
                 "or 'WORKLOAD_MANAGER_URL' environment variable."
             )
 
-        # Prioritize argument -> k8s service account token file
-        token_path = "/var/run/secrets/kubernetes.io/serviceaccount/token"
-        token = auth_token or read_token_from_file(token_path)
+        # Resolve auth: auth param > auth_token > k8s SA token file
+        if auth:
+            self._auth = auth
+        elif auth_token:
+            from agentcube.auth import TokenAuth
+            self._auth = TokenAuth(auth_token)
+        else:
+            token_path = "/var/run/secrets/kubernetes.io/serviceaccount/token"
+            token = read_token_from_file(token_path)
+            if token:
+                from agentcube.auth import TokenAuth
+                self._auth = TokenAuth(token)
+            else:
+                self._auth = None
+
         self.timeout = timeout
         self.connect_timeout = connect_timeout
 
@@ -70,8 +90,13 @@ def __init__(
         self.session.headers.update({
             "Content-Type": "application/json",
         })
-        if token:
-            self.session.headers["Authorization"] = f"Bearer {token}"
+
+    def _apply_auth(self, request_kwargs: dict) -> None:
+        """Add Authorization header from auth provider if available."""
+        if not self._auth:
+            return
+        headers = request_kwargs.setdefault("headers", {})
+        headers["Authorization"] = f"Bearer {self._auth.get_token()}"
 
     def create_session(
         self,
@@ -80,17 +105,6 @@ def create_session(
         metadata: Optional[Dict[str, Any]] = None,
         ttl: int = 3600,
     ) -> str:
-        """Create a new Code Interpreter session.
-
-        Args:
-            name: Name of the CodeInterpreter template (CRD name).
-            namespace: Kubernetes namespace.
-            metadata: Optional metadata.
-            ttl: Time to live (seconds).
-
-        Returns:
-            session_id (str): The ID of the created session.
-        """
         payload = {
             "name": name,
             "namespace": namespace,
@@ -102,11 +116,9 @@ def create_session(
         self.logger.debug(f"Creating session at {url} with payload: {payload}")
 
         try:
-            response = self.session.post(
-                url,
-                json=payload,
-                timeout=(self.connect_timeout, self.timeout)
-            )
+            kwargs = {"json": payload, "timeout": (self.connect_timeout, self.timeout)}
+            self._apply_auth(kwargs)
+            response = self.session.post(url, **kwargs)
             response.raise_for_status()
 
             data = response.json()
@@ -134,18 +146,53 @@ def delete_session(self, session_id: str) -> bool:
         self.logger.debug(f"Deleting session {session_id} at {url}")
 
         try:
-            response = self.session.delete(
-                url,
-                timeout=(self.connect_timeout, self.timeout)
-            )
+            kwargs: Dict[str, Any] = {"timeout": (self.connect_timeout, self.timeout)}
+            self._apply_auth(kwargs)
+            response = self.session.delete(url, **kwargs)
             if response.status_code == 404:
-                return True # Already gone
+                return True
             response.raise_for_status()
             return True
         except requests.exceptions.RequestException as e:
             self.logger.error(f"Failed to delete session {session_id}: {e}")
             return False
 
+    def delete_agent_runtime_session(self, session_id: str) -> Dict[str, Any]:
+        """
+        Delete an agent runtime session from the server.
+
+        Args:
+            session_id: The ID of the session to delete
+
+        Returns:
+            The response from the server (empty dict for 204)
+
+        Raises:
+            HTTPError: If the request fails
+        """
+        url = f"{self.base_url}/v1/agent-runtime/sessions/{session_id}"
+        self.logger.debug(f"Deleting agent runtime session {session_id} at {url}")
+
+        try:
+            kwargs: Dict[str, Any] = {"timeout": (self.connect_timeout, self.timeout)}
+            self._apply_auth(kwargs)
+            response = self.session.delete(url, **kwargs)
+            if response.status_code == 404:
+                return {}  # Already gone
+            response.raise_for_status()
+            # Handle empty response for 200 OK
+            if response.status_code == 200 and response.text.strip():
+                try:
+                    return response.json()
+                except ValueError:
+                    return {}
+            return {}
+        except requests.exceptions.RequestException as e:
+            self.logger.error(f"Failed to delete agent runtime session {session_id}: {e}")
+            if e.response is not None:
+                self.logger.error(f"Server response: {e.response.text}")
+            raise
+
     def close(self):
         """Close the underlying session and release connection pool resources."""
-        self.session.close()
+        self.session.close()