Skip to content

Add temporalio.contrib.pubsub module#1423

Draft
jssmith wants to merge 19 commits intomainfrom
contrib/pubsub
Draft

Add temporalio.contrib.pubsub module#1423
jssmith wants to merge 19 commits intomainfrom
contrib/pubsub

Conversation

@jssmith
Copy link
Copy Markdown
Contributor

@jssmith jssmith commented Apr 7, 2026

What was changed

Adds temporalio.contrib.pubsub, a reusable pub/sub primitive for streaming data out of Temporal workflows.

Why?

Streaming incremental results from long-running workflows (e.g., AI agent token streams, progress updates) is a common need with no built-in solution. This module provides a correct, reusable implementation so users don't have to roll their own poll/signal/dedup logic.

Checklist

  1. Closes — N/A (new contrib module, no existing issue)

  2. How was this tested:

    • 1100+ lines of pytest tests in tests/contrib/pubsub/test_pubsub.py
    • Covers: batching, flush safety, CAN serialization, replay guards, dedup (TTL pruning, truncation), offset-based resumption, max_batch_size, drain, and error handling
    • TLA+ model checking verifies dedup correctness (specs in temporalio/contrib/pubsub/verification/)
  3. Any docs updates needed?

    • Module includes README.md with usage examples and API reference
    • Design docs: DESIGN-v2.md, and addenda covering CAN, dedup, and topic semantics
    • No docs.temporal.io updates yet — can follow up once the API stabilizes

jssmith and others added 15 commits April 5, 2026 21:33
A workflow mixin (PubSubMixin) that turns any workflow into a pub/sub
broker. Activities and starters publish via batched signals; external
clients subscribe via long-poll updates exposed as an async iterator.

Key design decisions:
- Payloads are opaque bytes for cross-language compatibility
- Topics are plain strings, no hierarchy or prefix matching
- Global monotonic offsets (not per-topic) for simple continuation
- Batching built into PubSubClient with Nagle-like timer + priority flush
- Structured concurrency: no fire-and-forget tasks, trio-compatible
- Continue-as-new support: drain_pubsub() + get_pubsub_state() + validator
  to cleanly drain polls, plus follow_continues on the subscriber side

Module layout:
  _types.py  — PubSubItem, PublishInput, PollInput, PollResult, PubSubState
  _mixin.py  — PubSubMixin (signal, update, query handlers)
  _client.py — PubSubClient (batcher, async iterator, CAN resilience)

9 E2E integration tests covering: activity publish + subscribe, topic
filtering, offset-based replay, interleaved workflow/activity publish,
priority flush, iterator cancellation, context manager flush, concurrent
subscribers, and mixin coexistence with application signals/queries.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
PubSubState is now a Pydantic model so it survives serialization through
Pydantic-based data converters when embedded in Any-typed fields. Without
this, continue-as-new would fail with "'dict' object has no attribute 'log'"
because Pydantic deserializes Any fields as plain dicts.

Added two CAN tests:
- test_continue_as_new_any_typed_fails: documents that Any-typed fields
  lose PubSubState type information (negative test)
- test_continue_as_new_properly_typed: verifies CAN works with properly
  typed PubSubState | None fields

Simplified subscribe() exception handling: removed the broad except
Exception clause that tried _follow_continue_as_new() on every error.
Now only catches WorkflowUpdateRPCTimeoutOrCancelledError for CAN follow.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
README.md: usage-oriented documentation covering workflow mixin, activity
publishing, subscribing, continue-as-new, and cross-language protocol.

flush() safety: items are now removed from the buffer only after the
signal succeeds. Previously, buffer.clear() ran before the signal,
losing items on failure. Added test_flush_retains_items_on_signal_failure.

init_pubsub() guard: publish() and _pubsub_publish signal handler now
check for initialization and raise a clear RuntimeError instead of a
cryptic AttributeError.

PubSubClient.for_workflow() factory: preferred constructor that takes a
Client + workflow_id. Enables follow_continues in subscribe() without
accessing private WorkflowHandle._client. The handle-based constructor
remains for simple cases that don't need CAN following.

activity_pubsub_client() now uses for_workflow() internally with proper
keyword-only typed arguments instead of **kwargs: object.

CAN test timing: replaced asyncio.sleep(2) with assert_eq_eventually
polling for a different run_id, matching sdk-python test patterns.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
_pubsub_poll and _pubsub_offset now call _check_initialized() for a
clear RuntimeError instead of cryptic AttributeError when init_pubsub()
is forgotten.

