From 819216ca3ff1a4c6b33377a9dd1dbbbbb0ce1e6f Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 20:23:36 -0700 Subject: [PATCH] Added generic calls from Python to a worker extension. Lang code must reach the extension that Core calls, in the same native library. A call by method name with encoded input keeps the bridge free of code for one extension. A failed call raises ExtensionError with the gRPC code and details. --- temporalio/bridge/extension.py | 20 ++++++++ temporalio/bridge/src/extension.rs | 41 +++++++++++++-- temporalio/bridge/src/lib.rs | 4 ++ temporalio/worker/__init__.py | 3 +- temporalio/worker/_extension.py | 80 ++++++++++++++++++++++++++++++ tests/worker/test_extension.py | 38 ++++++++++++-- 6 files changed, 178 insertions(+), 8 deletions(-) diff --git a/temporalio/bridge/extension.py b/temporalio/bridge/extension.py index 5c4516077..b966552cd 100644 --- a/temporalio/bridge/extension.py +++ b/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) diff --git a/temporalio/bridge/src/extension.rs b/temporalio/bridge/src/extension.rs index cce72dce1..495876550 100644 --- a/temporalio/bridge/src/extension.rs +++ b/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, + 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) -> PyResult> { + 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, + ) -> PyResult> { + 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. diff --git a/temporalio/bridge/src/lib.rs b/temporalio/bridge/src/lib.rs index ae2eed586..1a0f33f0b 100644 --- a/temporalio/bridge/src/lib.rs +++ b/temporalio/bridge/src/lib.rs @@ -70,6 +70,10 @@ fn temporal_sdk_bridge(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_function(wrap_pyfunction!(new_worker, m)?)?; m.add_function(wrap_pyfunction!(new_replay_worker, m)?)?; + m.add( + "ExtensionCallError", + py.get_type::(), + )?; m.add_class::()?; m.add_function(wrap_pyfunction!(new_worker_extension, m)?)?; diff --git a/temporalio/worker/__init__.py b/temporalio/worker/__init__.py index 061175df8..b36c6920b 100644 --- a/temporalio/worker/__init__.py +++ b/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", diff --git a/temporalio/worker/_extension.py b/temporalio/worker/_extension.py index cbf218491..d387e358f 100644 --- a/temporalio/worker/_extension.py +++ b/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 diff --git a/tests/worker/test_extension.py b/tests/worker/test_extension.py index 1b9ca9ac7..bc8ad9dab 100644 --- a/tests/worker/test_extension.py +++ b/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