From eebdda25a4477ab4f27752e23b4dd5edba885f24 Mon Sep 17 00:00:00 2001 From: vivek41-glitch Date: Wed, 8 Jul 2026 00:30:22 +0530 Subject: [PATCH 1/4] feat: add stop() method for AgentRuntimeClient session deletion Signed-off-by: vivek41-glitch --- sdk-python/agentcube/agent_runtime.py | 35 +++++- sdk-python/agentcube/clients/control_plane.py | 113 ++++-------------- 2 files changed, 60 insertions(+), 88 deletions(-) diff --git a/sdk-python/agentcube/agent_runtime.py b/sdk-python/agentcube/agent_runtime.py index d5322178d..c67a61357 100644 --- a/sdk-python/agentcube/agent_runtime.py +++ b/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,24 @@ 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") + ) + self.auth_token = auth_token or os.getenv("AGENTCUBE_AUTH_TOKEN") + from agentcube.clients.control_plane import ControlPlaneClient + self._control_plane = ControlPlaneClient( + workload_manager_url=self.workload_manager_url, + auth_token=self.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,7 +89,7 @@ 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: if not self.session_id: @@ -91,3 +110,17 @@ def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None) -> An def close(self) -> None: if self.dp_client: self.dp_client.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}") + except Exception as e: + self.logger.warning(f"Error deleting AgentRuntime session: {e}") \ No newline at end of file diff --git a/sdk-python/agentcube/clients/control_plane.py b/sdk-python/agentcube/clients/control_plane.py index 48bee355c..b9aadea09 100644 --- a/sdk-python/agentcube/clients/control_plane.py +++ b/sdk-python/agentcube/clients/control_plane.py @@ -34,44 +34,7 @@ def __init__( pool_connections: int = 10, pool_maxsize: int = 10, ): - """Initialize the Control Plane client. - - Args: - workload_manager_url: URL of the WorkloadManager service. - auth_token: Kubernetes Service Account Token for authentication. - timeout: Default request timeout in seconds (default: 120). - 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). - """ - # Prioritize argument -> env var - self.base_url = workload_manager_url or os.getenv("WORKLOAD_MANAGER_URL") - if not self.base_url: - raise ValueError( - "Workload Manager URL must be provided via 'workload_manager_url' argument " - "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) - self.timeout = timeout - self.connect_timeout = connect_timeout - - self.logger = get_logger(f"{__name__}.ControlPlaneClient") - - # Create session with connection pooling using shared utility - self.session = create_session( - pool_connections=pool_connections, - pool_maxsize=pool_maxsize, - ) - - # Set default headers - self.session.headers.update({ - "Content-Type": "application/json", - }) - if token: - self.session.headers["Authorization"] = f"Bearer {token}" + # ... existing __init__ code ... def create_session( self, @@ -80,58 +43,26 @@ def create_session( metadata: Optional[Dict[str, Any]] = None, ttl: int = 3600, ) -> str: - """Create a new Code Interpreter session. + # ... existing create_session code ... - Args: - name: Name of the CodeInterpreter template (CRD name). - namespace: Kubernetes namespace. - metadata: Optional metadata. - ttl: Time to live (seconds). + def delete_session(self, session_id: str) -> bool: + # ... existing delete_session code ... - Returns: - session_id (str): The ID of the created session. + def delete_agent_runtime_session(self, session_id: str) -> Dict[str, Any]: """ - payload = { - "name": name, - "namespace": namespace, - "ttl": ttl, - "metadata": metadata or {} - } - - url = f"{self.base_url}/v1/code-interpreter" - 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) - ) - response.raise_for_status() - - data = response.json() - if "sessionId" not in data or not data["sessionId"]: - self.logger.error("Response JSON missing 'sessionId' in create_session response.") - self.logger.debug(f"Full response data: {data}") - raise ValueError("Failed to create session: 'sessionId' missing from response") - return data["sessionId"] - except requests.exceptions.RequestException as e: - self.logger.error(f"Failed to create session: {e}") - if e.response is not None: - self.logger.error(f"Server response: {e.response.text}") - raise - - def delete_session(self, session_id: str) -> bool: - """Delete a Code Interpreter session. + Delete an agent runtime session from the server. Args: - session_id: The session ID to delete. + session_id: The ID of the session to delete Returns: - True if deleted successfully (or didn't exist), False on failure. + The response from the server (empty dict for 204) + + Raises: + HTTPError: If the request fails """ - url = f"{self.base_url}/v1/code-interpreter/sessions/{session_id}" - self.logger.debug(f"Deleting session {session_id} at {url}") + url = f"{self.base_url}/v1/agent-runtime/sessions/{session_id}" + self.logger.debug(f"Deleting agent runtime session {session_id} at {url}") try: response = self.session.delete( @@ -139,13 +70,21 @@ def delete_session(self, session_id: str) -> bool: timeout=(self.connect_timeout, self.timeout) ) if response.status_code == 404: - return True # Already gone + return {} # Already gone response.raise_for_status() - return True + # 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 session {session_id}: {e}") - return False + 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() \ No newline at end of file From 320ce3a9329b334f339990011fb0c927e64f3f9a Mon Sep 17 00:00:00 2001 From: vivek41-glitch Date: Wed, 8 Jul 2026 01:25:09 +0530 Subject: [PATCH 2/4] Merge branch main into fix/agent-runtime-delete-session --- sdk-python/agentcube/agent_runtime.py | 15 ++- sdk-python/agentcube/clients/control_plane.py | 124 ++++++++++++++++-- 2 files changed, 128 insertions(+), 11 deletions(-) diff --git a/sdk-python/agentcube/agent_runtime.py b/sdk-python/agentcube/agent_runtime.py index c67a61357..7068ef9de 100644 --- a/sdk-python/agentcube/agent_runtime.py +++ b/sdk-python/agentcube/agent_runtime.py @@ -17,6 +17,7 @@ from typing import Any, Dict, Optional from requests.exceptions import JSONDecodeError +from agentcube.auth import AuthProvider from agentcube.clients.agent_runtime_data_plane import AgentRuntimeDataPlaneClient from agentcube.utils.log import get_logger @@ -33,6 +34,7 @@ def __init__( connect_timeout: float = 5.0, workload_manager_url: Optional[str] = None, auth_token: Optional[str] = None, + auth: Optional[AuthProvider] = None, ): self.agent_name = agent_name self.namespace = namespace @@ -42,6 +44,11 @@ def __init__( level = logging.DEBUG if verbose else logging.INFO self.logger = get_logger(__name__, level=level) + self._auth = auth + if not self._auth and auth_token: + from agentcube.auth import TokenAuth + self._auth = TokenAuth(auth_token) + router_url = router_url or os.getenv("ROUTER_URL") if not router_url: raise ValueError( @@ -56,11 +63,11 @@ def __init__( workload_manager_url or os.getenv("AGENTCUBE_WORKLOAD_MANAGER_URL") ) - self.auth_token = auth_token or os.getenv("AGENTCUBE_AUTH_TOKEN") from agentcube.clients.control_plane import ControlPlaneClient self._control_plane = ControlPlaneClient( workload_manager_url=self.workload_manager_url, - auth_token=self.auth_token + auth_token=auth_token or os.getenv("AGENTCUBE_AUTH_TOKEN"), + auth=self._auth, ) else: self._control_plane = None @@ -74,6 +81,7 @@ def __init__( agent_name=self.agent_name, timeout=self.timeout, connect_timeout=self.connect_timeout, + auth=self._auth, ) if verbose: self.dp_client.logger.setLevel(logging.DEBUG) @@ -91,7 +99,7 @@ def __enter__(self): def __exit__(self, exc_type, exc_val, exc_tb): 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") @@ -99,6 +107,7 @@ def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None) -> An session_id=self.session_id, payload=payload, timeout=timeout, + path=path, ) resp.raise_for_status() diff --git a/sdk-python/agentcube/clients/control_plane.py b/sdk-python/agentcube/clients/control_plane.py index b9aadea09..bd9514eb0 100644 --- a/sdk-python/agentcube/clients/control_plane.py +++ b/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,8 +39,64 @@ def __init__( connect_timeout: float = 5.0, pool_connections: int = 10, pool_maxsize: int = 10, + auth: Optional["AuthProvider"] = None, ): - # ... existing __init__ code ... + """Initialize the Control Plane client. + + Args: + workload_manager_url: URL of the WorkloadManager service. + auth_token: Kubernetes Service Account Token for authentication. + timeout: Default request timeout in seconds (default: 120). + 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") + if not self.base_url: + raise ValueError( + "Workload Manager URL must be provided via 'workload_manager_url' argument " + "or 'WORKLOAD_MANAGER_URL' environment variable." + ) + + # 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 + + self.logger = get_logger(f"{__name__}.ControlPlaneClient") + + # Create session with connection pooling using shared utility + self.session = create_session( + pool_connections=pool_connections, + pool_maxsize=pool_maxsize, + ) + + # Set default headers + self.session.headers.update({ + "Content-Type": "application/json", + }) + + 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, @@ -43,10 +105,57 @@ def create_session( metadata: Optional[Dict[str, Any]] = None, ttl: int = 3600, ) -> str: - # ... existing create_session code ... + payload = { + "name": name, + "namespace": namespace, + "ttl": ttl, + "metadata": metadata or {} + } + + url = f"{self.base_url}/v1/code-interpreter" + self.logger.debug(f"Creating session at {url} with payload: {payload}") + + try: + 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() + if "sessionId" not in data or not data["sessionId"]: + self.logger.error("Response JSON missing 'sessionId' in create_session response.") + self.logger.debug(f"Full response data: {data}") + raise ValueError("Failed to create session: 'sessionId' missing from response") + return data["sessionId"] + except requests.exceptions.RequestException as e: + self.logger.error(f"Failed to create session: {e}") + if e.response is not None: + self.logger.error(f"Server response: {e.response.text}") + raise def delete_session(self, session_id: str) -> bool: - # ... existing delete_session code ... + """Delete a Code Interpreter session. + + Args: + session_id: The session ID to delete. + + Returns: + True if deleted successfully (or didn't exist), False on failure. + """ + url = f"{self.base_url}/v1/code-interpreter/sessions/{session_id}" + self.logger.debug(f"Deleting 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 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]: """ @@ -65,10 +174,9 @@ def delete_agent_runtime_session(self, session_id: str) -> Dict[str, Any]: self.logger.debug(f"Deleting agent runtime 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 {} # Already gone response.raise_for_status() From 6d37331c61971b06b254e06b4f024c43775f78d1 Mon Sep 17 00:00:00 2001 From: vivek41-glitch Date: Wed, 8 Jul 2026 01:34:31 +0530 Subject: [PATCH 3/4] fix: simplify auth handling and add stop() method Signed-off-by: vivek41-glitch --- sdk-python/agentcube/agent_runtime.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/sdk-python/agentcube/agent_runtime.py b/sdk-python/agentcube/agent_runtime.py index c67a61357..35c27cdfc 100644 --- a/sdk-python/agentcube/agent_runtime.py +++ b/sdk-python/agentcube/agent_runtime.py @@ -56,11 +56,10 @@ def __init__( workload_manager_url or os.getenv("AGENTCUBE_WORKLOAD_MANAGER_URL") ) - self.auth_token = auth_token or os.getenv("AGENTCUBE_AUTH_TOKEN") from agentcube.clients.control_plane import ControlPlaneClient self._control_plane = ControlPlaneClient( workload_manager_url=self.workload_manager_url, - auth_token=self.auth_token + auth_token=auth_token or os.getenv("AGENTCUBE_AUTH_TOKEN"), ) else: self._control_plane = None @@ -91,7 +90,7 @@ def __enter__(self): def __exit__(self, exc_type, exc_val, exc_tb): 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") @@ -99,6 +98,7 @@ def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None) -> An session_id=self.session_id, payload=payload, timeout=timeout, + path=path, ) resp.raise_for_status() @@ -110,6 +110,8 @@ 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.""" @@ -117,10 +119,11 @@ def stop(self) -> None: 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}") \ No newline at end of file From becb8312b5e6f70a67e9b6512c95e5cfa1a114c0 Mon Sep 17 00:00:00 2001 From: vivek41-glitch Date: Wed, 8 Jul 2026 01:37:17 +0530 Subject: [PATCH 4/4] fix: remove auth module dependency, add stop() method Signed-off-by: vivek41-glitch --- sdk-python/agentcube/agent_runtime.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/sdk-python/agentcube/agent_runtime.py b/sdk-python/agentcube/agent_runtime.py index 520d022c6..35c27cdfc 100644 --- a/sdk-python/agentcube/agent_runtime.py +++ b/sdk-python/agentcube/agent_runtime.py @@ -17,7 +17,6 @@ from typing import Any, Dict, Optional from requests.exceptions import JSONDecodeError -from agentcube.auth import AuthProvider from agentcube.clients.agent_runtime_data_plane import AgentRuntimeDataPlaneClient from agentcube.utils.log import get_logger @@ -34,7 +33,6 @@ def __init__( connect_timeout: float = 5.0, workload_manager_url: Optional[str] = None, auth_token: Optional[str] = None, - auth: Optional[AuthProvider] = None, ): self.agent_name = agent_name self.namespace = namespace @@ -44,11 +42,6 @@ def __init__( level = logging.DEBUG if verbose else logging.INFO self.logger = get_logger(__name__, level=level) - self._auth = auth - if not self._auth and auth_token: - from agentcube.auth import TokenAuth - self._auth = TokenAuth(auth_token) - router_url = router_url or os.getenv("ROUTER_URL") if not router_url: raise ValueError( @@ -66,7 +59,7 @@ def __init__( from agentcube.clients.control_plane import ControlPlaneClient self._control_plane = ControlPlaneClient( workload_manager_url=self.workload_manager_url, - auth_token=self.auth_token + auth_token=auth_token or os.getenv("AGENTCUBE_AUTH_TOKEN"), ) else: self._control_plane = None @@ -80,7 +73,6 @@ def __init__( agent_name=self.agent_name, timeout=self.timeout, connect_timeout=self.connect_timeout, - auth=self._auth, ) if verbose: self.dp_client.logger.setLevel(logging.DEBUG)