Replace bespoke replace_staged_configs with general transaction mechanism

Adds staging_transaction() context manager to JobOrchestrator to batch multiple
staging operations with a single notification. This replaces the bespoke
replace_staged_configs() method with a more general, future-proof pattern.

Benefits:
- Handles any combination of staging operations, not just specific use cases
- New operations automatically benefit from batched notifications
- Follows standard Python context manager pattern
- Cleaner API with fewer methods to maintain

Changes:
- Add staging_transaction() context manager to JobOrchestrator
- Support nested transactions for the same workflow
- Reject attempts to nest transactions for different workflows
- Update _notify_staged_changed() to defer notifications during transactions
- Update WorkflowController.start_workflow() to use transaction
- Update WorkflowStatusWidget._on_remove_click() to use transaction
- Remove deprecated replace_staged_configs() method
- Replace 3 replace_staged_configs tests with 4 comprehensive transaction tests

This fixes the N+1 notification problem in both start_workflow (N+2 → 2
notifications) and _on_remove_click (N+1 → 1 notification), and prevents
similar issues in future code.

Original prompt: Have a look at the latest commit and related code. Think about whether JobOrchestrator should have a transaction mechanism instead of some bespoke methods.

Follow-up: Yes, please make a todo list and implement!

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Sonnet 4.5 <[email protected]>

Author: SimonHeybrock

src/ess/livedata/dashboard/job_orchestrator.py

@@ -14,6 +14,7 @@
 import logging
 import uuid
 from collections.abc import Callable, Mapping
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import NewType
 from uuid import UUID
@@ -138,6 +139,10 @@ def __init__(
             WidgetLifecycleSubscriptionId, WidgetLifecycleCallbacks
         ] = {}
 
+        # Transaction state for batching staging operations
+        self._transaction_workflow: WorkflowId | None = None
+        self._transaction_depth: int = 0
+
         # Load persisted configs
         self._load_configs_from_store()
 
@@ -224,6 +229,90 @@ def _load_configs_from_store(self) -> None:
             # Future subscribers are notified via subscribe_to_workflow(), which
             # checks for existing active jobs and notifies immediately.
 
+    @contextmanager
+    def staging_transaction(self, workflow_id: WorkflowId):
+        """
+        Context manager for batching staging operations.
+
+        All staging operations (clear_staged_configs, stage_config) within the
+        context will trigger only a single notification when the context exits.
+        This prevents N+1 UI rebuild issues when performing multiple staging
+        operations together.
+
+        Transactions can be nested - only the outermost transaction triggers
+        the notification on exit.
+
+        Parameters
+        ----------
+        workflow_id
+            The workflow to perform staging operations on.
+
+        Yields
+        ------
+        None
+
+        Raises
+        ------
+        ValueError
+            If attempting to nest transactions for different workflows.
+
+        Examples
+        --------
+        Replace all staged configs in one operation:
+
+            with orchestrator.staging_transaction(workflow_id):
+                orchestrator.clear_staged_configs(workflow_id)
+                for source in sources:
+                    orchestrator.stage_config(workflow_id, source, ...)
+            # Single notification sent here
+
+        Remove specific sources:
+
+            with orchestrator.staging_transaction(workflow_id):
+                staged = orchestrator.get_staged_config(workflow_id)
+                orchestrator.clear_staged_configs(workflow_id)
+                for source, config in staged.items():
+                    if source not in to_remove:
+                        orchestrator.stage_config(workflow_id, source, ...)
+            # Single notification sent here
+        """
+        # Enter transaction
+        if self._transaction_workflow is None:
+            self._transaction_workflow = workflow_id
+        elif self._transaction_workflow != workflow_id:
+            msg = (
+                f'Cannot nest transactions for different workflows: '
+                f'current={self._transaction_workflow}, new={workflow_id}'
+            )
+            raise ValueError(msg)
+
+        self._transaction_depth += 1
+        self._logger.debug(
+            'Entering staging transaction for workflow %s (depth=%d)',
+            workflow_id,
+            self._transaction_depth,
+        )
+
+        try:
+            yield
+        finally:
+            # Exit transaction
+            self._transaction_depth -= 1
+            self._logger.debug(
+                'Exiting staging transaction for workflow %s (depth=%d)',
+                workflow_id,
+                self._transaction_depth,
+            )
+
+            if self._transaction_depth == 0:
+                # Outermost transaction exiting - send notification
+                self._logger.debug(
+                    'Transaction complete for workflow %s, notifying subscribers',
+                    workflow_id,
+                )
+                self._transaction_workflow = None
+                self._notify_staged_changed(workflow_id)
+
     def clear_staged_configs(self, workflow_id: WorkflowId) -> None:
         """
         Clear all staged configs for a workflow.
@@ -263,38 +352,6 @@ def stage_config(
         )
         self._notify_staged_changed(workflow_id)
 
-    def replace_staged_configs(
-        self,
-        workflow_id: WorkflowId,
-        *,
-        configs: dict[SourceName, JobConfig],
-    ) -> None:
-        """
-        Replace all staged configs with new configs in a single operation.
-
-        This is more efficient than clear + multiple stage_config calls because
-        it triggers only a single notification to subscribers, reducing UI rebuilds.
-
-        Parameters
-        ----------
-        workflow_id
-            The workflow to configure.
-        configs
-            Dict mapping source names to their configs.
-        """
-        # Clear existing staged configs
-        self._workflows[workflow_id].staged_jobs.clear()
-
-        # Add all new configs (with copies to prevent external mutation)
-        for source_name, job_config in configs.items():
-            self._workflows[workflow_id].staged_jobs[source_name] = JobConfig(
-                params=job_config.params.copy(),
-                aux_source_names=job_config.aux_source_names.copy(),
-            )
-
-        # Single notification for the entire operation
-        self._notify_staged_changed(workflow_id)
-
     def commit_workflow(self, workflow_id: WorkflowId) -> list[JobId]:
         """
         Commit staged configs and start workflow.
