Fixes bug in sandbox_client

Author: igooch

.gitignore

@@ -9,6 +9,6 @@ bin/
 dist/
 release_assets/
 
-# Temporary folders
+# Ignore temporary folders
 tmp/
 temp/

clients/python/agentic-sandbox-client/agentic_sandbox/sandbox_client.py

@@ -39,8 +39,7 @@
 SANDBOX_PLURAL_NAME = "sandboxes"
 
 logging.basicConfig(level=logging.INFO,
-                    format='%(asctime)s - %(levelname)s - %(message)s',
-                    stream=sys.stdout)
+                    format='%(asctime)s - %(levelname)s - %(message)s', stream=sys.stdout)
 
 
 @dataclass
@@ -59,11 +58,8 @@ class SandboxClient:
     def __init__(
         self,
         template_name: str,
-        namespace: str = "default",  # Where Sandbox lives
-        gateway_name: str | None = None,  # Name of the Gateway
-        gateway_namespace: str = "default",  # Where Gateway lives
-        api_url: str | None = None,  # Allow custom URL (DNS or Localhost)
-        server_port: int = 8888,     # The port the runtime inside the sandbox listens on
+        namespace: str = "default",
+        server_port: int = 8888,
         sandbox_ready_timeout: int = 180,
         gateway_ready_timeout: int = 180,
         port_forward_ready_timeout: int = 30,
@@ -82,6 +78,9 @@ def __init__(
 
         self.claim_name: str | None = None
         self.sandbox_name: str | None = None
+        self.pod_name: str | None = None
+        self.base_url = f"http://127.0.0.1:{self.server_port}"
+        self.port_forward_process: subprocess.Popen | None = None
 
         try:
             config.load_incluster_config()
@@ -129,11 +128,40 @@ def _create_claim(self):
         )
 
     def _wait_for_sandbox_ready(self):
-        """Waits for the Sandbox custom resource to have a 'Ready' status."""
+        """
+        Waits for the SandboxClaim to be populated with the Sandbox name, and then
+        waits for the Sandbox custom resource to have a 'Ready' status condition.
+        This indicates that the underlying pod is running and has passed its checks.
+        """
         if not self.claim_name:
             raise RuntimeError(
-                "Cannot wait for sandbox; a sandboxclaim has not been created.")
+                "Cannot wait for sandbox, claim has not been created.")
+
+        # Watch for the SandboxClaim to be updated with the sandbox name
+        w_claim = watch.Watch()
+        for event in w_claim.stream(
+                self.custom_objects_api.list_namespaced_custom_object,
+                namespace=self.namespace,
+                group=CLAIM_API_GROUP,
+                version=CLAIM_API_VERSION,
+                plural=CLAIM_PLURAL_NAME,
+                field_selector=f"metadata.name={self.claim_name}",
+                timeout_seconds=self.sandbox_ready_timeout
+        ):
+            if event["type"] in ["ADDED", "MODIFIED"]:
+                claim_obj = event["object"]
+                status = claim_obj.get("status", {})
+                sandbox_status = status.get("sandbox", {})
+                if sandbox_status and sandbox_status.get("Name"):
+                    self.sandbox_name = sandbox_status.get("Name")
+                    w_claim.stop()
+                    break
+        else:
+            self.__exit__(None, None, None)  # Attempt cleanup
+            raise TimeoutError(
+                f"SandboxClaim did not become ready within {self.sandbox_ready_timeout} seconds.")
 
+        # Watch for the Sandbox to become ready
         w = watch.Watch()
         logging.info("Watching for Sandbox to become ready...")
         for event in w.stream(
@@ -142,56 +170,54 @@ def _wait_for_sandbox_ready(self):
             group=SANDBOX_API_GROUP,
             version=SANDBOX_API_VERSION,
             plural=SANDBOX_PLURAL_NAME,
-            field_selector=f"metadata.name={self.claim_name}",
+            field_selector=f"metadata.name={self.sandbox_name}",
             timeout_seconds=self.sandbox_ready_timeout
         ):
-            sandbox_object = event['object']
-            status = sandbox_object.get('status', {})
-            conditions = status.get('conditions', [])
-            is_ready = False
-            for cond in conditions:
-                if cond.get('type') == 'Ready' and cond.get('status') == 'True':
-                    is_ready = True
+            if event["type"] in ["ADDED", "MODIFIED"]:
+                sandbox_object = event['object']
+                status = sandbox_object.get('status', {})
+                conditions = status.get('conditions', [])
+                is_ready = False
+                for cond in conditions:
+                    if cond.get('type') == 'Ready' and cond.get('status') == 'True':
+                        is_ready = True
+                        break
+
+                if is_ready:
+                    annotations = sandbox_object.get(
+                        'metadata', {}).get('annotations', {})
+                    pod_name_annotation = "agents.x-k8s.io/pod-name"
+                    if pod_name_annotation in annotations:
+                        self.pod_name = annotations[pod_name_annotation]
+                        logging.info(
+                            f"Found pod name from annotation: {self.pod_name}")
+                    else:
+                        self.pod_name = self.sandbox_name
+                    w.stop()
+                    logging.info(f"Sandbox {self.sandbox_name} is ready.")
                     break
 
-            if is_ready:
-                self.sandbox_name = sandbox_object['metadata']['name']
-                w.stop()
-                logging.info(f"Sandbox {self.sandbox_name} is ready.")
-                break
-
-        if not self.sandbox_name:
+        if not self.pod_name:
             self.__exit__(None, None, None)
             raise TimeoutError(
-                f"Sandbox did not become ready within {self.sandbox_ready_timeout} seconds.")
-
-    def _get_free_port(self):
-        """Finds a free port on localhost."""
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            s.bind(('', 0))
-            return s.getsockname()[1]
+                f"Sandbox did not become ready or pod name could not be determined within {self.sandbox_ready_timeout} seconds.")
 
     def _start_and_wait_for_port_forward(self):
         """
         Starts 'kubectl port-forward' to the Router Service.
         This allows 'Dev Mode' without needing a public Gateway IP.
         """
