feat: generate System Nexus interceptor points

advanced/samples/python/wit/workflow_service/_system_nexus_interceptor.py

@@ -0,0 +1,97 @@
+# Generated by nexgen v0.2.2. DO NOT EDIT!
+
+from __future__ import annotations
+
+
+import abc
+import typing
+
+from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT
+
+from . import models
+
+if typing.TYPE_CHECKING:
+    import temporalio.workflow
+    from temporalio.worker._interceptor import StartNexusOperationInput
+
+
+__all__ = [
+    "_start_system_nexus_operation",
+    "_SystemNexusWorkflowOutboundInterceptorBase",
+    "_SystemNexusWorkflowOutboundInterceptorTerminal",
+]
+
+
+_InputT = typing.TypeVar("_InputT")
+_OutputT = typing.TypeVar("_OutputT")
+
+
+async def _start_system_nexus_operation(
+    interceptor: _SystemNexusWorkflowOutboundInterceptorBase,
+    input: StartNexusOperationInput[_InputT, _OutputT],
+) -> temporalio.workflow.NexusOperationHandle[_OutputT]:
+    if (
+        input.service == "temporal.api.workflowservice.v1.WorkflowService"
+        and input.operation_name == "SignalWithStartWorkflowExecution"
+    ):
+        typed_input = typing.cast(
+            "StartNexusOperationInput[models.SignalWithStartWorkflowRequest, models.SignalWithStartWorkflowResponse]",
+            input,
+        )
+        # The dispatch check above establishes that this operation's response type is _OutputT.
+        return typing.cast(
+            "temporalio.workflow.NexusOperationHandle[_OutputT]",
+            await interceptor.start_signal_with_start_workflow(typed_input.input),
+        )
+    raise ValueError(
+        f"unsupported System Nexus operation: {input.service}/{input.operation_name}"
+    )
+
+
+class _SystemNexusWorkflowOutboundInterceptorBase(abc.ABC):
+    @abc.abstractmethod
+    def _next_system_nexus_interceptor(
+        self,
+    ) -> _SystemNexusWorkflowOutboundInterceptorBase: ...
+
+    async def start_signal_with_start_workflow(
+        self, request: models.SignalWithStartWorkflowRequest
+    ) -> temporalio.workflow.NexusOperationHandle[
+        models.SignalWithStartWorkflowResponse
+    ]:
+        """Intercept the SignalWithStartWorkflow operation."""
+        return await self._next_system_nexus_interceptor().start_signal_with_start_workflow(
+            request
+        )
+
+
+class _SystemNexusWorkflowOutboundInterceptorTerminal(abc.ABC):
+    @abc.abstractmethod
+    async def _outbound_start_nexus_operation(
+        self,
+        input: StartNexusOperationInput[_InputT, _OutputT],
+    ) -> temporalio.workflow.NexusOperationHandle[_OutputT]: ...
+
+    async def start_signal_with_start_workflow(
+        self, request: models.SignalWithStartWorkflowRequest
+    ) -> temporalio.workflow.NexusOperationHandle[
+        models.SignalWithStartWorkflowResponse
+    ]:
+        from temporalio.worker._interceptor import StartNexusOperationInput
+        from temporalio.workflow import NexusOperationCancellationType
+
+        return await self._outbound_start_nexus_operation(
+            StartNexusOperationInput(
+                endpoint=TEMPORAL_SYSTEM_ENDPOINT,
+                service="temporal.api.workflowservice.v1.WorkflowService",
+                operation="SignalWithStartWorkflowExecution",
+                input=request,
+                output_type=models.SignalWithStartWorkflowResponse,
+                schedule_to_close_timeout=None,
+                schedule_to_start_timeout=None,
+                start_to_close_timeout=None,
+                cancellation_type=NexusOperationCancellationType.WAIT_COMPLETED,
+                headers=None,
+                summary=None,
+            )
+        )

src/generator/dotnet.rs

