Add changes required for deepagent-temporal

Author: pradithya

langgraph/temporal/activities.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import contextvars
 import itertools
 from collections.abc import Sequence
 from typing import Any, cast
@@ -35,6 +36,48 @@
 )
 from langgraph.temporal.converter import GraphRegistry
 
+# Generic extension point: context variable for collecting Child Workflow
+# requests during Activity execution. Middleware (e.g., TemporalSubAgentMiddleware)
+# appends requests here; the Activity collects them into NodeActivityOutput.
+_child_workflow_requests_var: contextvars.ContextVar[list[dict[str, Any]]] = (
+    contextvars.ContextVar("_child_workflow_requests")
+)
+
+# Worker-specific task queue for worker affinity.
+# Set at worker startup via `set_worker_task_queue()`. The
+# `get_available_task_queue` activity returns this value so the Workflow
+# can pin all subsequent Activities to this worker's queue.
+_worker_task_queue: str | None = None
+
+
+def set_worker_task_queue(queue: str) -> None:
+    """Set the worker-specific task queue name.
+
+    Called at worker startup. The `get_available_task_queue` activity
+    returns this value to the Workflow for worker affinity.
+    """
+    global _worker_task_queue  # noqa: PLW0603
+    _worker_task_queue = queue
+
+
+@activity.defn
+async def get_available_task_queue() -> str:
+    """Return this worker's unique task queue name.
+
+    The Workflow calls this activity on the shared distribution queue.
+    Whichever worker picks it up returns its own unique queue name. The
+    Workflow then dispatches all subsequent Activities to that queue,
+    achieving worker affinity (following the Temporal worker-specific
+    task queues pattern).
+    """
+    if _worker_task_queue is None:
+        raise ApplicationError(
+            "Worker task queue not configured. "
+            "Call set_worker_task_queue() at worker startup.",
+            non_retryable=True,
+        )
+    return _worker_task_queue
+
 
 def _make_counter() -> Any:
     """Create a simple callable counter using itertools.count."""
