Fix progress step recursion (home-assistant#153906)

Author: MartinHjelmare

homeassistant/data_entry_flow.py

@@ -645,12 +645,24 @@ class FlowHandler(Generic[_FlowContextT, _FlowResultT, _HandlerT]):
     __progress_task: asyncio.Task[Any] | None = None
     __no_progress_task_reported = False
     deprecated_show_progress = False
-    _progress_step_data: ProgressStepData[_FlowResultT] = {
-        "tasks": {},
-        "abort_reason": "",
-        "abort_description_placeholders": MappingProxyType({}),
-        "next_step_result": None,
-    }
+    __progress_step_data: ProgressStepData[_FlowResultT] | None = None
+
+    @property
+    def _progress_step_data(self) -> ProgressStepData[_FlowResultT]:
+        """Return progress step data.
+
+        A property is used instead of a simple attribute as derived classes
+        do not call super().__init__.
+        The property makes sure that the dict is initialized if needed.
+        """
+        if not self.__progress_step_data:
+            self.__progress_step_data = {
+                "tasks": {},
+                "abort_reason": "",
+                "abort_description_placeholders": MappingProxyType({}),
+                "next_step_result": None,
+            }
+        return self.__progress_step_data
 
     @property
     def source(self) -> str | None:
@@ -777,9 +789,10 @@ async def async_step__progress_step_abort(
         self, user_input: dict[str, Any] | None = None
     ) -> _FlowResultT:
         """Abort the flow."""
+        progress_step_data = self._progress_step_data
         return self.async_abort(
-            reason=self._progress_step_data["abort_reason"],
-            description_placeholders=self._progress_step_data[
+            reason=progress_step_data["abort_reason"],
+            description_placeholders=progress_step_data[
                 "abort_description_placeholders"
             ],
         )
@@ -795,14 +808,15 @@ async def async_step__progress_step_progress_done(
         without using async_show_progress_done.
         If no next step is set, abort the flow.
         """
-        if self._progress_step_data["next_step_result"] is None:
+        progress_step_data = self._progress_step_data
+        if (next_step_result := progress_step_data["next_step_result"]) is None:
             return self.async_abort(
-                reason=self._progress_step_data["abort_reason"],
-                description_placeholders=self._progress_step_data[
+                reason=progress_step_data["abort_reason"],
+                description_placeholders=progress_step_data[
                     "abort_description_placeholders"
                 ],
             )
-        return self._progress_step_data["next_step_result"]
+        return next_step_result
 
     @callback
     def async_external_step(
@@ -1021,48 +1035,55 @@ async def wrapper(
             self: FlowHandler[Any, ResultT], *args: P.args, **kwargs: P.kwargs
         ) -> ResultT:
             step_id = func.__name__.replace("async_step_", "")
-
+            progress_step_data = self._progress_step_data
             # Check if we have a progress task running
-            progress_task = self._progress_step_data["tasks"].get(step_id)
+            progress_task = progress_step_data["tasks"].get(step_id)
 
             if progress_task is None:
                 # First call - create and start the progress task
                 progress_task = self.hass.async_create_task(
                     func(self, *args, **kwargs),  # type: ignore[arg-type]
                     f"Progress step {step_id}",
                 )
-                self._progress_step_data["tasks"][step_id] = progress_task
-
-                if not progress_task.done():
-                    # Handle description placeholders
-                    placeholders = None
-                    if description_placeholders is not None:
-                        if callable(description_placeholders):
-                            placeholders = description_placeholders(self)
-                        else:
-                            placeholders = description_placeholders
-
-                    return self.async_show_progress(
-                        step_id=step_id,
-                        progress_action=step_id,
-                        progress_task=progress_task,
-                        description_placeholders=placeholders,
-                    )
+                progress_step_data["tasks"][step_id] = progress_task
+
+            if not progress_task.done():
+                # Handle description placeholders
+                placeholders = None
+                if description_placeholders is not None:
+                    if callable(description_placeholders):
+                        placeholders = description_placeholders(self)
+                    else:
+                        placeholders = description_placeholders
+
+                return self.async_show_progress(
+                    step_id=step_id,
+                    progress_action=step_id,
+                    progress_task=progress_task,
+                    description_placeholders=placeholders,
+                )
 
             # Task is done or this is a subsequent call
             try:
-                self._progress_step_data["next_step_result"] = await progress_task
+                progress_task_result = await progress_task
             except AbortFlow as err:
-                self._progress_step_data["abort_reason"] = err.reason
-                self._progress_step_data["abort_description_placeholders"] = (
+                progress_step_data["abort_reason"] = err.reason
+                progress_step_data["abort_description_placeholders"] = (
                     err.description_placeholders or {}
                 )
                 return self.async_show_progress_done(
                     next_step_id="_progress_step_abort"
                 )
             finally:
                 # Clean up task reference
-                self._progress_step_data["tasks"].pop(step_id, None)
+                progress_step_data["tasks"].pop(step_id, None)
+
+            # If the result type is FlowResultType.SHOW_PROGRESS_DONE
+            # an earlier show progress step has already been run and stored its result.
+            # In this case we should not overwrite the result,
+            # but just use the stored one.
+            if progress_task_result["type"] != FlowResultType.SHOW_PROGRESS_DONE:
+                progress_step_data["next_step_result"] = progress_task_result
 
             return self.async_show_progress_done(
                 next_step_id="_progress_step_progress_done"