README CAN example now includes the required imports (@DataClass,
workflow) and @workflow.init decorator.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The poll validator accesses _pubsub_draining, which would AttributeError
if init_pubsub() was never called. Added _check_initialized() guard.

Fixed PubSubState docstring: the field must be typed as PubSubState | None,
not Any. The old docstring incorrectly implied Any-typed fields would work.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
get_pubsub_state() and drain_pubsub() now call _check_initialized().
Previously drain_pubsub() could silently set _pubsub_draining on an
uninitialized instance, which init_pubsub() would then reset to False.

New tests:
- test_max_batch_size: verifies auto-flush when buffer reaches limit,
  using max_cached_workflows=0 to also test replay safety
- test_replay_safety: interleaved workflow/activity publish with
  max_cached_workflows=0, proving the mixin is determinism-safe

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Review comments (#@agent: annotations) capture design questions on:
- Topic offset model and information leakage (resolved: global offsets
  with BFF-layer containment, per NATS JetStream model)
- Exactly-once publish delivery (resolved: publisher ID + sequence number
  dedup, per Kafka producer model)
- Flush concurrency (resolved: asyncio.Lock with buffer swap)
- CAN follow behavior, poll rate limiting, activity context detection,
  validator purpose, pyright errors, API ergonomics

DESIGN-ADDENDUM-TOPICS.md: full exploration of per-topic vs global offsets
with industry survey (Kafka, Redis, NATS, PubNub, Google Pub/Sub,
RabbitMQ). Concludes global offsets are correct for workflow-scoped
pub/sub; leakage contained at BFF trust boundary.

DESIGN-ADDENDUM-DEDUP.md: exactly-once delivery via publisher ID +
monotonic sequence number. Workflow dedup state is dict[str, int],
bounded by publisher count. Buffer swap pattern with sequence reuse
on failure. PubSubState carries publisher_sequences through CAN.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Types:
- Remove offset from PubSubItem (global offset is now derived)
- Add publisher_id + sequence to PublishInput for exactly-once dedup
- Add base_offset + publisher_sequences to PubSubState for CAN
- Use Field(default_factory=...) for Pydantic mutable defaults

Mixin:
- Add _pubsub_base_offset for future log truncation support
- Add _pubsub_publisher_sequences for signal deduplication
- Dedup in signal handler: reject if sequence <= last seen
- Poll uses base_offset arithmetic for offset translation
- Class-body type declarations for basedpyright compatibility
- Validator docstring explaining drain/CAN interaction
- Module docstring gives specific init_pubsub() guidance

Client:
- asyncio.Lock + buffer swap for flush concurrency safety
- Publisher ID (uuid) + monotonic sequence for exactly-once delivery
- Sequence advances on failure to prevent data loss when new items
  merge with retry batch (found via Codex review)
- Remove follow_continues param — always follow CAN via describe()
- Configurable poll_interval (default 0.1s) for rate limiting
- Merge activity_pubsub_client() into for_workflow() with auto-detect
- _follow_continue_as_new is async with describe() check

Tests:
- New test_dedup_rejects_duplicate_signal
- Updated flush failure test for new sequence semantics
- All activities use PubSubClient.for_workflow()
- Remove PubSubItem.offset assertions
- poll_interval=0 in test helper for speed

Docs:
- DESIGN-v2.md: consolidated design doc superseding original + addenda
- README.md: updated API reference
- DESIGN-ADDENDUM-DEDUP.md: corrected flush failure semantics

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Rewrite the client-side dedup algorithm to match the formally verified
TLA+ protocol: failed flushes keep a separate _pending batch and retry
with the same sequence number. Only advance the confirmed sequence on
success. TLC proves NoDuplicates and OrderPreserved for the correct
algorithm, and finds duplicates in the old algorithm.

Add TTL-based pruning of publisher dedup entries during continue-as-new
(default 15 min). Add max_retry_duration (default 600s) to bound client
retries — must be less than publisher_ttl for safety. Both constraints
are formally verified in PubSubDedupTTL.tla.

Add truncate_pubsub() for explicit log prefix truncation. Add
publisher_last_seen timestamps for TTL tracking. Preserve legacy state
without timestamps during upgrade.

API changes: for_workflow→create, flush removed (use priority=True),
poll_interval→poll_cooldown, publisher ID shortened to 16 hex chars.

Includes TLA+ specs (correct, broken, inductive, multi-publisher TTL),
PROOF.md with per-action preservation arguments, scope and limitations.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
New analysis document evaluates whether publishing should use signals
or updates, examining Temporal's native dedup (Update ID per-run,
request_id for RPCs) vs the application-level (publisher_id, sequence)
protocol. Conclusion: app-level dedup is permanent for signals but
could be dropped for updates once temporal/temporal#6375 is fixed.
Non-blocking flush keeps signals as the right choice for streaming.

