Use target serialization context for system Nexus signal-with-start

scripts/gen_nexus_system_api.py

@@ -11,7 +11,7 @@
 
 base_dir = Path(__file__).parent.parent
 sys.path.insert(0, str(base_dir))
-wit_input_dir = (
+sdk_core_wit_input_dir = (
     base_dir
     / "temporalio"
     / "bridge"
@@ -22,8 +22,9 @@
     / "api_upstream"
     / "nexus"
 )
-wit_path = wit_input_dir / "workflow-service.wit"
-wit_deps_dir = wit_input_dir / "deps"
+# Temporary checked-in WIT used until sdk-core updates its API snapshot.
+wit_path = base_dir / "scripts" / "nexus_system_workflow_service.wit"
+wit_deps_dir = sdk_core_wit_input_dir / "deps"
 python_support_path = base_dir / "scripts" / "nex_gen_support.py"
 output_dir = base_dir / "temporalio" / "nexus" / "system" / "workflow_service"
 workflow_init_path = base_dir / "temporalio" / "workflow" / "__init__.py"
@@ -35,15 +36,28 @@
     / "v1"
     / "request_response.proto"
 )
+NEX_GEN_VERSION = "0.2.0"
 
 
 def nex_gen_command() -> list[str]:
     if bin_path := os.environ.get("NEX_GEN_BIN"):
         return [bin_path]
 
-    if shutil.which("nex-gen") is None:
-        subprocess.check_call(["cargo", "install", "--locked", "nex-gen", "--force"])
-    return ["nex-gen"]
+    if shutil.which("nexgen") is None:
+        subprocess.check_call(
+            [
+                "cargo",
+                "install",
+                "--locked",
+                "nex-gen",
+                "--version",
+                NEX_GEN_VERSION,
+                "--features",
+                "advanced",
+                "--force",
+            ]
+        )
+    return ["nexgen"]
 
 
 def build_descriptor_set(descriptor_path: Path) -> None:
@@ -114,13 +128,10 @@ def generate_nexus_system_api() -> None:
         subprocess.check_call(
             [
                 *command,
-                "generate",
-                "--lang",
                 "python",
-                "--input",
                 str(wit_path),
-                "--input",
                 str(wit_deps_dir),
+                "--native-api",
                 "--support-file",
                 str(python_support_path),
                 "--descriptors",

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
+from typing import cast, get_args, get_origin, get_type_hints
 
 import google.protobuf.message
 import nexusrpc
@@ -19,14 +19,15 @@
 from temporalio.bridge.proto.workflow_completion.workflow_completion_pb2 import (
     WorkflowActivationCompletion,
 )
+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" / "service.py"
-    )
+    module_path = base_dir / "temporalio" / "nexus" / "system" / "workflow_service"
     spec = spec_from_file_location(
-        "temporalio_nexus_system_workflow_service", module_path
+        "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}")
@@ -35,10 +36,14 @@ def discover_system_nexus_roots() -> list[Descriptor]:
     spec.loader.exec_module(module)
 
     roots: list[Descriptor] = []
-    for operation in vars(module.WorkflowService).values():
-        if not isinstance(operation, nexusrpc.Operation):
+    for annotation in get_type_hints(module._services.WorkflowService).values():
+        if get_origin(annotation) is not nexusrpc.Operation:
             continue
-        for proto_type in (operation.input_type, operation.output_type):
+        for operation_type in get_args(annotation):
+            converter = _get_transfer_type_converter(operation_type)
+            proto_type = (
+                converter.transfer_type if converter is not None else operation_type
+            )
             if isinstance(proto_type, type) and issubclass(
                 proto_type, google.protobuf.message.Message
             ):

scripts/nex_gen_support.py

@@ -10,6 +10,12 @@
 import temporalio.api.workflow.v1
 import temporalio.common
 import temporalio.converter
+import temporalio.nexus.system
+
+
+class SignalWithStartWorkflowModelRequest(typing.Protocol):
+    namespace: str
+    id: str
 
 
 def retry_policy_from_proto(
@@ -52,6 +58,12 @@ def workflow_type_to_proto(
     return common_pb2.WorkflowType(name=workflow_function_name(workflow_type))
 
 
+def workflow_type_from_proto(
+    proto: common_pb2.WorkflowType,
+) -> str:
+    return proto.name
+
+
 def task_queue_from_proto(
     proto: taskqueue_pb2.TaskQueue,
 ) -> str:
@@ -70,12 +82,33 @@ def workflow_namespace() -> str:
     return info().namespace
 
 
+def signal_with_start_workflow_serialization_context(
+    request: SignalWithStartWorkflowModelRequest,
+) -> temporalio.converter.WorkflowSerializationContext:
+    return temporalio.converter.WorkflowSerializationContext(
+        namespace=request.namespace,
+        workflow_id=request.id,
+    )
+
+
 def payloads_to_proto(
     values: collections.abc.Sequence[typing.Any],
 ) -> common_pb2.Payloads:
-    from temporalio.workflow import payload_converter
+    return (
+        temporalio.nexus.system._current_user_payload_converter().to_payloads_wrapper(
+            values
+        )
+    )
+
 
-    return payload_converter().to_payloads_wrapper(values)
+def payloads_from_proto(
+    proto: common_pb2.Payloads,
+) -> list[object]:
+    return list(
+        temporalio.nexus.system._current_user_payload_converter().from_payloads_wrapper(
+            proto
+        )
+    )
 
 
 def _clone_payload(payload: common_pb2.Payload) -> common_pb2.Payload:
@@ -87,20 +120,23 @@ def _clone_payload(payload: common_pb2.Payload) -> common_pb2.Payload:
 def _value_to_payload(value: object | common_pb2.Payload) -> common_pb2.Payload:
     if isinstance(value, common_pb2.Payload):
         return _clone_payload(value)
-    from temporalio.workflow import payload_converter
-
-    payloads = payload_converter().to_payloads_wrapper([value])
+    payloads = (
+        temporalio.nexus.system._current_user_payload_converter().to_payloads_wrapper(
+            [value]
+        )
+    )
     return _clone_payload(payloads.payloads[0])
 
 
 def _payload_to_value(payload: common_pb2.Payload) -> object:
     wrapper = common_pb2.Payloads()
     wrapper.payloads.add().CopyFrom(payload)
-    from temporalio.workflow import payload_converter
 
     return typing.cast(
         object,
-        payload_converter().from_payloads_wrapper(wrapper)[0],
+        temporalio.nexus.system._current_user_payload_converter().from_payloads_wrapper(
+            wrapper
+        )[0],
     )
 
 
@@ -177,6 +213,12 @@ def search_attributes_to_proto(
     return proto
 
 
+def search_attributes_from_proto(
+    proto: common_pb2.SearchAttributes,
+) -> temporalio.common.TypedSearchAttributes:
+    return temporalio.converter.decode_typed_search_attributes(proto)
+
+
 def priority_from_proto(
     proto: common_pb2.Priority,
 ) -> temporalio.common.Priority:
@@ -193,3 +235,25 @@ def versioning_override_to_proto(
     versioning_override: temporalio.common.VersioningOverride,
 ) -> temporalio.api.workflow.v1.VersioningOverride:
     return versioning_override._to_proto()  # pyright: ignore[reportPrivateUsage]
+
+
+def versioning_override_from_proto(
+    proto: temporalio.api.workflow.v1.VersioningOverride,
+) -> temporalio.common.VersioningOverride:
+    if proto.HasField("pinned") and proto.pinned.HasField("version"):
+        version = proto.pinned.version
+        return temporalio.common.PinnedVersioningOverride(
+            temporalio.common.WorkerDeploymentVersion(
+                deployment_name=version.deployment_name,
+                build_id=version.build_id,
+            )
+        )
+    if proto.pinned_version:
+        return temporalio.common.PinnedVersioningOverride(
+            temporalio.common.WorkerDeploymentVersion.from_canonical_string(
+                proto.pinned_version
+            )
+        )
+    if proto.auto_upgrade:
+        return temporalio.common.AutoUpgradeVersioningOverride()
+    raise ValueError("unknown versioning override proto shape")

scripts/nexus_system_workflow_service.wit

@@ -0,0 +1,128 @@
+package temporal:nexus@1.0.0;
+
+world system {
+  export workflow-service;
+}
+
+/// @nexus.endpoint "__temporal_system"
+/// @nexus.service-name "temporal.api.workflowservice.v1.WorkflowService"
+/// @nexus.namespace dotnet="Temporalio.Workflows"
+/// @nexus.operations-class dotnet="Workflow"
+/// @nexus.delay-load-temporalio-workflow
+/// @nexus.experimental
+interface workflow-service {
+  use nexus:temporal-types/model@1.0.0.{
+    duration,
+    memo,
+    payloads,
+    placeholder,
+    priority,
+    retry-policy,
+    search-attributes,
+    signal-function,
+    task-queue,
+    user-metadata,
+    versioning-override,
+    workflow-function,
+    workflow-id-conflict-policy,
+    workflow-id-reuse-policy,
+  };
+
+  /// @nexus.doc "Request fields for signaling a workflow, starting it first if needed."
+  /// @nexus.experimental
+  /// @nexus.proto "temporal.api.workflowservice.v1.SignalWithStartWorkflowExecutionRequest" typescript-import="@temporalio/proto"
+  record signal-with-start-workflow-request {
+    /// @nexus.doc
+    ///   python="Workflow type name or callable identifying the workflow to start."
+    ///   typescript="Workflow type name or workflow function identifying the workflow to start."
+    ///   dotnet="Workflow type name or workflow expression identifying the workflow to start."
+    /// @nexus.proto-field "workflow_type"
+    workflow: workflow-function,
+    /// @nexus.doc "Unique identifier for the workflow execution."
+    /// @nexus.proto-field "workflow_id"
+    id: string,
+    /// @nexus.doc "Task queue to run the workflow on."
+    task-queue: task-queue,
+    /// @nexus.doc
+    ///   python="Signal name or callable to send with the start request."
+    ///   typescript="Signal name or signal definition to send with the start request."
+    ///   dotnet="Signal name or signal expression to send with the start request."
+    /// @nexus.proto-field "signal_name"
+    signal: signal-function,
+    /// @nexus.doc "Total workflow execution timeout, including retries and continue-as-new."
+    /// @nexus.proto-field "workflow_execution_timeout"
+    execution-timeout: option<duration>,
+    /// @nexus.doc "Timeout of a single workflow run."
+    /// @nexus.proto-field "workflow_run_timeout"
+    run-timeout: option<duration>,
+    /// @nexus.doc "Timeout of a single workflow task."
+    /// @nexus.proto-field "workflow_task_timeout"
+    task-timeout: option<duration>,
+    /// @nexus.omit
+    identity: placeholder,
+    /// @nexus.omit
+    request-id: placeholder,
+    /// @nexus.doc "Behavior when a closed workflow with the same ID exists. Default is allow-duplicate."
+    /// @nexus.proto-field "workflow_id_reuse_policy"
+    /// @nexus.default "allow-duplicate"
+    id-reuse-policy: workflow-id-reuse-policy,
+    /// @nexus.doc "Behavior when a workflow is currently running with the same ID. Set to use-existing for idempotent deduplication on workflow ID. Cannot be set if id-reuse-policy is terminate-if-running."
+    /// @nexus.proto-field "workflow_id_conflict_policy"
+    id-conflict-policy: option<workflow-id-conflict-policy>,
+    /// @nexus.doc "Retry policy for the workflow."
+    retry-policy: option<retry-policy>,
+    /// @nexus.doc "Cron schedule for recurring workflow executions. See https://docs.temporal.io/cron-job."
+    cron-schedule: option<string>,
+    /// @nexus.doc "Memo for the workflow."
+    memo: option<memo>,
+    /// @nexus.doc "Typed search attributes for the workflow."
+    search-attributes: option<search-attributes>,
+    /// @nexus.doc "Priority of the workflow execution."
+    priority: option<priority>,
+    /// @nexus.doc "Override for workflow versioning behavior."
+    versioning-override: option<versioning-override>,
+    /// @nexus.doc "Amount of time to wait before starting the workflow. This does not work with cron-schedule."
+    /// @nexus.proto-field "workflow_start_delay"
+    start-delay: option<duration>,
+    user-metadata: option<user-metadata>,
+    /// @nexus.source python="workflow_namespace" typescript="workflowNamespace" dotnet="TemporalWorkflowContext.WorkflowNamespace"
+    namespace: string,
+    /// @nexus.omit
+    control: placeholder,
+    /// @nexus.omit
+    header: placeholder,
+    /// @nexus.omit
+    links: placeholder,
+    /// @nexus.omit
+    time-skipping-config: placeholder,
+  }
+
+  /// @nexus.experimental
+  /// @nexus.proto "temporal.api.workflowservice.v1.SignalWithStartWorkflowExecutionResponse" typescript-import="@temporalio/proto"
+  record signal-with-start-workflow-response {
+    run-id: option<string>,
+    started: option<bool>,
+    /// @nexus.omit
+    signal-link: placeholder,
+    /// @nexus.omit
+    first-execution-run-id: placeholder,
+  }
+
+  /// @nexus.doc
+  ///   "Signal a workflow, starting it first if needed."
+  ///   returns="A workflow handle to the started workflow."
+  /// @nexus.output-transform
+  ///   python-type="temporalio.workflow.ExternalWorkflowHandle[WorkflowResult]"
+  ///   python="temporalio.workflow.get_external_workflow_handle(request.id, run_id=result.run_id)"
+  ///   typescript-type="workflow.ExternalWorkflowHandle"
+  ///   typescript="workflow.getExternalWorkflowHandle(request.id, result.runId ?? undefined)"
+  ///   typescript-import="@temporalio/workflow"
+  ///   dotnet-type="Temporalio.Workflows.ExternalWorkflowHandle"
+  ///   dotnet="Temporalio.Workflows.Workflow.GetExternalWorkflowHandle(request.Id, result.RunId)"
+  /// @nexus.operation name="SignalWithStartWorkflowExecution"
+  /// @nexus.serialization-context python="signal_with_start_workflow_serialization_context"
+  /// @nexus.experimental
+  signal-with-start-workflow: func(
+    request: signal-with-start-workflow-request,
+  ) -> signal-with-start-workflow-response;
+}

temporalio/nexus/system/__init__.py

@@ -145,6 +145,20 @@ def _get_payload_converter(  # pyright: ignore[reportUnusedFunction]
     return _SystemNexusPayloadConverter(user_payload_converter)
 
 
+def _get_serialization_context(  # pyright: ignore[reportUnusedFunction]
+    service: str,
+    operation: str,
+    request: Any,
+) -> temporalio.converter.SerializationContext | None:
+    """Return the serialization context for a system Nexus operation."""
+    from .workflow_service import __nexus_operation_registry__
+
+    operation_info = __nexus_operation_registry__.get((service, operation))
+    if operation_info is None or operation_info.serialization_context is None:
+        return None
+    return operation_info.serialization_context(request)
+
+
 __all__ = [
     "TEMPORAL_SYSTEM_ENDPOINT",
     "is_system_endpoint",