diff --git a/sdk-python/agentcube/agent_runtime.py b/sdk-python/agentcube/agent_runtime.py index d5322178d..35c27cdfc 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,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,9 +88,9 @@ 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") @@ -80,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() @@ -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}") \ 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..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,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() \ No newline at end of file