samples: workflow_streams: race the LLM consumer with workflow result

If the LLM activity exhausts its retries (bad OPENAI_API_KEY,
provider outage, etc.), the workflow fails before the activity
publishes the `complete` terminator. The consumer's previous
async-for loop only exited on `complete`, so the script blocked
indefinitely on a terminator that would never arrive instead of
surfacing the workflow failure.

Wrap the subscriber in a `consume()` coroutine and run it through
the existing `race_with_workflow` helper (the same pattern
`run_publisher.py` uses): if the workflow finishes first the
subscriber gets cancelled and the workflow's exception propagates;
if the subscriber sees `complete` first, the helper waits for the
workflow result and returns it.

Found in a Codex code review of today's workflow_streams changes.

Author: jssmith

workflow_streams/run_llm.py

@@ -44,6 +44,7 @@
     TextComplete,
     TextDelta,
 )
+from workflow_streams.shared import race_with_workflow
 from workflow_streams.workflows.llm_workflow import LLMWorkflow
 
 # Long enough that you can comfortably kill the worker mid-stream and
@@ -93,34 +94,39 @@ async def main() -> None:
 
     stream = WorkflowStreamClient.create(client, workflow_id)
 
-    # Subscribe to all three topics on a single iterator.
-    # result_type=RawValue lets us dispatch on item.topic and decode
-    # against the right dataclass per topic.
-    async for item in stream.subscribe(
-        [TOPIC_DELTA, TOPIC_RETRY, TOPIC_COMPLETE],
-        result_type=RawValue,
-    ):
-        if item.topic == TOPIC_RETRY:
-            evt = converter.from_payload(item.data.payload, RetryEvent)
-            sys.stdout.write(ANSI_RESTORE_AND_CLEAR)
-            sys.stdout.flush()
-            print(f"[retry attempt {evt.attempt}] resetting output")
-            print()
-            sys.stdout.write(ANSI_SAVE)
-            sys.stdout.flush()
-        elif item.topic == TOPIC_DELTA:
-            delta = converter.from_payload(item.data.payload, TextDelta)
-            sys.stdout.write(delta.text)
-            sys.stdout.flush()
-        elif item.topic == TOPIC_COMPLETE:
-            # The full text is also in the payload (and returned by
-            # the workflow), but the consumer has already rendered it
-            # incrementally. Just terminate the line.
-            converter.from_payload(item.data.payload, TextComplete)
-            print()
-            break
-
-    result = await handle.result()
+    async def consume() -> None:
+        # Subscribe to all three topics on a single iterator.
+        # result_type=RawValue lets us dispatch on item.topic and
+        # decode against the right dataclass per topic.
+        async for item in stream.subscribe(
+            [TOPIC_DELTA, TOPIC_RETRY, TOPIC_COMPLETE],
+            result_type=RawValue,
+        ):
+            if item.topic == TOPIC_RETRY:
+                evt = converter.from_payload(item.data.payload, RetryEvent)
+                sys.stdout.write(ANSI_RESTORE_AND_CLEAR)
+                sys.stdout.flush()
+                print(f"[retry attempt {evt.attempt}] resetting output")
+                print()
+                sys.stdout.write(ANSI_SAVE)
+                sys.stdout.flush()
+            elif item.topic == TOPIC_DELTA:
+                delta = converter.from_payload(item.data.payload, TextDelta)
+                sys.stdout.write(delta.text)
+                sys.stdout.flush()
+            elif item.topic == TOPIC_COMPLETE:
+                # The full text is also in the payload (and returned
+                # by the workflow), but the consumer has already
+                # rendered it incrementally. Just terminate the line.
+                converter.from_payload(item.data.payload, TextComplete)
+                print()
+                return
+
+    # Race the subscriber against the workflow result so that if the
+    # activity exhausts its retries the workflow's failure surfaces
+    # here rather than leaving the subscriber blocked on a `complete`
+    # that will never arrive.
+    result = await race_with_workflow(consume(), handle)
     print(f"\n[workflow result: {len(result)} chars]")