fix(harness): heartbeat while auto_send delivers a turn

The provider helpers heartbeat once when a streamed turn starts, then hand
the stream to UnifiedEmitter.auto_send_turn, which never heartbeated. A
streamed OpenAI turn (or any harness turn delivered from an activity) that
outlasts a short heartbeat_timeout was timed out and retried while it was
still sending messages. auto_send now calls heartbeat_if_in_activity for
each event it delivers; temporalio throttles the actual RPCs, and it is a
no-op outside an activity.

New test: auto_send inside a temporalio ActivityEnvironment records one
heartbeat per event (none before this change).

Author: michaelxu2288

src/agentex/lib/core/harness/auto_send.py

@@ -6,6 +6,7 @@
 from datetime import datetime
 
 from agentex.types.text_delta import TextDelta
+from agentex.lib.utils.temporal import heartbeat_if_in_activity
 from agentex.types.text_content import TextContent
 from agentex.lib.core.harness.types import TurnUsage, TurnResult, StreamTaskMessage
 from agentex.lib.core.harness.tracer import SpanTracer
@@ -86,6 +87,7 @@ async def _close_all() -> None:
 
     try:
         async for event in events:
+            heartbeat_if_in_activity("auto send")
             if deriver is not None and tracer is not None:
                 for signal in deriver.observe(event):
                     await tracer.handle(signal)

tests/lib/core/harness/test_auto_send.py

@@ -478,3 +478,35 @@ async def test_auto_send_created_at_forwarded():
     await auto_send(_gen(events), task_id="task1", tracer=None, streaming=streaming, created_at=dt)
 
     assert all(ts == dt for ts in streaming.recorded_created_at)
+
+
+@pytest.mark.asyncio
+async def test_auto_send_heartbeats_per_event_inside_an_activity():
+    """A long streamed turn must keep heartbeating, or a short heartbeat_timeout
+    times out an activity that is still delivering messages."""
+    from temporalio import activity
+    from temporalio.testing import ActivityEnvironment
+
+    streaming = _FakeStreaming()
+    events = [
+        StreamTaskMessageStart(type="start", index=0, content=TextContent(type="text", author="agent", content="")),
+        StreamTaskMessageDelta(type="delta", index=0, delta=TextDelta(type="text", text_delta="a")),
+        StreamTaskMessageDelta(type="delta", index=0, delta=TextDelta(type="text", text_delta="b")),
+        StreamTaskMessageDone(type="done", index=0),
+    ]
+    beats: list[tuple[object, ...]] = []
+
+    @activity.defn(name="deliver_turn")
+    async def deliver_turn() -> None:
+        async def _source():
+            for e in events:
+                yield e
+
+        await auto_send(_source(), task_id="task1", tracer=None, streaming=streaming)
+
+    env = ActivityEnvironment()
+    env.on_heartbeat = lambda *details: beats.append(details)
+    await env.run(deliver_turn)
+
+    assert len(beats) == len(events)
+