[async worker] add configurable native UDF timeouts

docs/integration/integrations.md

@@ -77,6 +77,8 @@ UDFs can also produce **effects**: structured outputs that downstream systems ac
 
 UDFs that perform I/O (e.g. network calls or database reads) should subclass `AsyncUDFBase` when used in the async worker; see [`osprey_async_worker/src/osprey/async_worker/stdlib_udfs/async_mx_lookup.py`](https://github.com/roostorg/osprey/blob/main/osprey_async_worker/src/osprey/async_worker/stdlib_udfs/async_mx_lookup.py) for an example. Pure-computation UDFs like `TextContains` can be reused in both workers without modification.
 
+Native async UDFs time out after two seconds by default. Set `OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT` to change the process default. A UDF can set its own numeric `timeout` class attribute in seconds, which takes precedence over the process default. Values must be greater than zero and can't be NaN or infinity. Async UDFs must not suppress `asyncio.CancelledError`; catch it only to clean up, then re-raise it.
+
 ### Registering UDFs
 
 Return your UDF classes from the `register_udfs` hook; for example:

osprey_async_worker/Dockerfile

@@ -45,10 +45,11 @@ ADD osprey_rpc/pyproject.toml /osprey/osprey_rpc/pyproject.toml
 ADD osprey_worker/pyproject.toml /osprey/osprey_worker/pyproject.toml
 ADD osprey_async_worker/pyproject.toml /osprey/osprey_async_worker/pyproject.toml
 ADD example_plugins/pyproject.toml /osprey/example_plugins/pyproject.toml
+ADD example_atproto_plugins/pyproject.toml /osprey/example_atproto_plugins/pyproject.toml
 
 # Create minimal package structure required by uv
-RUN mkdir -p /osprey/osprey_worker /osprey/osprey_async_worker /osprey/osprey_rpc /osprey/example_plugins/src && \
-    touch /osprey/osprey_worker/__init__.py /osprey/osprey_async_worker/__init__.py /osprey/osprey_rpc/__init__.py /osprey/example_plugins/src/__init__.py
+RUN mkdir -p /osprey/osprey_worker /osprey/osprey_async_worker /osprey/osprey_rpc /osprey/example_plugins/src /osprey/example_atproto_plugins/src/atproto_plugin && \
+    touch /osprey/osprey_worker/__init__.py /osprey/osprey_async_worker/__init__.py /osprey/osprey_rpc/__init__.py /osprey/example_plugins/src/__init__.py /osprey/example_atproto_plugins/src/atproto_plugin/__init__.py
 
 # Install the full workspace, including osprey_async_worker (this is the one image
 # that does). This layer is cached when only source code changes.
@@ -65,6 +66,7 @@ ADD osprey_worker /osprey/osprey_worker
 ADD osprey_async_worker /osprey/osprey_async_worker
 ADD osprey_rpc /osprey/osprey_rpc
 ADD example_plugins /osprey/example_plugins
+ADD example_atproto_plugins /osprey/example_atproto_plugins
 
 COPY entrypoint.sh /osprey/entrypoint.sh
 

osprey_async_worker/src/osprey/async_worker/adaptor/interfaces.py

@@ -34,10 +34,15 @@ class AsyncUDFBase(UDFBase[Arguments, RValue]):
 
     The sync execute() raises so it can't accidentally be called in the
     async executor's sync path.
+
+    `timeout` sets the maximum runtime in seconds. Plugins may override it for
+    a UDF. `async_execute()` must not suppress `asyncio.CancelledError`; catch it
+    only to clean up, then re-raise it.
     """
 
     execute_async: ClassVar[bool] = True
     is_native_async: ClassVar[bool] = True
+    timeout: ClassVar[float] = 2.0
 
     def __init__(self, validation_context, arguments):
         super().__init__(validation_context, arguments)
@@ -74,10 +79,13 @@ class AsyncBatchableUDFBase(BatchableUDFBase[Arguments, RValue, BatchableArgumen
     """Native async batchable UDF base class.
 
     Same as AsyncUDFBase but for batchable UDFs. The async executor detects
-    these and awaits async_execute_batch() directly.
+    these and awaits async_execute_batch() directly. `timeout` sets the maximum
+    batch runtime in seconds. `async_execute_batch()` must not suppress
+    `asyncio.CancelledError`; catch it only to clean up, then re-raise it.
     """
 
     is_native_async: ClassVar[bool] = True
+    timeout: ClassVar[float] = 2.0
 
     def execute(self, execution_context: ExecutionContext, arguments: Arguments) -> RValue:
         raise RuntimeError(

osprey_async_worker/src/osprey/async_worker/adaptor/plugin_manager.py

@@ -7,13 +7,14 @@
 from __future__ import annotations
 
 import logging
+import math
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any, List, Type, cast
 
 import pluggy
 from osprey.async_worker.adaptor import hookspecs as async_hookspecs
 from osprey.async_worker.adaptor.constants import OSPREY_ASYNC_ADAPTOR
-from osprey.async_worker.adaptor.interfaces import AsyncBaseOutputSink
+from osprey.async_worker.adaptor.interfaces import AsyncBaseOutputSink, AsyncBatchableUDFBase, AsyncUDFBase
 from osprey.async_worker.sinks.sink.output_sink import AsyncMultiOutputSink
 from osprey.engine.ast_validator import ValidatorRegistry
 from osprey.engine.executor.udf_execution_helpers import HasHelper, UDFHelpers
@@ -24,6 +25,8 @@
 if TYPE_CHECKING:
     from osprey.worker.lib.config import Config
 
+OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT = 'OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT'
+
 hookimpl_osprey_async: pluggy.HookimplMarker = pluggy.HookimplMarker(OSPREY_ASYNC_ADAPTOR)
 
 plugin_manager = pluggy.PluginManager(OSPREY_ASYNC_ADAPTOR)
@@ -66,6 +69,17 @@ def _deduplicate_udfs(
     return deduplicated
 
 
+def _validate_udf_timeout(value: object, source: str) -> None:
+    """Validate a timeout value.
+
+    Raises ValueError naming the source and value when not finite or non-positive.
+    """
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise ValueError(f'{source} must be a positive finite number, got {value!r}')
+    if not math.isfinite(value) or value <= 0:
+        raise ValueError(f'{source} must be a positive finite number, got {value}')
+
+
 def bootstrap_async_udfs(config: 'Config | None' = None) -> tuple[UDFRegistry, UDFHelpers]:
     """Bootstrap UDFs from async plugins + stdlib.
 
@@ -76,16 +90,31 @@ def bootstrap_async_udfs(config: 'Config | None' = None) -> tuple[UDFRegistry, U
     `_async_stdlib_plugin`, which registers through the same `register_udfs`
     hook as third-party plugins. No sync fallbacks — all I/O UDFs must be
     native async.
+
+    Resolves the default UDF timeout from OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT config key,
+    defaulting to 2.0. Each call replaces the process-wide inherited default and
+    must finish before UDF execution starts. Validation completes before mutation.
     """
     from osprey.worker._stdlibplugin.udf_register import register_udfs as stdlib_register_udfs
 
+    # Resolve and validate the configured timeout before any mutations
+    resolved_timeout = 2.0 if config is None else config.get_float(OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT, 2.0)
+    _validate_udf_timeout(resolved_timeout, f'config[{OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT!r}]')
+
     load_all_async_plugins()
     udf_helpers = UDFHelpers()
 
     stdlib_udfs = list(stdlib_register_udfs())
     plugin_udfs = _flatten(plugin_manager.hook.register_udfs())
     all_udfs = _deduplicate_udfs(stdlib_udfs, plugin_udfs)
 
+    for udf in all_udfs:
+        if issubclass(udf, (AsyncUDFBase, AsyncBatchableUDFBase)):
+            _validate_udf_timeout(udf.timeout, f'{udf.__name__}.timeout')
+
+    AsyncUDFBase.timeout = resolved_timeout
+    AsyncBatchableUDFBase.timeout = resolved_timeout
+
     # Auto-register helpers for UDFs that extend HasHelper
     for udf in all_udfs:
         if issubclass(udf, HasHelper):

osprey_async_worker/src/osprey/async_worker/cli/main.py

@@ -50,6 +50,7 @@ def bootstrap_stdlib_engine(rules_path: str) -> tuple[AsyncOspreyEngine, UDFHelp
     """Bootstrap engine with only stdlib UDFs — no external plugins, no Postgres, no labels.
 
     This avoids loading example_plugins or any third-party plugins that require database connections.
+    It registers no native async UDFs, so native timeout configuration does not apply.
     """
     from osprey.engine.ast_validator import ValidatorRegistry
     from osprey.worker._stdlibplugin.udf_register import register_udfs as stdlib_register_udfs

osprey_async_worker/src/osprey/async_worker/executor.py

@@ -228,7 +228,13 @@ async def _execute_async_udf(
     context: ExecutionContext,
     error_info_: list[NodeErrorInfo],
 ) -> NodeResult:
-    """Execute a native async UDF. Awaited directly on the event loop."""
+    """Execute a native async UDF. Awaited directly on the event loop.
+
+    Timeout enforcement: asyncio.timeout() enforces udf.timeout only around
+    async_execute() after semaphore admission. TimeoutError is converted to
+    NodeErrorInfo through the existing exception handler, preserving current
+    failure handling and Sentry capture behavior without modification.
+    """
     async with semaphore:
         call_executor: CallExecutor = chain.executor  # type: ignore
         udf: AsyncUDFBase[Any, Any] = call_executor._udf  # type: ignore
@@ -239,7 +245,8 @@ async def _execute_async_udf(
         try:
             resolved_arguments = udf.resolve_arguments(context, call_executor)
             with metrics.timed('udf_execution_duration', tags=metric_tags, sample_rate=0.01):
-                result = await udf.async_execute(context, resolved_arguments)
+                async with asyncio.timeout(type(udf).timeout):
+                    result = await udf.async_execute(context, resolved_arguments)
             execution_result = Ok(udf.check_result_type(result))
         except Exception as e:
             if not isinstance(e, NodeFailurePropagationException):
@@ -259,15 +266,26 @@ async def _execute_async_batch(
     context: ExecutionContext,
     error_info_: list[NodeErrorInfo],
 ) -> Sequence[NodeResult]:
-    """Execute a batch of native async batchable UDFs."""
+    """Execute a batch of native async batchable UDFs.
+
+    Timeout enforcement: the shared batch deadline is computed as
+    max(udf.timeout for udf in udfs). asyncio.timeout() enforces this
+    deadline only around async_execute_batch() after semaphore admission.
+    TimeoutError is converted to NodeErrorInfo through the existing exception
+    handler, preserving current failure handling and Sentry capture behavior.
+    """
     async with semaphore:
         assert len(udfs) == len(nodes) == len(batchable_args)
         num_executions = len(udfs)
         metric_tags = _get_metric_tags(context, udfs[0])
 
+        # Compute the shared deadline using the maximum timeout from all UDFs
+        batch_timeout = max(type(udf).timeout for udf in udfs)
+
         try:
             with metrics.timed('udf_execution_batch_duration', tags=metric_tags, sample_rate=0.01):
-                results = await udfs[0].async_execute_batch(context, udfs, batchable_args)
+                async with asyncio.timeout(batch_timeout):
+                    results = await udfs[0].async_execute_batch(context, udfs, batchable_args)
             assert len(results) == num_executions
         except Exception as e:
             if not isinstance(e, NodeFailurePropagationException):
@@ -325,7 +343,7 @@ async def _enqueue_batches(
 
     Returns (remaining non-batched chains, dict of batch tasks -> chains).
     """
-    batch_chains: dict[tuple[type, str], list[tuple[DependencyChain, Any]]] = defaultdict(list)
+    batch_chains: dict[tuple[type, str, bool], list[tuple[DependencyChain, Any]]] = defaultdict(list)
     chains_to_remove: list[DependencyChain] = []
 
     for async_chain in ready_async:
@@ -337,11 +355,12 @@ async def _enqueue_batches(
         udf = call_executor._udf
 
         batch_type = udf.get_batchable_arguments_type()
+        is_native = isinstance(udf, AsyncBatchableUDFBase)
         try:
             resolved_arguments = udf.resolve_arguments(context, call_executor)
             batchable_arguments = udf.get_batchable_arguments(resolved_arguments)
             routing_key = udf.get_batch_routing_key(batchable_arguments)
-            batch_chains[(batch_type, routing_key)].append((async_chain, batchable_arguments))
+            batch_chains[(batch_type, routing_key, is_native)].append((async_chain, batchable_arguments))
         except Exception as e:
             if not isinstance(e, NodeFailurePropagationException):
                 error_infos.append(NodeErrorInfo(e, call_executor.node))

osprey_async_worker/src/osprey/async_worker/stdlib_udfs/async_mx_lookup.py

@@ -17,6 +17,7 @@
 from osprey.engine.stdlib.udfs.mx_lookup import MXLookup as SyncMXLookup
 
 _DNS_TIMEOUT = 5.0
+_DNS_TRIES = 3
 _resolver: aiodns.DNSResolver | None = None
 
 
@@ -25,14 +26,15 @@ def _get_resolver() -> aiodns.DNSResolver:
     global _resolver
     loop = asyncio.get_running_loop()
     if _resolver is None or _resolver.loop is not loop:
-        _resolver = aiodns.DNSResolver(timeout=_DNS_TIMEOUT, loop=loop)
+        _resolver = aiodns.DNSResolver(timeout=_DNS_TIMEOUT, tries=_DNS_TRIES, loop=loop)
     return _resolver
 
 
 class MXLookup(AsyncUDFBase[Arguments, str]):  # type: ignore[misc]
     """Async MXLookup — uses aiodns for non-blocking DNS resolution."""
 
     category = SyncMXLookup.category
+    timeout = (_DNS_TIMEOUT * _DNS_TRIES * 2) + 1.0
 
     @classmethod
     def _get_udf_base_args(cls):