[FROZEN][Bridge] Took standalone operation callbacks on a Worker.

.github/workflows/ci.yml

@@ -233,8 +233,13 @@ jobs:
       - run: uv add --dev --python 3.10 "googleapis-common-protos==1.70.0"
       - run: uv add --python 3.10 "protobuf<4"
       - run: uv sync --all-extras
-      - run: cargo install --locked nexgen --version 0.2.8 --features advanced --force
       - run: poe build-develop
+      # The generators build their pinned nexgen revision on first use, a full
+      # Rust build, so it's cached by the pin.
+      - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+        with:
+          path: .nexgen
+          key: nexgen-${{ hashFiles('scripts/gen_nexus_system_api.py') }}
       - run: poe gen-protos
       - name: Check generation unchanged
         run: |

scripts/gen_nexus_system_api.py

@@ -11,7 +11,8 @@
 
 base_dir = Path(__file__).parent.parent
 sys.path.insert(0, str(base_dir))
-wit_input_dir = (
+# The upstream WIT dependencies come from Core's API copy.
+wit_deps_dir = (
     base_dir
     / "temporalio"
     / "bridge"
@@ -21,10 +22,19 @@
     / "protos"
     / "api_upstream"
     / "nexus"
+    / "deps"
 )
+# The workflow service WIT with the stream notifier operations, and the notification service WIT,
+# come from the temporalio-streams crate, which carries the API they name.
+wit_input_dir = gen_protos.streams_proto_dir.parent / "nexus"
 wit_path = wit_input_dir / "workflow-service.wit"
-wit_deps_dir = wit_input_dir / "deps"
 output_dir = base_dir / "temporalio" / "nexus" / "system" / "workflow_service"
+# The service a Worker implements to take NexusHandler callbacks. The server calls
+# it, not a Workflow, so it gets plain definitions and no System Nexus caller.
+notification_wit_path = wit_input_dir / "notification-service.wit"
+notification_output_dir = (
+    base_dir / "temporalio" / "nexus" / "system" / "notification_service"
+)
 workflow_init_path = base_dir / "temporalio" / "workflow" / "__init__.py"
 workflowservice_request_response_proto = (
     gen_protos.api_proto_dir
@@ -34,28 +44,64 @@
     / "v1"
     / "request_response.proto"
 )
-NEX_GEN_VERSION = "0.2.8"
+notificationservice_request_response_proto = (
+    gen_protos.api_proto_dir
+    / "temporal"
+    / "api"
+    / "notificationservice"
+    / "v1"
+    / "request_response.proto"
+)
+workflowservice_stream_request_response_proto = (
+    gen_protos.streams_api_proto_dir
+    / "temporal"
+    / "api"
+    / "workflowservice"
+    / "v1"
+    / "stream_request_response.proto"
+)
+notificationservice_progress_proto = (
+    gen_protos.nexus_api_proto_dir
+    / "temporal"
+    / "api"
+    / "notificationservice"
+    / "v1"
+    / "progress.proto"
+)
+# The stream bindings need nexgen fixes that no release carries yet, so both generators
+# use this pinned revision of the nexgen fork. NEX_GEN_BIN points them at a binary built
+# some other way.
+NEX_GEN_REPOSITORY = "https://github.com/moetemp/nexgen"
+NEX_GEN_REVISION = "3c60ecc9987007d70f8d798d09373445b0219d54"
+nex_gen_root = base_dir / ".nexgen" / NEX_GEN_REVISION
 
 
 def nex_gen_command() -> list[str]:
     if bin_path := os.environ.get("NEX_GEN_BIN"):
         return [bin_path]
 
-    if shutil.which("nexgen") is None:
+    binary = nex_gen_root / "bin" / "nexgen"
+    if not binary.exists():
         subprocess.check_call(
             [
                 "cargo",
                 "install",
                 "--locked",
                 "nexgen",
-                "--version",
-                NEX_GEN_VERSION,
+                "--git",
+                NEX_GEN_REPOSITORY,
+                "--rev",
+                NEX_GEN_REVISION,
                 "--features",
                 "advanced",
-                "--force",
-            ]
+                "--root",
+                str(nex_gen_root),
+            ],
+            # The git CLI honours the user's URL rewrites and credentials,
+            # which cargo's built-in client does not.
+            env={**os.environ, "CARGO_NET_GIT_FETCH_WITH_CLI": "true"},
         )
