fix(temporal): send activity heartbeats from inside activities

heartbeat_if_in_workflow() only called activity.heartbeat() when
workflow.in_workflow() was true. activity.heartbeat() is only legal
inside an activity, where workflow.in_workflow() is false, so the guard
never passed where it mattered. Every ADK service that calls it (tasks,
messages, streaming, the LiteLLM/OpenAI/SGP providers, ACP, tracing,
templating) runs inside an activity and never heartbeated.

The shipped defaults hide this because each provider sets
heartbeat_timeout equal to start_to_close_timeout. An agent that sets a
shorter heartbeat_timeout for a long activity, the usual Temporal
pattern, gets ActivityTaskTimedOut (HEARTBEAT) on every attempt, and
activity cancellation is never delivered through the heartbeat channel.

Add heartbeat_if_in_activity(), which heartbeats when
activity.in_activity() is true and is a no-op everywhere else, and keep
heartbeat_if_in_workflow() as an alias so the existing call sites work
unchanged.

Verified with temporalio.testing.ActivityEnvironment: inside an activity
the old helper records no heartbeat, the new one (and the alias) records
one. tests/lib/test_temporal_utils.py, tests/lib/core/services and
tests/lib/core/temporal: 74 passed.

Author: michaelxu2288

src/agentex/lib/utils/temporal.py

@@ -12,11 +12,26 @@ def in_temporal_workflow():
         return False
 
 
-def heartbeat_if_in_workflow(heartbeat_name: str):
-    if in_temporal_workflow():
+def heartbeat_if_in_activity(heartbeat_name: str) -> None:
+    """Records a Temporal activity heartbeat when called from inside an activity.
+
+    ``activity.heartbeat`` is only legal inside an activity, so this stays a
+    silent no-op everywhere else: workflow code and sync agents reach the same
+    shared services and must not raise there.
+    """
+    if activity.in_activity():
         activity.heartbeat(heartbeat_name)
 
 
+def heartbeat_if_in_workflow(heartbeat_name: str) -> None:
+    """Deprecated alias for :func:`heartbeat_if_in_activity`.
+
+    Kept so the existing call sites keep working; prefer the activity-named
+    helper in new code.
+    """
+    heartbeat_if_in_activity(heartbeat_name)
+
+
 def workflow_now_if_in_workflow() -> datetime | None:
     # Returns Temporal's deterministic workflow clock when called from inside a
     # workflow, otherwise None. Used to stamp messages with a monotonic

tests/lib/test_temporal_utils.py

@@ -2,12 +2,18 @@
 
 from __future__ import annotations
 
+from typing import Any
 from datetime import datetime
 from unittest.mock import patch
 
+from temporalio import activity
+from temporalio.testing import ActivityEnvironment
+
 from agentex.lib.utils import temporal as _temporal_mod
 from agentex.lib.utils.temporal import (
     in_temporal_workflow,
+    heartbeat_if_in_activity,
+    heartbeat_if_in_workflow,
     workflow_now_if_in_workflow,
 )
 
@@ -18,6 +24,45 @@ def test_in_temporal_workflow_returns_false_outside_workflow() -> None:
     assert in_temporal_workflow() is False
 
 
+async def test_heartbeat_if_in_activity_heartbeats_inside_activity() -> None:
+    """The helper must reach Temporal's heartbeat channel from inside an activity."""
+    recorded: list[tuple[Any, ...]] = []
+
+    @activity.defn(name="heartbeating_activity")
+    async def heartbeating_activity() -> None:
+        heartbeat_if_in_activity("doing slow work")
+
+    env = ActivityEnvironment()
+    env.on_heartbeat = lambda *details: recorded.append(details)
+
+    await env.run(heartbeating_activity)
+
+    assert recorded == [("doing slow work",)]
+
+
+async def test_heartbeat_if_in_workflow_alias_heartbeats_inside_activity() -> None:
+    """The legacy name stays wired to the same behaviour for existing call sites."""
+    recorded: list[tuple[Any, ...]] = []
+
+    @activity.defn(name="legacy_heartbeating_activity")
+    async def legacy_heartbeating_activity() -> None:
+        heartbeat_if_in_workflow("doing slow work")
+
+    env = ActivityEnvironment()
+    env.on_heartbeat = lambda *details: recorded.append(details)
+
+    await env.run(legacy_heartbeating_activity)
+
+    assert recorded == [("doing slow work",)]
+
+
+def test_heartbeat_if_in_activity_is_noop_outside_activity() -> None:
+    """Sync agents and workflow code call the same helpers and must not raise there."""
+    assert activity.in_activity() is False
+    heartbeat_if_in_activity("doing slow work")
+    heartbeat_if_in_workflow("doing slow work")
+
+
 def test_workflow_now_if_in_workflow_returns_none_outside_workflow() -> None:
     assert workflow_now_if_in_workflow() is None