Streams API: combined branch for review

.github/actions/install-redis/action.yml

@@ -0,0 +1,25 @@
+name: Install Redis
+description: Install and start a Redis server for the tests of the Redis Stream Store.
+runs:
+  using: composite
+  steps:
+    - if: runner.os == 'Linux'
+      shell: bash
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y --no-install-recommends redis-server
+    - if: runner.os == 'macOS'
+      shell: bash
+      run: brew install redis
+    # Redis has no official Windows build, so the Redis Store tests skip there. The port is not the
+    # default one, since the Linux package starts its own server on 6379.
+    - if: runner.os != 'Windows'
+      shell: bash
+      run: |
+        redis-server --port 6380 --save "" --appendonly no --daemonize yes
+        for _ in $(seq 1 50); do
+          redis-cli -p 6380 ping && break
+          sleep 0.1
+        done
+        redis-cli -p 6380 ping
+        echo "STREAMS_REDIS_URL=redis://127.0.0.1:6380" >> "$GITHUB_ENV"

.github/workflows/ci.yml

@@ -58,6 +58,8 @@ jobs:
         if: ${{ matrix.clippyLinter }}
       - run: poe build-develop
       - run: poe lint
+      - name: Install and start Redis for the Redis Stream Store tests
+        uses: ./.github/actions/install-redis
       - run: mkdir junit-xml
       - run: poe test ${{matrix.pytestExtraArgs}} -s --junit-xml=junit-xml/${{ matrix.python }}--${{ matrix.os }}.xml
         timeout-minutes: 15

.gitmodules

@@ -1,3 +1,3 @@
 [submodule "sdk-core"]
 	path = temporalio/bridge/sdk-core
-	url = https://github.com/temporalio/sdk-rust.git
+	url = https://github.com/moetemp/sdk-rust.git

pyproject.toml

@@ -93,8 +93,8 @@ dev = [
 ]
 
 [tool.poe.tasks]
-build-develop = "uv run maturin develop --uv"
-build-develop-with-release = { cmd = "uv run maturin develop --release --uv" }
+build-develop = "uv run maturin develop --uv --features streams-redis"
+build-develop-with-release = { cmd = "uv run maturin develop --release --uv --features streams-redis" }
 # FIPS builds: swap the rustls stack onto aws-lc-rs FIPS mode (aws-lc-fips-sys), eliminating `ring`.
 build-develop-fips = { cmd = "uv run maturin develop --uv --no-default-features --features fips", env = { TEMPORALIO_FIPS = "1" } }
 build-wheel-fips = { cmd = "uv run maturin build --release --no-default-features --features fips", env = { TEMPORALIO_FIPS = "1" } }
@@ -153,18 +153,20 @@ filterwarnings = [
 before-all = "pip install protoc-wheel-0"
 build = "cp311-win_amd64 cp311-manylinux_x86_64 cp311-manylinux_aarch64 cp311-musllinux_x86_64 cp311-musllinux_aarch64 cp311-macosx_x86_64 cp311-macosx_arm64"
 build-verbosity = 1
+# The official wheels link the Stream runtime.
+environment = { MATURIN_PEP517_ARGS = "--features streams-redis" }
 
 [tool.cibuildwheel.macos]
-environment = { MACOSX_DEPLOYMENT_TARGET = "10.12" }
+environment = { MACOSX_DEPLOYMENT_TARGET = "10.12", MATURIN_PEP517_ARGS = "--features streams-redis" }
 
 [tool.cibuildwheel.linux]
 before-all = "sh scripts/cibuildwheel_before_all_linux.sh"
 before-build = "sh scripts/cibuildwheel_before_build_linux.sh"
-environment = { PATH = "$PATH:$HOME/.cargo/bin", CARGO_NET_GIT_FETCH_WITH_CLI = "true" }
+environment = { PATH = "$PATH:$HOME/.cargo/bin", CARGO_NET_GIT_FETCH_WITH_CLI = "true", MATURIN_PEP517_ARGS = "--features streams-redis" }
 
 [[tool.cibuildwheel.overrides]]
 select = "*musllinux*"
-environment = { PATH = "$PATH:$HOME/.cargo/bin", CARGO_NET_GIT_FETCH_WITH_CLI = "true", RUSTFLAGS = "-C target-feature=-crt-static" }
+environment = { PATH = "$PATH:$HOME/.cargo/bin", CARGO_NET_GIT_FETCH_WITH_CLI = "true", RUSTFLAGS = "-C target-feature=-crt-static", MATURIN_PEP517_ARGS = "--features streams-redis" }
 
 [tool.mypy]
 ignore_missing_imports = true
@@ -173,12 +175,17 @@ exclude = [
   'temporalio/api',
   'temporalio/bridge/proto',
   'temporalio/nexus/system/workflow_service',
+  'temporalio/nexus/system/stream_service',
 ]
 
 [[tool.mypy.overrides]]
 module = "temporalio.nexus.system.workflow_service.*"
 ignore_errors = true
 
+[[tool.mypy.overrides]]
+module = "temporalio.nexus.system.stream_service.*"
+ignore_errors = true
+
 [tool.pydocstyle]
 convention = "google"
 # https://github.com/PyCQA/pydocstyle/issues/363#issuecomment-625563088

scripts/gen_nexus_system_api.py

@@ -23,9 +23,14 @@
 )
 wit_deps_dir = wit_input_dir / "deps"
 workflow_output_dir = base_dir / "temporalio" / "nexus" / "system" / "workflow_service"
