ai-video-fullstack/backend/services/pipecat/call_lifecycle.py

"""Shared call termination timing for prompt tools and workflow end nodes."""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any
from uuid import uuid4

from loguru import logger
from pipecat.frames.frames import (
    BotStartedSpeakingFrame,
    BotStoppedSpeakingFrame,
    DataFrame,
    InterruptionFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.utils.time import time_now_iso8601


CALL_END_EVENT_BY_STATE = {
    "pending": "conversation_end_requested",
    "speaking": "conversation_end_speaking",
    "cancelled": "conversation_end_cancelled",
    "ended": "conversation_ended",
}


def call_end_trace_event(
    state: str,
    *,
    reason: str = "",
    interruptible: bool = False,
    voice_interruptible: bool = False,
) -> dict[str, Any]:
    """Build one shared record used by live calls and headless tests."""

    return {
        "type": "workflow-event",
        "eventId": f"wfe_{uuid4().hex[:20]}",
        "event": CALL_END_EVENT_BY_STATE.get(
            state,
            "conversation_end_state_changed",
        ),
        "timestamp": time_now_iso8601(),
        **({"reason": reason} if reason else {}),
        "interruptible": interruptible,
        "voiceInterruptible": voice_interruptible,
    }


class SpeechPlaybackCompletion:
    """Awaitable completed by its exact marker at the transport output."""

    def __init__(self, coordinator: CallEndCoordinator):
        self._coordinator = coordinator
        self._future: asyncio.Future[None] = (
            asyncio.get_running_loop().create_future()
        )
        self._queued = False

    def __await__(self):
        return self._future.__await__()

    def done(self) -> bool:
        return self._future.done()

    @property
    def queued(self) -> bool:
        return self._queued

    def mark_queued(self) -> None:
        self._queued = True

    async def mark_played(self) -> None:
        await self._coordinator.complete_tracked_speech(self)

    def _resolve(self) -> None:
        if not self._future.done():
            self._future.set_result(None)


@dataclass
class FixedSpeechPlaybackMarkerFrame(DataFrame):
    """Ordered frame that identifies one fixed utterance at transport output."""

    completion: SpeechPlaybackCompletion


def playback_marker_for(
    completion: Awaitable[None] | None,
) -> FixedSpeechPlaybackMarkerFrame | None:
    """Build a marker only for the production call-end coordinator."""
    if not isinstance(completion, SpeechPlaybackCompletion):
        return None
    return FixedSpeechPlaybackMarkerFrame(completion=completion)


CallEndStateHandler = Callable[[str, bool], Awaitable[None]]


class CallEndCoordinator:
    """Coordinate committed and user-cancellable conversation endings."""

    def __init__(
        self,
        queue_end: Callable[[str], Awaitable[None]],
        on_state_change: CallEndStateHandler | None = None,
    ):
        self._queue_end = queue_end
        self._on_state_change = on_state_change
        self._state_tasks: set[asyncio.Task[None]] = set()
        self._ending = False
        self._interruptible = False
        self._armed = False
        self._speaking = False
        self._speech_stopped = asyncio.Event()
        self._speech_stopped.set()
        self._response_speech_started = False
        self._tracked_speech_completions: set[SpeechPlaybackCompletion] = set()
        self._finish_after_tracked_speech = False
        self._finished = False
        self._reason = ""

    @property
    def ending(self) -> bool:
        return self._ending

    @property
    def accepts_input(self) -> bool:
        """Keep input open while an interruptible ending is still pending."""
        return not self._finished and (not self._ending or self._interruptible)

    @property
    def interruptible(self) -> bool:
        return self._ending and self._interruptible and not self._finished

    @property
    def speaking(self) -> bool:
        """Whether transport output is currently producing bot speech."""
        return self._speaking

    @property
    def reason(self) -> str:
        """Current requested ending reason for client feedback and records."""

        return self._reason

    async def wait_until_silent(self) -> None:
        """Wait for the transport-owned bot speech boundary."""
        await self._speech_stopped.wait()

    def begin(self, reason: str, *, interruptible: bool = False) -> None:
        if self._finished:
            return
        self._ending = True
        self._interruptible = interruptible
        self._armed = False
        self._finish_after_tracked_speech = False
        self._reason = reason or "completed"
        self._schedule_state("speaking" if self._speaking else "pending")

    def cancel_pending_end(self) -> bool:
        """Cancel one pending ending when its tool explicitly allows it."""
        if not self.interruptible:
            return False
        was_interruptible = self._interruptible
        self._ending = False
        self._interruptible = False
        self._armed = False
        self._finish_after_tracked_speech = False
        self._response_speech_started = False
        # Transport speech state remains authoritative: if audio is currently
        # playing, the interruption path will still emit BotStoppedSpeakingFrame.
        logger.info("用户打断结束语,已撤销待挂断状态")
        self._schedule_state("cancelled", interruptible=was_interruptible)
        return True

    def _schedule_state(
        self,
        state: str,
        *,
        interruptible: bool | None = None,
    ) -> None:
        if self._on_state_change is not None:
            task = asyncio.create_task(
                self._on_state_change(
                    state,
                    self.interruptible
                    if interruptible is None
                    else interruptible,
                )
            )
            self._state_tasks.add(task)
            task.add_done_callback(self._state_tasks.discard)

    def begin_response(self) -> None:
        """Start tracking speech produced by one LLM response."""
        self._response_speech_started = False

    def arm_after_speech(self) -> None:
        """Wait for the next observed bot speech to finish."""
        if not self._ending:
            return
        self._armed = True

    def track_speech(self) -> SpeechPlaybackCompletion:
        """Register fixed speech and return its transport completion signal."""
        completion = SpeechPlaybackCompletion(self)
        self._tracked_speech_completions.add(completion)
        return completion

    async def complete_tracked_speech(
        self,
        completion: SpeechPlaybackCompletion,
    ) -> None:
        """Complete one fixed utterance when its marker reaches output."""
        if completion not in self._tracked_speech_completions:
            return
        self._tracked_speech_completions.remove(completion)
        completion._resolve()
        if (
            self._finish_after_tracked_speech
            and not self._tracked_speech_completions
        ):
            logger.info("所有工作流结束语播报完毕,挂断通话")
            await self.finish()

    async def arm_after_tracked_speech(self) -> None:
        """Finish after every already queued fixed utterance has played."""
        if not self._ending:
            return
        self._finish_after_tracked_speech = True
        if not self._tracked_speech_completions:
            await self.finish()

    async def finish_after_current_speech(self, *, has_text: bool) -> None:
        """Finish now if speech is absent/done, otherwise wait for its stop."""
        if not self._ending:
            return
        if not has_text:
            await self.finish()
            return
        if self._response_speech_started and not self._speaking:
            await self.finish()
            return
        self._armed = True

    async def finish(self) -> None:
        if self._finished or not self._ending:
            return
        self._finished = True
        self._interruptible = False
        pending_states = tuple(self._state_tasks)
        if pending_states:
            await asyncio.gather(*pending_states, return_exceptions=True)
        await self._queue_end(self._reason)

    async def observe(self, frame) -> None:
        if isinstance(frame, InterruptionFrame):
            # Pipecat discards queued data frames on interruption, including
            # playback markers. Treat already queued fixed speech as stopped so
            # an interrupted Message cannot block a later EndNode forever.
            interrupted = tuple(
                completion
                for completion in self._tracked_speech_completions
                if completion.queued
            )
            for completion in interrupted:
                await self.complete_tracked_speech(completion)
        elif isinstance(frame, BotStartedSpeakingFrame):
            self._speaking = True
            self._speech_stopped.clear()
            self._response_speech_started = True
            if self._ending:
                await self._emit_state("speaking")
        elif isinstance(frame, BotStoppedSpeakingFrame) and self._speaking:
            self._speaking = False
            self._speech_stopped.set()
            if self._armed:
                logger.info("结束语播报完毕,挂断通话")
                await self.finish()

    async def _emit_state(self, state: str) -> None:
        if self._on_state_change is not None:
            await self._on_state_change(state, self.interruptible)


class EndCallAfterSpeechProcessor(FrameProcessor):
    def __init__(self, coordinator: CallEndCoordinator):
        super().__init__()
        self._coordinator = coordinator

    async def process_frame(self, frame, direction: FrameDirection):
        await super().process_frame(frame, direction)
        await self.push_frame(frame, direction)
        await self._coordinator.observe(frame)