From 9bc0300a90ecbefaf9fb5bb8f901ec3aa0ed300b Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 20:21:58 -0700 Subject: [PATCH 1/2] Added a Worker option that passes a worker extension to Core. Core calls a WorkerExtension around each Workflow Task, and only native code can build one. So the bridge builds an extension by its name, and the Worker passes it to Core. The Core pin moves to the Brook v3 c4 head, which has the trait. --- temporalio/bridge/extension.py | 28 ++++ .../bridge/proto/activity_task/__init__.py | 2 + .../proto/activity_task/activity_task_pb2.py | 34 ++-- .../proto/activity_task/activity_task_pb2.pyi | 49 ++++++ temporalio/bridge/sdk-core | 2 +- temporalio/bridge/src/extension.rs | 158 ++++++++++++++++++ temporalio/bridge/src/lib.rs | 12 ++ temporalio/bridge/src/worker.rs | 12 ++ temporalio/bridge/worker.py | 2 + temporalio/worker/__init__.py | 2 + temporalio/worker/_extension.py | 52 ++++++ temporalio/worker/_replayer.py | 1 + temporalio/worker/_worker.py | 12 ++ tests/worker/test_extension.py | 68 ++++++++ 14 files changed, 423 insertions(+), 11 deletions(-) create mode 100644 temporalio/bridge/extension.py create mode 100644 temporalio/bridge/src/extension.rs create mode 100644 temporalio/worker/_extension.py create mode 100644 tests/worker/test_extension.py diff --git a/temporalio/bridge/extension.py b/temporalio/bridge/extension.py new file mode 100644 index 0000000000..5c45160776 --- /dev/null +++ b/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 diff --git a/temporalio/bridge/proto/activity_task/__init__.py b/temporalio/bridge/proto/activity_task/__init__.py index 779ebf3535..9e69a0f894 100644 --- a/temporalio/bridge/proto/activity_task/__init__.py +++ b/temporalio/bridge/proto/activity_task/__init__.py @@ -4,6 +4,7 @@ ActivityTask, Cancel, Start, + WorkflowTaskAttempt, ) __all__ = [ @@ -12,4 +13,5 @@ "ActivityTask", "Cancel", "Start", + "WorkflowTaskAttempt", ] diff --git a/temporalio/bridge/proto/activity_task/activity_task_pb2.py b/temporalio/bridge/proto/activity_task/activity_task_pb2.py index 0e09839dc9..f269eb0cac 100644 --- a/temporalio/bridge/proto/activity_task/activity_task_pb2.py +++ b/temporalio/bridge/proto/activity_task/activity_task_pb2.py @@ -26,7 +26,7 @@ ) DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n3temporal/sdk/core/activity_task/activity_task.proto\x12\x15\x63oresdk.activity_task\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a$temporal/api/common/v1/message.proto\x1a%temporal/sdk/core/common/common.proto"\x8d\x01\n\x0c\x41\x63tivityTask\x12\x12\n\ntask_token\x18\x01 \x01(\x0c\x12-\n\x05start\x18\x03 \x01(\x0b\x32\x1c.coresdk.activity_task.StartH\x00\x12/\n\x06\x63\x61ncel\x18\x04 \x01(\x0b\x32\x1d.coresdk.activity_task.CancelH\x00\x42\t\n\x07variant"\xb1\x07\n\x05Start\x12\x1a\n\x12workflow_namespace\x18\x01 \x01(\t\x12\x15\n\rworkflow_type\x18\x02 \x01(\t\x12\x45\n\x12workflow_execution\x18\x03 \x01(\x0b\x32).temporal.api.common.v1.WorkflowExecution\x12\x13\n\x0b\x61\x63tivity_id\x18\x04 \x01(\t\x12\x15\n\ractivity_type\x18\x05 \x01(\t\x12\x45\n\rheader_fields\x18\x06 \x03(\x0b\x32..coresdk.activity_task.Start.HeaderFieldsEntry\x12.\n\x05input\x18\x07 \x03(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12:\n\x11heartbeat_details\x18\x08 \x03(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x32\n\x0escheduled_time\x18\t \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x42\n\x1e\x63urrent_attempt_scheduled_time\x18\n \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x30\n\x0cstarted_time\x18\x0b \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07\x61ttempt\x18\x0c \x01(\r\x12<\n\x19schedule_to_close_timeout\x18\r \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x39\n\x16start_to_close_timeout\x18\x0e \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x34\n\x11heartbeat_timeout\x18\x0f \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x39\n\x0cretry_policy\x18\x10 \x01(\x0b\x32#.temporal.api.common.v1.RetryPolicy\x12\x32\n\x08priority\x18\x12 \x01(\x0b\x32 .temporal.api.common.v1.Priority\x12\x10\n\x08is_local\x18\x11 \x01(\x08\x12\x0e\n\x06run_id\x18\x13 \x01(\t\x1aT\n\x11HeaderFieldsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12.\n\x05value\x18\x02 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload:\x02\x38\x01"\x8a\x01\n\x06\x43\x61ncel\x12;\n\x06reason\x18\x01 \x01(\x0e\x32+.coresdk.activity_task.ActivityCancelReason\x12\x43\n\x07\x64\x65tails\x18\x02 \x01(\x0b\x32\x32.coresdk.activity_task.ActivityCancellationDetails"\xa0\x01\n\x1b\x41\x63tivityCancellationDetails\x12\x14\n\x0cis_not_found\x18\x01 \x01(\x08\x12\x14\n\x0cis_cancelled\x18\x02 \x01(\x08\x12\x11\n\tis_paused\x18\x03 \x01(\x08\x12\x14\n\x0cis_timed_out\x18\x04 \x01(\x08\x12\x1a\n\x12is_worker_shutdown\x18\x05 \x01(\x08\x12\x10\n\x08is_reset\x18\x06 \x01(\x08*o\n\x14\x41\x63tivityCancelReason\x12\r\n\tNOT_FOUND\x10\x00\x12\r\n\tCANCELLED\x10\x01\x12\r\n\tTIMED_OUT\x10\x02\x12\x13\n\x0fWORKER_SHUTDOWN\x10\x03\x12\n\n\x06PAUSED\x10\x04\x12\t\n\x05RESET\x10\x05\x42\x32\xea\x02/Temporalio::Internal::Bridge::Api::ActivityTaskb\x06proto3' + b'\n3temporal/sdk/core/activity_task/activity_task.proto\x12\x15\x63oresdk.activity_task\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a$temporal/api/common/v1/message.proto\x1a%temporal/sdk/core/common/common.proto"\x8d\x01\n\x0c\x41\x63tivityTask\x12\x12\n\ntask_token\x18\x01 \x01(\x0c\x12-\n\x05start\x18\x03 \x01(\x0b\x32\x1c.coresdk.activity_task.StartH\x00\x12/\n\x06\x63\x61ncel\x18\x04 \x01(\x0b\x32\x1d.coresdk.activity_task.CancelH\x00\x42\t\n\x07variant"\xf4\x07\n\x05Start\x12\x1a\n\x12workflow_namespace\x18\x01 \x01(\t\x12\x15\n\rworkflow_type\x18\x02 \x01(\t\x12\x45\n\x12workflow_execution\x18\x03 \x01(\x0b\x32).temporal.api.common.v1.WorkflowExecution\x12\x13\n\x0b\x61\x63tivity_id\x18\x04 \x01(\t\x12\x15\n\ractivity_type\x18\x05 \x01(\t\x12\x45\n\rheader_fields\x18\x06 \x03(\x0b\x32..coresdk.activity_task.Start.HeaderFieldsEntry\x12.\n\x05input\x18\x07 \x03(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12:\n\x11heartbeat_details\x18\x08 \x03(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x32\n\x0escheduled_time\x18\t \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x42\n\x1e\x63urrent_attempt_scheduled_time\x18\n \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x30\n\x0cstarted_time\x18\x0b \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07\x61ttempt\x18\x0c \x01(\r\x12<\n\x19schedule_to_close_timeout\x18\r \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x39\n\x16start_to_close_timeout\x18\x0e \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x34\n\x11heartbeat_timeout\x18\x0f \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x39\n\x0cretry_policy\x18\x10 \x01(\x0b\x32#.temporal.api.common.v1.RetryPolicy\x12\x32\n\x08priority\x18\x12 \x01(\x0b\x32 .temporal.api.common.v1.Priority\x12\x10\n\x08is_local\x18\x11 \x01(\x08\x12\x0e\n\x06run_id\x18\x13 \x01(\t\x12\x41\n\rworkflow_task\x18\x14 \x01(\x0b\x32*.coresdk.activity_task.WorkflowTaskAttempt\x1aT\n\x11HeaderFieldsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12.\n\x05value\x18\x02 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload:\x02\x38\x01"r\n\x13WorkflowTaskAttempt\x12\x18\n\x10started_event_id\x18\x01 \x01(\x03\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\r\x12\x30\n\x0cstarted_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp"\x8a\x01\n\x06\x43\x61ncel\x12;\n\x06reason\x18\x01 \x01(\x0e\x32+.coresdk.activity_task.ActivityCancelReason\x12\x43\n\x07\x64\x65tails\x18\x02 \x01(\x0b\x32\x32.coresdk.activity_task.ActivityCancellationDetails"\xa0\x01\n\x1b\x41\x63tivityCancellationDetails\x12\x14\n\x0cis_not_found\x18\x01 \x01(\x08\x12\x14\n\x0cis_cancelled\x18\x02 \x01(\x08\x12\x11\n\tis_paused\x18\x03 \x01(\x08\x12\x14\n\x0cis_timed_out\x18\x04 \x01(\x08\x12\x1a\n\x12is_worker_shutdown\x18\x05 \x01(\x08\x12\x10\n\x08is_reset\x18\x06 \x01(\x08*o\n\x14\x41\x63tivityCancelReason\x12\r\n\tNOT_FOUND\x10\x00\x12\r\n\tCANCELLED\x10\x01\x12\r\n\tTIMED_OUT\x10\x02\x12\x13\n\x0fWORKER_SHUTDOWN\x10\x03\x12\n\n\x06PAUSED\x10\x04\x12\t\n\x05RESET\x10\x05\x42\x32\xea\x02/Temporalio::Internal::Bridge::Api::ActivityTaskb\x06proto3' ) _ACTIVITYCANCELREASON = DESCRIPTOR.enum_types_by_name["ActivityCancelReason"] @@ -42,6 +42,7 @@ _ACTIVITYTASK = DESCRIPTOR.message_types_by_name["ActivityTask"] _START = DESCRIPTOR.message_types_by_name["Start"] _START_HEADERFIELDSENTRY = _START.nested_types_by_name["HeaderFieldsEntry"] +_WORKFLOWTASKATTEMPT = DESCRIPTOR.message_types_by_name["WorkflowTaskAttempt"] _CANCEL = DESCRIPTOR.message_types_by_name["Cancel"] _ACTIVITYCANCELLATIONDETAILS = DESCRIPTOR.message_types_by_name[ "ActivityCancellationDetails" @@ -78,6 +79,17 @@ _sym_db.RegisterMessage(Start) _sym_db.RegisterMessage(Start.HeaderFieldsEntry) +WorkflowTaskAttempt = _reflection.GeneratedProtocolMessageType( + "WorkflowTaskAttempt", + (_message.Message,), + { + "DESCRIPTOR": _WORKFLOWTASKATTEMPT, + "__module__": "temporal.sdk.core.activity_task.activity_task_pb2", + # @@protoc_insertion_point(class_scope:coresdk.activity_task.WorkflowTaskAttempt) + }, +) +_sym_db.RegisterMessage(WorkflowTaskAttempt) + Cancel = _reflection.GeneratedProtocolMessageType( "Cancel", (_message.Message,), @@ -107,16 +119,18 @@ ) _START_HEADERFIELDSENTRY._options = None _START_HEADERFIELDSENTRY._serialized_options = b"8\001" - _ACTIVITYCANCELREASON._serialized_start = 1616 - _ACTIVITYCANCELREASON._serialized_end = 1727 + _ACTIVITYCANCELREASON._serialized_start = 1799 + _ACTIVITYCANCELREASON._serialized_end = 1910 _ACTIVITYTASK._serialized_start = 221 _ACTIVITYTASK._serialized_end = 362 _START._serialized_start = 365 - _START._serialized_end = 1310 - _START_HEADERFIELDSENTRY._serialized_start = 1226 - _START_HEADERFIELDSENTRY._serialized_end = 1310 - _CANCEL._serialized_start = 1313 - _CANCEL._serialized_end = 1451 - _ACTIVITYCANCELLATIONDETAILS._serialized_start = 1454 - _ACTIVITYCANCELLATIONDETAILS._serialized_end = 1614 + _START._serialized_end = 1377 + _START_HEADERFIELDSENTRY._serialized_start = 1293 + _START_HEADERFIELDSENTRY._serialized_end = 1377 + _WORKFLOWTASKATTEMPT._serialized_start = 1379 + _WORKFLOWTASKATTEMPT._serialized_end = 1493 + _CANCEL._serialized_start = 1496 + _CANCEL._serialized_end = 1634 + _ACTIVITYCANCELLATIONDETAILS._serialized_start = 1637 + _ACTIVITYCANCELLATIONDETAILS._serialized_end = 1797 # @@protoc_insertion_point(module_scope) diff --git a/temporalio/bridge/proto/activity_task/activity_task_pb2.pyi b/temporalio/bridge/proto/activity_task/activity_task_pb2.pyi index 1a2f8c7d2f..dead4d805e 100644 --- a/temporalio/bridge/proto/activity_task/activity_task_pb2.pyi +++ b/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,6 +330,8 @@ class Start(google.protobuf.message.Message): b"workflow_execution", "workflow_namespace", b"workflow_namespace", + "workflow_task", + b"workflow_task", "workflow_type", b"workflow_type", ], @@ -328,6 +339,44 @@ class Start(google.protobuf.message.Message): 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""" diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index 0a5d646279..d05afbadca 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit 0a5d646279824adcf07eeaa52d9179dcd934ab46 +Subproject commit d05afbadcae537e32cc51439cda177c5e5433228 diff --git a/temporalio/bridge/src/extension.rs b/temporalio/bridge/src/extension.rs new file mode 100644 index 0000000000..cce72dce16 --- /dev/null +++ b/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, +} + +#[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, +) -> PyResult { + 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 = 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, + } + + impl TestExtension { + /// The config is empty, or `fail_before_complete` to fail every workflow task. + pub(super) fn new(config: &[u8]) -> PyResult { + 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> { + Some(self.interceptor.clone()) + } + + fn call(&self, method: &str, input: &[u8]) -> Result, 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, + ) -> Result, 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); + } + } + } +} diff --git a/temporalio/bridge/src/lib.rs b/temporalio/bridge/src/lib.rs index a4fda26db6..ae2eed5860 100644 --- a/temporalio/bridge/src/lib.rs +++ b/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::()?; m.add_function(wrap_pyfunction!(new_worker, m)?)?; m.add_function(wrap_pyfunction!(new_replay_worker, m)?)?; + m.add_class::()?; + 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, +) -> PyResult { + extension::new_worker_extension(runtime_ref, name, config) +} + #[pyfunction] fn new_replay_worker<'a>( py: Python<'a>, diff --git a/temporalio/bridge/src/worker.rs b/temporalio/bridge/src/worker.rs index e74698e944..f838a0384e 100644 --- a/temporalio/bridge/src/worker.rs +++ b/temporalio/bridge/src/worker.rs @@ -26,6 +26,7 @@ use tokio_stream::wrappers::ReceiverStream; use tracing::error; use crate::client; +use crate::extension::WorkerExtensionRef; use crate::runtime; pyo3::create_exception!(temporal_sdk_bridge, PollShutdownError, PyException); @@ -65,6 +66,16 @@ pub struct WorkerConfig { plugins: Vec, storage_drivers: HashSet, disable_payload_error_limit: bool, + #[pyo3(from_py_with = extract_extension)] + extension: Option, +} + +/// Python wraps the extension, so the config has the wrapper. +fn extract_extension(obj: &Bound<'_, PyAny>) -> PyResult> { + if obj.is_none() { + return Ok(None); + } + obj.getattr("_ref")?.extract().map(Some).map_err(Into::into) } #[derive(FromPyObject)] @@ -780,6 +791,7 @@ fn convert_worker_config( .collect::>(), ) .disable_payload_error_limit(conf.disable_payload_error_limit) + .maybe_extension(conf.extension.map(|ext| ext.extension)) .build() .map_err(|err| PyValueError::new_err(format!("Invalid worker config: {err}"))) } diff --git a/temporalio/bridge/worker.py b/temporalio/bridge/worker.py index 4b7f55d092..0243e3627a 100644 --- a/temporalio/bridge/worker.py +++ b/temporalio/bridge/worker.py @@ -12,6 +12,7 @@ ) import temporalio.bridge.client +import temporalio.bridge.extension import temporalio.bridge.proto import temporalio.bridge.proto.activity_task import temporalio.bridge.proto.nexus @@ -60,6 +61,7 @@ class WorkerConfig: plugins: Sequence[str] storage_drivers: set[str] disable_payload_error_limit: bool + extension: temporalio.bridge.extension.WorkerExtension | None @dataclass diff --git a/temporalio/worker/__init__.py b/temporalio/worker/__init__.py index 4f6efe68c3..061175df8f 100644 --- a/temporalio/worker/__init__.py +++ b/temporalio/worker/__init__.py @@ -2,6 +2,7 @@ from ..common import WorkerDeploymentVersion from ._activity import SharedHeartbeatSender, SharedStateManager +from ._extension import WorkerExtension from ._interceptor import ( ActivityInboundInterceptor, ActivityOutboundInterceptor, @@ -79,6 +80,7 @@ "PollerBehaviorSimpleMaximum", "PollerBehaviorAutoscaling", "PatchActivationInput", + "WorkerExtension", # Interceptor base classes "Interceptor", "ActivityInboundInterceptor", diff --git a/temporalio/worker/_extension.py b/temporalio/worker/_extension.py new file mode 100644 index 0000000000..cbf218491c --- /dev/null +++ b/temporalio/worker/_extension.py @@ -0,0 +1,52 @@ +"""Native extensions that a Worker passes to Core.""" + +from __future__ import annotations + +import temporalio.bridge.extension +import temporalio.runtime + + +class WorkerExtension: + """A native extension that the SDK's bridge build links, which a Worker passes to Core. + + Core calls the extension at fixed points of each Workflow Task. A plugin builds the extension + and passes it in the Worker's ``extension`` option. The Worker must use the extension's + runtime. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + name: str, + config: bytes = b"", + *, + runtime: temporalio.runtime.Runtime | None = None, + ) -> None: + """Build the extension with a name that the bridge build knows. + + Args: + name: The name of the extension in the bridge build. + config: The extension's encoded config. Each extension has its own format. + runtime: The runtime the extension runs on. Defaults to + :py:meth:`temporalio.runtime.Runtime.default`. + + Raises: + ValueError: The bridge build has no extension with this name, or the config is not + valid. + """ + self._runtime = runtime or temporalio.runtime.Runtime.default() + self._bridge_extension = temporalio.bridge.extension.WorkerExtension( + self._runtime._core_runtime, name, config + ) + + @property + def name(self) -> str: + """Name the extension gives itself.""" + return self._bridge_extension.name + + @property + def runtime(self) -> temporalio.runtime.Runtime: + """Runtime the extension runs on.""" + return self._runtime diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index b3eb1a4d12..036cc42053 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -326,6 +326,7 @@ def on_eviction_hook( nonsticky_to_sticky_poll_ratio=1, no_remote_activities=True, disable_payload_error_limit=True, + extension=None, task_types=temporalio.bridge.worker.WorkerTaskTypes( enable_workflows=True, enable_local_activities=False, diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 3c6effd80e..adb638feea 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -31,6 +31,7 @@ ) from ._activity import SharedStateManager, _ActivityWorker +from ._extension import WorkerExtension from ._interceptor import Interceptor from ._nexus import _NexusWorker from ._plugin import Plugin @@ -152,6 +153,7 @@ def __init__( ), disable_payload_error_limit: bool = False, max_workflow_task_external_storage_concurrency: int = _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY, + extension: WorkerExtension | None = None, ) -> None: """Create a worker to process workflows and/or activities. @@ -345,6 +347,10 @@ def __init__( Please report any issues you encounter with this setting or if you feel the default should be changed. WARNING: This setting is experimental. + extension: A native extension that Core calls at fixed points of each + Workflow Task. It must use the same runtime as the client. Plugins + usually set it. + WARNING: This setting is experimental. """ config = WorkerConfig( @@ -393,6 +399,7 @@ def __init__( nexus_task_poller_behavior=nexus_task_poller_behavior, disable_payload_error_limit=disable_payload_error_limit, max_workflow_task_external_storage_concurrency=max_workflow_task_external_storage_concurrency, + extension=extension, ) plugins_from_client = cast( @@ -495,6 +502,9 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf or bridge_client.config.runtime or temporalio.runtime.Runtime.default() ) + extension = config.get("extension") + if extension is not None and extension.runtime is not self._runtime: + raise ValueError("Worker extension must use the same runtime as the client") activities = config.get("activities") if activities: _warn_if_activity_executor_max_workers_is_inconsistent(config) @@ -707,6 +717,7 @@ def check_activity(activity: str): disable_payload_error_limit=config.get( "disable_payload_error_limit", False ), + extension=extension._bridge_extension if extension else None, ), ) @@ -1030,6 +1041,7 @@ class WorkerConfig(TypedDict, total=False): nexus_task_poller_behavior: PollerBehavior disable_payload_error_limit: bool max_workflow_task_external_storage_concurrency: int + extension: WorkerExtension | None def _warn_if_activity_executor_max_workers_is_inconsistent( diff --git a/tests/worker/test_extension.py b/tests/worker/test_extension.py new file mode 100644 index 0000000000..1b9ca9ac7a --- /dev/null +++ b/tests/worker/test_extension.py @@ -0,0 +1,68 @@ +import uuid + +import pytest + +from temporalio import workflow +from temporalio.client import Client +from temporalio.runtime import Runtime, TelemetryConfig +from temporalio.worker import WorkerExtension +from tests.helpers import assert_task_fail_eventually, new_worker + +# Debug bridge builds link this extension for the SDK's tests. +TEST_EXTENSION = "temporalio-test" + + +def load_test_extension(config: bytes = b"") -> WorkerExtension: + try: + return WorkerExtension(TEST_EXTENSION, config) + except ValueError: + pytest.skip("Bridge build has no test extension") + + +@workflow.defn +class SimpleWorkflow: + @workflow.run + async def run(self) -> str: + return "done" + + +def test_unknown_extension_name(): + with pytest.raises(ValueError, match="no worker extension named not-linked"): + WorkerExtension("not-linked") + + +def test_extension_name(): + assert load_test_extension().name == TEST_EXTENSION + + +async def test_extension_lets_workflow_task_complete(client: Client): + async with new_worker( + client, SimpleWorkflow, extension=load_test_extension() + ) as worker: + result = await client.execute_workflow( + SimpleWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == "done" + + +async def test_extension_failure_fails_workflow_task(client: Client): + extension = load_test_extension(b"fail_before_complete") + async with new_worker(client, SimpleWorkflow, extension=extension) as worker: + handle = await client.start_workflow( + SimpleWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await assert_task_fail_eventually( + handle, message_contains="test extension refused the task" + ) + + +async def test_extension_must_use_client_runtime(client: Client): + load_test_extension() + other = Runtime(telemetry=TelemetryConfig()) + extension = WorkerExtension(TEST_EXTENSION, runtime=other) + with pytest.raises(ValueError, match="same runtime as the client"): + new_worker(client, SimpleWorkflow, extension=extension) From b86b6794d21d146bd778f81b21793fab49f9aa89 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 21:11:58 -0700 Subject: [PATCH 2/2] Moved the Core pin to the final c4 head. Lane C's review fixes moved c4. The crate pins the same sdk-rust commit, and the two must match. --- temporalio/bridge/sdk-core | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index d05afbadca..3a9203ac7c 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit d05afbadcae537e32cc51439cda177c5e5433228 +Subproject commit 3a9203ac7cb15171097d7f58e80b648f517ea658