Merge remote-tracking branch 'origin/main' into fix-swallowed-cancel

Author: yegorske50

.gitignore

@@ -2,6 +2,7 @@
 __pycache__
 /build
 /dist
+temporalio/bridge/libtemporal_sdk_bridge.dylib.dSYM/
 temporalio/bridge/target/
 temporalio/bridge/temporal_sdk_bridge*
 /tests/helpers/golangserver/golangserver

README.md

@@ -2059,28 +2059,33 @@ The environment is now ready to develop in.
 
 #### Testing
 
-To execute tests:
+To execute tests (in parallel if possible):
 
 ```bash
 poe test
 ```
 
-`poe test` spreads tests across multiple worker processes by default. If you
-need a serial run for debugging, invoke pytest directly:
+To execute tests serially:
 
 ```bash
 uv run pytest
 ```
 
-This runs against [Temporalite](https://github.com/temporalio/temporalite). To run against the time-skipping test
-server, pass `--workflow-environment time-skipping`. To run against the `default` namespace of an already-running
-server, pass the `host:port` to `--workflow-environment`. Can also use regular pytest arguments. For example, here's how
-to run a single test with debug logs on the console:
+To execute a single test:
 
 ```bash
 poe test -s --log-cli-level=DEBUG -k test_sync_activity_thread_cancel_caught
 ```
 
+**Temporal Server**
+
+- Tests that use the workflow test environment run against the [Temporal CLI dev server](https://docs.temporal.io/cli#start-dev-server).
+- By default, workflow-environment tests automatically start a local dev server.
+- On first run, the dev server binary may be downloaded so network access is required if no server is currently running.
+- To run workflow-environment tests against the time-skipping test server, pass `--workflow-environment time-skipping`. 
+- To run workflow-environment tests against the `default` namespace of an already-running server, pass the `host:port` to `--workflow-environment`.
+- Unit tests that do not use the workflow environment do not start a dev server.
+
 #### Proto Generation and Testing
 
 If you have docker available, run

temporalio/contrib/google_adk_agents/_model.py

@@ -1,4 +1,4 @@
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Callable
 from datetime import timedelta
 
 from google.adk.models import BaseLlm, LLMRegistry
@@ -40,20 +40,37 @@ class TemporalModel(BaseLlm):
     """A Temporal-based LLM model that executes model invocations as activities."""
 
     def __init__(
-        self, model_name: str, activity_config: ActivityConfig | None = None
+        self,
+        model_name: str,
+        activity_config: ActivityConfig | None = None,
+        *,
+        summary_fn: Callable[[LlmRequest], str | None] | None = None,
     ) -> None:
         """Initialize the TemporalModel.
 
         Args:
             model_name: The name of the model to use.
             activity_config: Configuration options for the activity execution.
+            summary_fn: Optional callable that receives the LlmRequest and
+                returns a summary string (or None) for the activity. Must be
+                deterministic as it is called during workflow execution. If
+                the callable raises, the exception will propagate and fail
+                the workflow task.
+
+        Raises:
+            ValueError: If both ``ActivityConfig["summary"]`` and ``summary_fn`` are set.
         """
         super().__init__(model=model_name)
         self._model_name = model_name
+        self._summary_fn = summary_fn
         self._activity_config = ActivityConfig(
             start_to_close_timeout=timedelta(seconds=60)
         )
-        if activity_config:
+        if activity_config is not None:
+            if summary_fn is not None and activity_config.get("summary") is not None:
+                raise ValueError(
+                    "Cannot specify both ActivityConfig 'summary' and 'summary_fn'"
+                )
             self._activity_config.update(activity_config)
 
     async def generate_content_async(
@@ -76,10 +93,20 @@ async def generate_content_async(
                 yield response
             return
 
+        config = self._activity_config.copy()
+        if self._summary_fn is not None:
+            summary = self._summary_fn(llm_request)
+            if summary is not None:
+                config["summary"] = summary
+        elif "summary" not in config:
+            if llm_request.config and llm_request.config.labels:
+                agent_name = llm_request.config.labels.get("adk_agent_name")
+                if agent_name:
+                    config["summary"] = agent_name
         responses = await workflow.execute_activity(
             invoke_model,
             args=[llm_request],
-            **self._activity_config,
+            **config,
         )
         for response in responses:
             yield response

temporalio/contrib/opentelemetry/_interceptor.py

@@ -341,6 +341,60 @@ async def start_update_with_start_workflow(
 
             return await super().start_update_with_start_workflow(input)
 
+    async def start_activity(
+        self, input: temporalio.client.StartActivityInput
+    ) -> temporalio.client.ActivityHandle[Any]:
+        with self.root._start_as_current_span(
+            f"StartActivity:{input.activity_type}",
+            attributes={
+                "temporalActivityID": input.id,
+                "temporalActivityType": input.activity_type,
+            },
+            input_with_headers=input,
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            return await super().start_activity(input)
+
+    async def cancel_activity(
+        self, input: temporalio.client.CancelActivityInput
+    ) -> None:
+        with self.root._start_as_current_span(
+            "CancelActivity",
+            attributes={"temporalActivityID": input.activity_id},
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            return await super().cancel_activity(input)
+
+    async def terminate_activity(
+        self, input: temporalio.client.TerminateActivityInput
+    ) -> None:
+        with self.root._start_as_current_span(
+            "TerminateActivity",
+            attributes={"temporalActivityID": input.activity_id},
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            return await super().terminate_activity(input)
+
+    async def describe_activity(
+        self, input: temporalio.client.DescribeActivityInput
+    ) -> temporalio.client.ActivityExecutionDescription:
+        with self.root._start_as_current_span(
+            "DescribeActivity",
+            attributes={"temporalActivityID": input.activity_id},
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            return await super().describe_activity(input)
+
+    async def count_activities(
+        self, input: temporalio.client.CountActivitiesInput
+    ) -> temporalio.client.ActivityExecutionCount:
+        with self.root._start_as_current_span(
+            "CountActivities",
+            attributes={},
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            return await super().count_activities(input)
+
 
 class _TracingActivityInboundInterceptor(temporalio.worker.ActivityInboundInterceptor):
     def __init__(

temporalio/contrib/opentelemetry/_otel_interceptor.py

@@ -303,6 +303,70 @@ async def start_update_with_start_workflow(
             )
             return await super().start_update_with_start_workflow(input)
 
+    async def start_activity(
+        self, input: temporalio.client.StartActivityInput
+    ) -> temporalio.client.ActivityHandle[Any]:
+        with _maybe_span(
+            get_tracer(__name__),
+            f"StartActivity:{input.activity_type}",
+            add_temporal_spans=self._add_temporal_spans,
+            attributes={
+                "temporalActivityID": input.id,
+                "temporalActivityType": input.activity_type,
+            },
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            input.headers = _context_to_headers(input.headers)
+            return await super().start_activity(input)
+
+    async def cancel_activity(
+        self, input: temporalio.client.CancelActivityInput
+    ) -> None:
+        with _maybe_span(
+            get_tracer(__name__),
+            "CancelActivity",
+            add_temporal_spans=self._add_temporal_spans,
+            attributes={"temporalActivityID": input.activity_id},
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            return await super().cancel_activity(input)
+
+    async def terminate_activity(
+        self, input: temporalio.client.TerminateActivityInput
+    ) -> None:
+        with _maybe_span(
+            get_tracer(__name__),
+            "TerminateActivity",
+            add_temporal_spans=self._add_temporal_spans,
+            attributes={"temporalActivityID": input.activity_id},
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            return await super().terminate_activity(input)
+
+    async def describe_activity(
+        self, input: temporalio.client.DescribeActivityInput
+    ) -> temporalio.client.ActivityExecutionDescription:
+        with _maybe_span(
+            get_tracer(__name__),
+            "DescribeActivity",
+            add_temporal_spans=self._add_temporal_spans,
+            attributes={"temporalActivityID": input.activity_id},
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            return await super().describe_activity(input)
+
+    async def count_activities(
+        self, input: temporalio.client.CountActivitiesInput
+    ) -> temporalio.client.ActivityExecutionCount:
+        with _maybe_span(
+            get_tracer(__name__),
+            "CountActivities",
+            add_temporal_spans=self._add_temporal_spans,
+            attributes={},
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+        ):
+            return await super().count_activities(input)
+
 
 class _TracingActivityInboundInterceptor(temporalio.worker.ActivityInboundInterceptor):
     def __init__(

temporalio/converter/_extstore.py

@@ -38,11 +38,17 @@ class StorageOperationMetrics:
     total_duration: timedelta = dataclasses.field(default_factory=timedelta)
     """Wall-clock time spent on external storage operations."""
 
-    def record_batch(self, count: int, size: int, duration: timedelta) -> None:
+    driver_names: set[str] = dataclasses.field(default_factory=set)
+    """Names of the drivers that participated in the operations."""
+
+    def record_batch(
+        self, count: int, size: int, duration: timedelta, driver_names: set[str]
+    ) -> None:
         """Record metrics from a batch of storage operations."""
         self.payload_count += count
         self.total_size += size
         self.total_duration += duration
+        self.driver_names.update(driver_names)
 
     @contextlib.contextmanager
     def track(self) -> Generator[Self, None, None]:
@@ -362,7 +368,7 @@ async def _store_payload(self, payload: Payload) -> Payload:
             )
         reference_payload.external_payloads.add().size_bytes = external_size
 
-        ExternalStorage._record_metrics(1, external_size, start_time)
+        ExternalStorage._record_metrics(1, external_size, start_time, {driver.name()})
 
         return reference_payload
 
@@ -407,6 +413,7 @@ async def _store_payload_sequence(
 
         external_count = 0
         external_size = 0
+        driver_names: set[str] = set()
         for (driver, indexed_payloads), claims in zip(driver_group_list, all_claims):
             indices = [idx for idx, _ in indexed_payloads]
             sizes = [p.ByteSize() for _, p in indexed_payloads]
@@ -428,8 +435,11 @@ async def _store_payload_sequence(
                 external_size += sizes[i]
 
             external_count += len(claims)
+            driver_names.add(driver.name())
 
-        ExternalStorage._record_metrics(external_count, external_size, start_time)
+        ExternalStorage._record_metrics(
+            external_count, external_size, start_time, driver_names
+        )
 
         return results
 
@@ -452,7 +462,9 @@ async def _retrieve_payload(self, payload: Payload) -> Payload:
 
         stored_payload = stored_payloads[0]
 
-        ExternalStorage._record_metrics(1, stored_payload.ByteSize(), start_time)
+        ExternalStorage._record_metrics(
+            1, stored_payload.ByteSize(), start_time, {driver.name()}
+        )
 
         return stored_payload
 
@@ -501,6 +513,7 @@ async def _retrieve_payload_sequence(
 
         external_count = 0
         external_size = 0
+        driver_names: set[str] = set()
         for (driver, indexed_claims), stored_payloads in zip(
             driver_claim_list, all_stored
         ):
@@ -517,14 +530,17 @@ async def _retrieve_payload_sequence(
                 external_size += stored_payload.ByteSize()
 
             external_count += len(stored_payloads)
+            driver_names.add(driver.name())
 
         retrieve_indices = sorted(stored_by_index.keys())
         stored_list = [stored_by_index[idx] for idx in retrieve_indices]
 
         for i, retrieved_payload in enumerate(stored_list):
             results[retrieve_indices[i]] = retrieved_payload
 
-        ExternalStorage._record_metrics(external_count, external_size, start_time)
+        ExternalStorage._record_metrics(
+            external_count, external_size, start_time, driver_names
+        )
 
         return results
 
@@ -545,9 +561,14 @@ def _validate_payload_length(
             )
 
     @staticmethod
-    def _record_metrics(count: int, size: int, start_time: float):
+    def _record_metrics(
+        count: int, size: int, start_time: float, driver_names: set[str]
+    ):
         metrics = _current_storage_metrics.get()
         if metrics is not None:
             metrics.record_batch(
-                count, size, timedelta(seconds=time.monotonic() - start_time)
+                count,
+                size,
+                timedelta(seconds=time.monotonic() - start_time),
+                driver_names,
             )