fix: bridge model.generate() to agenerate() for custom columns in async engine

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

@@ -12,7 +12,7 @@
 
 import data_designer.lazy_heavy_imports as lazy
 from data_designer.config.column_configs import CustomColumnConfig, GenerationStrategy
-from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.generators.base import _SYNC_BRIDGE_TIMEOUT, ColumnGenerator
 from data_designer.engine.column_generators.utils.errors import CustomColumnGenerationError
 from data_designer.logging import LOG_INDENT
 
@@ -22,6 +22,57 @@
 logger = logging.getLogger(__name__)
 
 
+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``.
+
+    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]:
+        facade = object.__getattribute__(self, "_facade")
+        try:
+            return facade.generate(*args, **kwargs)
+        except RuntimeError as exc:
+            if str(exc) != self._SYNC_CLIENT_ERROR:
+                raise
+
+        # We're in a worker thread (asyncio.to_thread) with no running loop.
+        # Guard against accidental use from the event loop itself (would deadlock).
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            pass  # No running loop - safe to bridge
+        else:
+            raise RuntimeError(
+                "model.generate() is not available in async engine mode from the event loop. "
+                "Use 'await model.agenerate()' in async custom columns."
+            )
+
+        from data_designer.engine.dataset_builders.utils.async_concurrency import ensure_async_engine_loop
+
+        loop = ensure_async_engine_loop()
+        future = asyncio.run_coroutine_threadsafe(facade.agenerate(*args, **kwargs), loop)
+        return future.result(timeout=_SYNC_BRIDGE_TIMEOUT)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(object.__getattribute__(self, "_facade"), name)
+
+    def __repr__(self) -> str:
+        return f"_AsyncBridgedModelFacade({object.__getattribute__(self, '_facade')!r})"
+
+
 class CustomColumnGenerator(ColumnGenerator[CustomColumnConfig]):
     """Column generator that uses a user-provided callable function.
 
@@ -277,9 +328,13 @@ def _invoke_generator_function(self, data: dict | pd.DataFrame) -> dict | pd.Dat
             return self.config.generator_function(data, self.config.generator_params, models)
 
     def _build_models_dict(self) -> dict[str, Any]:
-        """Build a dict of ModelFacade instances from model_aliases."""
+        """Build a dict of ModelFacade instances from model_aliases.
+
+        Facades are wrapped in ``_AsyncBridgedModelFacade`` so that sync custom
+        columns can call ``model.generate()`` transparently under the async engine.
+        """
         return {
-            alias: self.resource_provider.model_registry.get_model(model_alias=alias)
+            alias: _AsyncBridgedModelFacade(self.resource_provider.model_registry.get_model(model_alias=alias))
             for alias in self.config.model_aliases
         }
 

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

@@ -467,3 +467,105 @@ def df_func(df: pd.DataFrame) -> pd.DataFrame:
     gen = _create_test_generator(name="result", generator_function=df_func)
     with pytest.raises(CustomColumnGenerationError, match="first parameter must be 'row', got 'df'"):
         gen.generate({"input": 1})
+
+
+# Async model bridge tests
+
+
+class TestAsyncBridgedModelFacade:
+    """Tests for _AsyncBridgedModelFacade proxy used by custom columns with model access."""
+
+    def test_proxy_transparent_in_sync_mode(self, stub_resource_provider, stub_model_facade) -> None:
+        """Proxy passes through generate(), forwards attributes, and is used by _build_models_dict."""
+        from data_designer.engine.column_generators.generators.custom import _AsyncBridgedModelFacade
+
+        @custom_column_generator(required_columns=["input"], model_aliases=["test-model"])
+        def gen_with_model(row: dict, generator_params: SampleParams, models: dict) -> dict:
+            row["result"] = "ok"
+            return row
+
+        generator = _create_test_generator(
+            name="result",
+            generator_function=gen_with_model,
+            generator_params=SampleParams(),
+            resource_provider=stub_resource_provider,
+        )
+
+        models = generator._build_models_dict()
+        proxy = models["test-model"]
+        assert isinstance(proxy, _AsyncBridgedModelFacade)
+
+        # generate() passes through to the underlying facade (positional and keyword args)
+        result, _ = proxy.generate("test", parser=str)
+        assert result == "Generated summary text"
+        stub_model_facade.generate.assert_called_once_with("test", parser=str)
+
+        # Other attributes are forwarded
+        assert proxy.model_alias == "test_model"
+
+    def test_bridges_to_agenerate_on_sync_client_error(self) -> None:
+        """When sync generate() fails with an async/sync error, falls back to agenerate()."""
+        import asyncio
+        import threading
+        from unittest.mock import patch
+
+        from data_designer.engine.column_generators.generators.custom import _AsyncBridgedModelFacade
+
+        facade = Mock()
+        facade.generate.side_effect = RuntimeError("Sync methods are not available on an async-mode HttpModelClient.")
+
+        async def fake_agenerate(*args: Any, **kwargs: Any) -> tuple:
+            return ("async_result", list(args))
+
+        facade.agenerate = fake_agenerate
+        proxy = _AsyncBridgedModelFacade(facade)
+
+        engine_loop = asyncio.new_event_loop()
+        engine_thread = threading.Thread(target=engine_loop.run_forever, daemon=True)
+        engine_thread.start()
+
+        try:
+            with patch(
+                "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"])
+        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."""
+        from data_designer.engine.column_generators.generators.custom import _AsyncBridgedModelFacade
+
+        # ValueError - different type entirely
+        facade = Mock()
+        facade.generate.side_effect = ValueError("invalid prompt format")
+        proxy = _AsyncBridgedModelFacade(facade)
+        with pytest.raises(ValueError, match="invalid prompt format"):
+            proxy.generate(prompt="hello")
+
+        # RuntimeError with a different message - same type, not caught
+        facade = Mock()
+        facade.generate.side_effect = RuntimeError("connection timed out for async request")
+        proxy = _AsyncBridgedModelFacade(facade)
+        with pytest.raises(RuntimeError, match="connection timed out"):
+            proxy.generate(prompt="hello")
+
+    def test_deadlock_guard_on_event_loop(self) -> None:
+        """Raises a clear error instead of deadlocking when called from the event loop."""
+        import asyncio
+
+        from data_designer.engine.column_generators.generators.custom import _AsyncBridgedModelFacade
+
+        facade = Mock()
+        facade.generate.side_effect = RuntimeError("Sync methods are not available on an async-mode HttpModelClient.")
+        proxy = _AsyncBridgedModelFacade(facade)
+
+        async def call_from_loop() -> None:
+            proxy.generate(prompt="hello")
+
+        with pytest.raises(RuntimeError, match="Use 'await model.agenerate\\(\\)'"):
+            asyncio.run(call_from_loop())