@@ -123,6 +166,11 @@ def read_handler(
     try:
         activity.heartbeat(f"Starting node {input.node_name}")
 
+        # Initialize the child workflow requests context variable so that
+        # middleware (e.g., TemporalSubAgentMiddleware) can append requests
+        # during node execution.
+        _child_workflow_requests_var.set([])
+
         # Set the runnable config context so that interrupt() and other
         # functions that call get_config() can find it.  This is normally
         # done by the Runnable invoke/ainvoke tracing path, but Activities
@@ -161,6 +209,9 @@ def read_handler(
         # Filter out UntrackedValue writes
         serializable_writes = _filter_untracked_writes(all_writes, graph)
 
+        # Collect any Child Workflow requests stored during execution
+        child_wf_requests = _child_workflow_requests_var.get([])
+
         return NodeActivityOutput(
             node_name=input.node_name,
             writes=serializable_writes,
@@ -169,6 +220,7 @@ def read_handler(
             push_sends=push_sends if push_sends else None,
             command=command_data,
             custom_data=custom_data if custom_data else None,
+            child_workflow_requests=(child_wf_requests if child_wf_requests else None),
         )
     except GraphInterrupt as exc:
         # Node called interrupt() - propagate to Workflow

langgraph/temporal/config.py

@@ -26,17 +26,36 @@ class ActivityOptions:
     schedule_to_close_timeout: timedelta | None = None
 
 
+@dataclass
+class SubAgentConfig:
+    """Configuration for sub-agent Child Workflow dispatch.
+
+    Attributes:
+        task_queue: Task queue for sub-agent Child Workflows.
+        sticky_task_queue: Sticky task queue for sub-agent worker affinity.
+        execution_timeout_seconds: Maximum execution time for sub-agent
+            Child Workflows in seconds.
+    """
+
+    task_queue: str | None = None
+    sticky_task_queue: str | None = None
+    execution_timeout_seconds: float = 1800.0
+
+
 @dataclass
 class RestoredState:
     """State carried across continue-as-new boundaries.
 
     Attributes:
         checkpoint: The serialized Checkpoint TypedDict.
         step: The current step counter.
+        sticky_task_queue: Sticky task queue name for worker affinity,
+            preserved across continue-as-new.
     """
 
     checkpoint: dict[str, Any]
     step: int
+    sticky_task_queue: str | None = None
 
 
 @dataclass
@@ -74,6 +93,16 @@ class WorkflowInput:
         node_task_queues: Per-node task queue overrides.
         node_activity_options: Per-node Activity configuration (serialized).
         node_retry_policies: Per-node retry policy configuration.
+        sticky_task_queue: Sticky task queue for worker affinity. When set,
+            overrides per-node task queue routing for all Activities.
+            Can be set explicitly or discovered at runtime via
+            `get_available_task_queue` when `use_worker_affinity` is True.
+        use_worker_affinity: When True, the Workflow calls the
+            `get_available_task_queue` activity at startup to discover a
+            worker-specific queue. All subsequent Activities are dispatched
+            to that queue. The discovered queue is forwarded on
+            continue-as-new via `sticky_task_queue`.
+        subagent_config: Configuration for sub-agent Child Workflow dispatch.
     """
 
     graph_definition_ref: str
@@ -85,6 +114,9 @@ class WorkflowInput:
     node_task_queues: dict[str, str] | None = None
     node_activity_options: dict[str, ActivityOptions] | None = None
     node_retry_policies: dict[str, RetryPolicyConfig] | None = None
+    sticky_task_queue: str | None = None
+    use_worker_affinity: bool = False
+    subagent_config: SubAgentConfig | None = None
 
 
 @dataclass
@@ -138,6 +170,8 @@ class NodeActivityOutput:
         interrupts: Interrupt payloads if node called interrupt().
         push_sends: Dynamic Send objects emitted during execution.
         command: Command metadata if node returned a Command.
+        child_workflow_requests: Requests for Child Workflow dispatch
+            (e.g., sub-agent invocations collected via context variable).
     """
 
     node_name: str
@@ -148,6 +182,7 @@ class NodeActivityOutput:
     push_sends: list[dict[str, Any]] | None = None
     command: dict[str, Any] | None = None
     custom_data: list[Any] | None = None
+    child_workflow_requests: list[dict[str, Any]] | None = None
 
 
 @dataclass

langgraph/temporal/graph.py

@@ -20,6 +20,7 @@
     ActivityOptions,
     RetryPolicyConfig,
     StateUpdatePayload,
+    SubAgentConfig,
     WorkflowInput,
     WorkflowOutput,
 )
@@ -122,6 +123,27 @@ def _build_workflow_input(
                     max_attempts=getattr(policy, "max_attempts", 0),
                 )
 
+        # Read sticky_task_queue, use_worker_affinity, and subagent_config
+        sticky_task_queue: str | None = None
+        use_worker_affinity: bool = False
+        subagent_config: SubAgentConfig | None = None
+        if config and "configurable" in config:
+            sticky_task_queue = config["configurable"].get("sticky_task_queue")
+            use_worker_affinity = bool(
+                config["configurable"].get("use_worker_affinity", False)
+            )
+            raw_subagent = config["configurable"].get("subagent_config")
+            if isinstance(raw_subagent, SubAgentConfig):
+                subagent_config = raw_subagent
+            elif isinstance(raw_subagent, dict):
+                subagent_config = SubAgentConfig(
+                    task_queue=raw_subagent.get("task_queue"),
+                    sticky_task_queue=raw_subagent.get("sticky_task_queue"),
+                    execution_timeout_seconds=raw_subagent.get(
+                        "execution_timeout_seconds", 1800.0
+                    ),
+                )
+
         return WorkflowInput(
             graph_definition_ref=self._graph_ref,
             input_data=input if isinstance(input, dict) else {"__root__": input},
@@ -131,6 +153,9 @@ def _build_workflow_input(
             node_task_queues=self.node_task_queues if self.node_task_queues else None,
             node_activity_options=serialized_options,
             node_retry_policies=serialized_retry,
+            sticky_task_queue=sticky_task_queue,
+            use_worker_affinity=use_worker_affinity,
+            subagent_config=subagent_config,
         )
 
     async def ainvoke(

langgraph/temporal/worker.py

@@ -6,6 +6,9 @@
 
 from __future__ import annotations
 
+import uuid
+from pathlib import Path
+from types import TracebackType
 from typing import Any
 
 from langgraph.pregel import Pregel
@@ -16,40 +19,145 @@
     dynamic_execute_node,
     evaluate_conditional_edge,
     execute_node,
+    get_available_task_queue,
+    set_worker_task_queue,
 )
 from langgraph.temporal.converter import GraphRegistry
 from langgraph.temporal.workflow import LangGraphWorkflow
 
 
+class WorkerGroup:
+    """Manages multiple Temporal Workers as a single async context manager.
+
+    Used for worker-affinity mode where two workers are needed:
+    one on the shared queue (Workflows + discovery Activity) and one on
+    a worker-specific queue (node execution Activities).
+    """
+
+    def __init__(self, workers: list[Worker]) -> None:
+        self._workers = workers
+
+    async def __aenter__(self) -> WorkerGroup:
+        for w in self._workers:
+            await w.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        for w in reversed(self._workers):
+            await w.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def run(self) -> None:
+        """Run all workers. Blocks until shutdown."""
+        import asyncio
+
+        await asyncio.gather(*[w.run() for w in self._workers])
+
+
+def _resolve_worker_queue(
+    task_queue: str,
+    queue_file: Path | str | None,
+) -> str:
+    """Resolve the worker-specific queue name.
+
+    If `queue_file` is provided and exists, read the persisted queue name
+    (so a restarted worker re-registers on the same queue). Otherwise
+    generate a new one and persist it if a path was given.
+    """
+    if queue_file is not None:
+        path = Path(queue_file)
+        if path.exists():
+            stored = path.read_text().strip()
+            if stored:
+                return stored
+        # Generate and persist
+        queue = f"{task_queue}-worker-{uuid.uuid4().hex[:12]}"
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(queue)
+        return queue
+
+    return f"{task_queue}-worker-{uuid.uuid4().hex[:12]}"
+
+
 def create_worker(
     graph: Pregel,
     client: TemporalClient,
     task_queue: str = "langgraph-default",
+    *,
+    use_worker_affinity: bool = False,
+    worker_queue_file: Path | str | None = None,
     **kwargs: Any,
-) -> Worker:
+) -> Worker | WorkerGroup:
     """Create a Temporal Worker configured for a LangGraph graph.
 
     Registers the `LangGraphWorkflow` as a Workflow and `execute_node` /
     `evaluate_conditional_edge` as Activities. The graph is registered in
     the `GraphRegistry` for Activity-side lookup.
 
+    When `use_worker_affinity` is True, returns a `WorkerGroup` with two
+    workers following the Temporal worker-specific task queues pattern:
+    - A shared worker on `task_queue` (Workflows + discovery Activity)
+    - A worker-specific worker on a unique queue (node Activities)
+
     Args:
         graph: A compiled Pregel graph instance.
         client: A Temporal client instance.
         task_queue: The task queue to listen on.
+        use_worker_affinity: When True, create a dual-worker setup for
+            worker-specific task queue affinity.
+        worker_queue_file: Path to persist the worker-specific queue name.
+            On restart, the worker re-registers on the same queue so that
+            in-flight Activities resume on this worker. If None, a new
+            queue name is generated each time (no restart recovery).
         **kwargs: Additional Worker configuration (e.g.,
             `max_concurrent_activities`, `max_concurrent_workflow_tasks`).
 
     Returns:
-        A configured Temporal Worker instance.
+        A configured Temporal Worker or WorkerGroup instance.
     """
     # Ensure graph is registered
     GraphRegistry.get_instance().register(graph)
 
-    return Worker(
+    if not use_worker_affinity:
+        return Worker(
+            client,
+            task_queue=task_queue,
+            workflows=[LangGraphWorkflow],
+            activities=[
+                execute_node,
+                dynamic_execute_node,
+                evaluate_conditional_edge,
+            ],
+            **kwargs,
+        )
+
+    # Worker-affinity mode: two workers following the Temporal
+    # worker-specific task queues pattern.
+    worker_specific_queue = _resolve_worker_queue(task_queue, worker_queue_file)
+    set_worker_task_queue(worker_specific_queue)
+
+    # Worker 1: shared queue — Workflows + get_available_task_queue
+    shared_worker = Worker(
         client,
         task_queue=task_queue,
         workflows=[LangGraphWorkflow],
-        activities=[execute_node, dynamic_execute_node, evaluate_conditional_edge],
+        activities=[get_available_task_queue],
         **kwargs,
     )
+
+    # Worker 2: worker-specific queue — node execution Activities
+    specific_worker = Worker(
+        client,
+        task_queue=worker_specific_queue,
+        activities=[
+            execute_node,
+            dynamic_execute_node,
+            evaluate_conditional_edge,
+        ],
+    )
+
+    return WorkerGroup([shared_worker, specific_worker])