fix: catch genai.ClientError when sandbox is missing

This fixes an issue where AgentEngineSandboxCodeExecutor catches the wrong exception class when attempting to recover from externally-deleted sandboxes.

Fixes #5480

Co-authored-by: Amaad Martin <amaadmartin@google.com>
PiperOrigin-RevId: 908965545

Author: Amaad Martin

src/google/adk/code_executors/agent_engine_sandbox_code_executor.py

@@ -131,6 +131,7 @@ def execute_code(
     sandbox_name = self.sandbox_resource_name
     if self.sandbox_resource_name is None:
       from google.api_core import exceptions
+      from google.genai import errors as genai_errors
       from vertexai import types
 
       # use sandbox name stored in session if available.
@@ -148,6 +149,11 @@ def execute_code(
             create_new_sandbox = True
         except exceptions.NotFound:
           create_new_sandbox = True
+        except genai_errors.ClientError as exc:
+          if exc.code == 404:
+            create_new_sandbox = True
+          else:
+            raise
 
       if create_new_sandbox:
         # Create a new sandbox and assign it to sandbox_name.

tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py

@@ -155,7 +155,76 @@ def test_execute_code_recreates_sandbox_when_get_returns_none(
     mock_json_output = MagicMock()
     mock_json_output.mime_type = "application/json"
     mock_json_output.data = json.dumps(
-        {"stdout": "recreated sandbox run", "stderr": ""}
+        {"msg_out": "recreated sandbox run", "msg_err": ""}
+    ).encode("utf-8")
+    mock_json_output.metadata = None
+    mock_response.outputs = [mock_json_output]
+    mock_api_client.agent_engines.sandboxes.execute_code.return_value = (
+        mock_response
+    )
+
+    # Execute using agent_engine_resource_name so a sandbox can be created
+    executor = AgentEngineSandboxCodeExecutor(
+        agent_engine_resource_name=(
+            "projects/123/locations/us-central1/reasoningEngines/456"
+        )
+    )
+    code_input = CodeExecutionInput(code='print("hello world")')
+    result = executor.execute_code(mock_invocation_context, code_input)
+
+    # Assert get was called for the existing sandbox
+    mock_api_client.agent_engines.sandboxes.get.assert_called_once_with(
+        name=existing_sandbox_name
+    )
+
+    # Assert create was called and session updated with new sandbox
+    mock_api_client.agent_engines.sandboxes.create.assert_called_once()
+    assert (
+        mock_invocation_context.session.state["sandbox_name"]
+        == created_sandbox_name
+    )
+
+    # Assert execute_code used the created sandbox name
+    mock_api_client.agent_engines.sandboxes.execute_code.assert_called_once_with(
+        name=created_sandbox_name,
+        input_data={"code": 'print("hello world")'},
+    )
+
+  @patch("vertexai.Client")
+  def test_execute_code_recreates_sandbox_when_get_raises_client_error(
+      self,
+      mock_vertexai_client,
+      mock_invocation_context,
+  ):
+    # Setup Mocks
+    mock_api_client = MagicMock()
+    mock_vertexai_client.return_value = mock_api_client
+
+    # Existing sandbox name stored in session
+    existing_sandbox_name = "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/old"
+    mock_invocation_context.session.state = {
+        "sandbox_name": existing_sandbox_name
+    }
+
+    # Mock get to raise ClientError with code 404
+    from google.genai.errors import ClientError
+
+    mock_api_client.agent_engines.sandboxes.get.side_effect = ClientError(
+        code=404, response_json={"message": "Not Found"}
+    )
+
+    # Mock create operation to return a new sandbox resource name
+    operation_mock = MagicMock()
+    created_sandbox_name = "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789"
+    operation_mock.response.name = created_sandbox_name
+    mock_api_client.agent_engines.sandboxes.create.return_value = operation_mock
+
+    # Mock execute_code response
+    mock_response = MagicMock()
+    mock_json_output = MagicMock()
+    mock_json_output.mime_type = "application/json"
+    mock_json_output.data = json.dumps(
+        {"msg_out": "recreated sandbox run", "msg_err": ""}
     ).encode("utf-8")
     mock_json_output.metadata = None
     mock_response.outputs = [mock_json_output]