388 lines
13 KiB
Python
388 lines
13 KiB
Python
import asyncio
|
|
import logging
|
|
import random
|
|
from dataclasses import dataclass
|
|
|
|
from atproto import Client, SessionEvent
|
|
from atproto import models as at_models
|
|
|
|
from .dedupe import DedupeTracker
|
|
from .session_store import load_session, save_session
|
|
from .types import BlueskyAccount, InboundDM, InboundMention
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class BlueskyClientHandle:
|
|
client: Client
|
|
dm_client: Client
|
|
me_did: str
|
|
me_handle: str
|
|
account_id: str
|
|
|
|
|
|
@dataclass
|
|
class BlueskyPollHandle:
|
|
task: asyncio.Task
|
|
dedupe: DedupeTracker
|
|
|
|
|
|
class BlueskyGatewayError(Exception):
|
|
pass
|
|
|
|
|
|
class BlueskyGateway:
|
|
def __init__(self):
|
|
self._active_clients: dict[str, BlueskyClientHandle] = {}
|
|
self._poll_handles: dict[str, BlueskyPollHandle] = {}
|
|
self._notif_tasks: dict[str, asyncio.Task] = {}
|
|
|
|
async def _create_or_restore_client(self, account: BlueskyAccount, account_id: str) -> BlueskyClientHandle:
|
|
client = Client()
|
|
|
|
session_str = account.session_string or load_session(account_id)
|
|
if session_str:
|
|
try:
|
|
client.login(session_string=session_str)
|
|
me = client.me
|
|
logger.info(
|
|
"Bluesky session restored for %s (%s)",
|
|
account_id,
|
|
me.handle,
|
|
)
|
|
|
|
client.on_session_change = lambda event, session: (
|
|
save_session(account_id, client.export_session_string())
|
|
if event in (SessionEvent.CREATE, SessionEvent.REFRESH)
|
|
else None
|
|
)
|
|
|
|
self._label_as_bot(client, me.did)
|
|
|
|
return BlueskyClientHandle(
|
|
client=client,
|
|
dm_client=client.with_bsky_chat_proxy(),
|
|
me_did=me.did,
|
|
me_handle=me.handle,
|
|
account_id=account_id,
|
|
)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Bluesky session restore failed for %s: %s, will re-login",
|
|
account_id,
|
|
e,
|
|
)
|
|
|
|
client.login(account.handle, account.app_password)
|
|
me = client.me
|
|
logger.info("Bluesky login success for %s (%s)", account_id, me.handle)
|
|
|
|
session_str = client.export_session_string()
|
|
save_session(account_id, session_str)
|
|
|
|
client.on_session_change = lambda event, session: (
|
|
save_session(account_id, client.export_session_string())
|
|
if event in (SessionEvent.CREATE, SessionEvent.REFRESH)
|
|
else None
|
|
)
|
|
|
|
self._label_as_bot(client, me.did)
|
|
|
|
return BlueskyClientHandle(
|
|
client=client,
|
|
dm_client=client.with_bsky_chat_proxy(),
|
|
me_did=me.did,
|
|
me_handle=me.handle,
|
|
account_id=account_id,
|
|
)
|
|
|
|
def _label_as_bot(self, client: Client, did: str):
|
|
try:
|
|
profile = client.get_profile(actor=did)
|
|
existing_labels = []
|
|
if profile.labels:
|
|
existing_labels = [
|
|
at_models.ComAtprotoLabelDefs.SelfLabel(val=label.val)
|
|
for label in profile.labels
|
|
if hasattr(label, "val")
|
|
]
|
|
if any(label.val == "bot" for label in existing_labels):
|
|
return
|
|
|
|
new_labels = existing_labels + [at_models.ComAtprotoLabelDefs.SelfLabel(val="bot")]
|
|
client.app.bsky.actor.profile.put(
|
|
did,
|
|
at_models.AppBskyActorProfile.Record(
|
|
display_name=profile.display_name or "",
|
|
description=profile.description or "",
|
|
avatar=profile.avatar,
|
|
banner=profile.banner,
|
|
labels=at_models.ComAtprotoLabelDefs.SelfLabels(values=new_labels),
|
|
),
|
|
)
|
|
logger.info("Bluesky bot self-label applied for %s", did)
|
|
except Exception as e:
|
|
logger.warning("Failed to apply bot self-label: %s", e)
|
|
|
|
async def start(
|
|
self,
|
|
account_id: str,
|
|
account: BlueskyAccount,
|
|
on_dm: callable,
|
|
on_mention: callable,
|
|
authorize_sender: callable,
|
|
) -> BlueskyClientHandle:
|
|
if not account.is_configured:
|
|
raise BlueskyGatewayError(f"Bluesky account '{account_id}' not configured")
|
|
|
|
handle = await self._create_or_restore_client(account, account_id)
|
|
self._active_clients[account_id] = handle
|
|
|
|
dedupe = DedupeTracker()
|
|
task = asyncio.create_task(
|
|
self._poll_dm_loop(
|
|
account_id,
|
|
handle,
|
|
on_dm,
|
|
authorize_sender,
|
|
dedupe,
|
|
)
|
|
)
|
|
self._poll_handles[account_id] = BlueskyPollHandle(task=task, dedupe=dedupe)
|
|
|
|
if account.enable_notifications:
|
|
notif_task = asyncio.create_task(
|
|
self._poll_notifications_loop(
|
|
account_id,
|
|
handle,
|
|
on_mention,
|
|
dedupe,
|
|
)
|
|
)
|
|
self._notif_tasks[account_id] = notif_task
|
|
|
|
return handle
|
|
|
|
async def stop(self, account_id: str):
|
|
handle = self._poll_handles.pop(account_id, None)
|
|
if handle:
|
|
handle.task.cancel()
|
|
try:
|
|
await handle.task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
notif_task = self._notif_tasks.pop(account_id, None)
|
|
if notif_task:
|
|
notif_task.cancel()
|
|
try:
|
|
await notif_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
self._active_clients.pop(account_id, None)
|
|
logger.info("Bluesky gateway stopped for account %s", account_id)
|
|
|
|
def get_client(self, account_id: str) -> BlueskyClientHandle | None:
|
|
return self._active_clients.get(account_id)
|
|
|
|
async def stop_all(self):
|
|
for account_id in list(self._active_clients.keys()):
|
|
await self.stop(account_id)
|
|
|
|
async def _poll_dm_loop(
|
|
self,
|
|
account_id: str,
|
|
handle: BlueskyClientHandle,
|
|
on_dm: callable,
|
|
authorize_sender: callable,
|
|
dedupe: DedupeTracker,
|
|
interval: float = 5.0,
|
|
):
|
|
dm_client = handle.dm_client
|
|
cursor: str | None = None
|
|
consecutive_errors = 0
|
|
max_interval = 120.0
|
|
|
|
while True:
|
|
try:
|
|
params = {}
|
|
if cursor:
|
|
params["cursor"] = cursor
|
|
log = dm_client.chat.bsky.convo.get_log(params=params)
|
|
cursor = log.cursor
|
|
for item in log.logs:
|
|
self._process_log_item(
|
|
item,
|
|
handle,
|
|
on_dm,
|
|
authorize_sender,
|
|
dedupe,
|
|
)
|
|
consecutive_errors = 0
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
consecutive_errors += 1
|
|
delay = min(interval * (2**consecutive_errors), max_interval)
|
|
delay *= random.uniform(0.8, 1.2)
|
|
logger.error(
|
|
"Bluesky DM poll error for %s (attempt %d, retry in %.1fs): %s",
|
|
account_id,
|
|
consecutive_errors,
|
|
delay,
|
|
e,
|
|
)
|
|
await asyncio.sleep(delay)
|
|
continue
|
|
await asyncio.sleep(interval)
|
|
|
|
def _process_log_item(
|
|
self,
|
|
item,
|
|
handle: BlueskyClientHandle,
|
|
on_dm: callable,
|
|
authorize_sender: callable,
|
|
dedupe: DedupeTracker,
|
|
):
|
|
is_create_message = isinstance(item, at_models.ChatBskyConvoGetLog.LogCreateMessage)
|
|
is_add_reaction = isinstance(item, at_models.ChatBskyConvoGetLog.LogAddReaction)
|
|
is_remove_reaction = isinstance(item, at_models.ChatBskyConvoGetLog.LogRemoveReaction)
|
|
|
|
if is_create_message:
|
|
msg = item.message
|
|
mid = msg.id
|
|
|
|
if dedupe.has(mid):
|
|
return
|
|
|
|
if msg.sender.did == handle.me_did:
|
|
return
|
|
|
|
dm = InboundDM(
|
|
message_id=mid,
|
|
convo_id=item.convo_id,
|
|
sender_did=msg.sender.did,
|
|
sender_handle=msg.sender.handle or msg.sender.did,
|
|
text=msg.text,
|
|
sent_at=msg.sent_at,
|
|
rev=msg.rev,
|
|
)
|
|
|
|
asyncio.create_task(self._dispatch_dm(dm, handle, on_dm, authorize_sender))
|
|
elif is_add_reaction or is_remove_reaction:
|
|
mid = f"{item.convo_id}:{item.message_id}:{item.reaction.value}"
|
|
if dedupe.has(mid):
|
|
return
|
|
if item.reaction.sender.did == handle.me_did:
|
|
return
|
|
asyncio.create_task(self._dispatch_reaction(item, on_dm, is_add_reaction))
|
|
|
|
async def _dispatch_dm(
|
|
self,
|
|
dm: InboundDM,
|
|
handle: BlueskyClientHandle,
|
|
on_dm: callable,
|
|
authorize_sender: callable,
|
|
):
|
|
try:
|
|
if not await authorize_sender(dm.sender_did, dm.text):
|
|
return
|
|
await on_dm(dm)
|
|
try:
|
|
handle.dm_client.chat.bsky.convo.update_read(
|
|
at_models.ChatBskyConvoUpdateRead.Data(
|
|
convo_id=dm.convo_id,
|
|
message_id=dm.message_id,
|
|
)
|
|
)
|
|
except Exception:
|
|
pass
|
|
except Exception as e:
|
|
logger.error("DM dispatch error: %s", e)
|
|
|
|
async def _dispatch_reaction(self, item, on_dm, is_add):
|
|
|
|
reaction_dm = InboundDM(
|
|
message_id=item.message_id,
|
|
convo_id=item.convo_id,
|
|
sender_did=item.reaction.sender.did,
|
|
sender_handle=item.reaction.sender.handle or item.reaction.sender.did,
|
|
text=f"[reaction:{'add' if is_add else 'remove'}:{item.reaction.value}]",
|
|
sent_at=item.reaction.created_at,
|
|
rev="",
|
|
)
|
|
try:
|
|
await on_dm(reaction_dm)
|
|
except Exception as e:
|
|
logger.error("Reaction dispatch error: %s", e)
|
|
|
|
async def _poll_notifications_loop(
|
|
self,
|
|
account_id: str,
|
|
handle: BlueskyClientHandle,
|
|
on_mention: callable,
|
|
dedupe: DedupeTracker,
|
|
interval: float = 15.0,
|
|
):
|
|
client = handle.client
|
|
consecutive_errors = 0
|
|
max_interval = 120.0
|
|
|
|
while True:
|
|
try:
|
|
resp = client.app.bsky.notification.list_notifications(params={"limit": 50})
|
|
for notif in resp.notifications:
|
|
if notif.reason not in ("mention", "reply", "like", "repost", "follow", "quote"):
|
|
continue
|
|
|
|
if dedupe.has(notif.uri):
|
|
continue
|
|
|
|
mention = InboundMention(
|
|
uri=notif.uri,
|
|
cid=notif.cid,
|
|
author_did=notif.author.did,
|
|
author_handle=notif.author.handle or notif.author.did,
|
|
text=getattr(notif.record, "text", f"[{notif.reason}]"),
|
|
reason=notif.reason,
|
|
indexed_at=notif.indexed_at,
|
|
)
|
|
|
|
try:
|
|
thread = client.get_post_thread(uri=notif.uri)
|
|
if thread.thread.reply:
|
|
mention.root_uri = thread.thread.reply.root.uri
|
|
mention.root_cid = thread.thread.reply.root.cid
|
|
mention.parent_uri = thread.thread.reply.parent.uri
|
|
mention.parent_cid = thread.thread.reply.parent.cid
|
|
except Exception:
|
|
pass
|
|
|
|
asyncio.create_task(self._dispatch_mention(mention, on_mention))
|
|
|
|
consecutive_errors = 0
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
consecutive_errors += 1
|
|
delay = min(interval * (2**consecutive_errors), max_interval)
|
|
delay *= random.uniform(0.8, 1.2)
|
|
logger.error(
|
|
"Bluesky notification poll error for %s (attempt %d, retry in %.1fs): %s",
|
|
account_id,
|
|
consecutive_errors,
|
|
delay,
|
|
e,
|
|
)
|
|
await asyncio.sleep(delay)
|
|
continue
|
|
await asyncio.sleep(interval)
|
|
|
|
async def _dispatch_mention(self, mention: InboundMention, on_mention: callable):
|
|
try:
|
|
await on_mention(mention)
|
|
except Exception as e:
|
|
logger.error("Mention dispatch error: %s", e)
|