-    return ["nexgen"]
+    return [str(binary)]
 
 
 def build_descriptor_set(descriptor_path: Path) -> None:
@@ -64,10 +110,15 @@ def build_descriptor_set(descriptor_path: Path) -> None:
             sys.executable,
             "-mgrpc_tools.protoc",
             f"--proto_path={gen_protos.api_proto_dir}",
+            f"--proto_path={gen_protos.streams_api_proto_dir}",
+            f"--proto_path={gen_protos.nexus_api_proto_dir}",
             f"--proto_path={gen_protos.proto_dir}",
             "--include_imports",
             f"--descriptor_set_out={descriptor_path}",
             str(workflowservice_request_response_proto),
+            str(notificationservice_request_response_proto),
+            str(workflowservice_stream_request_response_proto),
+            str(notificationservice_progress_proto),
         ]
     )
 
@@ -113,6 +164,8 @@ def generate_nexus_system_api() -> None:
         raise RuntimeError(f"missing WIT source: {wit_path}")
     if not wit_deps_dir.exists():
         raise RuntimeError(f"missing WIT dependency directory: {wit_deps_dir}")
+    if not notification_wit_path.exists():
+        raise RuntimeError(f"missing WIT source: {notification_wit_path}")
 
     with tempfile.TemporaryDirectory(dir=base_dir) as temp_dir:
         descriptor_path = Path(temp_dir) / "temporal_api.bin"
@@ -138,6 +191,22 @@ def generate_nexus_system_api() -> None:
             ]
         )
 
+        shutil.rmtree(notification_output_dir, ignore_errors=True)
+        subprocess.check_call(
+            [
+                *command,
+                "python",
+                str(notification_wit_path),
+                str(wit_deps_dir),
+                "--support-package",
+                "temporalio.nexus.system._support",
+                "--descriptors",
+                str(descriptor_path),
+                "--output",
+                str(notification_output_dir),
+            ]
+        )
+
     (output_dir.parent / "__init__.py").touch()
     generate_workflow_exports()
     subprocess.check_call(
@@ -150,6 +219,7 @@ def generate_nexus_system_api() -> None:
             "I",
             "--fix",
             str(output_dir),
+            str(notification_output_dir),
             str(workflow_init_path),
         ]
     )
@@ -160,6 +230,7 @@ def generate_nexus_system_api() -> None:
             "ruff",
             "format",
             str(output_dir),
+            str(notification_output_dir),
             str(workflow_init_path),
         ]
     )

scripts/gen_payload_visitor.py

@@ -2,7 +2,7 @@
 import sys
 from importlib.util import module_from_spec, spec_from_file_location
 from pathlib import Path
-from typing import cast, get_args, get_origin, get_type_hints
+from typing import Any, cast, get_args, get_origin, get_type_hints
 
 import google.protobuf.message
 import nexusrpc
@@ -22,10 +22,10 @@
 from temporalio.converter._payload_converter import _get_transfer_type_converter
 
 
-def discover_system_nexus_roots() -> list[Descriptor]:
-    module_path = base_dir / "temporalio" / "nexus" / "system" / "workflow_service"
+def load_system_package(name: str) -> Any:
+    module_path = base_dir / "temporalio" / "nexus" / "system" / name
     spec = spec_from_file_location(
-        "temporalio_nexus_system_workflow_service",
+        f"temporalio_nexus_system_{name}",
         module_path / "__init__.py",
         submodule_search_locations=[str(module_path)],
     )
@@ -34,9 +34,23 @@ def discover_system_nexus_roots() -> list[Descriptor]:
     module = module_from_spec(spec)
     sys.modules[spec.name] = module
     spec.loader.exec_module(module)
+    return module
+
 
