Fixes bug in sandbox_client that assumed pod name matched the sandbox name (#150)

* Fixes bug in sandbox_client

* Removes unnecessary sandboxclaim watch

* Addresses review comments

* Rebase to main

* Adds tests for the pod name

Author: igooch

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

@@ -38,6 +38,8 @@
 SANDBOX_API_VERSION = "v1alpha1"
 SANDBOX_PLURAL_NAME = "sandboxes"
 
+POD_NAME_ANNOTATION = "agents.x-k8s.io/pod-name"
+
 logging.basicConfig(level=logging.INFO,
                     format='%(asctime)s - %(levelname)s - %(message)s',
                     stream=sys.stdout)
@@ -82,6 +84,8 @@ def __init__(
 
         self.claim_name: str | None = None
         self.sandbox_name: str | None = None
+        self.pod_name: str | None = None
+        self.annotations: dict | None = None
 
         try:
             config.load_incluster_config()
@@ -145,25 +149,41 @@ def _wait_for_sandbox_ready(self):
             field_selector=f"metadata.name={self.claim_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
-                    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:
-            self.__exit__(None, None, None)
-            raise TimeoutError(
-                f"Sandbox did not become ready within {self.sandbox_ready_timeout} seconds.")
+            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:
+                    metadata = sandbox_object.get(
+                        "metadata", {})
+                    self.sandbox_name = metadata.get(
+                        "name")
+                    if not self.sandbox_name:
+                        raise RuntimeError(
+                            "Could not determine sandbox name from sandbox object.")
+                    logging.info(f"Sandbox {self.sandbox_name} is ready.")
+
+                    self.annotations = sandbox_object.get(
+                        'metadata', {}).get('annotations', {})
+                    pod_name = self.annotations.get(POD_NAME_ANNOTATION)
+                    if pod_name:
+                        self.pod_name = pod_name
+                        logging.info(
+                            f"Found pod name from annotation: {self.pod_name}")
+                    else:
+                        self.pod_name = self.sandbox_name
+                    w.stop()
+                    return
+
+        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."""

clients/python/agentic-sandbox-client/test_client.py

@@ -16,14 +16,17 @@
 import asyncio
 from agentic_sandbox import SandboxClient
 
+POD_NAME_ANNOTATION = "agents.x-k8s.io/pod-name"
+
 
 async def main(template_name: str, gateway_name: str | None, api_url: str | None, namespace: str, server_port: int):
     """
     Tests the Sandbox client by creating a sandbox, running a command,
     and then cleaning up.
     """
 
-    print(f"--- Starting Sandbox Client Test (Namespace: {namespace}, Port: {server_port}) ---")
+    print(
+        f"--- Starting Sandbox Client Test (Namespace: {namespace}, Port: {server_port}) ---")
     if gateway_name:
         print(f"Mode: Gateway Discovery ({gateway_name})")
     elif api_url:
@@ -41,6 +44,20 @@ async def main(template_name: str, gateway_name: str | None, api_url: str | None
             server_port=server_port
         ) as sandbox:
 
+            print("\n--- Testing Pod Name Discovery ---")
+            assert sandbox.annotations is not None, "Sandbox annotations were not stored on the client"
+
+            pod_name_annotation = sandbox.annotations.get(POD_NAME_ANNOTATION)
+
+            if pod_name_annotation:
+                print(f"Found pod name from annotation: {pod_name_annotation}")
+                assert sandbox.pod_name == pod_name_annotation, f"Expected pod_name to be '{pod_name_annotation}', but got '{sandbox.pod_name}'"
+                print("--- Pod Name Discovery Test Passed (Annotation) ---")
+            else:
+                print("Pod name annotation not found, falling back to sandbox name.")
+                assert sandbox.pod_name == sandbox.sandbox_name, f"Expected pod_name to be '{sandbox.sandbox_name}', but got '{sandbox.pod_name}'"
+                print("--- Pod Name Discovery Test Passed (Fallback) ---")
+
             print("\n--- Testing Command Execution ---")
             command_to_run = "echo 'Hello from the sandbox!'"
             print(f"Executing command: '{command_to_run}'")
@@ -97,24 +114,27 @@ async def main(template_name: str, gateway_name: str | None, api_url: str | None
         default="python-sandbox-template",
         help="The name of the sandbox template to use for the test."
     )
-    
+
     # Default is None to allow testing the Port-Forward fallback
     parser.add_argument(
         "--gateway-name",
-        default=None, 
+        default=None,
         help="The name of the Gateway resource. If omitted, defaults to local port-forward mode."
     )
-    
-    parser.add_argument("--api-url", help="Direct URL to router (e.g. http://localhost:8080)", default=None)
-    parser.add_argument("--namespace", default="default", help="Namespace to create sandbox in")
-    parser.add_argument("--server-port", type=int, default=8888, help="Port the sandbox container listens on")
-    
+
+    parser.add_argument(
+        "--api-url", help="Direct URL to router (e.g. http://localhost:8080)", default=None)
+    parser.add_argument("--namespace", default="default",
+                        help="Namespace to create sandbox in")
+    parser.add_argument("--server-port", type=int, default=8888,
+                        help="Port the sandbox container listens on")
+
     args = parser.parse_args()
-    
+
     asyncio.run(main(
         template_name=args.template_name,
         gateway_name=args.gateway_name,
         api_url=args.api_url,
         namespace=args.namespace,
         server_port=args.server_port
-    ))
+    ))