+stream_output_dir = base_dir / "temporalio" / "nexus" / "system" / "stream_service"
 notification_output_dir = base_dir / "temporalio" / "nexus" / "notifications"
 workflow_init_path = base_dir / "temporalio" / "workflow" / "__init__.py"
 NEX_GEN_VERSION = "0.2.8"
+# The Stream package needs a newer nexgen, which rejects the other WITs. One pin is enough
+# again once the api WITs and a nexgen release agree. No release has nexgen commit 9653a0f5,
+# the `main` of nexgen, yet, so NEX_GEN_STREAM_BIN points at a build of it.
+NEX_GEN_STREAM_VERSION = "0.2.12"
 proto_files = [
     gen_protos.api_proto_dir
     / "temporal"
@@ -36,10 +41,22 @@
     gen_protos.api_proto_dir
     / "temporal"
     / "api"
+    / "streamservice"
+    / "v1"
+    / "request_response.proto",
+    gen_protos.api_proto_dir
+    / "temporal"
+    / "api"
     / "notificationservice"
     / "v1"
     / "request_response.proto",
 ]
+# The Workflow functions of these generated packages are exported by temporalio.workflow.
+# They are in import order, so that sorting the imports keeps them between the markers.
+workflow_export_packages = [
+    ("temporalio.nexus.system.stream_service.workflow", stream_output_dir / "workflow"),
+    ("temporalio.nexus.system.workflow_service", workflow_output_dir),
+]
 
 
 def nex_gen_command() -> list[str]:
@@ -66,6 +83,20 @@ def nex_gen_command() -> list[str]:
     return ["nexgen"]
 
 