+def discover_system_nexus_roots() -> list[Descriptor]:
+    # A Worker calls the workflow service, and the server calls the
+    # notification service on a Worker. Both carry system payloads.
+    services = [
+        load_system_package("workflow_service")._services.WorkflowService,
+        load_system_package("notification_service").NotificationService,
+    ]
     roots: list[Descriptor] = []
-    for annotation in get_type_hints(module._services.WorkflowService).values():
+    annotations = [
+        annotation
+        for service in services
+        for annotation in get_type_hints(service).values()
+    ]
+    for annotation in annotations:
         if get_origin(annotation) is not nexusrpc.Operation:
             continue
         for operation_type in get_args(annotation):

temporalio/bridge/_visitor.py

@@ -650,3 +650,29 @@ async def _visit_temporal_api_workflowservice_v1_SignalWithStartWorkflowExecutio
         self, fs: VisitorFunctions, o: Any
     ) -> None:
         pass
+
+    async def _visit_temporal_api_notificationservice_v1_OnCompleteRequest(
+        self, fs: VisitorFunctions, o: Any
+    ):
+        if o.HasField("source_context"):
+            await self._visit_temporal_api_common_v1_Payload(fs, o.source_context)
+        if o.HasField("success"):
+            await self._visit_temporal_api_common_v1_Payload(fs, o.success)
+        elif o.HasField("failure"):
+            await self._visit_temporal_api_failure_v1_Failure(fs, o.failure)
+
+    async def _visit_temporal_api_notificationservice_v1_OnCompleteResponse(
+        self, fs: VisitorFunctions, o: Any
+    ) -> None:
+        pass
+
+    async def _visit_temporal_api_notificationservice_v1_OnProgressRequest(
+        self, fs: VisitorFunctions, o: Any
+    ):
+        if o.HasField("source_context"):
+            await self._visit_temporal_api_common_v1_Payload(fs, o.source_context)
+
+    async def _visit_temporal_api_notificationservice_v1_OnProgressResponse(
+        self, fs: VisitorFunctions, o: Any
+    ) -> None:
+        pass

temporalio/client/__init__.py

@@ -145,6 +145,7 @@
 )
 from ._nexus import (
     NexusClient,
+    NexusHandlerCallback,
     NexusOperationExecution,
     NexusOperationExecutionAsyncIterator,
     NexusOperationExecutionCancellationInfo,
@@ -246,6 +247,7 @@
     "AsyncActivityHandle",
     "ActivityHandle",
     "NexusClient",
+    "NexusHandlerCallback",
     "NexusOperationExecution",
     "NexusOperationExecutionAsyncIterator",
     "NexusOperationExecutionCancellationInfo",

temporalio/client/_impl.py

@@ -1610,6 +1610,18 @@ async def start_nexus_operation(
             for k, v in input.headers.items():
                 req.nexus_header[k] = v
 
+        for callback in input.completion_callbacks:
+            handler = temporalio.api.common.v1.Callback.NexusHandler(
+                task_queue_name=callback.task_queue,
+                service=callback.service,
+                operation=callback.operation,
+                progress_operation=callback.progress_operation,
+            )
+            if callback.source_context is not None:
+                [context] = await data_converter.encode([callback.source_context])
+                handler.source_context.CopyFrom(context)
+            req.completion_callbacks.add(nexus_handler=handler)
+
         resp: temporalio.api.workflowservice.v1.StartNexusOperationExecutionResponse
         try:
             resp = await self._client.workflow_service.start_nexus_operation_execution(

temporalio/client/_interceptor.py

@@ -31,6 +31,7 @@
         AsyncActivityIDReference,
     )
     from ._nexus import (
+        NexusHandlerCallback,
         NexusOperationExecutionAsyncIterator,
         NexusOperationExecutionCount,
         NexusOperationExecutionDescription,
@@ -605,6 +606,7 @@ class StartNexusOperationInput:
     headers: Mapping[str, str]
     rpc_metadata: Mapping[str, str | bytes]
     rpc_timeout: timedelta | None
+    completion_callbacks: Sequence[NexusHandlerCallback] = ()
 
 
 @dataclass