@@ -2257,22 +2257,21 @@ impl ExternalModelBackend for DotNetExternalModels {
 pub(crate) fn generate(
     tree: &crate::spec::ApiSpecTree<PlannedFamily>,
     support: &crate::SupportFiles,
-    mode: GenerationMode,
 ) -> Result<GeneratedFiles> {
     match &tree.root {
         ApiSpecNode::Leaf(leaf) => {
             let support_fragments = support_fragments_for_plan(&leaf.spec, support);
-            generate_leaf(&leaf.spec, &support_fragments, mode)
+            generate_leaf(&leaf.spec, &support_fragments)
         }
-        ApiSpecNode::Branch(branch) => generate_tree(branch, support, mode),
+        ApiSpecNode::Branch(branch) => generate_tree(branch, support),
     }
 }
 
 fn generate_leaf(
     api_plan: &PlannedSpec,
     support_fragments: &[SupportFragmentSpec],
-    mode: GenerationMode,
 ) -> Result<GeneratedFiles> {
+    let mode = crate::nexgen_config::current().mode;
     let support_namespace = dotnet_support_namespace(support_fragments)?;
     let generator = ApiPlanner::new(api_plan, support_namespace.as_deref())?;
     validate_dotnet_support_references(
@@ -2322,12 +2321,11 @@ fn generate_leaf(
 fn generate_tree(
     branch: &ApiSpecBranch<PlannedFamily>,
     support: &crate::SupportFiles,
-    mode: GenerationMode,
 ) -> Result<GeneratedFiles> {
     let mut files = BTreeMap::new();
     let mut warnings = Vec::new();
     for node in branch.children.values() {
-        generate_tree_node(node, support, mode, &mut files, &mut warnings)?;
+        generate_tree_node(node, support, &mut files, &mut warnings)?;
     }
     Ok(GeneratedFiles {
         layout: crate::generator::GeneratedOutputLayout::Directory,
@@ -2339,14 +2337,13 @@ fn generate_tree(
 fn generate_tree_node(
     node: &ApiSpecNode<PlannedFamily>,
     support: &crate::SupportFiles,
-    mode: GenerationMode,
     files: &mut BTreeMap<PathBuf, String>,
     warnings: &mut Vec<String>,
 ) -> Result<()> {
     match node {
         ApiSpecNode::Leaf(leaf) => {
             let support_fragments = support_fragments_for_plan(&leaf.spec, support);
-            let generated = generate_leaf(&leaf.spec, &support_fragments, mode)?;
+            let generated = generate_leaf(&leaf.spec, &support_fragments)?;
             warnings.extend(generated.warnings);
             let prefix = leaf.module_path.to_path_buf();
             for (path, contents) in generated.files {
@@ -2356,7 +2353,7 @@ fn generate_tree_node(
         }
         ApiSpecNode::Branch(branch) => {
             for node in branch.children.values() {
-                generate_tree_node(node, support, mode, files, warnings)?;
+                generate_tree_node(node, support, files, warnings)?;
             }
             Ok(())
         }

src/generator/go.rs

@@ -625,9 +625,8 @@ pub(crate) fn generate(
     api_plan: &PlannedSpec,
     support_fragments: &[SupportFragmentSpec],
     options: &GoOptions,
-    mode: GenerationMode,
 ) -> Result<GeneratedFiles> {
-    generate_in_tree(api_plan, support_fragments, options, mode, &[])
+    generate_in_tree(api_plan, support_fragments, options, &[])
 }
 
 /// Generates one input file, given every JSON model the whole generate closure
@@ -643,10 +642,9 @@ fn generate_in_tree(
     api_plan: &PlannedSpec,
     support_fragments: &[SupportFragmentSpec],
     options: &GoOptions,
-    mode: GenerationMode,
     tree_models: &[PlannedJsonType],
 ) -> Result<GeneratedFiles> {
-    ApiPlanner::new(api_plan, options, mode, tree_models)?.generate(support_fragments)
+    ApiPlanner::new(api_plan, options, tree_models)?.generate(support_fragments)
 }
 
 /// Every JSON model declared anywhere in the generate closure.
@@ -662,11 +660,10 @@ pub(crate) fn generate_tree(
     tree: &ApiSpecTree<PlannedFamily>,
     support: &SupportFiles,
     options: &GoOptions,
-    mode: GenerationMode,
 ) -> Result<GeneratedFiles> {
     match &tree.root {
-        ApiSpecNode::Leaf(leaf) => generate_single_leaf(leaf, support, options, mode),
-        ApiSpecNode::Branch(branch) => generate_branch_tree(branch, support, options, mode),
+        ApiSpecNode::Leaf(leaf) => generate_single_leaf(leaf, support, options),
+        ApiSpecNode::Branch(branch) => generate_branch_tree(branch, support, options),
     }
 }
 
@@ -678,9 +675,8 @@ fn generate_single_leaf(
     leaf: &ApiSpecLeaf<PlannedFamily>,
     support: &SupportFiles,
     options: &GoOptions,
-    mode: GenerationMode,
 ) -> Result<GeneratedFiles> {
-    let mut generated = generate(&leaf.spec, &support.fragments, options, mode)?;
+    let mut generated = generate(&leaf.spec, &support.fragments, options)?;
     if go_tree_has_json_models(&[leaf]) {
         let package_name = GoPackageContext::new(&leaf.spec, options)?.package_name;
         insert_generated_file(
@@ -696,7 +692,6 @@ fn generate_branch_tree(
     branch: &ApiSpecBranch<PlannedFamily>,
     support: &SupportFiles,
     options: &GoOptions,
-    mode: GenerationMode,
 ) -> Result<GeneratedFiles> {
     let mut leaves = Vec::new();
     collect_leaf_specs(branch, &mut leaves);
@@ -729,7 +724,7 @@ fn generate_branch_tree(
         // references to sibling files are unqualified within the one package.
         let mut leaf_spec = leaf.spec.clone();
         leaf_spec.module_path = root.clone();
-        let generated = generate_in_tree(&leaf_spec, &[], options, mode, &tree_models)?;
+        let generated = generate_in_tree(&leaf_spec, &[], options, &tree_models)?;
         warnings.extend(generated.warnings);
 
         // `generate` names the single JSON member file `<package>.go`; re-key it
@@ -885,9 +880,11 @@ fn plan_uses_json_models(api_plan: &PlannedSpec) -> bool {
 }
 
 impl GoExternalModels {
-    fn new(api_plan: &PlannedSpec, package: GoPackageContext, mode: GenerationMode) -> Self {
+    fn new(api_plan: &PlannedSpec, package: GoPackageContext) -> Self {
         if plan_uses_json_models(api_plan) {
-            Self::Json(json::ModelBackend::new(mode == GenerationMode::NativeApi))
+            Self::Json(json::ModelBackend::new(
+                crate::nexgen_config::current().mode == GenerationMode::NativeApi,
+            ))
         } else {
             Self::Proto(proto::ModelBackend::new(package))
         }
@@ -1006,7 +1003,6 @@ impl GoExternalModels {
 
 struct ApiPlanner<'a> {
     api_plan: &'a PlannedSpec,
-    mode: GenerationMode,
     package: GoPackageContext,
     imports: BTreeSet<String>,
     external_models: GoExternalModels,
@@ -1020,7 +1016,6 @@ impl<'a> ApiPlanner<'a> {
     fn new(
         api_plan: &'a PlannedSpec,
         options: &GoOptions,
-        mode: GenerationMode,
         tree_models: &[PlannedJsonType],
     ) -> Result<Self> {
         let package = GoPackageContext::new(api_plan, options)?;
@@ -1034,13 +1029,12 @@ impl<'a> ApiPlanner<'a> {
         collect_imports_from_plan(api_plan, &mut imports);
         imports.retain(|import_path| !package.is_self_import(import_path));
 
-        let mut external_models = GoExternalModels::new(api_plan, package.clone(), mode);
+        let mut external_models = GoExternalModels::new(api_plan, package.clone());
         external_models.adopt_tree_models(tree_models);
         external_models.prepare(api_plan)?;
 
         Ok(Self {
             api_plan,
-            mode,
             package,
             imports,
             external_models,
@@ -1156,7 +1150,7 @@ impl<'a> ApiPlanner<'a> {
             self.imports
                 .insert("github.com/nexus-rpc/sdk-go/nexus".to_string());
         }
-        if self.mode == GenerationMode::NativeApi
+        if crate::nexgen_config::current().mode == GenerationMode::NativeApi
             && self.external_models.renders_operation_references()
             && !self.api_plan.services.is_empty()
             && !self.package.is_self_import("go.temporal.io/sdk/workflow")

src/generator/java.rs

@@ -6,7 +6,7 @@ use heck::ToLowerCamelCase;
 use crate::error::{Error, Result};
 use crate::generator::json_schema::java as java_json;
 use crate::generator::json_schema::java::JavaContext;
-use crate::generator::{GeneratedFiles, GeneratedOutputLayout, GenerationMode};
+use crate::generator::{GeneratedFiles, GeneratedOutputLayout};
 use crate::planning::{PlannedFamily, PlannedJsonType};
 use crate::spec::{ApiSpecLeaf, ApiSpecNode};
 use crate::spec::{ExternalTypeSpec, ModulePath, ServiceSpec, TypeSpec};
@@ -17,7 +17,6 @@ const JAVA_FORMAT_LINE_LENGTH: usize = 88;
 pub(crate) fn generate(
     tree: &crate::spec::ApiSpecTree<PlannedFamily>,
     _support: &crate::SupportFiles,
-    _mode: GenerationMode,
     base_package: Option<&str>,
 ) -> Result<GeneratedFiles> {
     let base_package = base_package.unwrap_or(DEFAULT_PACKAGE);

src/generator/mod.rs

@@ -133,24 +133,27 @@ pub(crate) fn generate_files_for_tree_with_mode_and_options(
     mode: GenerationMode,
     options: GenerateFilesOptions,
 ) -> Result<GeneratedFiles> {
-    crate::compile_tree_to_files(language, tree, descriptors, support, mode, options)
+    let config = crate::nexgen_config::NexgenConfig {
+        mode,
+        ..crate::nexgen_config::current()
+    };
+    let _scope = crate::nexgen_config::scope(config);
+    crate::compile_tree_to_files(language, tree, descriptors, support, options)
 }
 
 pub(crate) fn generate_files_from_planned_tree(
     language: Language,
     tree: &ApiSpecTree<PlannedFamily>,
     support: &SupportFiles,
-    mode: GenerationMode,
     options: GenerateFilesOptions,
 ) -> Result<GeneratedFiles> {
+    let mode = crate::nexgen_config::current().mode;
     let mut generated = match language {
-        Language::Dotnet => dotnet::generate(tree, support, mode),
-        Language::Go => generate_go_tree(tree, support, mode, options),
-        Language::Java => java::generate(tree, support, mode, options.java_package_root.as_deref()),
-        Language::Python => python::generate(tree, support, mode),
-        Language::TypeScript => {
-            typescript::generate(tree, support, mode, options.ts_date_time_types)
-        }
+        Language::Dotnet => dotnet::generate(tree, support),
+        Language::Go => generate_go_tree(tree, support, options),
+        Language::Java => java::generate(tree, support, options.java_package_root.as_deref()),
+        Language::Python => python::generate(tree, support),
+        Language::TypeScript => typescript::generate(tree, support, options.ts_date_time_types),
         language => Err(Error::UnsupportedLanguage { language }),
     }?;
     generated.warnings = if mode == GenerationMode::NativeApi {
@@ -164,7 +167,6 @@ pub(crate) fn generate_files_from_planned_tree(
 fn generate_go_tree(
     tree: &ApiSpecTree<PlannedFamily>,
     support: &SupportFiles,
-    mode: GenerationMode,
     options: GenerateFilesOptions,
 ) -> Result<GeneratedFiles> {
     go::generate_tree(
@@ -174,7 +176,6 @@ fn generate_go_tree(
             output_dir_name: options.go_output_dir_name,
             ..go::GoOptions::default()
         },
-        mode,
     )
 }