fix(application-integration): fail on errored or stuck schema operations

src/google/adk/tools/application_integration_tool/clients/connections_client.py

@@ -41,6 +41,8 @@
     "integrations.mtls.googleapis.com"
 )
 _DEFAULT_REQUEST_TIMEOUT_SECONDS = 30
+# Upper bound on waiting for a schema long-running operation to finish.
+_DEFAULT_OPERATION_TIMEOUT_SECONDS = 300
 
 
 class _ServiceAccountCredentialsFactory(Protocol):
@@ -962,15 +964,29 @@ def _poll_operation(self, operation_id: str) -> Dict[str, Any]:
 
     Raises:
         PermissionError: If there are credential issues.
-        ValueError: If there's a request error.
+        ValueError: If there's a request error, or the operation finished with
+          an error.
+        TimeoutError: If the operation is not done within
+          `_DEFAULT_OPERATION_TIMEOUT_SECONDS`.
         Exception: For any other unexpected errors.
     """
-    operation_done: bool = False
-    operation_response: Dict[str, Any] = {}
-    while not operation_done:
-      get_operation_url = f"{self.connector_url}/v1/{operation_id}"
+    get_operation_url = f"{self.connector_url}/v1/{operation_id}"
+    deadline = time.monotonic() + _DEFAULT_OPERATION_TIMEOUT_SECONDS
+    while True:
       response = self._execute_api_call(get_operation_url)
       operation_response = self._response_json(response)
-      operation_done = bool(operation_response.get("done", False))
+      if operation_response.get("done", False):
+        break
+      if time.monotonic() >= deadline:
+        raise TimeoutError(
+            f"Operation {operation_id} did not finish within"
+            f" {_DEFAULT_OPERATION_TIMEOUT_SECONDS} seconds."
+        )
       time.sleep(1)
+    # A finished operation carries either `response` or `error`; without this
+    # check a failed lookup reads as an empty schema and yields no tools.
+    error = operation_response.get("error")
+    if error:
+      message = error.get("message") if isinstance(error, dict) else None
+      raise ValueError(f"Operation {operation_id} failed: {message or error}")
     return operation_response

tests/unittests/tools/application_integration_tool/clients/test_connections_client.py

@@ -360,6 +360,90 @@ def test_get_entity_schema_and_operations_execute_api_call_error(
       with pytest.raises(ValueError, match="Request error"):
         client.get_entity_schema_and_operations("entity1")
 
+  def test_get_entity_schema_and_operations_failed_operation(
+      self, project, location, connection_name
+  ):
+    """A failed operation must not read as an empty schema with no tools."""
+    credentials = {"email": "test@example.com"}
+    client = ConnectionsClient(project, location, connection_name, credentials)
+    mock_execute_response_initial = mock.MagicMock()
+    mock_execute_response_initial.json.return_value = {
+        "name": "operations/test_op"
+    }
+    mock_execute_response_poll_failed = mock.MagicMock()
+    mock_execute_response_poll_failed.json.return_value = {
+        "done": True,
+        "error": {"code": 5, "message": "Entity type 'entity1' not found."},
+    }
+
+    with mock.patch.object(
+        client,
+        "_execute_api_call",
+        side_effect=[
+            mock_execute_response_initial,
+            mock_execute_response_poll_failed,
+        ],
+    ):
+      with pytest.raises(
+          ValueError,
+          match=(
+              "Operation operations/test_op failed: Entity type 'entity1' not"
+              " found."
+          ),
+      ):
+        client.get_entity_schema_and_operations("entity1")
+
+  def test_poll_operation_times_out(self, project, location, connection_name):
+    credentials = {"email": "test@example.com"}
+    client = ConnectionsClient(project, location, connection_name, credentials)
+    mock_execute_response_pending = mock.MagicMock()
+    mock_execute_response_pending.json.return_value = {"done": False}
+
+    with (
+        mock.patch.object(
+            client,
+            "_execute_api_call",
+            return_value=mock_execute_response_pending,
+        ),
+        mock.patch(
+            "google.adk.tools.application_integration_tool.clients.connections_client.time"
+        ) as mock_time,
+    ):
+      mock_time.monotonic.side_effect = [0, 1, 301]
+      with pytest.raises(
+          TimeoutError,
+          match="Operation operations/test_op did not finish within 300",
+      ):
+        client._poll_operation("operations/test_op")
+      assert client._execute_api_call.call_count == 2
+      mock_time.sleep.assert_called_once_with(1)
+
+  def test_poll_operation_does_not_sleep_once_done(
+      self, project, location, connection_name
+  ):
+    credentials = {"email": "test@example.com"}
+    client = ConnectionsClient(project, location, connection_name, credentials)
+    mock_execute_response_done = mock.MagicMock()
+    mock_execute_response_done.json.return_value = {
+        "done": True,
+        "response": {"jsonSchema": {}},
+    }
+
+    with (
+        mock.patch.object(
+            client, "_execute_api_call", return_value=mock_execute_response_done
+        ),
+        mock.patch(
+            "google.adk.tools.application_integration_tool.clients.connections_client.time"
+        ) as mock_time,
+    ):
+      mock_time.monotonic.return_value = 0
+      assert client._poll_operation("operations/test_op") == {
+          "done": True,
+          "response": {"jsonSchema": {}},
+      }
+      mock_time.sleep.assert_not_called()
+
   def test_get_action_schema_success(
       self, project, location, connection_name, mock_credentials
   ):