fix: address review feedback - typed exception, timeout cleanup, kwargs test

- Introduce SyncClientUnavailableError so the facade catches by type
  instead of matching error strings (review comment #1)
- Add future.cancel() + logger.warning() on timeout to match the
  _run_coroutine_sync pattern in base.py (review comment #2)
- Assert kwargs forwarding in the async bridge test (review comment #4)

Author: andreatgretel

packages/data-designer-engine/src/data_designer/engine/column_generators/generators/custom.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
 import inspect
 import logging
 from typing import TYPE_CHECKING, Any
@@ -26,27 +27,26 @@ class _AsyncBridgedModelFacade:
     """Proxy that bridges ``model.generate()`` to ``model.agenerate()`` in async engine mode.
 
     When a sync custom column runs inside ``asyncio.to_thread`` under the async engine,
-    the sync HTTP client is unavailable. This proxy intercepts the resulting error and
-    schedules ``agenerate()`` on the engine's persistent event loop via
-    ``run_coroutine_threadsafe``.
+    the sync HTTP client is unavailable. This proxy intercepts the resulting
+    ``SyncClientUnavailableError`` and schedules ``agenerate()`` on the engine's persistent
+    event loop via ``run_coroutine_threadsafe``.
 
     All other attributes are forwarded to the underlying facade unchanged.
     """
 
-    _SYNC_CLIENT_ERROR = "Sync methods are not available on an async-mode HttpModelClient."
-
     __slots__ = ("_facade",)
 
     def __init__(self, facade: Any) -> None:
         object.__setattr__(self, "_facade", facade)
 
     def generate(self, *args: Any, **kwargs: Any) -> tuple[Any, list]:
+        from data_designer.engine.models.clients.errors import SyncClientUnavailableError
+
         facade = object.__getattribute__(self, "_facade")
         try:
             return facade.generate(*args, **kwargs)
-        except RuntimeError as exc:
-            if str(exc) != self._SYNC_CLIENT_ERROR:
-                raise
+        except SyncClientUnavailableError:
+            pass  # Fall through to async bridge
 
         # We're in a worker thread (asyncio.to_thread) with no running loop.
         # Guard against accidental use from the event loop itself (would deadlock).
@@ -64,7 +64,12 @@ def generate(self, *args: Any, **kwargs: Any) -> tuple[Any, list]:
 
         loop = ensure_async_engine_loop()
         future = asyncio.run_coroutine_threadsafe(facade.agenerate(*args, **kwargs), loop)
-        return future.result(timeout=_SYNC_BRIDGE_TIMEOUT)
+        try:
+            return future.result(timeout=_SYNC_BRIDGE_TIMEOUT)
+        except concurrent.futures.TimeoutError as exc:
+            future.cancel()
+            logger.warning("Async model bridge timed out after %ss; coroutine cancelled", _SYNC_BRIDGE_TIMEOUT)
+            raise TimeoutError(f"model.generate() bridge timed out after {_SYNC_BRIDGE_TIMEOUT}s") from exc
 
     def __getattr__(self, name: str) -> Any:
         return getattr(object.__getattribute__(self, "_facade"), name)

packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/http_model_client.py

@@ -14,7 +14,7 @@
     resolve_timeout,
     wrap_transport_error,
 )
-from data_designer.engine.models.clients.errors import map_http_error_to_provider_error
+from data_designer.engine.models.clients.errors import SyncClientUnavailableError, map_http_error_to_provider_error
 from data_designer.engine.models.clients.retry import RetryConfig, RetryTransport, create_retry_transport
 
 if TYPE_CHECKING:
@@ -96,7 +96,7 @@ def _build_headers(self, extra_headers: dict[str, str]) -> dict[str, str]:
 
     def _get_sync_client(self) -> httpx.Client:
         if self._mode != ClientConcurrencyMode.SYNC:
-            raise RuntimeError("Sync methods are not available on an async-mode HttpModelClient.")
+            raise SyncClientUnavailableError("Sync methods are not available on an async-mode HttpModelClient.")
         with self._init_lock:
             if self._closed:
                 raise RuntimeError("Model client is closed.")

packages/data-designer-engine/src/data_designer/engine/models/clients/errors.py

@@ -29,6 +29,10 @@ class ProviderErrorKind(str, Enum):
     UNSUPPORTED_CAPABILITY = "unsupported_capability"
 
 
+class SyncClientUnavailableError(RuntimeError):
+    """Raised when sync methods are called on an async-mode HttpModelClient."""
+
+
 class ProviderError(Exception):
     def __init__(
         self,

packages/data-designer-engine/tests/engine/column_generators/generators/test_custom.py

@@ -509,12 +509,15 @@ def test_bridges_to_agenerate_on_sync_client_error(self) -> None:
         from unittest.mock import patch
 
         from data_designer.engine.column_generators.generators.custom import _AsyncBridgedModelFacade
+        from data_designer.engine.models.clients.errors import SyncClientUnavailableError
 
         facade = Mock()
-        facade.generate.side_effect = RuntimeError("Sync methods are not available on an async-mode HttpModelClient.")
+        facade.generate.side_effect = SyncClientUnavailableError(
+            "Sync methods are not available on an async-mode HttpModelClient."
+        )
 
         async def fake_agenerate(*args: Any, **kwargs: Any) -> tuple:
-            return ("async_result", list(args))
+            return ("async_result", list(args), kwargs)
 
         facade.agenerate = fake_agenerate
         proxy = _AsyncBridgedModelFacade(facade)
@@ -528,15 +531,14 @@ async def fake_agenerate(*args: Any, **kwargs: Any) -> tuple:
                 "data_designer.engine.dataset_builders.utils.async_concurrency.ensure_async_engine_loop",
                 return_value=engine_loop,
             ):
-                # Positional prompt arg is forwarded to agenerate
                 result = proxy.generate("hello", parser=str)
-            assert result == ("async_result", ["hello"])
+            assert result == ("async_result", ["hello"], {"parser": str})
         finally:
             engine_loop.call_soon_threadsafe(engine_loop.stop)
             engine_thread.join(timeout=5)
 
     def test_non_client_mode_errors_propagate(self) -> None:
-        """Only the specific HttpModelClient sync-mode RuntimeError triggers bridging."""
+        """Only SyncClientUnavailableError triggers bridging; other errors propagate."""
         from data_designer.engine.column_generators.generators.custom import _AsyncBridgedModelFacade
 
         # ValueError - different type entirely
@@ -546,7 +548,7 @@ def test_non_client_mode_errors_propagate(self) -> None:
         with pytest.raises(ValueError, match="invalid prompt format"):
             proxy.generate(prompt="hello")
 
-        # RuntimeError with a different message - same type, not caught
+        # RuntimeError - same base type as SyncClientUnavailableError, but not caught
         facade = Mock()
         facade.generate.side_effect = RuntimeError("connection timed out for async request")
         proxy = _AsyncBridgedModelFacade(facade)
@@ -558,9 +560,12 @@ def test_deadlock_guard_on_event_loop(self) -> None:
         import asyncio
 
         from data_designer.engine.column_generators.generators.custom import _AsyncBridgedModelFacade
+        from data_designer.engine.models.clients.errors import SyncClientUnavailableError
 
         facade = Mock()
-        facade.generate.side_effect = RuntimeError("Sync methods are not available on an async-mode HttpModelClient.")
+        facade.generate.side_effect = SyncClientUnavailableError(
+            "Sync methods are not available on an async-mode HttpModelClient."
+        )
         proxy = _AsyncBridgedModelFacade(facade)
 
         async def call_from_loop() -> None: