237 lines
9.5 KiB
Python
237 lines
9.5 KiB
Python
import asyncio
|
|
import json
|
|
import secrets
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Awaitable, Callable
|
|
|
|
from src.push_subscription_store import PushSubscriptionStore
|
|
from src.push_endpoint_policy import UnsafePushEndpoint, validate_public_push_endpoint
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PushConfiguration:
|
|
public_key: str
|
|
private_key: str
|
|
subject: str
|
|
|
|
@property
|
|
def enabled(self) -> bool:
|
|
return bool(self.public_key and self.private_key and self.subject)
|
|
|
|
|
|
async def send_web_push(
|
|
subscription: dict, payload: str, configuration: PushConfiguration
|
|
) -> None:
|
|
import requests
|
|
from pywebpush import webpush
|
|
|
|
session = requests.Session()
|
|
session.max_redirects = 0
|
|
|
|
await asyncio.to_thread(
|
|
webpush,
|
|
subscription_info=subscription,
|
|
data=payload,
|
|
vapid_private_key=configuration.private_key,
|
|
vapid_claims={"sub": configuration.subject},
|
|
ttl=300,
|
|
timeout=10,
|
|
requests_session=session,
|
|
)
|
|
|
|
|
|
async def dispatch_unread_updates(
|
|
store: PushSubscriptionStore,
|
|
configuration: PushConfiguration,
|
|
unread: Callable[[], Awaitable[dict]],
|
|
send: Callable[[dict, str], Awaitable[None]] | None = None,
|
|
*,
|
|
session_active: Callable[[str], Awaitable[bool]] | None = None,
|
|
lease_seconds: float = 60.0,
|
|
send_timeout_seconds: float = 10.0,
|
|
max_concurrency: int = 8,
|
|
max_individual_notifications: int = 3,
|
|
endpoint_validator: Callable[[str], Awaitable[str]] | None = None,
|
|
) -> int:
|
|
if not configuration.enabled:
|
|
return 0
|
|
owner = secrets.token_urlsafe(18)
|
|
acquired = await asyncio.to_thread(
|
|
store.acquire_dispatch_lease,
|
|
owner,
|
|
now=time.time(),
|
|
lease_seconds=lease_seconds,
|
|
)
|
|
if not acquired:
|
|
return 0
|
|
try:
|
|
page = await unread()
|
|
if page.get("complete") is False:
|
|
return 0
|
|
thread_revisions = {
|
|
int(item["id"]): str(item.get("updated_at") or "")
|
|
for item in page.get("items", [])
|
|
if isinstance(item, dict) and str(item.get("id", "")).isdigit()
|
|
}
|
|
await asyncio.to_thread(store.reconcile_unread, thread_revisions)
|
|
deliveries = await asyncio.to_thread(store.claim_unseen, thread_revisions)
|
|
if session_active is not None:
|
|
try:
|
|
authorized = [await session_active(item.session_id) for item in deliveries]
|
|
except Exception:
|
|
# Authorization state is mandatory for delivery. Preserve subscriptions
|
|
# so a temporary registry failure can be retried safely.
|
|
return 0
|
|
for delivery, active in zip(deliveries, authorized):
|
|
if not active:
|
|
await asyncio.to_thread(store.delete_session, delivery.session_id)
|
|
deliveries = [
|
|
delivery
|
|
for delivery, active in zip(deliveries, authorized)
|
|
if active
|
|
]
|
|
semaphore = asyncio.Semaphore(max(1, max_concurrency))
|
|
ownership_lost = asyncio.Event()
|
|
|
|
async def dispatch_device(delivery) -> int:
|
|
async with semaphore:
|
|
if ownership_lost.is_set():
|
|
return 0
|
|
validate_endpoint = endpoint_validator
|
|
if validate_endpoint is None and send is None:
|
|
validate_endpoint = validate_public_push_endpoint
|
|
try:
|
|
if validate_endpoint is not None:
|
|
await validate_endpoint(delivery.subscription["endpoint"])
|
|
except UnsafePushEndpoint:
|
|
await asyncio.to_thread(store.delete_session, delivery.session_id)
|
|
return 0
|
|
count = 0
|
|
digest_pending = set(delivery.digest_revisions)
|
|
new_revisions = tuple(
|
|
thread_revision
|
|
for thread_revision in delivery.thread_revisions
|
|
if thread_revision not in digest_pending
|
|
)
|
|
individual_revisions = new_revisions[
|
|
:max(0, max_individual_notifications)
|
|
]
|
|
overflow_revisions = (
|
|
delivery.digest_revisions
|
|
+ new_revisions[len(individual_revisions):]
|
|
)
|
|
can_send_digest = True
|
|
for thread_id, revision in individual_revisions:
|
|
still_owner = await asyncio.to_thread(
|
|
store.acquire_dispatch_lease,
|
|
owner,
|
|
now=time.time(),
|
|
lease_seconds=lease_seconds,
|
|
)
|
|
if not still_owner:
|
|
ownership_lost.set()
|
|
return count
|
|
payload = json.dumps(
|
|
{
|
|
"title": "New work update",
|
|
"body": "Tap to review it in Stackchain.",
|
|
"route": f"#/my-work/update/{thread_id}",
|
|
"tag": f"stackchain-update-{thread_id}",
|
|
"notification_id": thread_id,
|
|
},
|
|
separators=(",", ":"),
|
|
)
|
|
try:
|
|
if send is None:
|
|
operation = send_web_push(
|
|
delivery.subscription, payload, configuration
|
|
)
|
|
else:
|
|
operation = send(delivery.subscription, payload)
|
|
await asyncio.wait_for(operation, timeout=send_timeout_seconds)
|
|
except Exception as error:
|
|
status = getattr(
|
|
getattr(error, "response", None), "status_code", None
|
|
)
|
|
if status in {404, 410}:
|
|
await asyncio.to_thread(
|
|
store.delete_session, delivery.session_id
|
|
)
|
|
can_send_digest = False
|
|
# Leave this device's transient failures unseen for a later
|
|
# poll instead of paying the endpoint deadline repeatedly.
|
|
break
|
|
await asyncio.to_thread(
|
|
store.mark_delivered,
|
|
delivery.session_id,
|
|
((thread_id, revision),),
|
|
)
|
|
count += 1
|
|
if overflow_revisions and can_send_digest:
|
|
still_owner = await asyncio.to_thread(
|
|
store.acquire_dispatch_lease,
|
|
owner,
|
|
now=time.time(),
|
|
lease_seconds=lease_seconds,
|
|
)
|
|
if not still_owner:
|
|
ownership_lost.set()
|
|
return count
|
|
payload = json.dumps(
|
|
{
|
|
"title": f"{len(overflow_revisions)} new work updates",
|
|
"body": "Tap to review them in Stackchain.",
|
|
"route": "#/my-work/updates",
|
|
"tag": "stackchain-update-digest",
|
|
"update_count": len(overflow_revisions),
|
|
},
|
|
separators=(",", ":"),
|
|
)
|
|
try:
|
|
if send is None:
|
|
operation = send_web_push(
|
|
delivery.subscription, payload, configuration
|
|
)
|
|
else:
|
|
operation = send(delivery.subscription, payload)
|
|
await asyncio.wait_for(operation, timeout=send_timeout_seconds)
|
|
except Exception as error:
|
|
status = getattr(
|
|
getattr(error, "response", None), "status_code", None
|
|
)
|
|
if status in {404, 410}:
|
|
await asyncio.to_thread(
|
|
store.delete_session, delivery.session_id
|
|
)
|
|
else:
|
|
await asyncio.to_thread(
|
|
store.mark_digest_pending,
|
|
delivery.session_id,
|
|
overflow_revisions,
|
|
)
|
|
return count
|
|
await asyncio.to_thread(
|
|
store.mark_delivered,
|
|
delivery.session_id,
|
|
overflow_revisions,
|
|
)
|
|
count += 1
|
|
return count
|
|
|
|
async def dispatch_device_safely(delivery) -> int:
|
|
try:
|
|
return await dispatch_device(delivery)
|
|
except Exception:
|
|
# Device-specific validation, persistence, or provider failures
|
|
# must not cancel healthy siblings. Leave any uncheckpointed
|
|
# revisions unseen so a later poll can retry them.
|
|
return 0
|
|
|
|
counts = await asyncio.gather(
|
|
*(dispatch_device_safely(delivery) for delivery in deliveries)
|
|
)
|
|
return sum(counts)
|
|
finally:
|
|
await asyncio.to_thread(store.release_dispatch_lease, owner)
|