[Brook] Added generic calls from Python to a worker extension.

temporalio/bridge/extension.py

@@ -10,6 +10,10 @@
 import temporalio.bridge.runtime
 import temporalio.bridge.temporal_sdk_bridge
 
+ExtensionCallError: type[Exception] = (
+    temporalio.bridge.temporal_sdk_bridge.ExtensionCallError
+)
+
 
 class WorkerExtension:
     """An extension that the bridge builds by its name."""
@@ -26,3 +30,19 @@ def __init__(
     def name(self) -> str:
         """Name the extension gives itself."""
         return self._ref.name
+
+    def call(self, method: str, input: bytes) -> bytes:
+        """Call ``method`` and wait for its output.
+
+        Raises:
+            ExtensionCallError: The call failed, with the gRPC code, message and details as args.
+        """
+        return self._ref.call(method, input)
+
+    async def call_async(self, method: str, input: bytes) -> bytes:
+        """Call ``method`` on the runtime and await its output.
+
+        Raises:
+            ExtensionCallError: The call failed, with the gRPC code, message and details as args.
+        """
+        return await self._ref.call_async(method, input)

temporalio/bridge/src/extension.rs

@@ -2,18 +2,22 @@
 
 use std::sync::Arc;
 
-use pyo3::exceptions::PyValueError;
+use pyo3::exceptions::{PyException, PyValueError};
 use pyo3::prelude::*;
-use temporalio_sdk_core::WorkerExtension;
+use temporalio_sdk_core::{ExtensionError, WorkerExtension};
 
 use crate::runtime;
 
+// Its args are the gRPC code, the message and the encoded details, as with `RPCError`.
+pyo3::create_exception!(temporal_sdk_bridge, ExtensionCallError, PyException);
+
 /// An extension that a Worker config can carry. Python keeps it to reach the same instance that
 /// Core calls.
 #[pyclass(frozen, from_py_object)]
 #[derive(Clone)]
 pub struct WorkerExtensionRef {
     pub(crate) extension: Arc<dyn WorkerExtension>,
+    runtime: runtime::Runtime,
 }
 
 #[pymethods]
@@ -22,6 +26,37 @@ impl WorkerExtensionRef {
     fn name(&self) -> String {
         self.extension.name().to_string()
     }
+
+    /// Calls `method` and waits for its output. Workflow code can't await, so this is for code
+    /// that runs between the activation and its completion.
+    fn call(&self, py: Python<'_>, method: &str, input: Vec<u8>) -> PyResult<Vec<u8>> {
+        self.runtime
+            .assert_same_process("call a worker extension")?;
+        let extension = self.extension.clone();
+        py.detach(|| extension.call(method, &input))
+            .map_err(call_error)
+    }
+
+    fn call_async<'p>(
+        &self,
+        py: Python<'p>,
+        method: String,
+        input: Vec<u8>,
+    ) -> PyResult<Bound<'p, PyAny>> {
+        self.runtime
+            .assert_same_process("call a worker extension")?;
+        let extension = self.extension.clone();
+        self.runtime.future_into_py(py, async move {
+            extension
+                .call_async(method, input)
+                .await
+                .map_err(call_error)
+        })
+    }
+}
+
+fn call_error(err: ExtensionError) -> PyErr {
+    ExtensionCallError::new_err((err.code as u32, err.message, err.details))
 }
 
 /// Builds the extension with `name` from its encoded config. Only the extensions that this build
@@ -46,7 +81,7 @@ pub fn new_worker_extension(
             )))
         }
     };
-    Ok(WorkerExtensionRef { extension })
+    Ok(WorkerExtensionRef { extension, runtime })
 }
 
 /// An extension for the SDK's own tests. Release builds leave it out.

temporalio/bridge/src/lib.rs

