[Brook] Added a Worker option that passes a worker extension to Core.

temporalio/bridge/extension.py

@@ -0,0 +1,28 @@
+"""Worker extensions that this bridge build links. (unstable)
+
+Nothing in this module should be considered stable. The API may change.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import temporalio.bridge.runtime
+import temporalio.bridge.temporal_sdk_bridge
+
+
+class WorkerExtension:
+    """An extension that the bridge builds by its name."""
+
+    def __init__(
+        self, runtime: temporalio.bridge.runtime.Runtime, name: str, config: bytes
+    ) -> None:
+        """Build the extension with ``name`` from its encoded config."""
+        self._ref: Any = temporalio.bridge.temporal_sdk_bridge.new_worker_extension(
+            runtime._ref, name, config
+        )
+
+    @property
+    def name(self) -> str:
+        """Name the extension gives itself."""
+        return self._ref.name

temporalio/bridge/proto/activity_task/__init__.py

@@ -4,6 +4,7 @@
     ActivityTask,
     Cancel,
     Start,
+    WorkflowTaskAttempt,
 )
 
 __all__ = [
@@ -12,4 +13,5 @@
     "ActivityTask",
     "Cancel",
     "Start",
+    "WorkflowTaskAttempt",
 ]

temporalio/bridge/proto/activity_task/activity_task_pb2.pyi

@@ -160,6 +160,7 @@ class Start(google.protobuf.message.Message):
     PRIORITY_FIELD_NUMBER: builtins.int
     IS_LOCAL_FIELD_NUMBER: builtins.int
     RUN_ID_FIELD_NUMBER: builtins.int
+    WORKFLOW_TASK_FIELD_NUMBER: builtins.int
     workflow_namespace: builtins.str
     """The namespace the workflow lives in"""
     workflow_type: builtins.str
@@ -227,6 +228,11 @@ class Start(google.protobuf.message.Message):
     """
     run_id: builtins.str
     """Run ID of this activity execution. Only set for standalone activities."""
+    @property
+    def workflow_task(self) -> global___WorkflowTaskAttempt:
+        """For a local activity, the workflow task in which Core dispatched this execution. Retries
+        that Core runs without a new task keep it. Unset for other activities.
+        """
     def __init__(
         self,
         *,
@@ -258,6 +264,7 @@ class Start(google.protobuf.message.Message):
         priority: temporalio.api.common.v1.message_pb2.Priority | None = ...,
         is_local: builtins.bool = ...,
         run_id: builtins.str = ...,
+        workflow_task: global___WorkflowTaskAttempt | None = ...,
     ) -> None: ...
     def HasField(
         self,
@@ -280,6 +287,8 @@ class Start(google.protobuf.message.Message):
             b"started_time",
             "workflow_execution",
             b"workflow_execution",
+            "workflow_task",
+            b"workflow_task",
         ],
     ) -> builtins.bool: ...
     def ClearField(
@@ -321,13 +330,53 @@ class Start(google.protobuf.message.Message):
             b"workflow_execution",
             "workflow_namespace",
             b"workflow_namespace",
+            "workflow_task",
+            b"workflow_task",
             "workflow_type",
             b"workflow_type",
         ],
     ) -> None: ...
 
 global___Start = Start
 
+class WorkflowTaskAttempt(google.protobuf.message.Message):
+    """One attempt of a workflow task, with the values from the poll response that started it."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    STARTED_EVENT_ID_FIELD_NUMBER: builtins.int
+    ATTEMPT_FIELD_NUMBER: builtins.int
+    STARTED_TIME_FIELD_NUMBER: builtins.int
+    started_event_id: builtins.int
+    """The id of the task's WorkflowTaskStarted event."""
+    attempt: builtins.int
+    @property
+    def started_time(self) -> google.protobuf.timestamp_pb2.Timestamp:
+        """When the server started the task."""
+    def __init__(
+        self,
+        *,
+        started_event_id: builtins.int = ...,
+        attempt: builtins.int = ...,
+        started_time: google.protobuf.timestamp_pb2.Timestamp | None = ...,
+    ) -> None: ...
+    def HasField(
+        self, field_name: typing_extensions.Literal["started_time", b"started_time"]
+    ) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing_extensions.Literal[
+            "attempt",
+            b"attempt",
+            "started_event_id",
+            b"started_event_id",
+            "started_time",
+            b"started_time",
+        ],
+    ) -> None: ...
+
+global___WorkflowTaskAttempt = WorkflowTaskAttempt
+
 class Cancel(google.protobuf.message.Message):
     """Attempt to cancel a running activity"""
 

temporalio/bridge/src/extension.rs

