mirror of
https://github.com/andrewyng/openworker.git
synced 2026-09-04 16:42:35 +00:00
OpenWorker: initial import
Imported from andrewyng/aisuite@1b4bbf303e (contents of its platform/ directory, hoisted to the repo root). Development history prior to this commit lives in that repository. Co-authored-by: Devika <devikaverma11@gmail.com>
This commit is contained in:
Rohit C Prasad
@@ -0,0 +1,528 @@
|
||||
"""Managed-relay inbound adapter — the cloud-relay alternative to Socket Mode.
|
||||
|
||||
The desktop offers the user two ways to receive Slack:
|
||||
- **Socket Mode** (`SlackAdapter`): manual bot + app tokens, one workspace, a
|
||||
direct WebSocket to Slack. No cloud involved.
|
||||
- **Managed relay** (`SlackRelayAdapter`, here): "Add to Slack" OAuth, no tokens
|
||||
typed, *many* workspaces, events pushed from OpenWorker Cloud over one
|
||||
authenticated WebSocket. Replies still go desktop → Slack Web API directly
|
||||
with the per-team bot token (the relay is inbound-only).
|
||||
|
||||
Both register on the gateway as platform ``slack`` and produce the same
|
||||
``MessageEvent``/``InteractionEvent`` — downstream code doesn't care which mode
|
||||
delivered a message. Managed-relay reply handles are **team-qualified**
|
||||
(``slack:T…/C…``) so multi-workspace replies pick the right token (see
|
||||
``slack_addr``).
|
||||
|
||||
The socket transport is injectable so the frame-handling logic is tested with a
|
||||
fake relay (no live WebSocket); the default transport is a thin ``websockets``
|
||||
client, lazy-imported like the Socket-Mode SDK.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Awaitable, Callable, Optional, Protocol
|
||||
|
||||
from .adapters import _SLACK_MENTION_RE, slack_event_to_event
|
||||
from .base import BasePlatformAdapter, InteractionEvent, SendResult, SessionSource
|
||||
from .senders import _send_slack, _send_slack_interactive
|
||||
from .slack_addr import qualify
|
||||
|
||||
logger = logging.getLogger("coworker.connectors")
|
||||
|
||||
|
||||
class RelayTransport(Protocol):
|
||||
"""One live connection to the cloud relay. Implementations lazy-import their
|
||||
WebSocket library; the frame contract is decoded JSON dicts."""
|
||||
|
||||
async def open(self) -> None: ...
|
||||
async def recv(self) -> Optional[dict]:
|
||||
"""Next frame, or None when the connection has closed."""
|
||||
...
|
||||
|
||||
async def close(self) -> None: ...
|
||||
|
||||
|
||||
TransportFactory = Callable[[], RelayTransport]
|
||||
|
||||
# Slack errors that mean the BOT TOKEN is dead (uninstalled/revoked/suspended) —
|
||||
# distinct from transient network or method errors, which say nothing about it.
|
||||
_TOKEN_ERRORS = frozenset({"invalid_auth", "account_inactive", "token_revoked"})
|
||||
TokenProvider = Callable[[], str] # returns the current cloud sign-in JWT
|
||||
# team_id, channel, count -> list of raw Slack message dicts (newest last)
|
||||
HistoryFetcher = Callable[[str, str, int], Awaitable[list[dict]]]
|
||||
|
||||
|
||||
class RelayHub:
|
||||
"""The ONE desktop↔cloud relay socket, shared by every provider adapter.
|
||||
|
||||
The cloud pushes all of a user's events down a single authenticated WS;
|
||||
frames fan out here by their `provider` tag (slack / github / …). Owns the
|
||||
transport, the read loop, and the reconnect watchdog — adapters own only
|
||||
their provider's frame handling. Extracted from SlackRelayAdapter when
|
||||
GitHub became the second relay provider (github-relay-spec §8)."""
|
||||
|
||||
_RECONNECT_DELAY = 2.0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
relay_url: str,
|
||||
token_provider: TokenProvider,
|
||||
*,
|
||||
transport_factory: Optional[TransportFactory] = None,
|
||||
reconnect_delay: Optional[float] = None,
|
||||
) -> None:
|
||||
self.relay_url = relay_url
|
||||
self._token_provider = token_provider
|
||||
self._transport_factory = transport_factory or self._default_transport_factory
|
||||
self._reconnect_delay = (
|
||||
reconnect_delay if reconnect_delay is not None else self._RECONNECT_DELAY
|
||||
)
|
||||
self._handlers: dict[str, Callable[[dict], Awaitable[None]]] = {}
|
||||
self._transport: Optional[RelayTransport] = None
|
||||
self._task: Optional[asyncio.Task] = None
|
||||
self._closing = False
|
||||
self._connections = 0 # total successful opens; reconnects == connections-1
|
||||
self._connected = False # the desktop↔relay socket is open RIGHT NOW
|
||||
self._dispatched = 0 # frames dispatched (observable for tests)
|
||||
self.last_error: str = "" # last connect/reconnect failure ("" once healthy)
|
||||
self._progress = asyncio.Event()
|
||||
|
||||
def register(
|
||||
self, provider: str, handler: Callable[[dict], Awaitable[None]]
|
||||
) -> None:
|
||||
self._handlers[provider] = handler
|
||||
|
||||
async def release(self, provider: str) -> None:
|
||||
"""An adapter is done; the socket closes when the last one leaves."""
|
||||
self._handlers.pop(provider, None)
|
||||
if not self._handlers:
|
||||
await self.stop()
|
||||
|
||||
# -- lifecycle -----------------------------------------------------------
|
||||
async def start(self) -> bool:
|
||||
"""Open the socket (idempotent — the second adapter joins the running
|
||||
loop). True when the socket is up or already running."""
|
||||
if self._task is not None and not self._task.done():
|
||||
return True
|
||||
self._closing = False
|
||||
self._transport = self._transport_factory()
|
||||
try:
|
||||
await self._transport.open()
|
||||
except Exception as exc:
|
||||
logger.exception("relay connect failed")
|
||||
self.last_error = str(exc) or type(exc).__name__
|
||||
return False
|
||||
self._connections = 1
|
||||
self._connected = True
|
||||
self.last_error = ""
|
||||
self._task = asyncio.create_task(self._run())
|
||||
return True
|
||||
|
||||
async def _run(self) -> None:
|
||||
"""Read frames; on a dropped connection, reconnect (fresh transport) —
|
||||
the relay's own watchdog analogue on the desktop side."""
|
||||
while not self._closing:
|
||||
try:
|
||||
frame = await self._transport.recv() if self._transport else None
|
||||
except Exception:
|
||||
logger.exception("relay recv error")
|
||||
frame = None
|
||||
if frame is not None:
|
||||
handler = self._handlers.get(frame.get("provider") or "slack")
|
||||
if handler is not None:
|
||||
try:
|
||||
await handler(frame)
|
||||
except Exception:
|
||||
logger.exception("relay frame dispatch failed")
|
||||
self._dispatched += 1
|
||||
self._progress.set()
|
||||
continue
|
||||
# Connection closed → reconnect unless we're shutting down.
|
||||
self._connected = False
|
||||
if self._closing:
|
||||
break
|
||||
await self._reconnect()
|
||||
|
||||
async def _reconnect(self) -> None:
|
||||
try:
|
||||
await asyncio.sleep(self._reconnect_delay)
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
if self._closing:
|
||||
return
|
||||
self._transport = self._transport_factory()
|
||||
try:
|
||||
await self._transport.open()
|
||||
self._connections += 1
|
||||
self._connected = True
|
||||
self.last_error = ""
|
||||
logger.info("relay reconnected (#%d)", self._connections - 1)
|
||||
except Exception as exc:
|
||||
self.last_error = str(exc) or type(exc).__name__
|
||||
logger.exception("relay reconnect failed — will retry")
|
||||
|
||||
async def stop(self) -> None:
|
||||
self._closing = True
|
||||
self._connected = False
|
||||
if self._transport is not None:
|
||||
try:
|
||||
await self._transport.close()
|
||||
except Exception:
|
||||
pass
|
||||
if self._task is not None:
|
||||
self._task.cancel()
|
||||
self._task = None
|
||||
|
||||
@property
|
||||
def reconnects(self) -> int:
|
||||
return max(0, self._connections - 1)
|
||||
|
||||
def state(self) -> str:
|
||||
if self._connected:
|
||||
return "live"
|
||||
if self._task is not None and not self._closing:
|
||||
return "reconnecting"
|
||||
return "offline"
|
||||
|
||||
async def wait_dispatched(self, at_least: int, timeout: float = 2.0) -> None:
|
||||
"""Test helper: wait until at least N frames have been dispatched."""
|
||||
loop = asyncio.get_event_loop()
|
||||
deadline = loop.time() + timeout
|
||||
while self._dispatched < at_least:
|
||||
self._progress.clear()
|
||||
remaining = deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
raise TimeoutError(
|
||||
f"only {self._dispatched} frames dispatched (< {at_least})"
|
||||
)
|
||||
try:
|
||||
await asyncio.wait_for(self._progress.wait(), timeout=remaining)
|
||||
except asyncio.TimeoutError:
|
||||
raise TimeoutError(
|
||||
f"only {self._dispatched} frames dispatched (< {at_least})"
|
||||
)
|
||||
|
||||
# -- default transport ---------------------------------------------------
|
||||
def _default_transport_factory(self) -> RelayTransport:
|
||||
return _WebSocketsTransport(self.relay_url, self._token_provider)
|
||||
|
||||
|
||||
class SlackRelayAdapter(BasePlatformAdapter):
|
||||
platform = "slack"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
relay_url: str,
|
||||
token_provider: TokenProvider,
|
||||
*,
|
||||
teams: Optional[dict[str, dict[str, Any]]] = None,
|
||||
transport_factory: Optional[TransportFactory] = None,
|
||||
history_fetcher: Optional[HistoryFetcher] = None,
|
||||
reconnect_delay: Optional[float] = None,
|
||||
hub: Optional[RelayHub] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.relay_url = relay_url
|
||||
# A shared hub arrives when several relay providers coexist; standalone
|
||||
# construction (tests, single-provider setups) builds its own.
|
||||
self._hub = hub or RelayHub(
|
||||
relay_url,
|
||||
token_provider,
|
||||
transport_factory=transport_factory,
|
||||
reconnect_delay=reconnect_delay,
|
||||
)
|
||||
# team_id -> {"bot_token", "bot_user_id"}. Mutable: a `revoked` frame or a
|
||||
# new install updates it.
|
||||
self._teams: dict[str, dict[str, Any]] = dict(teams or {})
|
||||
self._history_fetcher = history_fetcher
|
||||
self.last_event_at: Optional[float] = None # last Slack event delivered
|
||||
# Name resolution caches, keyed PER WORKSPACE — a U…/C… id only means
|
||||
# something inside its team, and resolution uses that team's bot token.
|
||||
self._names: dict[str, dict[str, str]] = {} # team_id -> {uid: name}
|
||||
self._channels: dict[str, dict[str, str]] = {} # team_id -> {cid: name}
|
||||
|
||||
# -- lifecycle -----------------------------------------------------------
|
||||
async def connect(self) -> bool:
|
||||
self._hub.register(self.platform, self._dispatch)
|
||||
ok = await self._hub.start()
|
||||
if ok:
|
||||
logger.info(
|
||||
"slack adapter connected (managed relay), %d team(s)", len(self._teams)
|
||||
)
|
||||
return ok
|
||||
|
||||
async def disconnect(self) -> None:
|
||||
await self._hub.release(self.platform)
|
||||
|
||||
@property
|
||||
def reconnects(self) -> int:
|
||||
return self._hub.reconnects
|
||||
|
||||
@property
|
||||
def last_error(self) -> str:
|
||||
return self._hub.last_error
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
"""Health snapshot for the GUI: the desktop↔relay socket state plus each
|
||||
workspace's bot-token health. Says nothing about Slack↔cloud — the desktop
|
||||
can't observe that leg, and event silence is not an outage."""
|
||||
return {
|
||||
"state": self._hub.state(),
|
||||
"reconnects": self._hub.reconnects,
|
||||
"last_event_at": self.last_event_at,
|
||||
"last_error": self._hub.last_error,
|
||||
"teams": {
|
||||
tid: {"token_ok": bool(info.get("token_ok", True))}
|
||||
for tid, info in self._teams.items()
|
||||
},
|
||||
}
|
||||
|
||||
async def wait_dispatched(self, at_least: int, timeout: float = 2.0) -> None:
|
||||
await self._hub.wait_dispatched(at_least, timeout)
|
||||
|
||||
# -- team registry -------------------------------------------------------
|
||||
def set_team(
|
||||
self, team_id: str, bot_token: str, bot_user_id: Optional[str] = None
|
||||
) -> None:
|
||||
self._teams[team_id] = {"bot_token": bot_token, "bot_user_id": bot_user_id}
|
||||
|
||||
def _bot_user_id(self, team_id: str) -> Optional[str]:
|
||||
return (self._teams.get(team_id) or {}).get("bot_user_id")
|
||||
|
||||
def _bot_token(self, team_id: str) -> Optional[str]:
|
||||
return (self._teams.get(team_id) or {}).get("bot_token")
|
||||
|
||||
# -- frame dispatch ------------------------------------------------------
|
||||
async def _dispatch(self, frame: dict) -> None:
|
||||
kind = frame.get("kind")
|
||||
if kind == "missed":
|
||||
await self._on_missed(frame)
|
||||
return
|
||||
if kind == "revoked":
|
||||
self._teams.pop(frame.get("team_id", ""), None)
|
||||
logger.info("slack relay team %s revoked — dropped", frame.get("team_id"))
|
||||
return
|
||||
if kind == "interactivity":
|
||||
await self._on_interactivity(frame)
|
||||
return
|
||||
# A routed Slack event.
|
||||
await self._on_event(frame)
|
||||
|
||||
async def _on_event(self, frame: dict) -> None:
|
||||
await self._dispatch_slack_event(
|
||||
frame.get("team_id", ""), frame.get("event") or {}
|
||||
)
|
||||
|
||||
async def _dispatch_slack_event(self, team_id: str, event: dict) -> None:
|
||||
"""Map a raw Slack event → MessageEvent, resolve display names via the
|
||||
per-team bot token, team-qualify the reply handle, and dispatch."""
|
||||
self.last_event_at = time.time()
|
||||
mapped = slack_event_to_event(event, self._bot_user_id(team_id))
|
||||
if mapped is None:
|
||||
return
|
||||
channel = mapped.source.chat_id # bare channel id before qualification
|
||||
# Resolve friendly names with THIS workspace's bot token (cached per team),
|
||||
# mirroring the Socket-Mode adapter — so cards read "@ocw"/"Rohit"/"#ocw-test"
|
||||
# not raw U…/C… ids. Best-effort: ids fall through on failure.
|
||||
if not mapped.source.user_name:
|
||||
mapped.source.user_name = await self._display_name(
|
||||
team_id, mapped.source.user_id
|
||||
)
|
||||
if not mapped.source.chat_name:
|
||||
mapped.source.chat_name = await self._channel_name(team_id, channel)
|
||||
mapped.text = await self._resolve_mentions(team_id, mapped.text)
|
||||
# Team-qualify the reply handle so multi-workspace replies pick the right
|
||||
# per-team token.
|
||||
mapped.source.chat_id = qualify(team_id, channel)
|
||||
mapped.source.team_id = team_id
|
||||
await self.handle_message(mapped)
|
||||
|
||||
async def _on_interactivity(self, frame: dict) -> None:
|
||||
interaction = frame.get("interaction") or {}
|
||||
actions = interaction.get("actions") or [{}]
|
||||
value = actions[0].get("value", "")
|
||||
user = interaction.get("user") or {}
|
||||
team_id = frame.get("team_id", "")
|
||||
channel = (interaction.get("channel") or {}).get("id", "")
|
||||
ts = (interaction.get("message") or {}).get("ts")
|
||||
await self.handle_interaction(
|
||||
InteractionEvent(
|
||||
platform="slack",
|
||||
chat_id=qualify(team_id, channel),
|
||||
message_id=ts,
|
||||
value=str(value),
|
||||
user_name=user.get("username") or user.get("name"),
|
||||
)
|
||||
)
|
||||
|
||||
async def _on_missed(self, frame: dict) -> None:
|
||||
"""A nudge: content was dropped (offline > TTL / overflow). Pull the
|
||||
recent channel history ourselves via the per-team bot token and replay
|
||||
the missed messages (spec §7 channel-context / nudge)."""
|
||||
team_id = frame.get("team_id", "")
|
||||
channel = frame.get("channel", "")
|
||||
count = int(frame.get("count", 0)) or 1
|
||||
if self._history_fetcher is None or not channel:
|
||||
return
|
||||
try:
|
||||
messages = await self._history_fetcher(team_id, channel, count)
|
||||
except Exception:
|
||||
logger.exception("relay nudge history fetch failed")
|
||||
return
|
||||
for raw in messages:
|
||||
await self._dispatch_slack_event(team_id, {**raw, "channel": channel})
|
||||
|
||||
def _note_token_health(self, team_id: str, error: Optional[str]) -> None:
|
||||
"""Record what a Web API call said about the team's bot token: success
|
||||
proves it live; a token-class error marks it dead; anything else —
|
||||
network trouble, channel_not_found — says nothing, so changes nothing."""
|
||||
info = self._teams.get(team_id)
|
||||
if info is None:
|
||||
return
|
||||
if error is None:
|
||||
info["token_ok"] = True
|
||||
elif error in _TOKEN_ERRORS:
|
||||
info["token_ok"] = False
|
||||
|
||||
# -- name resolution (per workspace, via that team's bot token) ----------
|
||||
async def _slack_get(
|
||||
self, team_id: str, method: str, params: dict
|
||||
) -> Optional[dict]:
|
||||
"""Call a Slack Web API read method with the team's bot token. Best-effort
|
||||
(None on any failure). `SLACK_API_URL` redirects to the fake in tests."""
|
||||
import httpx
|
||||
|
||||
token = self._bot_token(team_id)
|
||||
if not token:
|
||||
return None
|
||||
base = os.environ.get("SLACK_API_URL", "https://slack.com/api/")
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=15) as http:
|
||||
resp = await http.get(
|
||||
base + method,
|
||||
params=params,
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
data = resp.json()
|
||||
except Exception:
|
||||
return None
|
||||
self._note_token_health(team_id, None if data.get("ok") else data.get("error"))
|
||||
return data if data.get("ok") else None
|
||||
|
||||
async def _display_name(self, team_id: str, uid: Optional[str]) -> Optional[str]:
|
||||
if not uid:
|
||||
return None
|
||||
cache = self._names.setdefault(team_id, {})
|
||||
if uid in cache:
|
||||
return cache[uid]
|
||||
data = await self._slack_get(team_id, "users.info", {"user": uid})
|
||||
u = (data or {}).get("user") or {}
|
||||
prof = u.get("profile") or {}
|
||||
name = (
|
||||
prof.get("display_name")
|
||||
or prof.get("real_name")
|
||||
or u.get("real_name")
|
||||
or u.get("name")
|
||||
)
|
||||
if name:
|
||||
cache[uid] = name
|
||||
return name
|
||||
|
||||
async def _channel_name(self, team_id: str, cid: Optional[str]) -> Optional[str]:
|
||||
if not cid:
|
||||
return None
|
||||
cache = self._channels.setdefault(team_id, {})
|
||||
if cid in cache:
|
||||
return cache[cid]
|
||||
data = await self._slack_get(team_id, "conversations.info", {"channel": cid})
|
||||
chan = (data or {}).get("channel") or {}
|
||||
name = chan.get("name") or chan.get("name_normalized")
|
||||
if name:
|
||||
cache[cid] = name
|
||||
return name
|
||||
|
||||
async def _resolve_mentions(self, team_id: str, text: str) -> str:
|
||||
"""Rewrite `<@U…>` tokens to `@display-name` (cached). Best-effort."""
|
||||
out = text
|
||||
for uid in set(_SLACK_MENTION_RE.findall(text or "")):
|
||||
name = await self._display_name(team_id, uid)
|
||||
if name:
|
||||
out = re.sub(rf"<@{re.escape(uid)}(?:\|[^>]*)?>", f"@{name}", out)
|
||||
return out
|
||||
|
||||
# -- outbound ------------------------------------------------------------
|
||||
async def send(
|
||||
self, chat_id: str, text: str, *, thread_id: Optional[str] = None
|
||||
) -> SendResult:
|
||||
"""Reply directly via the Slack Web API with the per-team bot token."""
|
||||
from .slack_addr import split
|
||||
|
||||
team_id, _channel = split(chat_id)
|
||||
token = self._bot_token(team_id or "")
|
||||
if not token:
|
||||
return SendResult(False, error=f"no bot token for team {team_id}")
|
||||
result = await asyncio.to_thread(_send_slack, token, chat_id, text, thread_id)
|
||||
self._note_token_health(team_id or "", None if result.ok else result.error)
|
||||
return result
|
||||
|
||||
async def send_interactive(
|
||||
self, chat_id: str, text: str, buttons, *, thread_id: Optional[str] = None
|
||||
) -> SendResult:
|
||||
from .slack_addr import split
|
||||
|
||||
team_id, _channel = split(chat_id)
|
||||
token = self._bot_token(team_id or "")
|
||||
if not token:
|
||||
return SendResult(False, error=f"no bot token for team {team_id}")
|
||||
result = await asyncio.to_thread(
|
||||
_send_slack_interactive, token, chat_id, text, buttons, thread_id
|
||||
)
|
||||
self._note_token_health(team_id or "", None if result.ok else result.error)
|
||||
return result
|
||||
|
||||
|
||||
class _WebSocketsTransport:
|
||||
"""Real transport: an authenticated `websockets` client. Sends the cloud
|
||||
sign-in JWT in the Authorization header (the relay's $connect authorizer)."""
|
||||
|
||||
def __init__(self, url: str, token_provider: TokenProvider) -> None:
|
||||
self._url = url
|
||||
self._token_provider = token_provider
|
||||
self._ws = None
|
||||
|
||||
async def open(self) -> None:
|
||||
import websockets # lazy: optional extra
|
||||
|
||||
token = self._token_provider()
|
||||
self._ws = await websockets.connect(
|
||||
self._url, additional_headers={"Authorization": f"Bearer {token}"}
|
||||
)
|
||||
|
||||
async def recv(self) -> Optional[dict]:
|
||||
import websockets
|
||||
|
||||
if self._ws is None:
|
||||
return None
|
||||
try:
|
||||
raw = await self._ws.recv()
|
||||
except websockets.ConnectionClosed:
|
||||
return None
|
||||
try:
|
||||
return json.loads(raw)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._ws is not None:
|
||||
try:
|
||||
await self._ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._ws = None
|
||||
Reference in New Issue
Block a user