mirror of
https://github.com/dograh-hq/dograh.git
synced 2026-07-25 12:01:04 +02:00
fix: address concurrency review findings (campaign-scoped counter, cleanup hardening, webrtc run validation)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
parent
bd69382a90
commit
6f977cf6d9
15 changed files with 594 additions and 50 deletions
|
|
@ -150,7 +150,11 @@ COUNTRY_CODES = {
|
|||
"IE": "353", # Ireland
|
||||
}
|
||||
|
||||
DEFAULT_ORG_CONCURRENCY_LIMIT = int(os.getenv("DEFAULT_ORG_CONCURRENCY_LIMIT", "10"))
|
||||
# Floor at 1 so a misconfigured env var (0 or negative) can't silently block
|
||||
# every call in the deployment.
|
||||
DEFAULT_ORG_CONCURRENCY_LIMIT = max(
|
||||
1, int(os.getenv("DEFAULT_ORG_CONCURRENCY_LIMIT", "10"))
|
||||
)
|
||||
DEFAULT_CAMPAIGN_RETRY_CONFIG = {
|
||||
"enabled": True,
|
||||
"max_retries": 1,
|
||||
|
|
|
|||
|
|
@ -721,6 +721,16 @@ async def signaling_websocket(
|
|||
logger.warning(f"workflow run {workflow_run_id} not found for user {user.id}")
|
||||
raise HTTPException(status_code=400, detail="Bad workflow_run_id")
|
||||
|
||||
# The URL's workflow_id drives org resolution, quota, and concurrency
|
||||
# accounting downstream — reject a workflow_id that doesn't own this run
|
||||
# so a caller can't charge their call to an unrelated workflow/org.
|
||||
if workflow_run.workflow_id != workflow_id:
|
||||
logger.warning(
|
||||
f"workflow run {workflow_run_id} does not belong to workflow "
|
||||
f"{workflow_id} (belongs to {workflow_run.workflow_id})"
|
||||
)
|
||||
raise HTTPException(status_code=400, detail="Bad workflow_run_id")
|
||||
|
||||
await signaling_manager.handle_websocket(
|
||||
websocket,
|
||||
workflow_id,
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ class CallConcurrencySlot:
|
|||
slot_id: str
|
||||
max_concurrent: int
|
||||
source: str
|
||||
scope_key: str | None = None
|
||||
|
||||
|
||||
class CallConcurrencyLimitError(Exception):
|
||||
|
|
@ -76,19 +77,28 @@ class CallConcurrencyService:
|
|||
*,
|
||||
source: str,
|
||||
timeout: float = 0,
|
||||
max_concurrent_override: int | None = None,
|
||||
scope_key: str | None = None,
|
||||
scope_max_concurrent: int | None = None,
|
||||
retry_interval: float = 1,
|
||||
) -> CallConcurrencySlot:
|
||||
org_concurrent_limit = await self.get_org_concurrent_limit(organization_id)
|
||||
if max_concurrent_override is None:
|
||||
max_concurrent = org_concurrent_limit
|
||||
else:
|
||||
max_concurrent = min(int(max_concurrent_override), org_concurrent_limit)
|
||||
"""Acquire a slot in the org-wide concurrency counter.
|
||||
|
||||
``scope_key``/``scope_max_concurrent`` additionally bound a secondary
|
||||
counter (e.g. ``campaign:<id>``) so a source can cap its own
|
||||
concurrency without measuring — or being starved by — unrelated calls
|
||||
in the same org.
|
||||
"""
|
||||
max_concurrent = await self.get_org_concurrent_limit(organization_id)
|
||||
if scope_max_concurrent is not None:
|
||||
scope_max_concurrent = int(scope_max_concurrent)
|
||||
|
||||
wait_start = time.time()
|
||||
while True:
|
||||
acquisition = await rate_limiter.try_acquire_concurrent_slot_details(
|
||||
organization_id, max_concurrent
|
||||
organization_id,
|
||||
max_concurrent,
|
||||
scope_key=scope_key,
|
||||
scope_max_concurrent=scope_max_concurrent,
|
||||
)
|
||||
if acquisition:
|
||||
logger.info(
|
||||
|
|
@ -102,15 +112,21 @@ class CallConcurrencyService:
|
|||
slot_id=acquisition.slot_id,
|
||||
max_concurrent=max_concurrent,
|
||||
source=source,
|
||||
scope_key=scope_key,
|
||||
)
|
||||
|
||||
wait_time = time.time() - wait_start
|
||||
if wait_time >= timeout:
|
||||
current_count = await rate_limiter.get_concurrent_count(organization_id)
|
||||
scope_note = (
|
||||
f", scope={scope_key} (limit={scope_max_concurrent})"
|
||||
if scope_key
|
||||
else ""
|
||||
)
|
||||
logger.warning(
|
||||
f"Concurrent call limit reached for org {organization_id}: "
|
||||
f"source={source}, active_calls={current_count}/{max_concurrent}, "
|
||||
f"waited={wait_time:.1f}s"
|
||||
f"source={source}, active_calls={current_count}/{max_concurrent}"
|
||||
f"{scope_note}, waited={wait_time:.1f}s"
|
||||
)
|
||||
raise CallConcurrencyLimitError(
|
||||
organization_id=organization_id,
|
||||
|
|
@ -132,6 +148,7 @@ class CallConcurrencyService:
|
|||
workflow_run_id,
|
||||
slot.organization_id,
|
||||
slot.slot_id,
|
||||
scope_key=slot.scope_key,
|
||||
)
|
||||
if stored:
|
||||
return
|
||||
|
|
@ -146,36 +163,66 @@ class CallConcurrencyService:
|
|||
*,
|
||||
source: str,
|
||||
timeout: float = 0,
|
||||
max_concurrent_override: int | None = None,
|
||||
scope_key: str | None = None,
|
||||
scope_max_concurrent: int | None = None,
|
||||
retry_interval: float = 1,
|
||||
) -> CallConcurrencySlot:
|
||||
slot = await self.acquire_org_slot(
|
||||
organization_id,
|
||||
source=source,
|
||||
timeout=timeout,
|
||||
max_concurrent_override=max_concurrent_override,
|
||||
scope_key=scope_key,
|
||||
scope_max_concurrent=scope_max_concurrent,
|
||||
retry_interval=retry_interval,
|
||||
)
|
||||
await self.bind_workflow_run(slot, workflow_run_id)
|
||||
return slot
|
||||
|
||||
async def unregister_active_call(self, workflow_run_id: int) -> bool:
|
||||
return await self.release_workflow_run_slot(workflow_run_id)
|
||||
"""Release the run's slot without ever raising.
|
||||
|
||||
Callers invoke this from ``finally`` blocks during pipeline/socket
|
||||
teardown; a cleanup failure must not mask the original exception.
|
||||
The slot mapping survives a failed release, so a later cleanup path
|
||||
(status callback, StasisEnd) or the Redis stale timeout recovers it.
|
||||
"""
|
||||
try:
|
||||
return await self.release_workflow_run_slot(workflow_run_id)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to release concurrent call slot for workflow run "
|
||||
f"{workflow_run_id}: {e}"
|
||||
)
|
||||
return False
|
||||
|
||||
async def release_slot(self, slot: CallConcurrencySlot | None) -> bool:
|
||||
if slot is None:
|
||||
return False
|
||||
return await rate_limiter.release_concurrent_slot(
|
||||
slot.organization_id, slot.slot_id
|
||||
released = await rate_limiter.release_concurrent_slot(
|
||||
slot.organization_id, slot.slot_id, scope_key=slot.scope_key
|
||||
)
|
||||
return bool(released)
|
||||
|
||||
async def release_workflow_run_slot(self, workflow_run_id: int) -> bool:
|
||||
mapping = await rate_limiter.get_workflow_slot_mapping(workflow_run_id)
|
||||
if not mapping:
|
||||
return False
|
||||
|
||||
org_id, slot_id = mapping
|
||||
released = await rate_limiter.release_concurrent_slot(org_id, slot_id)
|
||||
org_id, slot_id, scope_key = mapping
|
||||
released = await rate_limiter.release_concurrent_slot(
|
||||
org_id, slot_id, scope_key=scope_key
|
||||
)
|
||||
if released is None:
|
||||
# Redis error while releasing — keep the mapping so a later
|
||||
# cleanup path can retry instead of orphaning a live slot until
|
||||
# the stale timeout.
|
||||
logger.warning(
|
||||
f"Failed to release concurrent slot for workflow run "
|
||||
f"{workflow_run_id}; keeping mapping for retry"
|
||||
)
|
||||
return False
|
||||
await rate_limiter.delete_workflow_slot_mapping(workflow_run_id)
|
||||
if released:
|
||||
logger.info(f"Released concurrent slot for workflow run {workflow_run_id}")
|
||||
|
|
|
|||
|
|
@ -468,7 +468,11 @@ class CampaignCallDispatcher:
|
|||
Raises:
|
||||
ConcurrentSlotAcquisitionError: If slot cannot be acquired within timeout
|
||||
"""
|
||||
# Check for campaign-level max_concurrency in orchestrator_metadata
|
||||
# Check for campaign-level max_concurrency in orchestrator_metadata.
|
||||
# It caps this campaign's own concurrent calls via a campaign-scoped
|
||||
# counter — the org-wide limit still applies on top, but calls from
|
||||
# other sources (WebRTC, inbound, other campaigns) don't count
|
||||
# against the campaign's cap.
|
||||
campaign_max_concurrency = None
|
||||
if campaign.orchestrator_metadata:
|
||||
campaign_max_concurrency = campaign.orchestrator_metadata.get(
|
||||
|
|
@ -480,7 +484,12 @@ class CampaignCallDispatcher:
|
|||
organization_id,
|
||||
source=f"campaign:{campaign.id}",
|
||||
timeout=timeout,
|
||||
max_concurrent_override=campaign_max_concurrency,
|
||||
scope_key=(
|
||||
f"campaign:{campaign.id}"
|
||||
if campaign_max_concurrency is not None
|
||||
else None
|
||||
),
|
||||
scope_max_concurrent=campaign_max_concurrency,
|
||||
retry_interval=1,
|
||||
)
|
||||
except CallConcurrencyLimitError as e:
|
||||
|
|
|
|||
|
|
@ -113,41 +113,65 @@ class RateLimiter:
|
|||
return acquisition.slot_id if acquisition else None
|
||||
|
||||
async def try_acquire_concurrent_slot_details(
|
||||
self, organization_id: int, max_concurrent: int = 20
|
||||
self,
|
||||
organization_id: int,
|
||||
max_concurrent: int = 20,
|
||||
*,
|
||||
scope_key: str | None = None,
|
||||
scope_max_concurrent: int | None = None,
|
||||
) -> Optional[ConcurrentSlotAcquisition]:
|
||||
"""
|
||||
Try to acquire a concurrent call slot.
|
||||
Returns the slot_id and post-acquire active count if successful,
|
||||
or None if the limit is reached.
|
||||
|
||||
When ``scope_key``/``scope_max_concurrent`` are provided, the slot is
|
||||
also registered in a secondary counter (``concurrent_calls:<scope_key>``,
|
||||
e.g. ``campaign:<id>``) and acquisition additionally requires that
|
||||
counter to be below ``scope_max_concurrent``. Both counters are
|
||||
updated atomically. The scope-scoped slot must be released with the
|
||||
same ``scope_key``.
|
||||
"""
|
||||
redis_client = await self._get_redis()
|
||||
|
||||
concurrent_key = f"concurrent_calls:{organization_id}"
|
||||
scope_concurrent_key = f"concurrent_calls:{scope_key}" if scope_key else ""
|
||||
now = time.time()
|
||||
stale_cutoff = now - self.stale_call_timeout
|
||||
|
||||
# Lua script for atomic operation
|
||||
# Lua script for atomic operation across the org counter and the
|
||||
# optional scope counter (empty scope key = org-only acquisition).
|
||||
lua_script = """
|
||||
local key = KEYS[1]
|
||||
local scope_key = KEYS[2]
|
||||
local now = tonumber(ARGV[1])
|
||||
local max_concurrent = tonumber(ARGV[2])
|
||||
local stale_cutoff = tonumber(ARGV[3])
|
||||
local slot_id = ARGV[4]
|
||||
|
||||
-- Remove stale entries (older than 30 minutes)
|
||||
local scope_max_concurrent = tonumber(ARGV[5])
|
||||
|
||||
-- Remove stale entries (older than the stale-call timeout)
|
||||
redis.call('ZREMRANGEBYSCORE', key, 0, stale_cutoff)
|
||||
|
||||
|
||||
-- Get current count
|
||||
local current_count = redis.call('ZCARD', key)
|
||||
|
||||
if current_count < max_concurrent then
|
||||
-- Add new slot
|
||||
redis.call('ZADD', key, now, slot_id)
|
||||
redis.call('EXPIRE', key, 3600) -- Expire after 1 hour
|
||||
return {slot_id, current_count + 1}
|
||||
else
|
||||
|
||||
if current_count >= max_concurrent then
|
||||
return nil
|
||||
end
|
||||
|
||||
if scope_key ~= '' then
|
||||
redis.call('ZREMRANGEBYSCORE', scope_key, 0, stale_cutoff)
|
||||
if redis.call('ZCARD', scope_key) >= scope_max_concurrent then
|
||||
return nil
|
||||
end
|
||||
redis.call('ZADD', scope_key, now, slot_id)
|
||||
redis.call('EXPIRE', scope_key, 3600)
|
||||
end
|
||||
|
||||
redis.call('ZADD', key, now, slot_id)
|
||||
redis.call('EXPIRE', key, 3600) -- Expire after 1 hour
|
||||
return {slot_id, current_count + 1}
|
||||
"""
|
||||
|
||||
# Generate unique slot ID (timestamp + random component)
|
||||
|
|
@ -156,12 +180,14 @@ class RateLimiter:
|
|||
try:
|
||||
result = await redis_client.eval(
|
||||
lua_script,
|
||||
1,
|
||||
2,
|
||||
concurrent_key,
|
||||
scope_concurrent_key,
|
||||
now,
|
||||
max_concurrent,
|
||||
stale_cutoff,
|
||||
slot_id,
|
||||
scope_max_concurrent if scope_max_concurrent is not None else 0,
|
||||
)
|
||||
if not result:
|
||||
return None
|
||||
|
|
@ -175,10 +201,17 @@ class RateLimiter:
|
|||
logger.error(f"Concurrent limiter error: {e}")
|
||||
return None
|
||||
|
||||
async def release_concurrent_slot(self, organization_id: int, slot_id: str) -> bool:
|
||||
async def release_concurrent_slot(
|
||||
self,
|
||||
organization_id: int,
|
||||
slot_id: str,
|
||||
scope_key: str | None = None,
|
||||
) -> bool | None:
|
||||
"""
|
||||
Release a concurrent call slot.
|
||||
Returns True if slot was released, False otherwise.
|
||||
Release a concurrent call slot (and its scope counter entry, if any).
|
||||
Returns True if the slot was released, False if it was already gone
|
||||
(released/stale-expired), or None on a Redis error — callers that
|
||||
track cleanup state should keep it around for retry when None.
|
||||
"""
|
||||
if not slot_id:
|
||||
return False
|
||||
|
|
@ -188,6 +221,8 @@ class RateLimiter:
|
|||
|
||||
try:
|
||||
removed = await redis_client.zrem(concurrent_key, slot_id)
|
||||
if scope_key:
|
||||
await redis_client.zrem(f"concurrent_calls:{scope_key}", slot_id)
|
||||
if removed:
|
||||
logger.debug(
|
||||
f"Released concurrent slot {slot_id} for org {organization_id}"
|
||||
|
|
@ -195,7 +230,7 @@ class RateLimiter:
|
|||
return bool(removed)
|
||||
except Exception as e:
|
||||
logger.error(f"Error releasing concurrent slot: {e}")
|
||||
return False
|
||||
return None
|
||||
|
||||
async def get_concurrent_count(self, organization_id: int) -> int:
|
||||
"""
|
||||
|
|
@ -240,7 +275,11 @@ class RateLimiter:
|
|||
return False
|
||||
|
||||
async def store_workflow_slot_mapping_if_absent(
|
||||
self, workflow_run_id: int, organization_id: int, slot_id: str
|
||||
self,
|
||||
workflow_run_id: int,
|
||||
organization_id: int,
|
||||
slot_id: str,
|
||||
scope_key: str | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Store the workflow_run_id -> concurrent slot mapping only if no mapping
|
||||
|
|
@ -255,12 +294,16 @@ class RateLimiter:
|
|||
local org_id = ARGV[1]
|
||||
local slot_id = ARGV[2]
|
||||
local ttl = tonumber(ARGV[3])
|
||||
local scope_key = ARGV[4]
|
||||
|
||||
if redis.call('EXISTS', key) == 1 then
|
||||
return 0
|
||||
end
|
||||
|
||||
redis.call('HSET', key, 'org_id', org_id, 'slot_id', slot_id)
|
||||
if scope_key ~= '' then
|
||||
redis.call('HSET', key, 'scope_key', scope_key)
|
||||
end
|
||||
redis.call('EXPIRE', key, ttl)
|
||||
return 1
|
||||
"""
|
||||
|
|
@ -273,6 +316,7 @@ class RateLimiter:
|
|||
organization_id,
|
||||
slot_id,
|
||||
self.stale_call_timeout,
|
||||
scope_key or "",
|
||||
)
|
||||
return bool(stored)
|
||||
except Exception as e:
|
||||
|
|
@ -281,10 +325,11 @@ class RateLimiter:
|
|||
|
||||
async def get_workflow_slot_mapping(
|
||||
self, workflow_run_id: int
|
||||
) -> Optional[tuple[int, str]]:
|
||||
) -> Optional[tuple[int, str, str | None]]:
|
||||
"""
|
||||
Get the concurrent slot mapping for a workflow run.
|
||||
Returns (organization_id, slot_id) tuple or None if not found.
|
||||
Returns (organization_id, slot_id, scope_key) or None if not found;
|
||||
scope_key is None for slots acquired without a scope counter.
|
||||
"""
|
||||
redis_client = await self._get_redis()
|
||||
mapping_key = f"workflow_slot_mapping:{workflow_run_id}"
|
||||
|
|
@ -292,7 +337,11 @@ class RateLimiter:
|
|||
try:
|
||||
mapping = await redis_client.hgetall(mapping_key)
|
||||
if mapping and "org_id" in mapping and "slot_id" in mapping:
|
||||
return (int(mapping["org_id"]), mapping["slot_id"])
|
||||
return (
|
||||
int(mapping["org_id"]),
|
||||
mapping["slot_id"],
|
||||
mapping.get("scope_key") or None,
|
||||
)
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting workflow slot mapping: {e}")
|
||||
|
|
|
|||
|
|
@ -787,6 +787,14 @@ class ARIConnection:
|
|||
the bridge and both channels — like endConferenceOnExit.
|
||||
"""
|
||||
try:
|
||||
# Release the org concurrency slot. Normally the pipeline's own
|
||||
# teardown does this when the ext media websocket closes, but if
|
||||
# the pipeline never started (caller hung up before external
|
||||
# media connected, ext media creation failed, ...) this is the
|
||||
# only cleanup that runs before the Redis stale timeout. No-op
|
||||
# when the slot was already released.
|
||||
await call_concurrency.unregister_active_call(int(workflow_run_id))
|
||||
|
||||
workflow_run = await db_client.get_workflow_run_by_id(int(workflow_run_id))
|
||||
if not workflow_run or not workflow_run.gathered_context:
|
||||
logger.warning(
|
||||
|
|
|
|||
|
|
@ -365,6 +365,10 @@ class TelephonyProvider(ABC):
|
|||
the caller carries a provider stream protocol. ``organization_id`` is
|
||||
passed so providers can scope any config lookups to the workflow's org.
|
||||
Default raises so providers that haven't opted in fail loudly.
|
||||
|
||||
The route holds an org concurrency slot while this runs, so
|
||||
implementations must bound their pre-pipeline handshake reads with a
|
||||
timeout — an idle socket must not hold the slot indefinitely.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"Agent-stream not supported for provider {self.PROVIDER_NAME}"
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
Cloudonix implementation of the TelephonyProvider interface.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
|
@ -25,6 +26,11 @@ if TYPE_CHECKING:
|
|||
|
||||
CLOUDONIX_API_BASE_URL = "https://api.cloudonix.io"
|
||||
|
||||
# Cloudonix sends the connected/start handshake immediately after the media
|
||||
# stream opens. The agent-stream route holds an org concurrency slot while we
|
||||
# wait, so an idle socket must not be able to hold it indefinitely.
|
||||
AGENT_STREAM_HANDSHAKE_TIMEOUT_S = 10
|
||||
|
||||
|
||||
class CloudonixProvider(TelephonyProvider):
|
||||
"""
|
||||
|
|
@ -533,14 +539,31 @@ class CloudonixProvider(TelephonyProvider):
|
|||
from api.services.pipecat.run_pipeline import run_pipeline_telephony
|
||||
|
||||
try:
|
||||
first_msg = await websocket.receive_text()
|
||||
msg = json.loads(first_msg)
|
||||
if msg.get("event") != "connected":
|
||||
logger.error(f"Expected 'connected' event, got: {msg.get('event')}")
|
||||
await websocket.close(code=4400, reason="Expected connected event")
|
||||
try:
|
||||
first_msg = await asyncio.wait_for(
|
||||
websocket.receive_text(),
|
||||
timeout=AGENT_STREAM_HANDSHAKE_TIMEOUT_S,
|
||||
)
|
||||
msg = json.loads(first_msg)
|
||||
if msg.get("event") != "connected":
|
||||
logger.error(f"Expected 'connected' event, got: {msg.get('event')}")
|
||||
await websocket.close(code=4400, reason="Expected connected event")
|
||||
return
|
||||
|
||||
start_msg = json.loads(
|
||||
await asyncio.wait_for(
|
||||
websocket.receive_text(),
|
||||
timeout=AGENT_STREAM_HANDSHAKE_TIMEOUT_S,
|
||||
)
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
f"Cloudonix agent-stream handshake timed out for workflow_run "
|
||||
f"{workflow_run_id}"
|
||||
)
|
||||
await websocket.close(code=4408, reason="Handshake timeout")
|
||||
return
|
||||
|
||||
start_msg = json.loads(await websocket.receive_text())
|
||||
if start_msg.get("event") != "start":
|
||||
logger.error("Expected 'start' event second")
|
||||
await websocket.close(code=4400, reason="Expected start event")
|
||||
|
|
|
|||
|
|
@ -244,3 +244,34 @@ def test_parse_cloudonix_cdr_preserves_zero_billsec():
|
|||
)
|
||||
|
||||
assert req["duration"] == "0"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_stream_handshake_timeout_closes_socket(monkeypatch):
|
||||
"""An idle agent-stream socket holds an org concurrency slot, so the
|
||||
handshake read must be bounded rather than waiting forever."""
|
||||
import asyncio
|
||||
|
||||
from api.services.telephony.providers.cloudonix import provider as provider_module
|
||||
|
||||
monkeypatch.setattr(provider_module, "AGENT_STREAM_HANDSHAKE_TIMEOUT_S", 0.05)
|
||||
provider = CloudonixProvider({})
|
||||
|
||||
async def never_returns():
|
||||
await asyncio.Event().wait()
|
||||
|
||||
websocket = SimpleNamespace(receive_text=never_returns, close=AsyncMock())
|
||||
|
||||
await asyncio.wait_for(
|
||||
provider.handle_external_websocket(
|
||||
websocket,
|
||||
organization_id=1,
|
||||
workflow_id=2,
|
||||
user_id=3,
|
||||
workflow_run_id=4,
|
||||
params={},
|
||||
),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
websocket.close.assert_awaited_once_with(code=4408, reason="Handshake timeout")
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -7,6 +8,13 @@ from api.services.telephony import ari_manager
|
|||
from api.services.telephony.ari_manager import ARIConnection
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_call_concurrency(monkeypatch):
|
||||
fake = SimpleNamespace(unregister_active_call=AsyncMock(return_value=True))
|
||||
monkeypatch.setattr(ari_manager, "call_concurrency", fake)
|
||||
return fake
|
||||
|
||||
|
||||
class _FakeDbClient:
|
||||
def __init__(self, gathered_context):
|
||||
self.workflow_run = SimpleNamespace(gathered_context=gathered_context)
|
||||
|
|
@ -77,6 +85,7 @@ def _completed_transfer_context():
|
|||
@pytest.mark.asyncio
|
||||
async def test_completed_transfer_tears_down_destination_when_caller_leaves(
|
||||
monkeypatch,
|
||||
fake_call_concurrency,
|
||||
):
|
||||
fake_db = _FakeDbClient(_completed_transfer_context())
|
||||
monkeypatch.setattr(ari_manager, "db_client", fake_db)
|
||||
|
|
@ -90,11 +99,13 @@ async def test_completed_transfer_tears_down_destination_when_caller_leaves(
|
|||
assert "dest-chan" in conn.deleted_channel_runs
|
||||
assert conn.deleted_transfer_channel_mappings == ["dest-chan"]
|
||||
assert fake_db.workflow_run.gathered_context["transfer_state"] == "terminated"
|
||||
fake_call_concurrency.unregister_active_call.assert_awaited_with(123)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_transfer_tears_down_caller_when_destination_leaves(
|
||||
monkeypatch,
|
||||
fake_call_concurrency,
|
||||
):
|
||||
fake_db = _FakeDbClient(_completed_transfer_context())
|
||||
monkeypatch.setattr(ari_manager, "db_client", fake_db)
|
||||
|
|
@ -108,6 +119,31 @@ async def test_completed_transfer_tears_down_caller_when_destination_leaves(
|
|||
assert "dest-chan" in conn.deleted_channel_runs
|
||||
assert conn.deleted_transfer_channel_mappings == ["dest-chan"]
|
||||
assert fake_db.workflow_run.gathered_context["transfer_state"] == "terminated"
|
||||
fake_call_concurrency.unregister_active_call.assert_awaited_with(123)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stasis_end_releases_concurrency_slot_on_normal_teardown(
|
||||
monkeypatch,
|
||||
fake_call_concurrency,
|
||||
):
|
||||
"""StasisEnd must release the org slot even when the pipeline never ran
|
||||
(e.g. the caller hung up before external media connected)."""
|
||||
fake_db = _FakeDbClient(
|
||||
{
|
||||
"call_id": "caller-chan",
|
||||
"ext_channel_id": "ext-chan",
|
||||
"bridge_id": "bridge-1",
|
||||
}
|
||||
)
|
||||
monkeypatch.setattr(ari_manager, "db_client", fake_db)
|
||||
conn = _RecordingARIConnection()
|
||||
|
||||
await conn._handle_stasis_end("caller-chan", "123")
|
||||
|
||||
fake_call_concurrency.unregister_active_call.assert_awaited_once_with(123)
|
||||
assert conn.deleted_bridges == ["bridge-1"]
|
||||
assert conn.deleted_channels == ["ext-chan"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ async def test_acquire_org_slot_logs_post_acquire_count_and_limit():
|
|||
assert slot.max_concurrent == 10
|
||||
assert slot.source == "test_source"
|
||||
mock_rate_limiter.try_acquire_concurrent_slot_details.assert_awaited_once_with(
|
||||
199, 10
|
||||
199, 10, scope_key=None, scope_max_concurrent=None
|
||||
)
|
||||
mock_logger.info.assert_called_once()
|
||||
log_message = mock_logger.info.call_args.args[0]
|
||||
|
|
@ -67,3 +67,222 @@ async def test_acquire_org_slot_logs_warning_when_limit_reached():
|
|||
assert "Concurrent call limit reached for org 199" in log_message
|
||||
assert "source=test_source" in log_message
|
||||
assert "active_calls=12/10" in log_message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_org_slot_passes_scope_to_rate_limiter():
|
||||
service = CallConcurrencyService()
|
||||
|
||||
with (
|
||||
patch("api.services.call_concurrency.db_client") as mock_db,
|
||||
patch("api.services.call_concurrency.rate_limiter") as mock_rate_limiter,
|
||||
):
|
||||
mock_db.get_configuration = AsyncMock(return_value=None)
|
||||
mock_rate_limiter.try_acquire_concurrent_slot_details = AsyncMock(
|
||||
return_value=ConcurrentSlotAcquisition(slot_id="slot-123", active_count=1)
|
||||
)
|
||||
mock_rate_limiter.store_workflow_slot_mapping_if_absent = AsyncMock(
|
||||
return_value=True
|
||||
)
|
||||
|
||||
slot = await service.acquire_org_slot(
|
||||
199,
|
||||
source="campaign:42",
|
||||
scope_key="campaign:42",
|
||||
scope_max_concurrent=3,
|
||||
)
|
||||
await service.bind_workflow_run(slot, 501)
|
||||
|
||||
assert slot.scope_key == "campaign:42"
|
||||
mock_rate_limiter.try_acquire_concurrent_slot_details.assert_awaited_once_with(
|
||||
199, 10, scope_key="campaign:42", scope_max_concurrent=3
|
||||
)
|
||||
mock_rate_limiter.store_workflow_slot_mapping_if_absent.assert_awaited_once_with(
|
||||
501, 199, "slot-123", scope_key="campaign:42"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_workflow_run_slot_keeps_mapping_on_redis_error():
|
||||
service = CallConcurrencyService()
|
||||
|
||||
with patch("api.services.call_concurrency.rate_limiter") as mock_rate_limiter:
|
||||
mock_rate_limiter.get_workflow_slot_mapping = AsyncMock(
|
||||
return_value=(11, "slot-1", None)
|
||||
)
|
||||
# None = Redis error during release (vs False = slot already gone)
|
||||
mock_rate_limiter.release_concurrent_slot = AsyncMock(return_value=None)
|
||||
mock_rate_limiter.delete_workflow_slot_mapping = AsyncMock()
|
||||
|
||||
released = await service.release_workflow_run_slot(501)
|
||||
|
||||
assert released is False
|
||||
mock_rate_limiter.release_concurrent_slot.assert_awaited_once_with(
|
||||
11, "slot-1", scope_key=None
|
||||
)
|
||||
mock_rate_limiter.delete_workflow_slot_mapping.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_workflow_run_slot_deletes_mapping_when_slot_already_gone():
|
||||
service = CallConcurrencyService()
|
||||
|
||||
with patch("api.services.call_concurrency.rate_limiter") as mock_rate_limiter:
|
||||
mock_rate_limiter.get_workflow_slot_mapping = AsyncMock(
|
||||
return_value=(11, "slot-1", "campaign:42")
|
||||
)
|
||||
mock_rate_limiter.release_concurrent_slot = AsyncMock(return_value=False)
|
||||
mock_rate_limiter.delete_workflow_slot_mapping = AsyncMock(return_value=True)
|
||||
|
||||
released = await service.release_workflow_run_slot(501)
|
||||
|
||||
assert released is False
|
||||
mock_rate_limiter.release_concurrent_slot.assert_awaited_once_with(
|
||||
11, "slot-1", scope_key="campaign:42"
|
||||
)
|
||||
mock_rate_limiter.delete_workflow_slot_mapping.assert_awaited_once_with(501)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unregister_active_call_never_raises():
|
||||
service = CallConcurrencyService()
|
||||
|
||||
with patch("api.services.call_concurrency.rate_limiter") as mock_rate_limiter:
|
||||
mock_rate_limiter.get_workflow_slot_mapping = AsyncMock(
|
||||
side_effect=RuntimeError("redis down")
|
||||
)
|
||||
|
||||
released = await service.unregister_active_call(501)
|
||||
|
||||
assert released is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Redis integration tests for scoped (campaign-level) slot acquisition
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
import os # noqa: E402
|
||||
import uuid # noqa: E402
|
||||
|
||||
from api.services.campaign.rate_limiter import RateLimiter # noqa: E402
|
||||
|
||||
requires_redis = pytest.mark.skipif(
|
||||
"REDIS_URL" not in os.environ,
|
||||
reason="Requires Redis (set REDIS_URL via .env.test)",
|
||||
)
|
||||
|
||||
|
||||
def _unique_org_id() -> int:
|
||||
return uuid.uuid4().int % 10_000_000
|
||||
|
||||
|
||||
@requires_redis
|
||||
@pytest.mark.asyncio
|
||||
async def test_scoped_acquisition_enforces_scope_limit_independently_of_org():
|
||||
"""A campaign scope caps its own calls without measuring — or being
|
||||
starved by — other calls in the same org counter."""
|
||||
rl = RateLimiter()
|
||||
org_id = _unique_org_id()
|
||||
scope = f"campaign:{org_id}"
|
||||
org_key = f"concurrent_calls:{org_id}"
|
||||
scope_key_full = f"concurrent_calls:{scope}"
|
||||
redis_client = await rl._get_redis()
|
||||
|
||||
try:
|
||||
# Unscoped (e.g. WebRTC) calls fill part of the org counter.
|
||||
for _ in range(3):
|
||||
assert await rl.try_acquire_concurrent_slot_details(org_id, 10)
|
||||
|
||||
# Scope limit 2: two scoped acquisitions succeed...
|
||||
first = await rl.try_acquire_concurrent_slot_details(
|
||||
org_id, 10, scope_key=scope, scope_max_concurrent=2
|
||||
)
|
||||
second = await rl.try_acquire_concurrent_slot_details(
|
||||
org_id, 10, scope_key=scope, scope_max_concurrent=2
|
||||
)
|
||||
assert first and second
|
||||
|
||||
# ...the third is rejected by the scope even though the org has room.
|
||||
third = await rl.try_acquire_concurrent_slot_details(
|
||||
org_id, 10, scope_key=scope, scope_max_concurrent=2
|
||||
)
|
||||
assert third is None
|
||||
|
||||
# Unscoped calls are unaffected by the scope being full.
|
||||
assert await rl.try_acquire_concurrent_slot_details(org_id, 10)
|
||||
|
||||
# Releasing with the scope key frees both counters.
|
||||
released = await rl.release_concurrent_slot(
|
||||
org_id, first.slot_id, scope_key=scope
|
||||
)
|
||||
assert released is True
|
||||
assert await redis_client.zscore(org_key, first.slot_id) is None
|
||||
assert await redis_client.zscore(scope_key_full, first.slot_id) is None
|
||||
|
||||
# And the scope accepts a new call again.
|
||||
assert await rl.try_acquire_concurrent_slot_details(
|
||||
org_id, 10, scope_key=scope, scope_max_concurrent=2
|
||||
)
|
||||
finally:
|
||||
await redis_client.delete(org_key, scope_key_full)
|
||||
await rl.close()
|
||||
|
||||
|
||||
@requires_redis
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_limit_still_binds_scoped_acquisition():
|
||||
rl = RateLimiter()
|
||||
org_id = _unique_org_id()
|
||||
scope = f"campaign:{org_id}"
|
||||
org_key = f"concurrent_calls:{org_id}"
|
||||
scope_key_full = f"concurrent_calls:{scope}"
|
||||
redis_client = await rl._get_redis()
|
||||
|
||||
try:
|
||||
assert await rl.try_acquire_concurrent_slot_details(org_id, 1)
|
||||
|
||||
# Org counter is full, so the scoped acquire fails and must not
|
||||
# leave a phantom entry in the scope counter.
|
||||
rejected = await rl.try_acquire_concurrent_slot_details(
|
||||
org_id, 1, scope_key=scope, scope_max_concurrent=5
|
||||
)
|
||||
assert rejected is None
|
||||
assert await redis_client.zcard(scope_key_full) == 0
|
||||
finally:
|
||||
await redis_client.delete(org_key, scope_key_full)
|
||||
await rl.close()
|
||||
|
||||
|
||||
@requires_redis
|
||||
@pytest.mark.asyncio
|
||||
async def test_workflow_slot_mapping_round_trips_scope_key():
|
||||
rl = RateLimiter()
|
||||
run_id = _unique_org_id()
|
||||
mapping_key = f"workflow_slot_mapping:{run_id}"
|
||||
redis_client = await rl._get_redis()
|
||||
|
||||
try:
|
||||
stored = await rl.store_workflow_slot_mapping_if_absent(
|
||||
run_id, 11, "slot-1", scope_key="campaign:42"
|
||||
)
|
||||
assert stored is True
|
||||
assert await rl.get_workflow_slot_mapping(run_id) == (
|
||||
11,
|
||||
"slot-1",
|
||||
"campaign:42",
|
||||
)
|
||||
|
||||
# Unscoped mappings surface scope_key=None.
|
||||
run_id_2 = _unique_org_id()
|
||||
try:
|
||||
await rl.store_workflow_slot_mapping_if_absent(run_id_2, 11, "slot-2")
|
||||
assert await rl.get_workflow_slot_mapping(run_id_2) == (
|
||||
11,
|
||||
"slot-2",
|
||||
None,
|
||||
)
|
||||
finally:
|
||||
await redis_client.delete(f"workflow_slot_mapping:{run_id_2}")
|
||||
finally:
|
||||
await redis_client.delete(mapping_key)
|
||||
await rl.close()
|
||||
|
|
|
|||
|
|
@ -266,8 +266,9 @@ def mock_call_concurrency():
|
|||
return CallConcurrencySlot(
|
||||
organization_id=organization_id,
|
||||
slot_id=f"slot-{uuid.uuid4().hex[:8]}",
|
||||
max_concurrent=kwargs.get("max_concurrent_override") or 20,
|
||||
max_concurrent=20,
|
||||
source=source,
|
||||
scope_key=kwargs.get("scope_key"),
|
||||
)
|
||||
|
||||
with patch(
|
||||
|
|
@ -814,3 +815,49 @@ class TestProcessBatchEdgeCases:
|
|||
{"campaign_id": campaign_test_data.campaign_id},
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
|
||||
class TestAcquireConcurrentSlotScoping:
|
||||
"""Campaign max_concurrency must scope to the campaign, not the org counter."""
|
||||
|
||||
def _campaign(self, orchestrator_metadata):
|
||||
campaign = MagicMock()
|
||||
campaign.id = 42
|
||||
campaign.orchestrator_metadata = orchestrator_metadata
|
||||
return campaign
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_campaign_max_concurrency_uses_campaign_scope(
|
||||
self, mock_call_concurrency
|
||||
):
|
||||
dispatcher = CampaignCallDispatcher()
|
||||
campaign = self._campaign({"max_concurrency": 3})
|
||||
|
||||
await dispatcher.acquire_concurrent_slot(7, campaign, timeout=5)
|
||||
|
||||
mock_call_concurrency.acquire_org_slot.assert_awaited_once_with(
|
||||
7,
|
||||
source="campaign:42",
|
||||
timeout=5,
|
||||
scope_key="campaign:42",
|
||||
scope_max_concurrent=3,
|
||||
retry_interval=1,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_campaign_max_concurrency_skips_scope(
|
||||
self, mock_call_concurrency
|
||||
):
|
||||
dispatcher = CampaignCallDispatcher()
|
||||
campaign = self._campaign({})
|
||||
|
||||
await dispatcher.acquire_concurrent_slot(7, campaign, timeout=5)
|
||||
|
||||
mock_call_concurrency.acquire_org_slot.assert_awaited_once_with(
|
||||
7,
|
||||
source="campaign:42",
|
||||
timeout=5,
|
||||
scope_key=None,
|
||||
scope_max_concurrent=None,
|
||||
retry_interval=1,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -154,7 +154,8 @@ def test_workflow_uuid_route_uses_scoped_lookup_and_shared_execution():
|
|||
new=AsyncMock(return_value=("https://api.example.com", "wss://ignored")),
|
||||
),
|
||||
):
|
||||
mock_concurrency.acquire_org_slot = AsyncMock(return_value=object())
|
||||
slot = object()
|
||||
mock_concurrency.acquire_org_slot = AsyncMock(return_value=slot)
|
||||
mock_concurrency.bind_workflow_run = AsyncMock()
|
||||
mock_concurrency.release_workflow_run_slot = AsyncMock()
|
||||
mock_concurrency.release_slot = AsyncMock()
|
||||
|
|
@ -180,6 +181,12 @@ def test_workflow_uuid_route_uses_scoped_lookup_and_shared_execution():
|
|||
11,
|
||||
)
|
||||
assert not mock_db.get_agent_trigger_by_path.called
|
||||
mock_concurrency.acquire_org_slot.assert_awaited_once_with(
|
||||
workflow.organization_id,
|
||||
source="public_agent",
|
||||
timeout=0,
|
||||
)
|
||||
mock_concurrency.bind_workflow_run.assert_awaited_once_with(slot, 601)
|
||||
|
||||
create_kwargs = mock_db.create_workflow_run.await_args.kwargs
|
||||
assert create_kwargs["user_id"] == workflow.user_id
|
||||
|
|
|
|||
|
|
@ -115,3 +115,53 @@ async def test_public_embed_renegotiation_does_not_acquire_another_slot():
|
|||
mock_concurrency.acquire_org_slot.assert_not_called()
|
||||
pc.renegotiate.assert_awaited_once()
|
||||
assert ws.send_json.await_args.args[0]["type"] == "answer"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_signaling_websocket_rejects_run_not_owned_by_workflow():
|
||||
"""The URL workflow_id drives org/quota/concurrency accounting, so a
|
||||
workflow_id that doesn't own the run must be rejected before signaling."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from api.routes.webrtc_signaling import signaling_websocket
|
||||
|
||||
ws = _FakeWebSocket()
|
||||
user = SimpleNamespace(id=7)
|
||||
|
||||
with (
|
||||
patch("api.routes.webrtc_signaling.db_client") as mock_db,
|
||||
patch("api.routes.webrtc_signaling.signaling_manager") as mock_manager,
|
||||
):
|
||||
mock_db.get_workflow_run = AsyncMock(
|
||||
return_value=SimpleNamespace(id=501, workflow_id=99)
|
||||
)
|
||||
mock_manager.handle_websocket = AsyncMock()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await signaling_websocket(
|
||||
ws, workflow_id=33, workflow_run_id=501, user=user
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
mock_manager.handle_websocket.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_signaling_websocket_accepts_matching_workflow_and_run():
|
||||
from api.routes.webrtc_signaling import signaling_websocket
|
||||
|
||||
ws = _FakeWebSocket()
|
||||
user = SimpleNamespace(id=7)
|
||||
|
||||
with (
|
||||
patch("api.routes.webrtc_signaling.db_client") as mock_db,
|
||||
patch("api.routes.webrtc_signaling.signaling_manager") as mock_manager,
|
||||
):
|
||||
mock_db.get_workflow_run = AsyncMock(
|
||||
return_value=SimpleNamespace(id=501, workflow_id=33)
|
||||
)
|
||||
mock_manager.handle_websocket = AsyncMock()
|
||||
|
||||
await signaling_websocket(ws, workflow_id=33, workflow_run_id=501, user=user)
|
||||
|
||||
mock_manager.handle_websocket.assert_awaited_once()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue