diff --git a/api/services/pipecat/run_pipeline.py b/api/services/pipecat/run_pipeline.py index 8fb565cc..39718a1c 100644 --- a/api/services/pipecat/run_pipeline.py +++ b/api/services/pipecat/run_pipeline.py @@ -167,6 +167,34 @@ async def run_pipeline_telephony( user_id: int, call_id: str, transport_kwargs: dict, +) -> None: + """Run a pipeline for any telephony provider.""" + # Register before any async setup so deploy drains see calls that are still + # resolving DB/config/transport state. + register_active_call(workflow_run_id) + try: + await _run_pipeline_telephony_impl( + websocket, + provider_name=provider_name, + workflow_id=workflow_id, + workflow_run_id=workflow_run_id, + user_id=user_id, + call_id=call_id, + transport_kwargs=transport_kwargs, + ) + finally: + unregister_active_call(workflow_run_id) + + +async def _run_pipeline_telephony_impl( + websocket, + *, + provider_name: str, + workflow_id: int, + workflow_run_id: int, + user_id: int, + call_id: str, + transport_kwargs: dict, ) -> None: """Run a pipeline for any telephony provider. @@ -240,7 +268,7 @@ async def run_pipeline_telephony( ) try: - await _run_pipeline( + await _run_pipeline_impl( transport, workflow_id, workflow_run_id, @@ -264,6 +292,31 @@ async def run_pipeline_smallwebrtc( user_id: int, call_context_vars: dict = {}, user_provider_id: str | None = None, +) -> None: + """Run pipeline for WebRTC connections.""" + # Register before any async setup so deploy drains see calls that are still + # resolving DB/config/transport state. + register_active_call(workflow_run_id) + try: + await _run_pipeline_smallwebrtc_impl( + webrtc_connection, + workflow_id, + workflow_run_id, + user_id, + call_context_vars=call_context_vars, + user_provider_id=user_provider_id, + ) + finally: + unregister_active_call(workflow_run_id) + + +async def _run_pipeline_smallwebrtc_impl( + webrtc_connection: SmallWebRTCConnection, + workflow_id: int, + workflow_run_id: int, + user_id: int, + call_context_vars: dict = {}, + user_provider_id: str | None = None, ) -> None: """Run pipeline for WebRTC connections""" logger.debug( @@ -313,7 +366,7 @@ async def run_pipeline_smallwebrtc( ambient_noise_config, is_realtime=is_realtime, ) - await _run_pipeline( + await _run_pipeline_impl( transport, workflow_id, workflow_run_id, @@ -336,6 +389,35 @@ async def _run_pipeline( user_provider_id: str | None = None, workflow_run=None, resolved_user_config=None, +) -> None: + """Run the pipeline with active-call drain accounting.""" + register_active_call(workflow_run_id) + try: + await _run_pipeline_impl( + transport, + workflow_id, + workflow_run_id, + user_id, + call_context_vars=call_context_vars, + audio_config=audio_config, + user_provider_id=user_provider_id, + workflow_run=workflow_run, + resolved_user_config=resolved_user_config, + ) + finally: + unregister_active_call(workflow_run_id) + + +async def _run_pipeline_impl( + transport, + workflow_id: int, + workflow_run_id: int, + user_id: int, + call_context_vars: dict = {}, + audio_config: AudioConfig = None, + user_provider_id: str | None = None, + workflow_run=None, + resolved_user_config=None, ) -> None: """ Run the pipeline with the given transport and configuration @@ -895,12 +977,6 @@ async def _run_pipeline( register_audio_data_handler(audio_buffer, workflow_run_id, in_memory_audio_buffer) - # Mark this run active so a deploy can drain it before stopping this worker: - # the orchestrator polls GET /api/v1/health/active-calls and waits for zero - # before sending SIGTERM (uvicorn force-closes live call WebSockets on - # SIGTERM, so the wait has to come first). See api/services/pipecat/ - # active_calls.py and scripts/rolling_update.sh. - register_active_call(workflow_run_id) try: # Run the pipeline await run_pipeline_worker(task) @@ -908,7 +984,6 @@ async def _run_pipeline( except asyncio.CancelledError: logger.warning("Received CancelledError in _run_pipeline") finally: - unregister_active_call(workflow_run_id) # Close MCP sessions here, not in engine.cleanup(). The anyio cancel # scopes opened by MCPClient.start() in engine.initialize() are # task-affine; this finally runs in the same task as initialize(), diff --git a/api/tests/test_active_calls.py b/api/tests/test_active_calls.py index bb40b60a..40447906 100644 --- a/api/tests/test_active_calls.py +++ b/api/tests/test_active_calls.py @@ -6,11 +6,15 @@ worker. The guarantees that matter for draining: register/unregister are idempotent, and the count only reaches zero when every registered run is gone. """ +import asyncio + +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from api.routes import main as main_routes from api.services.pipecat import active_calls +from api.services.pipecat import run_pipeline as run_pipeline_module def setup_function(): @@ -72,6 +76,107 @@ def test_full_lifecycle_drains_to_zero(): assert active_calls.active_call_count() == 0 +@pytest.mark.asyncio +async def test_run_pipeline_counts_call_during_setup(monkeypatch): + entered_setup = asyncio.Event() + release_setup = asyncio.Event() + + async def fake_get_workflow_run(*args, **kwargs): + entered_setup.set() + await release_setup.wait() + raise RuntimeError("setup failed") + + monkeypatch.setattr( + run_pipeline_module.db_client, + "get_workflow_run", + fake_get_workflow_run, + ) + + task = asyncio.create_task( + run_pipeline_module._run_pipeline( + transport=object(), + workflow_id=1, + workflow_run_id=42, + user_id=7, + ) + ) + + await asyncio.wait_for(entered_setup.wait(), timeout=1.0) + assert active_calls.active_call_count() == 1 + + release_setup.set() + with pytest.raises(RuntimeError, match="setup failed"): + await asyncio.wait_for(task, timeout=1.0) + assert active_calls.active_call_count() == 0 + + +@pytest.mark.asyncio +async def test_webrtc_entrypoint_counts_call_during_setup(monkeypatch): + entered_setup = asyncio.Event() + release_setup = asyncio.Event() + + async def fake_get_workflow(*args, **kwargs): + entered_setup.set() + await release_setup.wait() + raise RuntimeError("setup failed") + + monkeypatch.setattr( + run_pipeline_module.db_client, "get_workflow", fake_get_workflow + ) + + task = asyncio.create_task( + run_pipeline_module.run_pipeline_smallwebrtc( + webrtc_connection=object(), + workflow_id=1, + workflow_run_id=43, + user_id=7, + ) + ) + + await asyncio.wait_for(entered_setup.wait(), timeout=1.0) + assert active_calls.active_call_count() == 1 + + release_setup.set() + with pytest.raises(RuntimeError, match="setup failed"): + await asyncio.wait_for(task, timeout=1.0) + assert active_calls.active_call_count() == 0 + + +@pytest.mark.asyncio +async def test_telephony_entrypoint_counts_call_during_setup(monkeypatch): + entered_setup = asyncio.Event() + release_setup = asyncio.Event() + + async def fake_get_workflow(*args, **kwargs): + entered_setup.set() + await release_setup.wait() + raise RuntimeError("setup failed") + + monkeypatch.setattr( + run_pipeline_module.db_client, "get_workflow", fake_get_workflow + ) + + task = asyncio.create_task( + run_pipeline_module.run_pipeline_telephony( + websocket=object(), + provider_name="twilio", + workflow_id=1, + workflow_run_id=44, + user_id=7, + call_id="call-1", + transport_kwargs={}, + ) + ) + + await asyncio.wait_for(entered_setup.wait(), timeout=1.0) + assert active_calls.active_call_count() == 1 + + release_setup.set() + with pytest.raises(RuntimeError, match="setup failed"): + await asyncio.wait_for(task, timeout=1.0) + assert active_calls.active_call_count() == 0 + + def test_active_calls_route_requires_configured_secret(monkeypatch): client = _make_active_calls_client(monkeypatch, configured_secret=None)