@@ -0,0 +1,158 @@
+//! Worker extensions that this bridge build links, so Python can pass one to Core by its name.
+
+use std::sync::Arc;
+
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use temporalio_sdk_core::WorkerExtension;
+
+use crate::runtime;
+
+/// 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>,
+}
+
+#[pymethods]
+impl WorkerExtensionRef {
+    #[getter]
+    fn name(&self) -> String {
+        self.extension.name().to_string()
+    }
+}
+
+/// Builds the extension with `name` from its encoded config. Only the extensions that this build
+/// links are known.
+// A release build links no extension yet, so it never reads the config.
+#[cfg_attr(not(debug_assertions), allow(unused_variables))]
+pub fn new_worker_extension(
+    runtime_ref: &runtime::RuntimeRef,
+    name: &str,
+    config: Vec<u8>,
+) -> PyResult<WorkerExtensionRef> {
+    let runtime = runtime_ref.runtime.clone();
+    runtime.assert_same_process("create a worker extension")?;
+    // An extension may start tasks on the runtime that it lives on.
+    let _guard = runtime.core.tokio_handle().enter();
+    let extension: Arc<dyn WorkerExtension> = match name {
+        #[cfg(debug_assertions)]
+        test_extension::NAME => Arc::new(test_extension::TestExtension::new(&config)?),
+        _ => {
+            return Err(PyValueError::new_err(format!(
+                "This build has no worker extension named {name}"
+            )))
+        }
+    };
+    Ok(WorkerExtensionRef { extension })
+}
+
+/// An extension for the SDK's own tests. Release builds leave it out.
+#[cfg(debug_assertions)]
+mod test_extension {
+    use std::sync::{
+        atomic::{AtomicU64, Ordering},
+        Arc,
+    };
+
+    use pyo3::exceptions::PyValueError;
+    use pyo3::PyResult;
+    use temporalio_sdk_core::{
+        ExtensionError, WorkerExtension, WorkflowTaskCompletionInfo, WorkflowTaskInterceptError,
+        WorkflowTaskInterceptor,
+    };
+
+    pub(super) const NAME: &str = "temporalio-test";
+
+    pub(super) struct TestExtension {
+        interceptor: Arc<CountingInterceptor>,
+    }
+
+    impl TestExtension {
+        /// The config is empty, or `fail_before_complete` to fail every workflow task.
+        pub(super) fn new(config: &[u8]) -> PyResult<Self> {
+            let fail_before_complete = match config {
+                b"" => false,
+                b"fail_before_complete" => true,
+                _ => return Err(PyValueError::new_err("Unknown test extension config")),
+            };
+            Ok(Self {
+                interceptor: Arc::new(CountingInterceptor {
+                    fail_before_complete,
+                    before: AtomicU64::new(0),
+                    after: AtomicU64::new(0),
+                    accepted: AtomicU64::new(0),
+                }),
+            })
+        }
+    }
+
+    #[async_trait::async_trait]
+    impl WorkerExtension for TestExtension {
+        fn name(&self) -> &str {
+            NAME
+        }
+
+        fn workflow_task_interceptor(&self) -> Option<Arc<dyn WorkflowTaskInterceptor>> {
+            Some(self.interceptor.clone())
+        }
+
+        fn call(&self, method: &str, input: &[u8]) -> Result<Vec<u8>, ExtensionError> {
+            match method {
+                "Echo" => Ok(input.to_vec()),
+                "Fail" => Err(ExtensionError::new(tonic::Code::NotFound, "test failure")
+                    .with_details(input.to_vec())),
+                // Comma separated, so a test reads it with no proto.
+                "Counts" => Ok(format!(
+                    "{},{},{}",
+                    self.interceptor.before.load(Ordering::SeqCst),
+                    self.interceptor.after.load(Ordering::SeqCst),
+                    self.interceptor.accepted.load(Ordering::SeqCst)
+                )
+                .into_bytes()),
+                _ => Err(ExtensionError::unknown_method(method)),
+            }
+        }
+
+        async fn call_async(
+            &self,
+            method: String,
+            input: Vec<u8>,
+        ) -> Result<Vec<u8>, ExtensionError> {
+            tokio::task::yield_now().await;
+            self.call(&method, &input)
+        }
+    }
+
+    struct CountingInterceptor {
+        fail_before_complete: bool,
+        before: AtomicU64,
+        after: AtomicU64,
+        accepted: AtomicU64,
+    }
+
+    #[async_trait::async_trait]
+    impl WorkflowTaskInterceptor for CountingInterceptor {
+        async fn before_complete(
+            &self,
+            _task: &WorkflowTaskCompletionInfo,
+        ) -> Result<(), WorkflowTaskInterceptError> {
+            self.before.fetch_add(1, Ordering::SeqCst);
+            if self.fail_before_complete {
+                return Err(WorkflowTaskInterceptError::new(
+                    "test extension refused the task",
+                ));
+            }
+            Ok(())
+        }
+
+        async fn after_complete(&self, _task: &WorkflowTaskCompletionInfo, accepted: bool) {
+            self.after.fetch_add(1, Ordering::SeqCst);
+            if accepted {
+                self.accepted.fetch_add(1, Ordering::SeqCst);
+            }
+        }
+    }
+}

temporalio/bridge/src/lib.rs

@@ -4,6 +4,7 @@ use pyo3::types::PyTuple;
 mod client;
 mod client_rpc_generated;
 mod envconfig;
+mod extension;
 mod metric;
 mod runtime;
 mod testing;
@@ -69,6 +70,8 @@ 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_class::<extension::WorkerExtensionRef>()?;
+    m.add_function(wrap_pyfunction!(new_worker_extension, m)?)?;
 
     // envconfig
     let envconfig_module = PyModule::new(py, "envconfig")?;
@@ -131,6 +134,15 @@ fn new_worker(
     worker::new_worker(runtime_ref, client, config)
 }
 
+#[pyfunction]
+fn new_worker_extension(
+    runtime_ref: &runtime::RuntimeRef,
+    name: &str,
+    config: Vec<u8>,
+) -> PyResult<extension::WorkerExtensionRef> {
+    extension::new_worker_extension(runtime_ref, name, config)
+}
+
 #[pyfunction]
 fn new_replay_worker<'a>(
     py: Python<'a>,