@@ -70,6 +70,10 @@ fn temporal_sdk_bridge(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<worker::LocalActivitySlotInfo>()?;
     m.add_function(wrap_pyfunction!(new_worker, m)?)?;
     m.add_function(wrap_pyfunction!(new_replay_worker, m)?)?;
+    m.add(
+        "ExtensionCallError",
+        py.get_type::<extension::ExtensionCallError>(),
+    )?;
     m.add_class::<extension::WorkerExtensionRef>()?;
     m.add_function(wrap_pyfunction!(new_worker_extension, m)?)?;
 

temporalio/worker/__init__.py

@@ -2,7 +2,7 @@
 
 from ..common import WorkerDeploymentVersion
 from ._activity import SharedHeartbeatSender, SharedStateManager
-from ._extension import WorkerExtension
+from ._extension import ExtensionError, WorkerExtension
 from ._interceptor import (
     ActivityInboundInterceptor,
     ActivityOutboundInterceptor,
@@ -81,6 +81,7 @@
     "PollerBehaviorAutoscaling",
     "PatchActivationInput",
     "WorkerExtension",
+    "ExtensionError",
     # Interceptor base classes
     "Interceptor",
     "ActivityInboundInterceptor",

temporalio/worker/_extension.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import temporalio.bridge.extension
+import temporalio.exceptions
 import temporalio.runtime
+import temporalio.service
 
 
 class WorkerExtension:
@@ -50,3 +52,81 @@ def name(self) -> str:
     def runtime(self) -> temporalio.runtime.Runtime:
         """Runtime the extension runs on."""
         return self._runtime
+
+    def call(self, method: str, input: bytes) -> bytes:
+        """Call a method of the extension and wait for its output.
+
+        This blocks the thread, so call it only for methods that the extension documents as
+        fast.
+
+        Args:
+            method: The method name. Each extension has its own methods.
+            input: The encoded input.
+
+        Returns:
+            The encoded output.
+
+        Raises:
+            ExtensionError: The extension refused the call.
+        """
+        try:
+            return self._bridge_extension.call(method, input)
+        except temporalio.bridge.extension.ExtensionCallError as err:
+            raise ExtensionError._from_bridge(err) from None
+
+    async def call_async(self, method: str, input: bytes) -> bytes:
+        """Call a method of the extension on its runtime and await the output.
+
+        Args:
+            method: The method name. Each extension has its own methods.
+            input: The encoded input.
+
+        Returns:
+            The encoded output.
+
+        Raises:
+            ExtensionError: The extension refused the call.
+        """
+        try:
+            return await self._bridge_extension.call_async(method, input)
+        except temporalio.bridge.extension.ExtensionCallError as err:
+            raise ExtensionError._from_bridge(err) from None
+
+
+class ExtensionError(temporalio.exceptions.TemporalError):
+    """Error that a :py:class:`WorkerExtension` call raises.
+
+    It has the fields of a gRPC status, so callers can treat it like a failed service call.
+
+    .. warning::
+        This API is experimental.
+    """
+
+    def __init__(
+        self, message: str, code: temporalio.service.RPCStatusCode, details: bytes
+    ) -> None:
+        """Create an extension error."""
+        super().__init__(message)
+        self._message = message
+        self._code = code
+        self._details = details
+
+    @staticmethod
+    def _from_bridge(err: Exception) -> ExtensionError:
+        code, message, details = err.args
+        return ExtensionError(message, temporalio.service.RPCStatusCode(code), details)
+
+    @property
+    def message(self) -> str:
+        """Message from the extension."""
+        return self._message
+
+    @property
+    def code(self) -> temporalio.service.RPCStatusCode:
+        """Status code of the error, as in gRPC."""
+        return self._code
+
+    @property
+    def details(self) -> bytes:
+        """Encoded error details. Each extension has its own format. Empty if none."""
+        return self._details

tests/worker/test_extension.py

@@ -5,7 +5,8 @@
 from temporalio import workflow
 from temporalio.client import Client
 from temporalio.runtime import Runtime, TelemetryConfig
-from temporalio.worker import WorkerExtension
+from temporalio.service import RPCStatusCode
+from temporalio.worker import ExtensionError, WorkerExtension
 from tests.helpers import assert_task_fail_eventually, new_worker
 
 # Debug bridge builds link this extension for the SDK's tests.
@@ -36,15 +37,16 @@ def test_extension_name():
 
 
 async def test_extension_lets_workflow_task_complete(client: Client):
-    async with new_worker(
-        client, SimpleWorkflow, extension=load_test_extension()
-    ) as worker:
+    worker_extension = load_test_extension()
+    async with new_worker(client, SimpleWorkflow, extension=worker_extension) as worker:
         result = await client.execute_workflow(
             SimpleWorkflow.run,
             id=f"wf-{uuid.uuid4()}",
             task_queue=worker.task_queue,
         )
     assert result == "done"
+    before, after, accepted = map(int, extension_counts(worker_extension))
+    assert before >= 1 and after == before and accepted == before
 
 
 async def test_extension_failure_fails_workflow_task(client: Client):
@@ -66,3 +68,31 @@ async def test_extension_must_use_client_runtime(client: Client):
     extension = WorkerExtension(TEST_EXTENSION, runtime=other)
     with pytest.raises(ValueError, match="same runtime as the client"):
         new_worker(client, SimpleWorkflow, extension=extension)
+
+
+def extension_counts(extension: WorkerExtension) -> list[str]:
+    return extension.call("Counts", b"").decode().split(",")
+
+
+async def test_extension_calls():
+    extension = load_test_extension()
+    assert extension.call("Echo", b"abc") == b"abc"
+    assert await extension.call_async("Echo", b"xyz") == b"xyz"
+    assert extension_counts(extension) == ["0", "0", "0"]
+
+
+async def test_extension_call_errors():
+    extension = load_test_extension()
+    with pytest.raises(ExtensionError) as sync_err:
+        extension.call("Fail", b"details")
+    assert sync_err.value.code == RPCStatusCode.NOT_FOUND
+    assert sync_err.value.message == "test failure"
+    assert sync_err.value.details == b"details"
+    with pytest.raises(ExtensionError) as async_err:
+        await extension.call_async("Fail", b"")
+    assert async_err.value.code == RPCStatusCode.NOT_FOUND
+    assert async_err.value.details == b""
+    with pytest.raises(ExtensionError) as unknown_err:
+        await extension.call_async("Missing", b"")
+    assert unknown_err.value.code == RPCStatusCode.UNIMPLEMENTED
+    assert "Missing" in unknown_err.value.message