+def nex_gen_stream_command() -> list[str]:
+    bin_path = os.environ.get("NEX_GEN_STREAM_BIN")
+    if not bin_path:
+        raise RuntimeError(
+            "NEX_GEN_STREAM_BIN must point at a nexgen build of commit 9653a0f5"
+        )
+    version = subprocess.check_output([bin_path, "--version"], text=True).strip()
+    if version != f"nexgen {NEX_GEN_STREAM_VERSION}":
+        raise RuntimeError(
+            f"NEX_GEN_STREAM_BIN is {version}, not nexgen {NEX_GEN_STREAM_VERSION}"
+        )
+    return [bin_path]
+
+
 def build_descriptor_set(descriptor_path: Path) -> None:
     subprocess.check_call(
         [
@@ -85,14 +116,15 @@ def generate_package(
     output_dir: Path,
     descriptor_path: Path,
     *,
-    native_api: bool,
+    flags: list[str],
 ) -> None:
-    args = [*command, "python", str(wit_input_dir / wit_name), str(wit_deps_dir)]
-    if native_api:
-        args.append("--native-api")
     subprocess.check_call(
         [
-            *args,
+            *command,
+            "python",
+            str(wit_input_dir / wit_name),
+            str(wit_deps_dir),
+            *flags,
             "--system-nexus",
             "--support-package",
             "temporalio.nexus.system._support",
@@ -132,17 +164,26 @@ def replace_marker_block(
 
 
 def generate_workflow_exports() -> None:
-    exports = generated_exports(workflow_output_dir / "__init__.py")
+    package_exports = [
+        (module, generated_exports(package_dir / "__init__.py"))
+        for module, package_dir in workflow_export_packages
+    ]
     import_block = [
         "# BEGIN GENERATED NEXUS SYSTEM EXPORTS\n",
-        "from temporalio.nexus.system.workflow_service import (\n",
-        *[f"    {export},\n" for export in exports],
-        ")\n",
+        *[
+            line
+            for module, exports in package_exports
+            for line in [
+                f"from {module} import (\n",
+                *[f"    {export},\n" for export in exports],
+                ")\n",
+            ]
+        ],
         "# END GENERATED NEXUS SYSTEM EXPORTS\n",
     ]
     all_block = [
         "    # BEGIN GENERATED NEXUS SYSTEM __ALL__\n",
-        *[f'    "{export}",\n' for export in exports],
+        *[f'    "{export}",\n' for _, exports in package_exports for export in exports],
         "    # END GENERATED NEXUS SYSTEM __ALL__\n",
     ]
     content = workflow_init_path.read_text()
@@ -162,6 +203,7 @@ def generate_workflow_exports() -> None:
 def generate_nexus_system_api() -> None:
     required_paths = [
         wit_input_dir / "workflow-service.wit",
+        wit_input_dir / "stream-service.wit",
         wit_input_dir / "notification-service.wit",
         wit_deps_dir,
         *proto_files,
@@ -173,23 +215,31 @@ def generate_nexus_system_api() -> None:
         descriptor_path = Path(temp_dir) / "temporal_api.bin"
         build_descriptor_set(descriptor_path)
         command = nex_gen_command()
-        for wit_name, output_dir, native_api in (
-            ("workflow-service.wit", workflow_output_dir, True),
-            ("notification-service.wit", notification_output_dir, False),
+        stream_command = nex_gen_stream_command()
+        for wit_name, output_dir, package_command, flags in (
+            ("workflow-service.wit", workflow_output_dir, command, ["--native-api"]),
+            (
+                "stream-service.wit",
+                stream_output_dir,
+                stream_command,
+                ["--targets", "workflow,client"],
+            ),
+            ("notification-service.wit", notification_output_dir, command, []),
         ):
             shutil.rmtree(output_dir, ignore_errors=True)
             output_dir.parent.mkdir(parents=True, exist_ok=True)
             generate_package(
-                command,
+                package_command,
                 wit_name,
                 output_dir,
                 descriptor_path,
-                native_api=native_api,
+                flags=flags,
             )
     workflow_output_dir.parent.joinpath("__init__.py").touch()
     generate_workflow_exports()
     format_paths = [
         str(workflow_output_dir),
+        str(stream_output_dir),
         str(notification_output_dir),
         str(workflow_init_path),
     ]

scripts/gen_payload_visitor.py

@@ -24,20 +24,30 @@
 
 
 def discover_system_nexus_roots() -> list[Descriptor]:
-    module_path = base_dir / "temporalio" / "nexus" / "system" / "workflow_service"
-    spec = spec_from_file_location(
-        "temporalio_nexus_system_workflow_service",
-        module_path / "__init__.py",
-        submodule_search_locations=[str(module_path)],
-    )
-    if spec is None or spec.loader is None:
-        raise RuntimeError(f"Cannot load generated system service from {module_path}")
-    module = module_from_spec(spec)
-    sys.modules[spec.name] = module
-    spec.loader.exec_module(module)
+    services: list[type] = []
+    system_path = base_dir / "temporalio" / "nexus" / "system"
+    for module_path in sorted(p.parent for p in system_path.glob("*/services.py")):
+        spec = spec_from_file_location(
+            f"temporalio_nexus_system_{module_path.name}",
+            module_path / "__init__.py",
+            submodule_search_locations=[str(module_path)],
+        )
+        if spec is None or spec.loader is None:
+            raise RuntimeError(
+                f"Cannot load generated system service from {module_path}"
+            )
+        module = module_from_spec(spec)
+        sys.modules[spec.name] = module
+        spec.loader.exec_module(module)
+        services.extend(
+            value
+            for value in vars(module.services).values()
+            if isinstance(value, type) and nexusrpc.get_service_definition(value)
+        )
 
+    annotations = [a for service in services for a in get_type_hints(service).values()]
     roots: list[Descriptor] = []
-    for annotation in get_type_hints(module._services.WorkflowService).values():
+    for annotation in annotations:
         if get_origin(annotation) is not nexusrpc.Operation:
             continue
         for operation_type in get_args(annotation):