-        local_port = self._get_free_port()
-
-        # Assumes the router service name from sandbox_router.yaml
-        router_svc = "svc/sandbox-router-svc"
-
+        if not self.pod_name:
+            raise RuntimeError(
+                "Cannot start port-forwarding, sandbox pod name is not known.")
         logging.info(
-            f"Starting Dev Mode tunnel: localhost:{local_port} -> {router_svc}:8080...")
-
+            f"Starting port-forwarding for sandbox {self.sandbox_name} with sandbox pod {self.pod_name}...")
         self.port_forward_process = subprocess.Popen(
             [
                 "kubectl", "port-forward",
-                router_svc,
-                # Tunnel to Router (8080), not Sandbox (8888)
-                f"{local_port}:8080",
-                # The router lives in the Gateway/Default NS
-                "-n", self.gateway_namespace
+                f"pod/{self.pod_name}",
+                f"{self.server_port}:{self.server_port}",
+                "-n", self.namespace
             ],
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE
@@ -206,54 +232,16 @@ def _start_and_wait_for_port_forward(self):
                     f"Tunnel crashed: {stderr.decode(errors='ignore')}")
 
             try:
-                # Connect to localhost
-                with socket.create_connection(("127.0.0.1", local_port), timeout=0.1):
-                    self.base_url = f"http://127.0.0.1:{local_port}"
+                with socket.create_connection(("127.0.0.1", self.server_port), timeout=0.1):
                     logging.info(
-                        f"Dev Mode ready. Tunneled to Router at {self.base_url}")
-                    # No need for huge sleeps; the Router service is stable.
-                    time.sleep(0.5)
+                        f"Port-forwarding is ready on port {self.server_port}.")
                     return
             except (socket.timeout, ConnectionRefusedError):
                 time.sleep(0.5)
 
         self.__exit__(None, None, None)
-        raise TimeoutError("Failed to establish tunnel to Router Service.")
-
-    def _wait_for_gateway_ip(self):
-        """Waits for the Gateway to be assigned an external IP."""
-        # Check if we already have a manually provided URL
-        if self.base_url:
-            logging.info(f"Using configured API URL: {self.base_url}")
-            return
-
-        logging.info(
-            f"Waiting for Gateway '{self.gateway_name}' in namespace '{self.gateway_namespace}'...")
-
-        w = watch.Watch()
-        for event in w.stream(
-            func=self.custom_objects_api.list_namespaced_custom_object,
-            namespace=self.gateway_namespace, group=GATEWAY_API_GROUP,
-            version=GATEWAY_API_VERSION, plural=GATEWAY_PLURAL,
-            field_selector=f"metadata.name={self.gateway_name}",
-            timeout_seconds=self.gateway_ready_timeout,
-        ):
-            if event["type"] in ["ADDED", "MODIFIED"]:
-                gateway_object = event['object']
-                status = gateway_object.get('status', {})
-                addresses = status.get('addresses', [])
-                if addresses:
-                    ip_address = addresses[0].get('value')
-                    if ip_address:
-                        self.base_url = f"http://{ip_address}"
-                        logging.info(
-                            f"Gateway is ready. Base URL set to: {self.base_url}")
-                        w.stop()
-                        return
-
-        if not self.base_url:
-            raise TimeoutError(
-                f"Gateway '{self.gateway_name}' in namespace '{self.gateway_namespace}' did not get an IP within {self.gateway_ready_timeout} seconds.")
+        raise TimeoutError(
+            f"Port-forwarding did not become ready within {self.port_forward_ready_timeout} seconds.")
 
     def __enter__(self) -> 'SandboxClient':
         self._create_claim()
@@ -320,22 +308,15 @@ def _request(self, method: str, endpoint: str, **kwargs) -> requests.Response:
         headers["X-Sandbox-Port"] = str(self.server_port)
         kwargs["headers"] = headers
 
+        url = f"{self.base_url}/{endpoint}"
         try:
             response = self.session.request(method, url, **kwargs)
             response.raise_for_status()
             return response
         except requests.exceptions.RequestException as e:
-            # Check if port-forward died DURING request
-            if self.port_forward_process and self.port_forward_process.poll() is not None:
-                stdout, stderr = self.port_forward_process.communicate()
-                raise RuntimeError(
-                    f"Kubectl Port-Forward crashed DURING request!\n"
-                    f"Stderr: {stderr.decode(errors='ignore')}"
-                ) from e
-
-            logging.error(f"Request to gateway router failed: {e}")
+            logging.error(f"Request to sandbox failed: {e}")
             raise RuntimeError(
-                f"Failed to communicate with the sandbox via the gateway at {url}.") from e
+                f"Failed to communicate with the sandbox at {url}.") from e
 
     def run(self, command: str, timeout: int = 60) -> ExecutionResult:
         payload = {"command": command}
@@ -359,6 +340,10 @@ def write(self, path: str, content: bytes | str, timeout: int = 60):
         self._request("POST", "upload", files=files_payload, timeout=timeout)
         logging.info(f"File '{filename}' uploaded successfully.")
 
-    def read(self, path: str, timeout: int = 60) -> bytes:
-        response = self._request("GET", f"download/{path}", timeout=timeout)
+    def read(self, path: str) -> bytes:
+        """
+        Downloads a file from the sandbox.
+        The base path for the download is the root of the sandbox's filesystem.
+        """
+        response = self._request("GET", f"download/{path}")
         return response.content