Refactor to create task on StartFrame, also cleanup

src/pipecat/processors/aggregators/dtmf_aggregator.py

@@ -14,6 +14,7 @@ from pipecat.frames.frames import (
     Frame,
     InputDTMFFrame,
     KeypadEntry,
+    StartFrame,
     TranscriptionFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
@@ -53,93 +54,73 @@ class DTMFAggregator(FrameProcessor):
         self._digit_event = asyncio.Event()
         self._aggregation_task: Optional[asyncio.Task] = None
 
-    def _create_aggregation_task(self) -> None:
-        """Creates the aggregation task if it hasn't been created yet."""
-        if not self._aggregation_task:
-            self._aggregation_task = self.create_task(self._aggregation_task_handler())
-
-    async def _stop(self) -> None:
-        """Stops and cleans up the aggregation task."""
-        if self._aggregation_task:
-            await self.cancel_task(self._aggregation_task)
-            self._aggregation_task = None
-
-    async def cleanup(self) -> None:
-        await super().cleanup()
-        if self._aggregation_task:
-            await self._stop()
-
     async def process_frame(self, frame: Frame, direction: FrameDirection) -> None:
         await super().process_frame(frame, direction)
 
-        if isinstance(frame, (EndFrame, CancelFrame)):
-            # Flush any pending aggregation before stopping
+        if isinstance(frame, StartFrame):
+            self._create_aggregation_task()
+            await self.push_frame(frame, direction)
+        elif isinstance(frame, (EndFrame, CancelFrame)):
             if self._aggregation:
                 await self._flush_aggregation()
-            await self._stop()
+            await self._stop_aggregation_task()
             await self.push_frame(frame, direction)
-            return
-
-        # Push all other frames downstream immediately
-        await self.push_frame(frame, direction)
-
-        if isinstance(frame, InputDTMFFrame):
+        elif isinstance(frame, InputDTMFFrame):
+            # Push the DTMF frame downstream first
+            await self.push_frame(frame, direction)
+            # Then handle it in order for the TranscriptionFrame to be emitted
+            # after the InputDTMFFrame
             await self._handle_dtmf_frame(frame)
+        else:
+            await self.push_frame(frame, direction)
 
-    async def _handle_dtmf_frame(
-        self,
-        frame: InputDTMFFrame,
-    ):
-        # Create task on first DTMF input if needed
-        if not self._aggregation_task:
-            self._create_aggregation_task()
+    async def _handle_dtmf_frame(self, frame: InputDTMFFrame):
+        """Handle DTMF input frame."""
+        is_first_digit = not self._aggregation
 
         digit_value = frame.button.value
         self._aggregation += digit_value
 
-        # If this is the first digit, send BotInterruptionFrame upstream
-        # But use a separate task to avoid interfering with current frame processing
-        if len(self._aggregation) == 1:
-            # Use create_task to avoid queue issues
-            self.create_task(self._send_interruption_frame())
+        # For first digit, schedule interruption in separate task
+        if is_first_digit:
+            asyncio.create_task(self._send_interruption_task())
 
         # Check for immediate flush conditions
         if frame.button == self._termination_digit:
             await self._flush_aggregation()
         else:
-            # Signal new digit received
+            # Signal digit received for timeout handling
             self._digit_event.set()
 
-    async def _send_interruption_frame(self):
-        """Send an interruption frame in a separate task to avoid queue issues.
-
-        This interruption frame allows for the user to interrupt the bot's speaking
-        with a keypress. Without this, the TranscriptionFrame generated is accompanied
-        by EmulatedUserStarted/StoppedSpeakingFrames, which the bot currently ignores.
-        This treats the keypress as an explicit input which the bot should respond to.
-        """
+    async def _send_interruption_task(self):
+        """Send interruption frame safely in a separate task."""
         try:
+            # Send the interruption frame
             await self.push_frame(BotInterruptionFrame(), FrameDirection.UPSTREAM)
         except Exception as e:
-            print(f"Error sending interruption frame: {e}")
+            # Log error but don't propagate
+            print(f"Error sending interruption: {e}")
+
+    def _create_aggregation_task(self) -> None:
+        """Creates the aggregation task if it hasn't been created yet."""
+        if not self._aggregation_task:
+            self._aggregation_task = self.create_task(self._aggregation_task_handler())
+
+    async def _stop_aggregation_task(self) -> None:
+        """Stops the aggregation task."""
+        if self._aggregation_task:
+            await self.cancel_task(self._aggregation_task)
+            self._aggregation_task = None
 
     async def _aggregation_task_handler(self):
         """Background task that handles timeout-based flushing."""
-        try:
-            while True:
-                try:
-                    await asyncio.wait_for(self._digit_event.wait(), timeout=self._idle_timeout)
-                except asyncio.TimeoutError:
-                    if self._aggregation:
-                        await self._flush_aggregation()
-                finally:
-                    self._digit_event.clear()
-        except asyncio.CancelledError:
-            # Task was cancelled - exit cleanly
-            pass
-        except Exception as e:
-            # Log unexpected errors but don't crash
-            print(f"Unexpected error in DTMF aggregation task: {e}")
+        while True:
+            try:
+                await asyncio.wait_for(self._digit_event.wait(), timeout=self._idle_timeout)
+                self._digit_event.clear()
+            except asyncio.TimeoutError:
+                if self._aggregation:
+                    await self._flush_aggregation()
 
     async def _flush_aggregation(self):
         """Flush the current aggregation as a TranscriptionFrame."""
@@ -147,16 +128,16 @@ class DTMFAggregator(FrameProcessor):
             return
 
         sequence = self._aggregation
-
-        # Create transcription with prefix for LLM context
         transcription_text = f"{self._prefix}{sequence}"
 
-        # Create and push transcription frame
         transcription_frame = TranscriptionFrame(
             text=transcription_text, user_id="", timestamp=time_now_iso8601()
         )
-
         await self.push_frame(transcription_frame)
 
-        # Reset aggregation
         self._aggregation = ""
+
+    async def cleanup(self) -> None:
+        """Clean up resources."""
+        await super().cleanup()
+        await self._stop_aggregation_task()