@@ -735,7 +792,21 @@ def unsubscribe_from_widget_lifecycle(
             )
 
     def _notify_staged_changed(self, workflow_id: WorkflowId) -> None:
-        """Notify all widget subscribers that staging area changed."""
+        """Notify all widget subscribers that staging area changed.
+
+        Notifications are deferred if currently in a staging transaction.
+        The transaction context manager will call this on exit.
+        """
+        # Skip notification if we're in a transaction - it will be sent on exit
+        if self._transaction_workflow is not None:
+            self._logger.debug(
+                'Deferring staged_changed notification for workflow %s '
+                '(in transaction, depth=%d)',
+                workflow_id,
+                self._transaction_depth,
+            )
+            return
+
         for subscription_id, callbacks in self._widget_subscriptions.items():
             if callbacks.on_staged_changed is not None:
                 try:

src/ess/livedata/dashboard/widgets/workflow_status_widget.py

@@ -799,18 +799,19 @@ def _on_remove_click(self, source_names: list[str]) -> None:
         """Handle remove button click - remove sources from staged config."""
         staged = self._orchestrator.get_staged_config(self._workflow_id)
 
-        # Remove the specified sources and re-stage the rest.
-        # Note: clear_staged_configs() and stage_config() both notify subscribers,
-        # so the widget will rebuild via _on_lifecycle_event callback.
-        self._orchestrator.clear_staged_configs(self._workflow_id)
-        for source_name, config in staged.items():
-            if source_name not in source_names:
-                self._orchestrator.stage_config(
-                    self._workflow_id,
-                    source_name=source_name,
-                    params=config.params,
-                    aux_source_names=config.aux_source_names,
-                )
+        # Remove the specified sources and re-stage the rest in a transaction.
+        # Transaction ensures only a single notification is sent, so the widget
+        # rebuilds once via _on_lifecycle_event callback.
+        with self._orchestrator.staging_transaction(self._workflow_id):
+            self._orchestrator.clear_staged_configs(self._workflow_id)
+            for source_name, config in staged.items():
+                if source_name not in source_names:
+                    self._orchestrator.stage_config(
+                        self._workflow_id,
+                        source_name=source_name,
+                        params=config.params,
+                        aux_source_names=config.aux_source_names,
+                    )
 
     def _on_reset_click(self) -> None:
         """Handle reset button click."""

src/ess/livedata/dashboard/workflow_controller.py

@@ -21,7 +21,7 @@
 from .configuration_adapter import ConfigurationState
 from .correlation_histogram import CorrelationHistogramController, make_workflow_spec
 from .data_service import DataService
-from .job_orchestrator import JobConfig, JobOrchestrator
+from .job_orchestrator import JobOrchestrator
 from .workflow_configuration_adapter import WorkflowConfigurationAdapter
 
 
@@ -133,15 +133,16 @@ def start_workflow(
         params_dict = config.model_dump(mode='json')
         aux_dict = aux_source_names.model_dump(mode='json') if aux_source_names else {}
 
-        # Build configs for all sources
-        configs = {
-            source_name: JobConfig(params=params_dict, aux_source_names=aux_dict)
-            for source_name in source_names
-        }
-
-        # Replace staged configs in a single operation (more efficient than
-        # clear + multiple stage_config calls, reduces UI rebuilds)
-        self._orchestrator.replace_staged_configs(workflow_id, configs=configs)
+        # Replace all staged configs in a transaction (single notification)
+        with self._orchestrator.staging_transaction(workflow_id):
+            self._orchestrator.clear_staged_configs(workflow_id)
+            for source_name in source_names:
+                self._orchestrator.stage_config(
+                    workflow_id,
+                    source_name=source_name,
+                    params=params_dict,
+                    aux_source_names=aux_dict,
+                )
 
         # Commit and start workflow
         return self._orchestrator.commit_workflow(workflow_id)