Updates DESIGN-v2.md section 6 to be precise about the two Temporal
guarantees that signal ordering relies on: sequential send order and
history-order handler invocation.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Analyzes deduplication through the end-to-end principle lens. Three
types of duplicates exist in the pipeline, each handled at the layer
that introduces them:

- Type A (duplicate LLM work): belongs at application layer — data
  escapes to consumers before the duplicate exists, so only the
  application can resolve it
- Type B (duplicate signal batches): belongs in pub/sub workflow —
  encapsulates transport details and is the only layer that can
  detect them correctly
- Type C (duplicate SSE delivery): belongs at BFF/browser layer

Concludes the (publisher_id, sequence) protocol is correctly placed.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
… design

Fill gaps identified during design review:
- Document why per-topic offsets were rejected (trust model, cursor
  portability, unjustified complexity) inline rather than only in historical
  addendum
- Expand BFF section with the four reconnection options considered and
  the decision to use SSE Last-Event-ID with BFF-assigned gapless IDs
- Add poll efficiency characteristics (O(new items) common case)
- Document BFF restart fallback (replay from turn start)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Wire types (PublishEntry, _WireItem, PollResult, PubSubState) encode
data as base64 strings for cross-language compatibility across all
Temporal SDKs. User-facing types (PubSubItem) use native bytes.

Conversion happens inside handlers:
- Signal handler decodes base64 → bytes on ingest
- Poll handler encodes bytes → base64 on response
- Client publish() accepts bytes, encodes for signal
- Client subscribe() decodes poll response, yields bytes

This means Go/Java/.NET ports get cross-language compat for free since
their JSON serializers encode byte[] as base64 by default.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
class PubSubState(BaseModel):
"""Serializable snapshot of pub/sub state for continue-as-new.

This is a Pydantic model (not a dataclass) so that Pydantic-based data
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this case it needs to somehow be clear that the pydantic data converter is required for this.

jssmith and others added 4 commits April 7, 2026 20:10
Remove the bounded poll wait from PubSubMixin and trim trailing
whitespace from types. Update DESIGN-v2.md with streaming plugin
rationale (no fencing needed, UI handles repeat delivery).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Add opt-in streaming code path to both agent framework plugins.
When enabled, the model activity calls the streaming LLM endpoint,
publishes TEXT_DELTA/THINKING_DELTA/TOOL_CALL_START events via
PubSubClient as a side channel, and returns the complete response
for the workflow to process (unchanged interface).

OpenAI Agents SDK:
- ModelActivityParameters.enable_streaming flag
- New invoke_model_activity_streaming method on ModelActivity
- ModelResponse reconstructed from ResponseCompletedEvent
- Uses @_auto_heartbeater for periodic heartbeats
- Routing in _temporal_model_stub (rejects local activities)

Google ADK:
- TemporalModel(streaming=True) constructor parameter
- New invoke_model_streaming activity using stream=True
- Registered in GoogleAdkPlugin

Both use batch_interval=0.1s for near-real-time token delivery.
No pubsub module changes needed.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The Pydantic BaseModel was introduced as a workaround for Any-typed fields
losing type information during continue-as-new serialization. The actual fix
is using concrete type annotations (PubSubState | None), which the default
data converter handles correctly for dataclasses — no Pydantic dependency
needed.

This removes the pydantic import from the pubsub contrib module entirely,
making it work out of the box with the default data converter. All 18 tests
pass, including both continue-as-new tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Implements DESIGN-ADDENDUM-ITEM-OFFSET.md. The poll handler now annotates
each item with its global offset (base_offset + position in log), enabling
subscribers to track fine-grained consumption progress for truncation.
This is needed for the voice-terminal agent where audio chunks must not be
truncated until actually played, not merely received.

- Add offset field to PubSubItem and _WireItem (default 0)
- Poll handler computes offset from base_offset + log_offset + enumerate index
- subscribe() passes wire_item.offset through to yielded PubSubItem
- Tests: per-item offsets, offsets with topic filtering, offsets after truncation

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants