fix: implement PR review

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(),

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)