griptape-nodes 0.55.1__py3-none-any.whl → 0.56.0__py3-none-any.whl

griptape_nodes/retained_mode/managers/workflow_manager.py

@@ -4,7 +4,6 @@ import ast
 import asyncio
 import logging
 import pickle
-import pkgutil
 import re
 from dataclasses import dataclass, field, fields, is_dataclass
 from datetime import UTC, datetime
@@ -24,7 +23,7 @@ from griptape_nodes.drivers.storage import StorageBackend
 from griptape_nodes.exe_types.core_types import ParameterTypeBuiltin
 from griptape_nodes.exe_types.flow import ControlFlow
 from griptape_nodes.exe_types.node_types import BaseNode, EndNode, StartNode
-from griptape_nodes.node_library.workflow_registry import Workflow, WorkflowMetadata, WorkflowRegistry
+from griptape_nodes.node_library.workflow_registry import Workflow, WorkflowMetadata, WorkflowRegistry, WorkflowShape
 from griptape_nodes.retained_mode.events.app_events import (
     GetEngineVersionRequest,
     GetEngineVersionResultSuccess,
@@ -100,6 +99,9 @@ from griptape_nodes.retained_mode.events.workflow_events import (
     RunWorkflowWithCurrentStateRequest,
     RunWorkflowWithCurrentStateResultFailure,
     RunWorkflowWithCurrentStateResultSuccess,
+    SaveWorkflowFileFromSerializedFlowRequest,
+    SaveWorkflowFileFromSerializedFlowResultFailure,
+    SaveWorkflowFileFromSerializedFlowResultSuccess,
     SaveWorkflowRequest,
     SaveWorkflowResultFailure,
     SaveWorkflowResultSuccess,
@@ -115,7 +117,6 @@ if TYPE_CHECKING:
     from types import TracebackType
 
     from griptape_nodes.exe_types.core_types import Parameter
-    from griptape_nodes.node_library.library_registry import LibraryNameAndVersion
     from griptape_nodes.retained_mode.events.base_events import ResultPayload
     from griptape_nodes.retained_mode.events.node_events import SetLockNodeStateRequest
     from griptape_nodes.retained_mode.managers.event_manager import EventManager
@@ -123,6 +124,10 @@ if TYPE_CHECKING:
 
 T = TypeVar("T")
 
+# Type aliases for workflow shape building
+ParameterShapeInfo = dict[str, Any]  # Parameter metadata dict from _convert_parameter_to_minimal_dict
+NodeParameterMap = dict[str, ParameterShapeInfo]  # {param_name: param_info}
+WorkflowShapeNodes = dict[str, NodeParameterMap]  # {node_name: {param_name: param_info}}
 
 logger = logging.getLogger("griptape_nodes")
 
@@ -269,6 +274,10 @@ class WorkflowManager:
             SaveWorkflowRequest,
             self.on_save_workflow_request,
         )
+        event_manager.assign_manager_to_request_type(
+            SaveWorkflowFileFromSerializedFlowRequest,
+            self.on_save_workflow_file_from_serialized_flow_request,
+        )
         event_manager.assign_manager_to_request_type(LoadWorkflowMetadata, self.on_load_workflow_metadata_request)
         event_manager.assign_manager_to_request_type(
             PublishWorkflowRequest,
@@ -1115,111 +1124,287 @@ class WorkflowManager:
 
         return True
 
-    def _gather_workflow_imports(self) -> list[str]:
-        """Gathers all the imports for the saved workflow file, specifically for the events."""
-        import_template = "from {} import *"
-        import_statements = []
-
-        from griptape_nodes.retained_mode import events as events_pkg
+    class WriteWorkflowFileResult(NamedTuple):
+        """Result of writing a workflow file."""
 
-        # Iterate over all modules in the events package
-        for _finder, module_name, _is_pkg in pkgutil.iter_modules(events_pkg.__path__, events_pkg.__name__ + "."):
-            if module_name.endswith("generate_request_payload_schemas"):
-                continue
-            import_statements.append(import_template.format(module_name))
+        success: bool
+        error_details: str
 
-        return import_statements
-
-    def _generate_workflow_file_contents_and_metadata(  # noqa: C901, PLR0912, PLR0915
-        self,
-        file_name: str,
-        creation_date: datetime,
-        image_path: str | None = None,
-        prior_workflow: Workflow | None = None,
-        custom_metadata: WorkflowMetadata | None = None,
-    ) -> tuple[str, WorkflowMetadata]:
-        """Generate the contents of a workflow file.
+    def _write_workflow_file(self, file_path: Path, content: str, file_name: str) -> WriteWorkflowFileResult:
+        """Write workflow content to file with proper validation and error handling.
 
         Args:
-            file_name: The name of the workflow file
-            creation_date: The creation date for the workflow
-            image_path: Optional; the path to an image to include in the workflow metadata
-            prior_workflow: Optional; existing workflow to preserve branch info from
-            custom_metadata: Optional; pre-constructed metadata to use instead of generating it
-                           from the current workflow state. When provided, this metadata will be
-                           used directly, allowing branch/merge operations to pass specific metadata.
+            file_path: Path where to write the file
+            content: Content to write
+            file_name: Name for error messages
 
         Returns:
-            A tuple of (workflow_file_contents, workflow_metadata)
-
-        Raises:
-            ValueError, TypeError: If workflow generation fails
+            WriteWorkflowFileResult with success status and error details if failed
         """
-        # Get the engine version.
-        engine_version_request = GetEngineVersionRequest()
-        engine_version_result = GriptapeNodes.handle_request(request=engine_version_request)
-        if not isinstance(engine_version_result, GetEngineVersionResultSuccess):
-            details = f"Failed getting the engine version for workflow '{file_name}'."
-            raise TypeError(details)
+        # Check disk space before any file system operations
+        config_manager = GriptapeNodes.ConfigManager()
+        min_space_gb = config_manager.get_config_value("minimum_disk_space_gb_workflows")
+        if not OSManager.check_available_disk_space(file_path.parent, min_space_gb):
+            error_msg = OSManager.format_disk_space_error(file_path.parent)
+            details = f"Attempted to save workflow '{file_name}' (requires {min_space_gb:.1f} GB). Failed due to insufficient disk space: {error_msg}"
+            return self.WriteWorkflowFileResult(success=False, error_details=details)
+
+        # Create directory structure
         try:
-            engine_version_success = cast("GetEngineVersionResultSuccess", engine_version_result)
-            engine_version = (
-                f"{engine_version_success.major}.{engine_version_success.minor}.{engine_version_success.patch}"
-            )
-        except Exception as err:
-            details = f"Failed getting the engine version for workflow '{file_name}': {err}"
-            raise ValueError(details) from err
+            file_path.parent.mkdir(parents=True, exist_ok=True)
+        except OSError as e:
+            details = f"Attempted to save workflow '{file_name}'. Failed when creating directory structure: {e}"
+            return self.WriteWorkflowFileResult(success=False, error_details=details)
 
-        # Keep track of all of the nodes we create and the generated variable names for them.
-        node_uuid_to_node_variable_name: dict[SerializedNodeCommands.NodeUUID, str] = {}
+        # Write the file content
+        try:
+            with file_path.open("w", encoding="utf-8") as file:
+                file.write(content)
+        except OSError as e:
+            details = f"Attempted to save workflow '{file_name}'. Failed when writing file content: {e}"
+            return self.WriteWorkflowFileResult(success=False, error_details=details)
 
-        # Keep track of each flow and node index we've created.
-        flow_creation_index = 0
+        return self.WriteWorkflowFileResult(success=True, error_details="")
 
-        # Serialize from the top.
+    def on_save_workflow_request(self, request: SaveWorkflowRequest) -> ResultPayload:  # noqa: C901, PLR0912, PLR0915
+        # Start with the file name provided; we may change it.
+        file_name = request.file_name
+
+        # See if we had an existing workflow for this.
+        prior_workflow = None
+        creation_date = None
+        if file_name and WorkflowRegistry.has_workflow_with_name(file_name):
+            # Get the metadata.
+            prior_workflow = WorkflowRegistry.get_workflow_by_name(file_name)
+            # We'll use its creation date.
+            creation_date = prior_workflow.metadata.creation_date
+        elif file_name:
+            # If no prior workflow exists for the new name, check if there's a current workflow
+            # context (e.g., during rename operations) to preserve metadata from
+            context_manager = GriptapeNodes.ContextManager()
+            if context_manager.has_current_workflow():
+                current_workflow_name = context_manager.get_current_workflow_name()
+                if current_workflow_name and WorkflowRegistry.has_workflow_with_name(current_workflow_name):
+                    prior_workflow = WorkflowRegistry.get_workflow_by_name(current_workflow_name)
+                    creation_date = prior_workflow.metadata.creation_date
+
+        if (creation_date is None) or (creation_date == WorkflowManager.EPOCH_START):
+            # Either a new workflow, or a backcompat situation.
+            creation_date = datetime.now(tz=UTC)
+
+        # Let's see if this is a template file; if so, re-route it as a copy in the customer's workflow directory.
+        if prior_workflow and prior_workflow.metadata.is_template:
+            # Aha! User is attempting to save a template. Create a differently-named file in their workspace.
+            # Find the first available file name that doesn't conflict.
+            curr_idx = 1
+            free_file_found = False
+            while not free_file_found:
+                # Composite a new candidate file name to test.
+                new_file_name = f"{file_name}_{curr_idx}"
+                new_file_name_with_extension = f"{new_file_name}.py"
+                new_file_full_path = GriptapeNodes.ConfigManager().workspace_path.joinpath(new_file_name_with_extension)
+                if new_file_full_path.exists():
+                    # Keep going.
+                    curr_idx += 1
+                else:
+                    free_file_found = True
+                    file_name = new_file_name
+
+        # Get file name stuff prepped.
+        # Use the existing registered file path if this is an existing workflow (not a template)
+        if prior_workflow and not prior_workflow.metadata.is_template:
+            # Use the existing registered file path
+            relative_file_path = prior_workflow.file_path
+            file_path = Path(WorkflowRegistry.get_complete_file_path(relative_file_path))
+            # Extract file name from the path for metadata generation
+            if not file_name:
+                file_name = prior_workflow.metadata.name
+        else:
+            # Create new path in workspace for new workflows or templates
+            if not file_name:
+                file_name = datetime.now(tz=UTC).strftime("%d.%m_%H.%M")
+            relative_file_path = f"{file_name}.py"
+            file_path = GriptapeNodes.ConfigManager().workspace_path.joinpath(relative_file_path)
+
+        # First, serialize the current workflow state
         top_level_flow_request = GetTopLevelFlowRequest()
         top_level_flow_result = GriptapeNodes.handle_request(top_level_flow_request)
         if not isinstance(top_level_flow_result, GetTopLevelFlowResultSuccess):
-            details = f"Failed when requesting to get top level flow for workflow '{file_name}'."
-            raise TypeError(details)
+            details = f"Attempted to save workflow '{relative_file_path}'. Failed when requesting top level flow."
+            return SaveWorkflowResultFailure(result_details=details)
+
         top_level_flow_name = top_level_flow_result.flow_name
         serialized_flow_request = SerializeFlowToCommandsRequest(
             flow_name=top_level_flow_name, include_create_flow_command=True
         )
         serialized_flow_result = GriptapeNodes.handle_request(serialized_flow_request)
         if not isinstance(serialized_flow_result, SerializeFlowToCommandsResultSuccess):
-            details = f"Failed when serializing flow for workflow '{file_name}'."
-            raise TypeError(details)
+            details = f"Attempted to save workflow '{relative_file_path}'. Failed when serializing flow."
+            return SaveWorkflowResultFailure(result_details=details)
+
         serialized_flow_commands = serialized_flow_result.serialized_flow_commands
 
-        # Use custom metadata if provided, otherwise generate it
-        if custom_metadata is not None:
-            workflow_metadata = custom_metadata
+        # Extract branched_from information if it exists
+        branched_from = None
+        if prior_workflow and prior_workflow.metadata.branched_from:
+            branched_from = prior_workflow.metadata.branched_from
+
+        # Extract workflow shape if possible
+        workflow_shape = None
+        try:
+            workflow_shape_dict = self.extract_workflow_shape(workflow_name=file_name)
+            workflow_shape = WorkflowShape(inputs=workflow_shape_dict["input"], outputs=workflow_shape_dict["output"])
+        except ValueError:
+            # If we can't extract workflow shape, continue without it
+            pass
+
+        # Use the standalone request to save the workflow file
+        save_file_request = SaveWorkflowFileFromSerializedFlowRequest(
+            serialized_flow_commands=serialized_flow_commands,
+            file_name=file_name,
+            creation_date=creation_date,
+            image_path=request.image_path,
+            execution_flow_name=top_level_flow_name,
+            branched_from=branched_from,
+            workflow_shape=workflow_shape,
+            file_path=str(file_path),
+        )
+        save_file_result = self.on_save_workflow_file_from_serialized_flow_request(save_file_request)
+
+        if not isinstance(save_file_result, SaveWorkflowFileFromSerializedFlowResultSuccess):
+            details = f"Attempted to save workflow '{relative_file_path}'. Failed during file generation: {save_file_result.result_details}"
+            return SaveWorkflowResultFailure(result_details=details)
+
+        # Use the metadata returned by the save operation
+        workflow_metadata = save_file_result.workflow_metadata
+
+        # save the created workflow as an entry in the JSON config file.
+        registered_workflows = WorkflowRegistry.list_workflows()
+        if file_name not in registered_workflows:
+            # Only add to config if it's in the workspace directory (not an external file)
+            if not Path(relative_file_path).is_absolute():
+                try:
+                    GriptapeNodes.ConfigManager().save_user_workflow_json(str(file_path))
+                except OSError as e:
+                    details = f"Attempted to save workflow '{file_name}'. Failed when saving configuration: {e}"
+                    return SaveWorkflowResultFailure(result_details=details)
+            WorkflowRegistry.generate_new_workflow(metadata=workflow_metadata, file_path=relative_file_path)
+        # Update existing workflow's metadata in the registry
+        existing_workflow = WorkflowRegistry.get_workflow_by_name(file_name)
+        existing_workflow.metadata = workflow_metadata
+        details = f"Successfully saved workflow to: {save_file_result.file_path}"
+        return SaveWorkflowResultSuccess(
+            file_path=save_file_result.file_path, result_details=ResultDetails(message=details, level=logging.INFO)
+        )
+
+    def on_save_workflow_file_from_serialized_flow_request(
+        self, request: SaveWorkflowFileFromSerializedFlowRequest
+    ) -> ResultPayload:
+        """Save a workflow file from serialized flow commands without registry overhead."""
+        # Determine file path
+        if request.file_path:
+            # Use provided file path
+            file_path = Path(request.file_path)
         else:
-            # Create the Workflow Metadata header.
-            workflows_referenced = None
-            if serialized_flow_commands.referenced_workflows:
-                workflows_referenced = list(serialized_flow_commands.referenced_workflows)
-
-            workflow_metadata = self._generate_workflow_metadata(
-                file_name=file_name,
-                engine_version=engine_version,
+            # Default to workspace path
+            relative_file_path = f"{request.file_name}.py"
+            file_path = GriptapeNodes.ConfigManager().workspace_path.joinpath(relative_file_path)
+
+        # Use provided creation date or default to current time
+        creation_date = request.creation_date
+        if creation_date is None:
+            creation_date = datetime.now(tz=UTC)
+
+        # Generate metadata from the serialized commands
+        try:
+            workflow_metadata = self._generate_workflow_metadata_from_commands(
+                serialized_flow_commands=request.serialized_flow_commands,
+                file_name=request.file_name,
                 creation_date=creation_date,
-                node_libraries_referenced=list(serialized_flow_commands.node_libraries_used),
-                workflows_referenced=workflows_referenced,
-                prior_workflow=prior_workflow,
+                image_path=request.image_path,
+                branched_from=request.branched_from,
+                workflow_shape=request.workflow_shape,
             )
-            if workflow_metadata is None:
-                details = f"Failed to generate metadata for workflow '{file_name}'."
-                raise ValueError(details)
+        except Exception as err:
+            details = f"Attempted to save workflow file '{request.file_name}' from serialized flow commands. Failed during metadata generation: {err}"
+            return SaveWorkflowFileFromSerializedFlowResultFailure(result_details=details)
 
-        # Set the image if provided
-        if image_path:
-            workflow_metadata.image = image_path
+        # Use provided execution flow name or default to file name
+        execution_flow_name = request.execution_flow_name
+        if execution_flow_name is None:
+            execution_flow_name = request.file_name
 
+        # Generate the workflow file content
+        try:
+            final_code_output = self._generate_workflow_file_content(
+                serialized_flow_commands=request.serialized_flow_commands,
+                workflow_metadata=workflow_metadata,
+                execution_flow_name=execution_flow_name,
+            )
+        except Exception as err:
+            details = f"Attempted to save workflow file '{request.file_name}' from serialized flow commands. Failed during content generation: {err}"
+            return SaveWorkflowFileFromSerializedFlowResultFailure(result_details=details)
+
+        # Write the workflow file
+        write_result = self._write_workflow_file(file_path, final_code_output, request.file_name)
+        if not write_result.success:
+            return SaveWorkflowFileFromSerializedFlowResultFailure(result_details=write_result.error_details)
+
+        details = f"Successfully saved workflow file at: {file_path}"
+        return SaveWorkflowFileFromSerializedFlowResultSuccess(
+            file_path=str(file_path),
+            workflow_metadata=workflow_metadata,
+            result_details=ResultDetails(message=details, level=logging.INFO),
+        )
+
+    def _generate_workflow_metadata_from_commands(  # noqa: PLR0913
+        self,
+        serialized_flow_commands: SerializedFlowCommands,
+        file_name: str,
+        creation_date: datetime,
+        image_path: str | None = None,
+        branched_from: str | None = None,
+        workflow_shape: WorkflowShape | None = None,
+    ) -> WorkflowMetadata:
+        """Generate workflow metadata from serialized commands."""
+        # Get the engine version
+        engine_version_request = GetEngineVersionRequest()
+        engine_version_result = GriptapeNodes.handle_request(request=engine_version_request)
+        if not isinstance(engine_version_result, GetEngineVersionResultSuccess):
+            details = f"Failed getting the engine version for workflow '{file_name}'."
+            raise TypeError(details)
+
+        engine_version_success = cast("GetEngineVersionResultSuccess", engine_version_result)
+        engine_version = f"{engine_version_success.major}.{engine_version_success.minor}.{engine_version_success.patch}"
+
+        # Create the Workflow Metadata header
+        workflows_referenced = None
+        if serialized_flow_commands.referenced_workflows:
+            workflows_referenced = list(serialized_flow_commands.referenced_workflows)
+
+        return WorkflowMetadata(
+            name=str(file_name),
+            schema_version=WorkflowMetadata.LATEST_SCHEMA_VERSION,
+            engine_version_created_with=engine_version,
+            node_libraries_referenced=list(serialized_flow_commands.node_libraries_used),
+            workflows_referenced=workflows_referenced,
+            creation_date=creation_date,
+            last_modified_date=datetime.now(tz=UTC),
+            branched_from=branched_from,
+            workflow_shape=workflow_shape,
+            image=image_path,
+        )
+
+    def _generate_workflow_file_content(  # noqa: PLR0912, C901
+        self,
+        serialized_flow_commands: SerializedFlowCommands,
+        workflow_metadata: WorkflowMetadata,
+        execution_flow_name: str,
+    ) -> str:
+        """Generate workflow file content from serialized commands and metadata."""
         metadata_block = self._generate_workflow_metadata_header(workflow_metadata=workflow_metadata)
         if metadata_block is None:
-            details = f"Failed to generate metadata block for workflow '{file_name}'."
+            details = f"Failed to generate metadata block for workflow '{workflow_metadata.name}'."
             raise ValueError(details)
 
         import_recorder = ImportRecorder()
@@ -1233,7 +1418,7 @@ class WorkflowManager:
         for node in prereq_code:
             ast_container.add_node(node)
 
-        # Generate unique values code AST node.
+        # Generate unique values code AST node
         unique_values_node = self._generate_unique_values_code(
             unique_parameter_uuid_to_values=serialized_flow_commands.unique_parameter_uuid_to_values,
             prefix="top_level",
@@ -1241,7 +1426,10 @@ class WorkflowManager:
         )
         ast_container.add_node(unique_values_node)
 
-        # See if this serialized flow has a flow initialization command; if it does, we'll need to insert that.
+        # Keep track of each flow and node index we've created
+        flow_creation_index = 0
+
+        # See if this serialized flow has a flow initialization command; if it does, we'll need to insert that
         flow_initialization_command = serialized_flow_commands.flow_initialization_command
 
         match flow_initialization_command:
@@ -1263,7 +1451,7 @@ class WorkflowManager:
                 # No initialization command, deserialize into current context
                 pass
 
-        # Generate assign flow context AST node, if we have any children commands.
+        # Generate assign flow context AST node, if we have any children commands
         # Skip content generation for referenced workflows - they should only have the import command
         is_referenced_workflow = isinstance(flow_initialization_command, ImportWorkflowAsReferencedSubFlowRequest)
         has_content_to_serialize = (
@@ -1275,12 +1463,15 @@ class WorkflowManager:
         )
 
         if not is_referenced_workflow and has_content_to_serialize:
+            # Keep track of all of the nodes we create and the generated variable names for them
+            node_uuid_to_node_variable_name: dict[SerializedNodeCommands.NodeUUID, str] = {}
+
             # Create the "with..." statement
             assign_flow_context_node = self._generate_assign_flow_context(
                 flow_initialization_command=flow_initialization_command, flow_creation_index=flow_creation_index
             )
 
-            # Generate nodes in flow AST node. This will create the node and apply all element modifiers.
+            # Generate nodes in flow AST node. This will create the node and apply all element modifiers
             nodes_in_flow = self._generate_nodes_in_flow(
                 serialized_flow_commands, import_recorder, node_uuid_to_node_variable_name
             )
@@ -1307,7 +1498,7 @@ class WorkflowManager:
                             )
                             assign_flow_context_node.body.append(cast("ast.stmt", sub_flow_import_node))
 
-            # Now generate the connection code and add it to the flow context.
+            # Now generate the connection code and add it to the flow context
             connection_asts = self._generate_connections_code(
                 serialized_connections=serialized_flow_commands.serialized_connections,
                 node_uuid_to_node_variable_name=node_uuid_to_node_variable_name,
@@ -1315,7 +1506,7 @@ class WorkflowManager:
             )
             assign_flow_context_node.body.extend(connection_asts)
 
-            # Now generate all the set parameter value code and add it to the flow context.
+            # Generate parameter values
             set_parameter_value_asts = self._generate_set_parameter_value_code(
                 set_parameter_value_commands=serialized_flow_commands.set_parameter_value_commands,
                 lock_commands=serialized_flow_commands.set_lock_commands_per_node,
@@ -1327,166 +1518,20 @@ class WorkflowManager:
 
             ast_container.add_node(assign_flow_context_node)
 
-        workflow_execution_code = (
-            self._generate_workflow_execution(
-                flow_name=top_level_flow_name,
-                import_recorder=import_recorder,
-            )
-            if top_level_flow_name
-            else None
+        # Generate workflow execution code
+        workflow_execution_code = self._generate_workflow_execution(
+            flow_name=execution_flow_name,
+            import_recorder=import_recorder,
+            workflow_metadata=workflow_metadata,
         )
         if workflow_execution_code is not None:
             for node in workflow_execution_code:
                 ast_container.add_node(node)
 
-            # TODO: https://github.com/griptape-ai/griptape-nodes/issues/1190 do child workflows
-
         # Generate final code from ASTContainer
         ast_output = "\n\n".join([ast.unparse(node) for node in ast_container.get_ast()])
         import_output = import_recorder.generate_imports()
-        final_code_output = f"{metadata_block}\n\n{import_output}\n\n{ast_output}\n"
-
-        return final_code_output, workflow_metadata
-
-    def on_save_workflow_request(self, request: SaveWorkflowRequest) -> ResultPayload:  # noqa: C901, PLR0912, PLR0915
-        # Start with the file name provided; we may change it.
-        file_name = request.file_name
-
-        # See if we had an existing workflow for this.
-        prior_workflow = None
-        creation_date = None
-        if file_name and WorkflowRegistry.has_workflow_with_name(file_name):
-            # Get the metadata.
-            prior_workflow = WorkflowRegistry.get_workflow_by_name(file_name)
-            # We'll use its creation date.
-            creation_date = prior_workflow.metadata.creation_date
-        elif file_name:
-            # If no prior workflow exists for the new name, check if there's a current workflow
-            # context (e.g., during rename operations) to preserve metadata from
-            context_manager = GriptapeNodes.ContextManager()
-            if context_manager.has_current_workflow():
-                current_workflow_name = context_manager.get_current_workflow_name()
-                if current_workflow_name and WorkflowRegistry.has_workflow_with_name(current_workflow_name):
-                    prior_workflow = WorkflowRegistry.get_workflow_by_name(current_workflow_name)
-                    creation_date = prior_workflow.metadata.creation_date
-
-        if (creation_date is None) or (creation_date == WorkflowManager.EPOCH_START):
-            # Either a new workflow, or a backcompat situation.
-            creation_date = datetime.now(tz=UTC)
-
-        # Let's see if this is a template file; if so, re-route it as a copy in the customer's workflow directory.
-        if prior_workflow and prior_workflow.metadata.is_template:
-            # Aha! User is attempting to save a template. Create a differently-named file in their workspace.
-            # Find the first available file name that doesn't conflict.
-            curr_idx = 1
-            free_file_found = False
-            while not free_file_found:
-                # Composite a new candidate file name to test.
-                new_file_name = f"{file_name}_{curr_idx}"
-                new_file_name_with_extension = f"{new_file_name}.py"
-                new_file_full_path = GriptapeNodes.ConfigManager().workspace_path.joinpath(new_file_name_with_extension)
-                if new_file_full_path.exists():
-                    # Keep going.
-                    curr_idx += 1
-                else:
-                    free_file_found = True
-                    file_name = new_file_name
-
-        # Get file name stuff prepped.
-        # Use the existing registered file path if this is an existing workflow (not a template)
-        if prior_workflow and not prior_workflow.metadata.is_template:
-            # Use the existing registered file path
-            relative_file_path = prior_workflow.file_path
-            file_path = Path(WorkflowRegistry.get_complete_file_path(relative_file_path))
-            # Extract file name from the path for metadata generation
-            if not file_name:
-                file_name = prior_workflow.metadata.name
-        else:
-            # Create new path in workspace for new workflows or templates
-            if not file_name:
-                file_name = datetime.now(tz=UTC).strftime("%d.%m_%H.%M")
-            relative_file_path = f"{file_name}.py"
-            file_path = GriptapeNodes.ConfigManager().workspace_path.joinpath(relative_file_path)
-
-        # Generate the workflow file contents
-        try:
-            final_code_output, workflow_metadata = self._generate_workflow_file_contents_and_metadata(
-                file_name=file_name,
-                creation_date=creation_date,
-                image_path=request.image_path,
-                prior_workflow=prior_workflow,
-            )
-        except Exception as err:
-            details = f"Attempted to save workflow '{relative_file_path}', but {err}"
-            return SaveWorkflowResultFailure(result_details=details)
-
-        # Create the pathing and write the file
-        try:
-            file_path.parent.mkdir(parents=True, exist_ok=True)
-        except OSError as e:
-            details = f"Attempted to save workflow '{file_name}'. Failed when creating directory: {e}"
-            return SaveWorkflowResultFailure(result_details=details)
-
-        # Check disk space before writing
-        config_manager = GriptapeNodes.ConfigManager()
-        min_space_gb = config_manager.get_config_value("minimum_disk_space_gb_workflows")
-        if not OSManager.check_available_disk_space(file_path.parent, min_space_gb):
-            error_msg = OSManager.format_disk_space_error(file_path.parent)
-            details = f"Attempted to save workflow '{file_name}' (requires {min_space_gb:.1f} GB). Failed: {error_msg}"
-            return SaveWorkflowResultFailure(result_details=details)
-
-        try:
-            with file_path.open("w", encoding="utf-8") as file:
-                file.write(final_code_output)
-        except OSError as e:
-            details = f"Attempted to save workflow '{file_name}'. Failed when writing file: {e}"
-            return SaveWorkflowResultFailure(result_details=details)
-
-        # save the created workflow as an entry in the JSON config file.
-        registered_workflows = WorkflowRegistry.list_workflows()
-        if file_name not in registered_workflows:
-            # Only add to config if it's in the workspace directory (not an external file)
-            if not Path(relative_file_path).is_absolute():
-                try:
-                    GriptapeNodes.ConfigManager().save_user_workflow_json(str(file_path))
-                except OSError as e:
-                    details = f"Attempted to save workflow '{file_name}'. Failed when saving configuration: {e}"
-                    return SaveWorkflowResultFailure(result_details=details)
-            WorkflowRegistry.generate_new_workflow(metadata=workflow_metadata, file_path=relative_file_path)
-        # Update existing workflow's metadata in the registry
-        existing_workflow = WorkflowRegistry.get_workflow_by_name(file_name)
-        existing_workflow.metadata = workflow_metadata
-        details = f"Successfully saved workflow to: {file_path}"
-        return SaveWorkflowResultSuccess(
-            file_path=str(file_path), result_details=ResultDetails(message=details, level=logging.INFO)
-        )
-
-    def _generate_workflow_metadata(  # noqa: PLR0913
-        self,
-        file_name: str,
-        engine_version: str,
-        creation_date: datetime,
-        node_libraries_referenced: list[LibraryNameAndVersion],
-        workflows_referenced: list[str] | None = None,
-        prior_workflow: Workflow | None = None,
-    ) -> WorkflowMetadata | None:
-        # Preserve branched workflow information if it exists
-        branched_from = None
-        if prior_workflow and prior_workflow.metadata.branched_from:
-            branched_from = prior_workflow.metadata.branched_from
-
-        workflow_metadata = WorkflowMetadata(
-            name=str(file_name),
-            schema_version=WorkflowMetadata.LATEST_SCHEMA_VERSION,
-            engine_version_created_with=engine_version,
-            node_libraries_referenced=node_libraries_referenced,
-            workflows_referenced=workflows_referenced,
-            creation_date=creation_date,
-            last_modified_date=datetime.now(tz=UTC),
-            branched_from=branched_from,
-        )
-
-        return workflow_metadata
+        return f"{metadata_block}\n\n{import_output}\n\n{ast_output}\n"
 
     def _replace_workflow_metadata_header(self, workflow_content: str, new_metadata: WorkflowMetadata) -> str | None:
         """Replace the metadata header in a workflow file with new metadata.
@@ -1516,11 +1561,13 @@ class WorkflowManager:
             toml_doc = tomlkit.document()
             toml_doc.add("dependencies", tomlkit.item([]))
             griptape_tool_table = tomlkit.table()
-            metadata_dict = workflow_metadata.model_dump()
+            # Strip out the Nones since TOML doesn't like those
+            # WorkflowShape is now serialized as JSON string by Pydantic field_serializer;
+            # this preserves the nil/null/None values that we WANT, but for all of the
+            # Python-related Nones, TOML will flip out if they are not stripped.
+            metadata_dict = workflow_metadata.model_dump(exclude_none=True)
             for key, value in metadata_dict.items():
-                # Strip out the Nones since TOML doesn't like those.
-                if value is not None:
-                    griptape_tool_table.add(key=key, value=value)
+                griptape_tool_table.add(key=key, value=value)
             toml_doc["tool"] = tomlkit.table()
             toml_doc["tool"]["griptape-nodes"] = griptape_tool_table  # type: ignore (this is the only way I could find to get tomlkit to do the dotted notation correctly)
         except Exception as err:
@@ -1545,14 +1592,20 @@ class WorkflowManager:
         self,
         flow_name: str,
         import_recorder: ImportRecorder,
+        workflow_metadata: WorkflowMetadata,
     ) -> list[ast.AST] | None:
         """Generates execute_workflow(...) and the __main__ guard."""
-        try:
-            workflow_shape = self.extract_workflow_shape(flow_name)
-        except ValueError:
+        # Use workflow shape from metadata if available, otherwise skip execution block
+        if workflow_metadata.workflow_shape is None:
             logger.info("Workflow shape does not have required Start or End Nodes. Skipping local execution block.")
             return None
 
+        # Convert WorkflowShape to dict format expected by the rest of the method
+        workflow_shape = {
+            "input": workflow_metadata.workflow_shape.inputs,
+            "output": workflow_metadata.workflow_shape.outputs,
+        }
+
         # === imports ===
         import_recorder.add_import("argparse")
         import_recorder.add_import("asyncio")
@@ -2608,7 +2661,7 @@ class WorkflowManager:
         )
 
         if flow_initialization_command is None:
-            # Construct AST for "GriptapeNodes.ContextManager().flow(GriptapeNodes.ContextManager().get_current_flow_name())"
+            # Construct AST for "GriptapeNodes.ContextManager().flow(GriptapeNodes.ContextManager().get_current_flow().flow_name)"
             flow_call = ast.Call(
                 func=ast.Attribute(
                     value=ast.Call(func=context_manager, args=[], keywords=[], lineno=1, col_offset=0),
@@ -2618,16 +2671,22 @@ class WorkflowManager:
                     col_offset=0,
                 ),
                 args=[
-                    ast.Call(
-                        func=ast.Attribute(
-                            value=ast.Call(func=context_manager, args=[], keywords=[], lineno=1, col_offset=0),
-                            attr="get_current_flow_name",
-                            ctx=ast.Load(),
+                    ast.Attribute(
+                        value=ast.Call(
+                            func=ast.Attribute(
+                                value=ast.Call(func=context_manager, args=[], keywords=[], lineno=1, col_offset=0),
+                                attr="get_current_flow",
+                                ctx=ast.Load(),
+                                lineno=1,
+                                col_offset=0,
+                            ),
+                            args=[],
+                            keywords=[],
                             lineno=1,
                             col_offset=0,
                         ),
-                        args=[],
-                        keywords=[],
+                        attr="flow_name",
+                        ctx=ast.Load(),
                         lineno=1,
                         col_offset=0,
                     )
@@ -3101,17 +3160,12 @@ class WorkflowManager:
         for node in nodes:
             for param in node.parameters:
                 # Expose only the parameters that are relevant for workflow input and output.
-                # Excluding list types as the individual parameters are exposed in the workflow shape.
-                # TODO (https://github.com/griptape-ai/griptape-nodes/issues/1090): This is a temporary solution until we know how to handle container types.
-                if param.type != ParameterTypeBuiltin.CONTROL_TYPE.value and not param.type.startswith("list"):
+                param_info = self.extract_parameter_shape_info(param, include_control_params=False)
+                if param_info is not None:
                     if node.name in workflow_shape[workflow_shape_type]:
-                        cast("dict", workflow_shape[workflow_shape_type][node.name])[param.name] = (
-                            self._convert_parameter_to_minimal_dict(param)
-                        )
+                        cast("dict", workflow_shape[workflow_shape_type][node.name])[param.name] = param_info
                     else:
-                        workflow_shape[workflow_shape_type][node.name] = {
-                            param.name: self._convert_parameter_to_minimal_dict(param)
-                        }
+                        workflow_shape[workflow_shape_type][node.name] = {param.name: param_info}
         return workflow_shape
 
     def extract_workflow_shape(self, workflow_name: str) -> dict[str, Any]:
@@ -3162,6 +3216,50 @@ class WorkflowManager:
 
         return workflow_shape
 
+    def extract_parameter_shape_info(
+        self, parameter: Parameter, *, include_control_params: bool = False
+    ) -> ParameterShapeInfo | None:
+        """Extract shape information from a parameter for workflow shape building.
+
+        Expose only the parameters that are relevant for workflow input and output.
+
+        Args:
+            parameter: The parameter to extract shape info from
+            include_control_params: Whether to include control type parameters (default: False)
+
+        Returns:
+            Parameter info dict if relevant for workflow shape, None if should be excluded
+        """
+        # TODO (https://github.com/griptape-ai/griptape-nodes/issues/1090): This is a temporary solution until we know how to handle container types.
+        # Always exclude list types until container type handling is implemented
+        if parameter.type.startswith("list"):
+            logger.warning(
+                "Skipping list parameter '%s' of type '%s' in workflow shape - container types not yet supported",
+                parameter.name,
+                parameter.type,
+            )
+            return None
+
+        # Conditionally exclude control types
+        if not include_control_params and parameter.type == ParameterTypeBuiltin.CONTROL_TYPE.value:
+            return None
+
+        return self._convert_parameter_to_minimal_dict(parameter)
+
+    def build_workflow_shape_from_parameter_info(
+        self, input_node_params: WorkflowShapeNodes, output_node_params: WorkflowShapeNodes
+    ) -> WorkflowShape:
+        """Build a WorkflowShape from collected parameter information.
+
+        Args:
+            input_node_params: Mapping of input node names to their parameter info
+            output_node_params: Mapping of output node names to their parameter info
+
+        Returns:
+            WorkflowShape object with inputs and outputs
+        """
+        return WorkflowShape(inputs=input_node_params, outputs=output_node_params)
+
     async def on_publish_workflow_request(self, request: PublishWorkflowRequest) -> ResultPayload:
         try:
             publisher_name = request.publisher_name