diff --git a/backend/apps/tools_lib/tools_lib.py b/backend/apps/tools_lib/tools_lib.py index 3edeb0d0..d811026f 100644 --- a/backend/apps/tools_lib/tools_lib.py +++ b/backend/apps/tools_lib/tools_lib.py @@ -326,6 +326,32 @@ def derive_mcp_config(tool: ToolDefinition) -> Optional[dict]: existing_pp = env.get("PYTHONPATH") or os.environ.get("PYTHONPATH", "") env["PYTHONPATH"] = (_project_root + os.pathsep + existing_pp) if existing_pp else _project_root + # Twitter MCP runs as a Python shim (`backend.apps.twitter_mcp_shim`) + # that forwards each tool call to the local backend's + # /api/twitter/* routes. The shim carries no twikit dependency — it + # only speaks HTTP — so we just give it the base URL and a path to + # the bearer-token file (re-read on every call so a backend + # restart that rotates the token can't strand the shim in 401). + if tool.name.lower() == "twitter" and config.get("type") == "stdio": + from backend.config.paths import AUTH_TOKEN_FILE + from backend.auth import get_auth_token + env = config.setdefault("env", {}) + env.setdefault( + "OPENSWARM_BASE_URL", + f"http://127.0.0.1:{os.environ.get('OPENSWARM_PORT', '8324')}", + ) + env.setdefault("OPENSWARM_AUTH_TOKEN_FILE", AUTH_TOKEN_FILE) + # Belt-and-suspenders: pass the current token via env too. The + # shim prefers the file but falls back to the env var if the + # file isn't readable (e.g. packaged-mode permission glitch). + env.setdefault("OPENSWARM_AUTH_TOKEN", get_auth_token()) + # Same PYTHONPATH trick the Discord shim uses — the subprocess + # needs to import `backend.apps.twitter_mcp_shim`, so we point + # PYTHONPATH at the project root (parent of backend/). + _project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + existing_pp = env.get("PYTHONPATH") or os.environ.get("PYTHONPATH", "") + env["PYTHONPATH"] = (_project_root + os.pathsep + existing_pp) if existing_pp else _project_root + # Microsoft 365 MCP: use a stable token cache path shared across process spawns if tool.name.lower() == "microsoft 365" and config.get("type") == "stdio": env = config.setdefault("env", {}) diff --git a/backend/apps/twitter/__init__.py b/backend/apps/twitter/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/twitter/_twikit_patches.py b/backend/apps/twitter/_twikit_patches.py new file mode 100644 index 00000000..fb241b9f --- /dev/null +++ b/backend/apps/twitter/_twikit_patches.py @@ -0,0 +1,674 @@ +"""Runtime patches for twikit 2.3.3. + +This module ships four independent workarounds for upstream gaps that +accumulated through 2026. All are gated by env vars so they can be disabled +once twikit releases real fixes. They stack: Patch 3 (TLS) gets us past +Cloudflare's pre-HTTP scoring, Patch 2 (headers) keeps the post-handshake +HTTP layer internally consistent, Patch 1 (transaction-id) lets the actual +authenticated request reach X's GraphQL API, Patch 4 (User parser) keeps +the *parsed* response from silently dropping every tweet whose author has +a profile shape twikit's strict-key access doesn't expect. + +Patch 1: x_client_transaction regex (default: ON) +------------------------------------------------- +X rotated the minified JS that twikit's ``ClientTransaction.get_indices`` +regex-scrapes for the ``x-client-transaction-id`` signing keys. twikit 2.3.3 +raises ``Exception("Couldn't get KEY_BYTE indices")`` on every authenticated +request (including ``client.login``) until a new release is cut. Tracked in +https://github.com/d60/twikit/issues/408 and an unmerged fix at +https://github.com/d60/twikit/pull/407. + +Specifically this: + +1. Replaces ``ON_DEMAND_FILE_REGEX``. The home page chunk map used to embed + ``'ondemand.s':''`` directly; the post-rotation build instead writes + ``,:"ondemand.s",...,:""``, so we have to find the chunk + index first and then look up its hash in a second pass. +2. Adds ``ON_DEMAND_HASH_PATTERN`` for that second pass. +3. Relaxes ``INDICES_REGEX`` from ``\\w{1}`` to ``\\w{1,2}``. The byte-array + variable name in the minified JS is now 1-2 chars (e.g. ``xx[12]``) where + it used to be exactly one (``x[12]``). +4. Replaces ``ClientTransaction.get_indices`` with the two-step extractor. + +Disable with ``OPENSWARM_TWITTER_DISABLE_TWIKIT_PATCH=1``. + +Patch 2: Cloudflare-friendly request headers (default: ON) +---------------------------------------------------------- +After Patch 1 lets the request reach X's edge, X's Cloudflare layer aggressively +challenges twikit's default identifiers: the shipped User-Agent claims Safari 17 +but the ``httpx`` TLS fingerprint and missing ``sec-ch-ua-*`` / ``sec-fetch-*`` +headers give it away as a bot. Login POSTs often come back with a +``403 Forbidden`` containing a Cloudflare "Sorry, you have been blocked" +interstitial. Same upstream PR proposes bumping the UA to Chrome 133 and adding +the modern Sec-* set. + +Specifically this: + +1. Sets a Chrome 133 User-Agent on every new ``Client`` instance that didn't + request a specific UA (we don't override user-supplied UAs). +2. Wraps ``Client._base_headers`` to merge in ``sec-ch-ua``, ``sec-ch-ua-mobile``, + ``sec-ch-ua-platform``, ``sec-fetch-dest``, ``sec-fetch-mode``, + ``sec-fetch-site``. + +This addresses the HTTP-level fingerprint but cannot fix the TLS-level one +on its own — Patch 3 below is what actually clears Cloudflare's pre-HTTP +JA3/JA4 scoring. Without Patch 3 this patch is largely cosmetic: Cloudflare +rejects the connection before the UA is read. + +Disable with ``OPENSWARM_TWITTER_DISABLE_TWIKIT_HEADER_PATCH=1``. + +Patch 3: TLS transport via curl-impersonate (default: ON) +--------------------------------------------------------- +Cloudflare scores the TLS ClientHello (JA3/JA4: cipher order, extensions, +GREASE values) and the HTTP/2 SETTINGS frame *before* it reads any HTTP +headers. Stock httpx → Python's OpenSSL binding → a recognizable +"library, not browser" fingerprint → 403 on every request to x.com, +regardless of Patch 2's headers. This is the actual blocker behind +``twikit.errors.Forbidden: status: 403 ... Sorry, you have been blocked`` +reported widely since late 2025 (twikit#396). + +We replace twikit's transport with ``httpx-curl-cffi``'s ``AsyncCurlTransport``, +which wraps curl-impersonate (Chrome's actual BoringSSL build under the hood). +The result is a byte-identical TLS ClientHello + HTTP/2 SETTINGS to a real +Chrome 133. + +Specifically this: + +1. Monkey-patches the ``AsyncClient`` symbol in ``twikit.client.client`` + (and ``twikit.guest.client``) so twikit's ``Client.__init__`` — + ``self.http = AsyncClient(proxy=proxy, **kwargs)`` — gets our wrapper that + injects ``transport=AsyncCurlTransport(impersonate="chrome133", ...)``. +2. Monkey-patches the ``AsyncHTTPTransport`` symbol in the same modules. + twikit's ``proxy`` setter runs unconditionally in ``__init__`` + (``self.proxy = proxy`` at line 109) and assigns + ``self.http._mounts = {URLPattern('all://'): AsyncHTTPTransport(proxy=url)}``. + Without this second patch, the mount silently overrides our default + transport the moment the setter fires — we'd be back on stock httpx + immediately after construction. + +We patch the *imported names* in twikit's namespace, not ``httpx.AsyncClient`` +itself, so only twikit's ``self.http`` lives behind curl-impersonate. FastAPI, +the agent SDK's outbound HTTP, ``anthropic_proxy``, etc., all stay on stock +httpx. + +This is the silver-bullet patch — the previous two are necessary but not +sufficient without this one. + +Disable with ``OPENSWARM_TWITTER_DISABLE_TWIKIT_TLS_PATCH=1``. + +Patch 4: User parser tolerance (default: ON) +-------------------------------------------- +Even when Patches 1–3 get us a valid authenticated response from X, +twikit's ``twikit.user.User.__init__`` still hard-accesses ~30 keys off +``data['legacy'][...]`` (e.g. ``legacy['entities']['description']['urls']``, +``legacy['fast_followers_count']``). X has been gradually omitting fields +from the legacy shape — accounts with empty bios drop +``entities.description.urls``, accounts without a pinned tweet drop +``pinned_tweet_ids_str``, etc. Any missing key raises ``KeyError`` mid- +construction. + +The damage compounds because ``twikit.tweet.tweet_from_data`` calls +``User(client, ...)`` for every tweet's author, and ``client.search_tweet`` +silently swallows ``KeyError`` per-item (``twikit/client/client.py:765``): + + try: + tweet = tweet_from_data(self, item) + except KeyError: + tweet = None + +So a single missing author-field turns the entire SearchResult into an +empty ``items`` list, with no log line and no exception surface. Same +class of failure hits ``client.user()`` (the smoke probe) directly: +``KeyError`` raises out of the call and our pool flips the account to +``needs_relogin`` even though the cookies are perfectly valid. + +Specifically this: + +1. Imports ``twikit.user`` and replaces ``User.__init__`` with a tolerant + version that mirrors the original field-by-field, but uses + ``data.get(...)`` and ``legacy.get(...)`` with type-appropriate defaults + (``''`` / ``None`` for strings, ``0`` for counts, ``False`` for bools, + ``[]`` for lists) instead of bare ``[key]`` lookups. +2. Keeps ``rest_id`` as the one hard access. A User without a rest_id is + genuinely unidentifiable; failing fast there is correct. + +After this patch, ``client.user()`` returns a real User on any +authenticated session, ``/verify`` reflects truth, the smoke probe stops +auto-quarantining accounts, and ``/search``/``/user/{id}/tweets``/etc. +stop silently dropping tweets at the parser layer. + +Disable with ``OPENSWARM_TWITTER_DISABLE_TWIKIT_USER_PATCH=1``. + +Operational notes +----------------- +- ``apply()`` is idempotent; safe to call from any number of import sites. +- If a patch itself fails (twikit not importable, ``httpx_curl_cffi`` missing, + internal class shape changed again), we log a warning and continue. The + SubApp will still 503 or raise login errors, but startup won't blow up. +- The smoke probe in ``twitter.twitter._smoke_probe`` will catch any + *further* X-side drift past these patches and audit-log ``smoke_fail``. +""" + +from __future__ import annotations + +import logging +import os +import re + +logger = logging.getLogger(__name__) + +# Chrome 133 on macOS. Matches the sec-ch-ua trio below so the fingerprint +# is internally consistent (mismatches between UA and sec-ch-ua are a strong +# bot signal). Update both together if you bump this. +_CHROME_UA = ( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/133.0.0.0 Safari/537.36" +) +_CHROME_SEC_HEADERS = { + "sec-ch-ua": '"Not(A:Brand";v="99", "Google Chrome";v="133", "Chromium";v="133"', + "sec-ch-ua-mobile": "?0", + "sec-ch-ua-platform": '"macOS"', + "sec-fetch-dest": "empty", + "sec-fetch-mode": "cors", + "sec-fetch-site": "same-origin", +} + +_APPLIED_TX = False +_APPLIED_HEADERS = False +_APPLIED_TLS = False +_APPLIED_USER = False + +# curl-impersonate target preference list for Patch 3. We probe the +# installed curl_cffi's BrowserType enum at apply time and pick the first +# of these that's actually supported, so a curl_cffi upgrade or +# downgrade can't strand us on a missing target name. +# +# Order: chrome133a (matches our Chrome 133 UA in `_CHROME_UA`) first, +# then graceful degradation through older but still-recent Chrome +# targets. The actual JA3/JA4 difference between adjacent Chrome +# majors is tiny — Cloudflare scoring doesn't materially distinguish +# Chrome 131 from 133. The UA-vs-impersonate mismatch one notch off is +# similarly insignificant; both still claim "modern Chrome." +# +# When you bump `_CHROME_UA` to a newer Chrome major, prepend the +# matching impersonate target here. Newest-Chrome-first ordering is +# the only invariant. +_IMPERSONATE_PREFERENCE: tuple[str, ...] = ( + "chrome146", + "chrome142", + "chrome136", + "chrome133a", + "chrome131", + "chrome124", + "chrome120", +) + + +def apply() -> bool: + """Apply all patches. Idempotent. + + Returns True if at least one patch is now in place, False if everything + was skipped or failed. + """ + tx_ok = _apply_transaction_patch() + headers_ok = _apply_header_patch() + tls_ok = _apply_tls_transport_patch() + user_ok = _apply_user_patch() + return tx_ok or headers_ok or tls_ok or user_ok + + +def _apply_transaction_patch() -> bool: + global _APPLIED_TX + if _APPLIED_TX: + return True + if os.environ.get("OPENSWARM_TWITTER_DISABLE_TWIKIT_PATCH"): + logger.info( + "twitter: skipping twikit x_client_transaction patch " + "(OPENSWARM_TWITTER_DISABLE_TWIKIT_PATCH set)" + ) + return False + + try: + from twikit.x_client_transaction import transaction as _tx + except ImportError as e: + logger.warning( + "twitter: cannot apply twikit x_client_transaction patch (%s)", e + ) + return False + + try: + _tx.ON_DEMAND_FILE_REGEX = re.compile( + r""",(\d+):["']ondemand\.s["']""", + flags=re.MULTILINE, + ) + _tx.ON_DEMAND_HASH_PATTERN = r',{}:"([0-9a-f]+)"' + _tx.INDICES_REGEX = re.compile( + r"""(\(\w{1,2}\[(\d{1,2})\],\s*16\))+""", + flags=re.MULTILINE, + ) + + async def get_indices(self, home_page_response, session, headers): + """Two-step extractor: locate the ``ondemand.s`` chunk index in + the home page, look up its hash, then scrape byte indices from + the resulting JS file. Raises the same exception as the original + on failure so callers (and our smoke probe) see identical + behavior when X drifts again. + """ + key_byte_indices: list[str] = [] + response = ( + self.validate_response(home_page_response) + or self.home_page_response + ) + response_str = str(response) + + on_demand_file = _tx.ON_DEMAND_FILE_REGEX.search(response_str) + if on_demand_file: + chunk_index = on_demand_file.group(1) + hash_regex = re.compile( + _tx.ON_DEMAND_HASH_PATTERN.format(chunk_index) + ) + hash_match = hash_regex.search(response_str) + if hash_match: + filename = hash_match.group(1) + on_demand_file_url = ( + "https://abs.twimg.com/responsive-web/client-web/" + f"ondemand.s.{filename}a.js" + ) + on_demand_file_response = await session.request( + method="GET", + url=on_demand_file_url, + headers=headers, + ) + for item in _tx.INDICES_REGEX.finditer( + str(on_demand_file_response.text) + ): + key_byte_indices.append(item.group(2)) + + if not key_byte_indices: + raise Exception("Couldn't get KEY_BYTE indices") + idxs = list(map(int, key_byte_indices)) + return idxs[0], idxs[1:] + + _tx.ClientTransaction.get_indices = get_indices + except Exception as e: + logger.warning( + "twitter: failed to apply twikit x_client_transaction patch (%s); " + "login and authenticated calls will likely fail with " + "'Couldn't get KEY_BYTE indices'", e, + ) + return False + + _APPLIED_TX = True + # WARNING level (not INFO) because the rest of the backend's + # logging config swallows INFO from non-uvicorn loggers — and + # "did my patches actually apply?" is exactly the question an + # operator needs to answer when twitter routes 403. Once-per- + # startup status line is cheap. + logger.warning( + "twitter: applied twikit x_client_transaction patch " + "(workaround for github.com/d60/twikit/issues/408)" + ) + return True + + +def _apply_header_patch() -> bool: + """Bump default UA + inject modern browser headers into ``Client._base_headers``. + + Only affects ``Client`` instances that didn't pass an explicit ``user_agent`` + kwarg — we never override a caller's intentional choice. + """ + global _APPLIED_HEADERS + if _APPLIED_HEADERS: + return True + if os.environ.get("OPENSWARM_TWITTER_DISABLE_TWIKIT_HEADER_PATCH"): + logger.info( + "twitter: skipping twikit header patch " + "(OPENSWARM_TWITTER_DISABLE_TWIKIT_HEADER_PATCH set)" + ) + return False + + try: + from twikit.client import client as _client_mod + except ImportError as e: + logger.warning("twitter: cannot apply twikit header patch (%s)", e) + return False + + try: + # Sentinel string burned into the shipped 2.3.3 wheel. If we ever see + # a different default we bail out — that means twikit updated UAs on + # its own and our override would be a regression. + _expected_default_ua = ( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 14_6_1) " + "AppleWebKit/605.1.15 (KHTML, like Gecko) " + "Version/17.5 Safari/605.1.15" + ) + + _original_init = _client_mod.Client.__init__ + + def _patched_init(self, *args, **kwargs): + _original_init(self, *args, **kwargs) + # Only swap if the caller didn't request a specific UA and the + # default is still what we expect. This lets a user override us + # via Client(user_agent="...") without us silently clobbering it. + user_supplied = kwargs.get("user_agent") is not None + if not user_supplied and self._user_agent == _expected_default_ua: + self._user_agent = _CHROME_UA + + _original_base_headers_prop = _client_mod.Client._base_headers + _original_base_headers_fget = _original_base_headers_prop.fget + + def _patched_base_headers(self): + headers = _original_base_headers_fget(self) + for k, v in _CHROME_SEC_HEADERS.items(): + headers.setdefault(k, v) + return headers + + _client_mod.Client.__init__ = _patched_init + _client_mod.Client._base_headers = property(_patched_base_headers) + except Exception as e: + logger.warning( + "twitter: failed to apply twikit header patch (%s); " + "Cloudflare blocks on login POSTs will be more likely", + e, + ) + return False + + _APPLIED_HEADERS = True + logger.warning( + "twitter: applied twikit header patch (Chrome 133 UA + sec-ch-ua/sec-fetch headers)" + ) + return True + + +def _apply_tls_transport_patch() -> bool: + """Replace twikit's httpx.AsyncClient TLS layer with curl-impersonate. + + See the module docstring's "Patch 3" section for the full rationale. + Summary: Cloudflare scores TLS ClientHello and HTTP/2 SETTINGS frames + before reading HTTP headers; stock httpx looks like a Python script + at the TLS layer and gets 403'd regardless of the Patch 2 headers. + Wrapping the transport in ``httpx_curl_cffi.AsyncCurlTransport`` makes + twikit's wire-level fingerprint match a real Chrome 133. + + We patch two names in twikit's module namespaces (NOT ``httpx`` itself + — keeps the rest of the backend on stock httpx): + + - ``AsyncClient`` — intercepts twikit's ``self.http = AsyncClient( + proxy=proxy, **kwargs)`` and injects ``transport=AsyncCurlTransport( + impersonate="chrome133", ...)``. We strip ``proxy=`` before passing + to the real ``httpx.AsyncClient`` because httpx errors when both + ``transport=`` and ``proxy=`` are supplied; the proxy travels with + the transport instead. + - ``AsyncHTTPTransport`` — intercepts twikit's + ``self.http._mounts = {URLPattern('all://'): AsyncHTTPTransport( + proxy=url)}`` (run unconditionally from ``Client.__init__`` via + ``self.proxy = proxy``). Without this second patch, the mount + clobbers our default transport and twikit silently falls back to + stock httpx the moment the proxy setter fires. + + Both ``twikit.client.client`` and ``twikit.guest.client`` import these + names directly, so we patch both namespaces. We don't currently use + ``GuestClient`` in OpenSwarm but it's cheap to cover for future use. + + Disable with ``OPENSWARM_TWITTER_DISABLE_TWIKIT_TLS_PATCH=1``. + """ + global _APPLIED_TLS + if _APPLIED_TLS: + return True + if os.environ.get("OPENSWARM_TWITTER_DISABLE_TWIKIT_TLS_PATCH"): + logger.info( + "twitter: skipping twikit TLS transport patch " + "(OPENSWARM_TWITTER_DISABLE_TWIKIT_TLS_PATCH set); " + "Cloudflare 403s on x.com are likely" + ) + return False + + try: + from httpx_curl_cffi import AsyncCurlTransport, CurlOpt + except ImportError as e: + logger.warning( + "twitter: httpx_curl_cffi not importable (%s); skipping TLS " + "transport patch. Cloudflare 403s on x.com are likely — " + "ensure `httpx-curl-cffi` and `curl-cffi` are installed.", + e, + ) + return False + + # Pick the best impersonate target available in the installed + # curl_cffi. Target names drift between curl_cffi versions + # (e.g. 0.15 ships `chrome133a` not `chrome133`), so a hardcoded + # string would be brittle. We consult the BrowserType enum and + # take the newest entry from `_IMPERSONATE_PREFERENCE` that's + # actually present. + try: + from curl_cffi.requests.impersonate import BrowserType + available = {bt.value for bt in BrowserType} + except ImportError as e: + logger.warning( + "twitter: curl_cffi BrowserType not importable (%s); skipping TLS patch", + e, + ) + return False + + impersonate_target = next( + (t for t in _IMPERSONATE_PREFERENCE if t in available), + None, + ) + if impersonate_target is None: + logger.warning( + "twitter: no preferred Chrome target found in curl_cffi " + "(available chrome targets: %s); skipping TLS patch", + sorted(t for t in available if "chrome" in t.lower()), + ) + return False + + # Patch every twikit module that imports the two httpx names directly. + # We treat this as a list rather than hardcoding the strings inline so + # adding a third namespace later (e.g. a new twikit subpackage) is a + # one-line change. + target_modules: list = [] + for mod_path in ("twikit.client.client", "twikit.guest.client"): + try: + import importlib + target_modules.append(importlib.import_module(mod_path)) + except ImportError as e: + # The guest client may not be present in every twikit build; + # main client must be. We log either way and continue with + # whatever we have. + logger.debug("twitter: TLS patch — %s not importable (%s)", mod_path, e) + + if not target_modules: + logger.warning( + "twitter: TLS patch found no twikit modules to patch; " + "package layout may have changed" + ) + return False + + try: + def _make_transport(proxy): + # `FRESH_CONNECT=True` is required by httpx-curl-cffi when + # issuing parallel async requests (see their README's + # "curl_cffi issues"). RateGate runs at most one twikit call + # per account at a time (concurrency semaphore), but the + # smoke probe + a tool call on the same Client can interleave + # briefly, so cheap insurance. + # + # `default_headers=True` is critical for Cloudflare. twikit's + # `Client.request()` adds `_base_headers` (which our Patch 2 + # enriches with sec-ch-ua/sec-fetch) only on some code paths + # — e.g. `V11Client.guest_activate` passes them, but + # `V11Client.onboarding_task` hardcodes a minimal 2-header + # dict (x-guest-token + Authorization) and inherits nothing + # else. A POST that presents a Chrome TLS+HTTP/2 fingerprint + # but goes out with no User-Agent, no Accept-Language, no + # sec-ch-ua-* is an obvious inconsistency that Cloudflare + # scores as bot and 403s. With default_headers=True, + # curl-impersonate fills in Chrome's standard browser + # headers for any name twikit didn't set explicitly — making + # the whole request internally consistent. twikit's own + # headers (Authorization, x-guest-token, x-csrf-token, + # X-Client-Transaction-Id, etc.) are preserved verbatim. + return AsyncCurlTransport( + impersonate=impersonate_target, + default_headers=True, + curl_options={CurlOpt.FRESH_CONNECT: True}, + proxy=proxy, + ) + + def _patched_async_client_factory(real_async_client): + def _patched(*args, **kwargs): + # Honor an explicitly-supplied transport (no caller in + # twikit does this today, but cheap to be polite — and + # makes the patch trivially testable: pass your own + # transport to bypass curl_cffi entirely). + if "transport" in kwargs: + return real_async_client(*args, **kwargs) + # httpx errors when both `transport=` and `proxy=` are + # given. Strip `proxy` from the AsyncClient kwargs and + # hand it to the transport instead. + proxy = kwargs.pop("proxy", None) + kwargs["transport"] = _make_transport(proxy) + return real_async_client(*args, **kwargs) + return _patched + + def _patched_http_transport_factory(): + def _patched(*args, **kwargs): + # Twikit only ever calls this as + # `AsyncHTTPTransport(proxy=url)` from the proxy setter. + # We discard any other transport kwargs (`verify=`, + # `cert=`, `http1=`, `http2=`, ...) on the assumption + # that curl-impersonate's Chrome 133 defaults are what + # we want — twikit never passes those anyway. + proxy = kwargs.pop("proxy", None) + return _make_transport(proxy) + return _patched + + for mod in target_modules: + real_async_client = getattr(mod, "AsyncClient", None) + if real_async_client is None: + logger.debug( + "twitter: TLS patch — %s has no AsyncClient name; skipping", + mod.__name__, + ) + continue + mod.AsyncClient = _patched_async_client_factory(real_async_client) + mod.AsyncHTTPTransport = _patched_http_transport_factory() + except Exception as e: + logger.warning( + "twitter: failed to apply twikit TLS transport patch (%s); " + "Cloudflare 403s likely until fixed", + e, + ) + return False + + _APPLIED_TLS = True + logger.warning( + "twitter: applied twikit TLS transport patch " + "(curl-impersonate %s; workaround for github.com/d60/twikit/issues/396)", + impersonate_target, + ) + return True + + +def _apply_user_patch() -> bool: + """Replace ``twikit.user.User.__init__`` with a missing-key-tolerant version. + + See the module docstring's "Patch 4" section for full rationale. + Summary: twikit's User constructor hard-accesses ~30 keys off + ``data['legacy']``. X is dropping fields from that shape, and the + KeyError propagates up to ``client.search_tweet``'s silent + ``except KeyError`` (``twikit/client/client.py:763-766``), which + swallows the whole tweet. End result: empty SearchResult ``items`` + with valid cursors, no log line. Same KeyError flow flips smoke-probe + accounts to ``needs_relogin`` even when cookies are valid. + + We mirror the original constructor field-for-field but use ``.get()`` + with type-appropriate defaults. ``rest_id`` stays hard — a User + without an id is genuinely unidentifiable and we want to know. + + Disable with ``OPENSWARM_TWITTER_DISABLE_TWIKIT_USER_PATCH=1``. + """ + global _APPLIED_USER + if _APPLIED_USER: + return True + if os.environ.get("OPENSWARM_TWITTER_DISABLE_TWIKIT_USER_PATCH"): + logger.info( + "twitter: skipping twikit User parser patch " + "(OPENSWARM_TWITTER_DISABLE_TWIKIT_USER_PATCH set); " + "tweet author parse failures will silently empty SearchResults" + ) + return False + + try: + from twikit import user as _user_mod + except ImportError as e: + logger.warning( + "twitter: cannot apply twikit User parser patch (%s); " + "tweet author parse failures will silently empty SearchResults", + e, + ) + return False + + try: + def _patched_init(self, client, data): + # `data` is the GraphQL result envelope: a dict with + # `rest_id`, `legacy`, `is_blue_verified`, etc. Anything + # missing gets a typed default; rest_id stays hard because + # an id-less User is meaningless. + self._client = client + + legacy = data.get('legacy') or {} + entities = legacy.get('entities') or {} + description_entities = entities.get('description') or {} + url_entities = entities.get('url') or {} + + self.id = data['rest_id'] + self.created_at = legacy.get('created_at') + self.name = legacy.get('name') + self.screen_name = legacy.get('screen_name') + self.profile_image_url = legacy.get('profile_image_url_https') + self.profile_banner_url = legacy.get('profile_banner_url') + self.url = legacy.get('url') + self.location = legacy.get('location') + self.description = legacy.get('description') + self.description_urls = description_entities.get('urls') or [] + self.urls = url_entities.get('urls') or [] + self.pinned_tweet_ids = legacy.get('pinned_tweet_ids_str') or [] + self.is_blue_verified = data.get('is_blue_verified', False) + self.verified = legacy.get('verified', False) + self.possibly_sensitive = legacy.get('possibly_sensitive', False) + self.can_dm = legacy.get('can_dm', False) + self.can_media_tag = legacy.get('can_media_tag', False) + self.want_retweets = legacy.get('want_retweets', False) + self.default_profile = legacy.get('default_profile', False) + self.default_profile_image = legacy.get('default_profile_image', False) + self.has_custom_timelines = legacy.get('has_custom_timelines', False) + self.followers_count = legacy.get('followers_count', 0) + self.fast_followers_count = legacy.get('fast_followers_count', 0) + self.normal_followers_count = legacy.get('normal_followers_count', 0) + self.following_count = legacy.get('friends_count', 0) + self.favourites_count = legacy.get('favourites_count', 0) + self.listed_count = legacy.get('listed_count', 0) + self.media_count = legacy.get('media_count', 0) + self.statuses_count = legacy.get('statuses_count', 0) + self.is_translator = legacy.get('is_translator', False) + self.translator_type = legacy.get('translator_type') + self.withheld_in_countries = legacy.get('withheld_in_countries') or [] + self.protected = legacy.get('protected', False) + + _user_mod.User.__init__ = _patched_init + except Exception as e: + logger.warning( + "twitter: failed to apply twikit User parser patch (%s); " + "tweet author parse failures will silently empty SearchResults", + e, + ) + return False + + _APPLIED_USER = True + logger.warning( + "twitter: applied twikit User parser patch " + "(tolerant .get() for ~30 legacy fields; unblocks search + verify)" + ) + return True diff --git a/backend/apps/twitter/cache.py b/backend/apps/twitter/cache.py new file mode 100644 index 00000000..ac5206d2 --- /dev/null +++ b/backend/apps/twitter/cache.py @@ -0,0 +1,139 @@ +"""TTL cache for twikit read responses. + +Caches successful (non-error) twikit responses keyed by +`(endpoint, normalized_args)`. Two reasons this matters: + +1. LLM agents replay tool calls — the same `get_user("openai")` shows up + many times in a typical conversation as the model thinks. A short TTL + absorbs the redundancy without staleness mattering for the bulk of + queries. +2. Cache hits skip the rate-limit bucket entirely, which is the single + biggest throughput win for the entire SubApp. + +Backed by sqlite so a process restart doesn't blow the budget — the +first agent call after a restart can still hit the cache instead of +twikit. The in-memory layer is just a write-through speedup; durability +lives in sqlite. + +Schema is owned by `persistence.py` (one connection per SubApp, +WAL mode, see there). This module only knows how to GET/SET against +it. + +We never cache errors. If `set()` is called with a non-OK payload (the +caller never should, but defensively) it's a no-op. +""" + +from __future__ import annotations + +import json +import logging +import sqlite3 +import threading +import time +from typing import Any + +logger = logging.getLogger(__name__) + + +def _normalize_key(key: tuple) -> str: + """Deterministic string form for a cache key tuple. + + Sort dict members so callers don't have to be careful about key + order in arg tuples. Falls back to `repr()` for anything not JSON- + serializable (shouldn't happen in our paths, but defensive). + """ + try: + return json.dumps(list(key), sort_keys=True, default=str) + except Exception: + return repr(key) + + +class TTLCache: + """SQLite-backed TTL cache with an in-process write-through layer. + + Single-process semantics: the in-memory dict is the authoritative + fast path. sqlite gets a copy so the next process start can + re-warm. We don't try to coordinate across processes — there's only + one backend process. + + Thread/async safety: sqlite connections aren't shareable across + threads safely; we wrap reads/writes in a Lock. The hot path stays + short (single SELECT or INSERT OR REPLACE) so contention is + negligible for our request volume. + """ + + def __init__(self, conn: sqlite3.Connection) -> None: + self._conn = conn + self._lock = threading.Lock() + self._mem: dict[str, tuple[float, Any]] = {} + # On startup, warm the in-memory dict from disk so cache hits are + # immediate after a restart. Drops expired entries while we're at + # it. + self._warm_from_disk() + + def _warm_from_disk(self) -> None: + now = time.time() + with self._lock: + cur = self._conn.execute("SELECT key, value_json, expires_at FROM twitter_cache") + for k, v_json, exp in cur.fetchall(): + if exp > now: + try: + self._mem[k] = (exp, json.loads(v_json)) + except json.JSONDecodeError: + continue + # Garbage-collect expired rows so disk doesn't grow unbounded. + self._conn.execute("DELETE FROM twitter_cache WHERE expires_at <= ?", (now,)) + self._conn.commit() + logger.info("twitter cache: warmed %d entries from disk", len(self._mem)) + + def get(self, key: tuple) -> Any | None: + nk = _normalize_key(key) + now = time.time() + # Hot path: in-memory hit. + entry = self._mem.get(nk) + if entry is not None: + exp, val = entry + if exp > now: + return val + # Expired — drop it so we don't return again. + self._mem.pop(nk, None) + return None + + def set(self, key: tuple, value: Any, *, ttl: int) -> None: + if ttl <= 0: + return + nk = _normalize_key(key) + expires_at = time.time() + ttl + self._mem[nk] = (expires_at, value) + try: + payload = json.dumps(value, default=str) + except Exception as e: + # If we can't serialize, just keep the in-memory copy and + # skip the disk write. The cache will still satisfy reads + # this session. + logger.warning("cache: serializer failed for %s (%s); keeping memory-only", key, e) + return + with self._lock: + self._conn.execute( + "INSERT OR REPLACE INTO twitter_cache (key, value_json, expires_at) VALUES (?, ?, ?)", + (nk, payload, expires_at), + ) + self._conn.commit() + + def invalidate(self, key: tuple) -> None: + nk = _normalize_key(key) + self._mem.pop(nk, None) + with self._lock: + self._conn.execute("DELETE FROM twitter_cache WHERE key = ?", (nk,)) + self._conn.commit() + + def clear(self) -> None: + self._mem.clear() + with self._lock: + self._conn.execute("DELETE FROM twitter_cache") + self._conn.commit() + + def stats(self) -> dict: + with self._lock: + (n_disk,) = self._conn.execute("SELECT COUNT(*) FROM twitter_cache").fetchone() + return {"in_memory": len(self._mem), "on_disk": int(n_disk)} diff --git a/backend/apps/twitter/import_cookies.py b/backend/apps/twitter/import_cookies.py new file mode 100644 index 00000000..eaa1a9a6 --- /dev/null +++ b/backend/apps/twitter/import_cookies.py @@ -0,0 +1,196 @@ +"""One-shot CLI for importing browser-extracted cookies into the pool. + +When Cloudflare's bot detection blocks ``POST /accounts/login`` (the +underlying ``httpx`` TLS fingerprint gets flagged even after the UA + +``sec-ch-ua-*`` header patches in :mod:`._twikit_patches`), the operator +can still get a working session by: + +1. Logging in via a real browser on the same machine. +2. Copying ``auth_token`` and ``ct0`` from DevTools (Application → + Cookies → ``https://x.com``). +3. Running this script with those two values. + +The script writes a twikit-format cookies file (twikit's +``save_cookies`` just dumps ``dict(self.http.cookies)`` to JSON, so any +JSON dict of name->value works) and registers a matching account record +in ``accounts.json``. The next backend restart will pick it up via +``_hydrate_pool``; the operator can then confirm with +``POST /api/twitter/accounts/{id}/verify``, which calls the much softer +``client.user()`` endpoint that historically survives Cloudflare even +when the login POST does not. + +This is intentionally a CLI script and not a route: accepting raw +session cookies over an HTTP endpoint widens the attack surface (any +request that can authenticate to the backend could exfiltrate a session +into the pool), and the import use case is rare enough that paying the +"restart the backend" cost is fine. + +Usage +----- + +:: + + python -m backend.apps.twitter.import_cookies \\ + --auth-token AAA... \\ + --ct0 BBB... \\ + --label "personal" \\ + [--handle myname] \\ + [--role primary|read_only] \\ + [--id ] + +If ``--id`` is omitted, a fresh uuid4 is generated. If ``--id`` matches +an existing account, that account's cookies are overwritten in place +(re-login path) and its label/handle are updated if supplied. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from uuid import uuid4 + +from backend.apps.twitter import persistence +from backend.apps.twitter.models import TwitterAccount + + +def _build_cookie_dict(auth_token: str, ct0: str, extra: dict | None = None) -> dict: + """Minimum cookie set that x.com's GraphQL endpoints accept. + + ``auth_token`` is the session secret; ``ct0`` is the CSRF token X + cross-checks against the ``x-csrf-token`` header twikit sends on + every authenticated request. Anything else (``guest_id``, ``kdt``, + ``att``) is nice-to-have for fingerprint consistency but not + required for the API to authorize the call. + """ + cookies = {"auth_token": auth_token, "ct0": ct0} + if extra: + cookies.update(extra) + return cookies + + +def _upsert_account_record( + account_id: str, + label: str, + handle: str | None, + role: str, +) -> TwitterAccount: + """Find or create the matching record in ``accounts.json``. + + We treat this as an upsert: re-running the script with the same + ``--id`` (e.g. to refresh cookies after they expire) updates the + label/handle on the existing row instead of duplicating it. State + is forced back to ``active`` so the next hydrate doesn't skip the + account; the smoke probe / verify will downgrade it if the imported + cookies are stale. + """ + accounts_raw = persistence.load_accounts() + for i, raw in enumerate(accounts_raw): + if raw.get("id") == account_id: + record = TwitterAccount(**raw) + record.label = label or record.label + if handle: + record.handle = handle + record.role = role # type: ignore[assignment] + record.state = "active" + record.last_error = None + record.last_verified_at = 0.0 # let /verify stamp this + accounts_raw[i] = record.model_dump() + persistence.save_accounts(accounts_raw) + return record + + record = TwitterAccount( + id=account_id, + label=label or (handle or "imported"), + handle=handle, + role=role, # type: ignore[arg-type] + ) + record.state = "active" + record.last_verified_at = 0.0 + accounts_raw.append(record.model_dump()) + persistence.save_accounts(accounts_raw) + return record + + +def _write_cookies(account_id: str, cookies: dict) -> str: + """Drop cookies file at the canonical path with mode 0600. + + Atomic via tmp-then-rename so a crash mid-write can't leave the + pool hydrating from a half-written JSON file at next startup. + """ + persistence.ensure_dirs() + final = persistence.cookies_path(account_id) + tmp = final + ".tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump(cookies, f) + os.replace(tmp, final) + persistence.chmod_cookies(final) + return final + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description=( + "Import browser-extracted x.com cookies into the OpenSwarm " + "Twitter account pool. Use when /accounts/login is being " + "Cloudflare-blocked but the same account works in a real " + "browser on the same machine." + ), + ) + parser.add_argument("--auth-token", required=True, help="auth_token cookie value") + parser.add_argument("--ct0", required=True, help="ct0 cookie value (CSRF)") + parser.add_argument("--label", default="", help="human-readable label for the account") + parser.add_argument("--handle", default=None, help="screen name (optional; /verify will fill this in)") + parser.add_argument( + "--role", + default="primary", + choices=("primary", "read_only"), + help="account role; defaults to primary", + ) + parser.add_argument( + "--id", + default=None, + help="reuse an existing account id (re-import path); omit to mint a fresh uuid", + ) + args = parser.parse_args(argv) + + auth_token = args.auth_token.strip() + ct0 = args.ct0.strip() + if not auth_token or not ct0: + print("error: --auth-token and --ct0 must both be non-empty", file=sys.stderr) + return 2 + + account_id = args.id or uuid4().hex + cookies = _build_cookie_dict(auth_token, ct0) + + cookie_path = _write_cookies(account_id, cookies) + record = _upsert_account_record( + account_id=account_id, + label=args.label, + handle=args.handle, + role=args.role, + ) + + print(json.dumps({ + "ok": True, + "id": record.id, + "label": record.label, + "handle": record.handle, + "role": record.role, + "state": record.state, + "cookies_path": cookie_path, + "imported_at": time.time(), + "next_steps": [ + "Restart the backend so _hydrate_pool picks up the new account.", + f"Call POST /api/twitter/accounts/{record.id}/verify to confirm " + "the cookies are live (uses client.user(), which is much softer " + "than the login POST).", + ], + }, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/backend/apps/twitter/models.py b/backend/apps/twitter/models.py new file mode 100644 index 00000000..a56e9e67 --- /dev/null +++ b/backend/apps/twitter/models.py @@ -0,0 +1,147 @@ +"""Pydantic types for the Twitter SubApp. + +`TwitterAccount` is the public-facing record (serialized via API). +Credentials never live in this model — passwords are accepted in +`LoginRequest`, used once for `client.login()`, then discarded; cookies +live on disk only. +""" + +from __future__ import annotations + +import time +from typing import Literal, Optional +from uuid import uuid4 + +from pydantic import BaseModel, Field, model_validator + + +# The state machine documented in the plan ("Account lifecycle" section). +AccountState = Literal["active", "locked", "needs_relogin", "suspended"] +AccountRole = Literal["primary", "read_only"] + + +class TwitterAccount(BaseModel): + """Public account record. Mirrors what's stored in accounts.json + minus secrets (no password ever; cookies live in a separate file). + """ + + id: str = Field(default_factory=lambda: uuid4().hex) + label: str = "" + handle: Optional[str] = None # @screen_name, set after first login + role: AccountRole = "primary" + state: AccountState = "active" + trust_multiplier: float = 0.4 # see Rate-limit semantics in the plan + proxy: Optional[str] = None + created_at: float = Field(default_factory=time.time) + last_verified_at: float = 0.0 + last_error: Optional[str] = None + + +class LoginRequest(BaseModel): + """Inbound login payload. + + twikit's `client.login()` is flexible about which of auth_info_1 / + auth_info_2 is the username/email/phone — we just forward both. The + password is here and only here; it's used to call `client.login()` + and then dropped on the floor (never written to disk, never logged). + """ + + auth_info_1: str + auth_info_2: Optional[str] = None + password: str + totp_secret: Optional[str] = None + label: Optional[str] = None + role: AccountRole = "primary" + + +class TrustUpdateRequest(BaseModel): + """PATCH /accounts/{id} body — currently just trust_multiplier. + + Constrained to [0, 1]. Upper bound prevents a stray decimal from + multiplying the budget by 10 and tripping a wave of 429s; lower + bound of 0 (inclusive) lets the operator pause an account + in-place without deleting it ("dial trust to 0 while debugging + why this account is hitting locks"). Bucket math treats + capacity=0 as "never refill," so a paused account stays in the + pool, keeps its cookies, but loses its turn in `pick()`. + """ + + trust_multiplier: float + + @model_validator(mode="after") + def _check_range(self) -> "TrustUpdateRequest": + if not (0.0 <= self.trust_multiplier <= 1.0): + raise ValueError("trust_multiplier must be in [0, 1]") + return self + + +class BucketSnapshot(BaseModel): + """One bucket's current state, for /health.""" + + endpoint: str + capacity: int + tokens: float + locked_until: float + seconds_until_available: float + + +class AccountHealth(BaseModel): + """Output of GET /accounts/{id}/health. + + Served from in-memory pool state — does NOT call twikit, so health + polling can be aggressive without burning rate budget. + """ + + id: str + label: str + handle: Optional[str] + state: AccountState + role: AccountRole + trust_multiplier: float + last_verified_at: float + last_error: Optional[str] + recent_429_count: int = 0 + buckets: list[BucketSnapshot] = [] + + +# --- Tool request schemas. Pydantic enforces enums + bounds before +# --- anything hits twikit; saves us defensive checks in the route. + +TweetProduct = Literal["Top", "Latest", "Media"] +UserTweetType = Literal["Tweets", "Replies", "Media", "Likes"] + + +class SearchRequest(BaseModel): + q: str + product: TweetProduct = "Latest" + count: int = Field(default=20, ge=1, le=50) + cursor: Optional[str] = None + + +class UserLookupRequest(BaseModel): + """One of handle or user_id must be provided, not both.""" + + handle: Optional[str] = None + user_id: Optional[str] = None + + @model_validator(mode="after") + def _exactly_one(self) -> "UserLookupRequest": + if bool(self.handle) == bool(self.user_id): + raise ValueError("specify exactly one of: handle, user_id") + return self + + +class UserTweetsRequest(BaseModel): + user_id: str + type: UserTweetType = "Tweets" + count: int = Field(default=20, ge=1, le=50) + cursor: Optional[str] = None + + +class TweetLookupRequest(BaseModel): + tweet_id: str + + +class TweetRepliesRequest(BaseModel): + tweet_id: str + cursor: Optional[str] = None diff --git a/backend/apps/twitter/persistence.py b/backend/apps/twitter/persistence.py new file mode 100644 index 00000000..d0949e67 --- /dev/null +++ b/backend/apps/twitter/persistence.py @@ -0,0 +1,288 @@ +"""On-disk state for the Twitter SubApp. + +Everything durable lives under `DATA_ROOT/twitter/`: + + twitter/ + accounts.json — account index (id, label, role, state, trust) + cookies/.json — twikit cookie jar per account, mode 0600 + state.sqlite — buckets, response cache, 429 audit log + +`accounts.json` is a small, human-editable file. The cookies dir is the +sensitive bit — we make the dir mode 0700 and each file mode 0600 (twikit +itself writes the JSON, so we chmod after). The sqlite file holds the +high-churn state that benefits from indexed lookups and atomic +transactions. + +This module knows nothing about twikit — it speaks only in plain dicts +and Bucket snapshots. The pool layer composes these primitives with +twikit.Client instances. +""" + +from __future__ import annotations + +import json +import logging +import os +import sqlite3 +import time + +from backend.config.paths import TWITTER_DIR + +logger = logging.getLogger(__name__) + +ACCOUNTS_PATH = os.path.join(TWITTER_DIR, "accounts.json") +COOKIES_DIR = os.path.join(TWITTER_DIR, "cookies") +STATE_DB_PATH = os.path.join(TWITTER_DIR, "state.sqlite") + + +# --------------------------------------------------------------------------- +# Filesystem setup +# --------------------------------------------------------------------------- + +def ensure_dirs() -> None: + """Create the twitter data dirs with restrictive permissions. + + Cookies are auth material — equivalent to bearer tokens for the + user's X session. The dir is mode 0700 so other local users on a + multi-user macOS box can't read them. (On single-user laptops this + is belt-and-suspenders, but free.) + """ + os.makedirs(TWITTER_DIR, exist_ok=True) + os.makedirs(COOKIES_DIR, exist_ok=True) + try: + os.chmod(COOKIES_DIR, 0o700) + except OSError as e: + logger.warning("twitter: chmod 0700 on cookies dir failed: %s", e) + + +def cookies_path(account_id: str) -> str: + return os.path.join(COOKIES_DIR, f"{account_id}.json") + + +def chmod_cookies(path: str) -> None: + """Lock down a cookies file to mode 0600 (twikit writes it 0644).""" + try: + os.chmod(path, 0o600) + except OSError as e: + logger.warning("twitter: chmod 0600 on %s failed: %s", path, e) + + +def delete_cookies(account_id: str) -> None: + path = cookies_path(account_id) + if os.path.exists(path): + try: + os.remove(path) + except OSError as e: + logger.warning("twitter: rm %s failed: %s", path, e) + + +# --------------------------------------------------------------------------- +# accounts.json (small enough to load whole on every read) +# --------------------------------------------------------------------------- + +def load_accounts() -> list[dict]: + """Return the list of account records, empty list on first run.""" + if not os.path.isfile(ACCOUNTS_PATH): + return [] + try: + with open(ACCOUNTS_PATH) as f: + data = json.load(f) + if isinstance(data, list): + return data + logger.warning("twitter: accounts.json wasn't a list, ignoring") + return [] + except (OSError, json.JSONDecodeError) as e: + logger.exception("twitter: accounts.json read failed: %s", e) + return [] + + +def save_accounts(accounts: list[dict]) -> None: + """Atomic write so a crash mid-rename can't leave us with an empty file.""" + ensure_dirs() + tmp = ACCOUNTS_PATH + ".tmp" + with open(tmp, "w") as f: + json.dump(accounts, f, indent=2, default=str) + os.replace(tmp, ACCOUNTS_PATH) + + +# --------------------------------------------------------------------------- +# state.sqlite — buckets, cache, audit +# --------------------------------------------------------------------------- + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS twitter_buckets ( + account_id TEXT NOT NULL, + endpoint TEXT NOT NULL, + capacity INTEGER NOT NULL, + tokens REAL NOT NULL, + locked_until REAL NOT NULL DEFAULT 0, + updated_at REAL NOT NULL, + PRIMARY KEY (account_id, endpoint) +); + +CREATE TABLE IF NOT EXISTS twitter_cache ( + key TEXT PRIMARY KEY, + value_json TEXT NOT NULL, + expires_at REAL NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_twitter_cache_expires_at + ON twitter_cache (expires_at); + +CREATE TABLE IF NOT EXISTS twitter_audit ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + ts REAL NOT NULL, + account_id TEXT NOT NULL, + endpoint TEXT NOT NULL, + event TEXT NOT NULL, + detail TEXT +); + +CREATE INDEX IF NOT EXISTS idx_twitter_audit_ts + ON twitter_audit (ts); +""" + + +def open_state_db() -> sqlite3.Connection: + """Open the state sqlite, applying schema migrations if needed. + + WAL mode means concurrent readers don't block the snapshot writer, + which we want because the route handlers read from `twitter_cache` + on every call while the lifespan task is also writing bucket + snapshots. + + Why `execute(...).fetchone()` for the PRAGMAs instead of + `executescript`: `PRAGMA journal_mode=WAL` is a query that *returns + a row* (the new mode), and `executescript` ignores result rows. + On some sqlite builds that's enough for the pragma to silently + not apply — using execute() + fetchone() forces the driver to + actually run it and consume the result. + + `check_same_thread=False` is kept because the backend currently + only touches sqlite from the asyncio event loop (single thread), + but `TTLCache` defensively uses a `threading.Lock` in case + something later offloads to a threadpool. If you add a sync + sqlite call from a worker thread, share that lock. + """ + ensure_dirs() + conn = sqlite3.connect(STATE_DB_PATH, check_same_thread=False) + conn.execute("PRAGMA journal_mode=WAL").fetchone() + conn.execute("PRAGMA synchronous=NORMAL").fetchone() + conn.executescript(_SCHEMA) + conn.commit() + try: + os.chmod(STATE_DB_PATH, 0o600) + except OSError: + pass + return conn + + +# --------------------------------------------------------------------------- +# Bucket snapshot read/write +# --------------------------------------------------------------------------- + +def load_buckets(conn: sqlite3.Connection, account_id: str) -> dict[str, dict]: + """Return {endpoint: snapshot_dict} for one account. + + Snapshot dict matches `Bucket.snapshot()` keys (capacity, tokens, + locked_until) so the pool can call `Bucket.restore(snap)` directly. + """ + cur = conn.execute( + "SELECT endpoint, capacity, tokens, locked_until " + "FROM twitter_buckets WHERE account_id = ?", + (account_id,), + ) + out: dict[str, dict] = {} + for endpoint, capacity, tokens, locked_until in cur.fetchall(): + out[endpoint] = { + "capacity": int(capacity), + "tokens": float(tokens), + "locked_until": float(locked_until), + } + return out + + +def save_bucket( + conn: sqlite3.Connection, + account_id: str, + endpoint: str, + snapshot: dict, +) -> None: + """Upsert one bucket. Called by the periodic snapshot loop, ~1Hz.""" + conn.execute( + """ + INSERT INTO twitter_buckets + (account_id, endpoint, capacity, tokens, locked_until, updated_at) + VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT(account_id, endpoint) DO UPDATE SET + capacity = excluded.capacity, + tokens = excluded.tokens, + locked_until = excluded.locked_until, + updated_at = excluded.updated_at + """, + ( + account_id, + endpoint, + int(snapshot.get("capacity", 1)), + float(snapshot.get("tokens", 0.0)), + float(snapshot.get("locked_until", 0.0)), + time.time(), + ), + ) + + +def delete_buckets_for(conn: sqlite3.Connection, account_id: str) -> None: + """Drop all bucket rows for an account (called on account delete).""" + conn.execute("DELETE FROM twitter_buckets WHERE account_id = ?", (account_id,)) + conn.commit() + + +# --------------------------------------------------------------------------- +# Audit log +# --------------------------------------------------------------------------- + +def audit( + conn: sqlite3.Connection, + account_id: str, + endpoint: str, + event: str, + detail: str | None = None, +) -> None: + """Append a row to the audit log. Cheap (no commit per-row in WAL). + + Events worth logging: `429`, `locked`, `suspended`, `login_ok`, + `login_fail`, `verify_ok`, `verify_fail`, `relogin`, `delete`. + The route layer keys off these when computing /health summaries. + """ + conn.execute( + "INSERT INTO twitter_audit (ts, account_id, endpoint, event, detail) " + "VALUES (?, ?, ?, ?, ?)", + (time.time(), account_id, endpoint, event, detail), + ) + + +def recent_429s(conn: sqlite3.Connection, account_id: str, since_s: float) -> int: + """How many 429 events did this account hit since `since_s` seconds ago? + + Surfaced via the /health endpoint so the operator can see whether + the trust_multiplier needs to come down. + """ + cutoff = time.time() - since_s + row = conn.execute( + "SELECT COUNT(*) FROM twitter_audit " + "WHERE account_id = ? AND event = '429' AND ts >= ?", + (account_id, cutoff), + ).fetchone() + return int(row[0]) + + +def trim_audit(conn: sqlite3.Connection, keep_days: int = 30) -> None: + """Drop audit rows older than `keep_days`. Called on lifespan startup. + + Audit data is purely for human inspection / health UI — nobody + depends on the full history. Keep it bounded so the sqlite file + doesn't drift toward "the size of the disk." + """ + cutoff = time.time() - keep_days * 86400 + conn.execute("DELETE FROM twitter_audit WHERE ts < ?", (cutoff,)) + conn.commit() diff --git a/backend/apps/twitter/pool.py b/backend/apps/twitter/pool.py new file mode 100644 index 00000000..bdd8b390 --- /dev/null +++ b/backend/apps/twitter/pool.py @@ -0,0 +1,397 @@ +"""ManagedAccount + AccountPool: the live state of each logged-in account. + +A `ManagedAccount` wraps a single `twikit.Client` together with: +- per-endpoint `Bucket` instances (rate budget) +- an `asyncio.Semaphore(1)` so we never run two twikit calls in parallel + on the same Client (cookies would race + X would flag the pattern) +- the mutable lifecycle state (`active` / `locked` / `needs_relogin` / + `suspended`) and the `trust_multiplier` knob + +The `AccountPool` is the registry that the routes go through. It +exposes: + +- `pick(endpoint)` — returns the active account whose bucket frees up + soonest. Returns `None` if no active accounts exist. +- `add(account, client)` — runtime hand-off from the login route. Sets + up buckets, kicks off persistence wiring. +- `remove(account_id)` — waits for any in-flight call (under the + semaphore, with a hard cap of `REMOVE_TIMEOUT_S`) then disposes the + Client. Returns False if the timeout fired (caller decides whether + to force or 503). +- `mark_locked` / `mark_suspended` / `mark_needs_relogin` / `record_429` + — side-effect hooks called by `RateGate` on twikit errors. + +Why module-level globals: this matches the rest of OpenSwarm's SubApp +pattern (`mcp_registry` uses module-level `_cache`, `_refresh_task`, +etc.). Routes import the pool object directly. Single-process backend, +so no cross-process coordination needed. +""" + +from __future__ import annotations + +import asyncio +import logging +import sqlite3 +import time +from dataclasses import dataclass, field +from typing import Optional + +from backend.apps.twitter import persistence +from backend.apps.twitter.models import TwitterAccount +from backend.apps.twitter.ratelimit import ( + DEFAULT_BUDGETS, + Bucket, +) + +logger = logging.getLogger(__name__) + +# How long `remove()` waits for an in-flight call before giving up and +# letting the route return 503. The semaphore is held only for the +# duration of a single twikit GraphQL call (a few seconds upper bound +# in normal operation); 5s is generous enough that we never preempt +# legitimate work and short enough that a deleted account doesn't +# stall the UI. +REMOVE_TIMEOUT_S: float = 5.0 + +# Same idea for the re-login path: when /accounts/login swaps a fresh +# twikit.Client onto an existing ManagedAccount, we want any in-flight +# tool call on the old Client to finish first. Otherwise the call sees +# its cookies replaced mid-request. Same budget as REMOVE_TIMEOUT_S +# (one twikit call's worth). +REPLACE_CLIENT_TIMEOUT_S: float = 5.0 + + +@dataclass +class ManagedAccount: + """One live account in the pool. + + `record` is the public-facing `TwitterAccount` we serialize via the + API. Mutable lifecycle state lives directly on `record.state`. The + Client + semaphore + buckets are runtime-only and never leave the + pool. + """ + + record: TwitterAccount + client: object # twikit.client.client.Client — duck-typed so tests don't need twikit + concurrency: asyncio.Semaphore = field(default_factory=lambda: asyncio.Semaphore(1)) + _buckets: dict[str, Bucket] = field(default_factory=dict) + + # ---- record passthrough conveniences ------------------------------ + @property + def id(self) -> str: + return self.record.id + + @property + def state(self) -> str: + return self.record.state + + @property + def role(self) -> str: + return self.record.role + + @property + def trust_multiplier(self) -> float: + return self.record.trust_multiplier + + # ---- buckets ------------------------------------------------------- + + def bucket(self, endpoint: str) -> Bucket: + """Lazy-create the bucket for this endpoint at trust-scaled capacity. + + Lazy creation matters because we don't pre-allocate every + endpoint at startup — a long-lived process might only ever + touch `search_tweet`, and creating buckets we don't use is + wasted state. + + Capacity floors to 0 (paused) rather than 1; the Bucket math + handles capacity=0 by returning a long sentinel wait so pick() + deprioritizes paused accounts but doesn't crash. + """ + if endpoint not in self._buckets: + base = DEFAULT_BUDGETS.get(endpoint, 20) + cap = max(0, int(round(base * self.trust_multiplier))) + self._buckets[endpoint] = Bucket(capacity=cap) + return self._buckets[endpoint] + + def restore_buckets(self, snapshots: dict[str, dict]) -> None: + """On startup, replace each Bucket with one restored from disk. + + Crash-safe: `Bucket.restore` clamps tokens to `capacity / 2` + regardless of what was on disk, so an unclean shutdown can't + leak budget into a post-restart burst. + """ + for endpoint, snap in snapshots.items(): + try: + self._buckets[endpoint] = Bucket.restore(snap) + except Exception as e: + logger.warning( + "twitter: bucket restore failed for %s/%s: %s", + self.id, + endpoint, + e, + ) + + def rescale_buckets(self) -> None: + """Re-apply `trust_multiplier` after a PATCH. + + Raising the multiplier widens the ceiling; existing `tokens` + carry over so the bucket starts under the new cap and refills + normally. Lowering the multiplier clamps `tokens` down to the + new ceiling so we can't burst on the next call. We do not + invent tokens out of thin air on a raise — refill is what + replenishes within the 15-min window. + + A `trust_multiplier` of 0 effectively pauses the account: the + capacity gets floored to 0 and `pick()` will skip endpoints + whose `time_until_available()` is huge (the bucket never + refills past zero). + """ + for endpoint, bucket in list(self._buckets.items()): + base = DEFAULT_BUDGETS.get(endpoint, 20) + new_cap = max(0, int(round(base * self.trust_multiplier))) + if new_cap == bucket.capacity: + continue + bucket.capacity = new_cap + if bucket.tokens > new_cap: + bucket.tokens = float(new_cap) + + +class AccountPool: + """Registry of `ManagedAccount` instances and pick/lifecycle ops.""" + + def __init__(self, conn: sqlite3.Connection) -> None: + self._conn = conn + self._accounts: dict[str, ManagedAccount] = {} + # Tracks remove-in-progress account IDs so concurrent pick()s + # don't return an account that's about to be disposed. + self._removing: set[str] = set() + self._lock = asyncio.Lock() + + # ---- introspection ------------------------------------------------- + + @property + def accounts(self) -> list[ManagedAccount]: + return list(self._accounts.values()) + + def get(self, account_id: str) -> Optional[ManagedAccount]: + return self._accounts.get(account_id) + + def by_handle(self, handle: str) -> Optional[ManagedAccount]: + """Find an account by @screen_name. Used by /login for re-login + on a stuck account (locked/needs_relogin) — same handle = same + ManagedAccount, we just refresh its cookies in place rather + than create a new entry.""" + h = handle.lstrip("@").lower() + for acct in self._accounts.values(): + if acct.record.handle and acct.record.handle.lower() == h: + return acct + return None + + # ---- picking ------------------------------------------------------- + + async def pick(self, endpoint: str) -> Optional[ManagedAccount]: + """Choose the active account with the soonest availability. + + Skips: + - accounts being removed (we'd race the disposal) + - any state != "active" + Among active accounts, breaks ties by `time_until_available()` + on this endpoint's bucket. With one account this is just "is + it active and not being removed." + """ + candidates = [ + a for a in self._accounts.values() + if a.id not in self._removing and a.state == "active" + ] + if not candidates: + return None + candidates.sort(key=lambda a: a.bucket(endpoint).time_until_available()) + return candidates[0] + + # ---- mutation ------------------------------------------------------ + + async def add(self, account: TwitterAccount, client: object) -> ManagedAccount: + """Register a freshly-logged-in account. Idempotent for re-login. + + If `account.id` is already in the pool, we replace the Client + in place. This is the re-login-after-lock path: the login route + matches by handle, finds the existing record, calls `login()` + on the existing Client (or a fresh one), then hands us the + result here. We keep the same `ManagedAccount` so the buckets' + runtime state survives a re-login (no reason to discard + partially-used budget). + + Re-login waits for any in-flight twikit call (under the + account's concurrency semaphore) before swapping the Client. + Without this, a request that was halfway through GraphQL would + come back to a fresh Client whose cookies don't match the one + it started on — symptom is mysterious deserialization errors + or session drift. Times out at REPLACE_CLIENT_TIMEOUT_S; if + the in-flight call is stuck, we swap anyway (the alternative + is stranding the re-login indefinitely). + """ + async with self._lock: + existing = self._accounts.get(account.id) + + if existing is not None: + # Drain the in-flight call BEFORE swapping the Client. + try: + await asyncio.wait_for( + existing.concurrency.acquire(), + timeout=REPLACE_CLIENT_TIMEOUT_S, + ) + try: + existing.record = account + existing.client = client + snaps = persistence.load_buckets(self._conn, account.id) + if snaps: + existing.restore_buckets(snaps) + finally: + existing.concurrency.release() + except asyncio.TimeoutError: + logger.warning( + "twitter: re-login swap timed out for %s after %.1fs; " + "swapping Client without semaphore (in-flight call " + "may see cookie drift)", + account.id, + REPLACE_CLIENT_TIMEOUT_S, + ) + existing.record = account + existing.client = client + snaps = persistence.load_buckets(self._conn, account.id) + if snaps: + existing.restore_buckets(snaps) + return existing + + async with self._lock: + # Race window check: someone else might have added in + # between our two lock acquisitions. Cheap to re-check. + already = self._accounts.get(account.id) + if already is not None: + return already + managed = ManagedAccount(record=account, client=client) + snaps = persistence.load_buckets(self._conn, account.id) + if snaps: + managed.restore_buckets(snaps) + self._accounts[account.id] = managed + return managed + + async def remove(self, account_id: str, *, wipe_buckets: bool = True) -> bool: + """Remove an account, awaiting any in-flight call. + + Returns True on clean removal, False if `REMOVE_TIMEOUT_S` fired + while waiting for the semaphore (in which case we evict from + the dict anyway and the in-flight call will see `CancelledError` + when it next yields). + + `wipe_buckets=True` is the normal path (DELETE endpoint); set + False if a caller wants to keep the budget state around (e.g. + replacing a logged-out account with a re-login). + """ + async with self._lock: + managed = self._accounts.pop(account_id, None) + if managed is None: + return True + self._removing.add(account_id) + + try: + # Acquire the semaphore so we know no twikit call is mid- + # flight before we let go of the Client. The shield + wait_for + # combo means: if the in-flight call wraps up within + # REMOVE_TIMEOUT_S, we get a clean acquire; if not, we + # bail and let the call error out on its own. + try: + await asyncio.wait_for( + managed.concurrency.acquire(), + timeout=REMOVE_TIMEOUT_S, + ) + managed.concurrency.release() + clean = True + except asyncio.TimeoutError: + logger.warning( + "twitter: remove(%s) timed out waiting for in-flight call", + account_id, + ) + clean = False + + if wipe_buckets: + persistence.delete_buckets_for(self._conn, account_id) + return clean + finally: + self._removing.discard(account_id) + + # ---- audit / commit (keep the routes out of pool._conn) ------------ + + def audit_lifecycle(self, account_id: str, event: str, detail: str | None = None) -> None: + """Write a `_lifecycle`-endpoint audit row (login_ok, delete, ...). + + Routes used to call `persistence.audit(pool._conn, ...)` directly, + which leaked the sqlite connection through. This is the same + thing with the endpoint fixed and the private kept private. + """ + persistence.audit(self._conn, account_id, "_lifecycle", event, detail) + + def commit(self) -> None: + """Flush the underlying audit/state writes. + + The snapshot loop commits ~1Hz, so most callers don't need + this — but the login/delete/verify routes do, because they + return to the user immediately and we want their audit rows + durable before the response goes out. + """ + self._conn.commit() + + def recent_429s(self, account_id: str, since_s: float) -> int: + """Count 429 audit rows in the last `since_s` seconds. Used by /health.""" + return persistence.recent_429s(self._conn, account_id, since_s=since_s) + + # ---- side-effect hooks (called by RateGate on twikit errors) ------- + + def record_429(self, account_id: str, endpoint: str) -> None: + persistence.audit(self._conn, account_id, endpoint, "429") + # No commit — the snapshot loop commits ~1Hz and audit writes + # ride along. If we crash before the commit we lose a few + # audit rows; nobody depends on those for correctness. + + def mark_locked(self, account_id: str, reason: str) -> None: + acct = self._accounts.get(account_id) + if acct is not None: + acct.record.state = "locked" + acct.record.last_error = reason + persistence.audit(self._conn, account_id, "_lifecycle", "locked", reason) + + def mark_suspended(self, account_id: str, reason: str) -> None: + acct = self._accounts.get(account_id) + if acct is not None: + acct.record.state = "suspended" + acct.record.last_error = reason + persistence.audit(self._conn, account_id, "_lifecycle", "suspended", reason) + + def mark_needs_relogin(self, account_id: str, reason: str) -> None: + acct = self._accounts.get(account_id) + if acct is not None: + acct.record.state = "needs_relogin" + acct.record.last_error = reason + persistence.audit(self._conn, account_id, "_lifecycle", "needs_relogin", reason) + + def mark_active(self, account_id: str) -> None: + """Called after a successful login or verify.""" + acct = self._accounts.get(account_id) + if acct is not None: + acct.record.state = "active" + acct.record.last_error = None + acct.record.last_verified_at = time.time() + + # ---- snapshotting -------------------------------------------------- + + def snapshot_all(self) -> None: + """Write every bucket's current state to sqlite. Idempotent. + + Called by the lifespan's periodic snapshot task (~1Hz) and on + shutdown. Cheap because sqlite is local + WAL mode; the heavy + lifting is just the UPSERTs. + """ + for managed in self._accounts.values(): + for endpoint, bucket in managed._buckets.items(): + persistence.save_bucket(self._conn, managed.id, endpoint, bucket.snapshot()) + self._conn.commit() diff --git a/backend/apps/twitter/ratelimit.py b/backend/apps/twitter/ratelimit.py new file mode 100644 index 00000000..b3cbd8de --- /dev/null +++ b/backend/apps/twitter/ratelimit.py @@ -0,0 +1,358 @@ +"""Token-bucket rate limiter for twikit calls. + +Twitter's internal GraphQL API rate-limits per-(account, endpoint) in a +rolling 15-minute window. Twikit doesn't track this proactively — it just +raises `twikit.errors.TooManyRequests` after the fact, with the server's +`x-rate-limit-reset` epoch parsed into `exc.rate_limit_reset`. That's +fine as a safety net, but for a multi-agent setup we want to *avoid* +ever hitting 429 in the first place: every 429 is a heuristic flag on +the user's account, and accumulating them is how accounts get locked. + +This module provides: + +- `Bucket` — a token bucket that refills continuously over a 15-min + window, with an extra `locked_until` field driven reactively by 429 + responses (when X tells us the bucket reset time, we trust the server + over our local accounting). +- `DEFAULT_BUDGETS` — community-observed per-endpoint ceilings. The + effective per-bucket capacity is `DEFAULT_BUDGETS[endpoint] * + account.trust_multiplier`, where the multiplier defaults to 0.4 for + primary accounts (keep way clear of the heuristic threshold) and is a + writable field the operator can ratchet up after observing no 429s. +- `RateGate` — orchestrator that combines a `TTLCache`, an `AccountPool`, + and per-bucket `acquire()` into a single async call that either returns + the serialized result, returns a structured "rate_limited" with + `retry_after_s`, or surfaces an error category for the route to handle. + +Concurrency: every code path that mutates `tokens` / `locked_until` runs +on the asyncio event loop, single-threaded, never awaits between a read +and a write. No lock needed. +""" + +from __future__ import annotations + +import asyncio +import logging +import math +import random +import time +from dataclasses import dataclass, field +from typing import Awaitable, Callable, Literal, TypeVar + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + +# Community-observed per-endpoint ceilings (calls per 15-minute window). +# Treat these as upper bounds; the effective cap is multiplied by the +# account's `trust_multiplier` (default 0.4 for primary). +# +# Source: community measurements on twscrape / twikit repos; X doesn't +# publish these. They're approximate and drift over time. The lifespan's +# smoke probe logs loudly if any single endpoint trips a 429 within the +# first hour after startup — that's the signal to bump caps down. +DEFAULT_BUDGETS: dict[str, int] = { + "search_tweet": 50, + "get_user_by_screen_name": 95, + "get_user_by_id": 95, + "get_user_tweets": 50, + "get_tweet_by_id": 150, # also covers tweet replies (same endpoint) + # Internal/admin endpoints. /verify and the startup smoke probe + # share this bucket. We used to set this to 10, which at + # trust_multiplier=0.4 left only 4 tokens/15min — meaning a single + # dev hot-reload cycle ate 25% of the verify window. X doesn't + # appear to rate-limit `client.user()` tightly so a roomier budget + # is safe. + "_self_user": 100, +} + +WINDOW_S: float = 15 * 60.0 + + +@dataclass +class Bucket: + """Continuous-refill token bucket for one (account, endpoint) pair. + + - `capacity` is the cap **after** trust_multiplier scaling. Callers + pass the already-scaled value; this class doesn't know about + multipliers. + - `tokens` is a float in `[0, capacity]`. It refills at + `capacity / WINDOW_S` per second. Restoring from a snapshot starts + lower (see persistence) to make crash-burst impossible. + - `locked_until` is a wall-clock UNIX timestamp set when twikit + raises `TooManyRequests`. While `now < locked_until`, no calls + go through regardless of `tokens` — we trust the server's reset + hint over our local clock. + """ + + capacity: int + tokens: float = field(init=False) + last_refill: float = field(default_factory=time.monotonic) + locked_until: float = 0.0 + + def __post_init__(self) -> None: + self.tokens = float(self.capacity) + + def _refill(self) -> None: + """Add tokens proportional to elapsed monotonic time.""" + now = time.monotonic() + elapsed = now - self.last_refill + if elapsed <= 0: + return + if self.capacity <= 0: + # Paused bucket. No refill, no division — just mark time + # and return so future calls don't compute against stale + # last_refill. + self.last_refill = now + return + rate = self.capacity / WINDOW_S + self.tokens = min(float(self.capacity), self.tokens + elapsed * rate) + self.last_refill = now + + def time_until_available(self) -> float: + """Seconds the caller would block in `acquire()` right now. + + Returns 0.0 if a call could go through immediately. Useful for + the route layer to decide between blocking (short wait) and + returning a 429 with `retry_after_s` (long wait). + + A capacity of 0 means the operator has paused this account + (trust_multiplier=0). The bucket never refills past zero, so + we return a long-but-finite sentinel: pick() ranks by this + value, so any other account will outrank a paused one. + """ + self._refill() + wall_wait = max(0.0, self.locked_until - time.time()) + token_wait = 0.0 + if self.tokens < 1.0: + if self.capacity <= 0: + # 24 hours — far above any block_ceiling_s we'd set, + # so RateGate will return rate_limited rather than + # waiting. No division by zero. + token_wait = 86400.0 + else: + deficit = 1.0 - self.tokens + token_wait = deficit * WINDOW_S / self.capacity + return max(token_wait, wall_wait) + + async def acquire(self) -> None: + """Block until one token is available, then consume it. + + Loops in case multiple coroutines race the same bucket — each + wakes from sleep, re-checks `time_until_available`, and decides + to sleep again or take the token. The +50–250ms jitter prevents + synchronized thundering-herd patterns that look bot-like to X. + """ + while True: + wait = self.time_until_available() + if wait <= 0: + self.tokens = max(0.0, self.tokens - 1.0) + return + await asyncio.sleep(wait + random.uniform(0.05, 0.25)) + + def mark_rate_limited(self, reset_at: float | None) -> None: + """React to a 429 from twikit by trusting the server's reset. + + `reset_at` is an absolute UNIX timestamp (the contents of + `x-rate-limit-reset`). We add a few seconds of jitter to keep + post-cooldown traffic from hitting the wall in lockstep with + other clients on the same account. + """ + now = time.time() + if reset_at is None or reset_at <= now: + # Server didn't tell us a reset time (or it's in the past). + # Cool down for one full window to be safe. + reset_at = now + WINDOW_S + self.locked_until = reset_at + random.uniform(1.0, 5.0) + self.tokens = 0.0 + + def snapshot(self) -> dict: + """Serializable representation for persistence.""" + self._refill() + return { + "capacity": self.capacity, + "tokens": self.tokens, + "locked_until": self.locked_until, + } + + @classmethod + def restore(cls, snap: dict) -> "Bucket": + """Reconstruct from `snapshot()` with crash-safe headroom. + + On startup we deliberately restore `tokens` at `min(saved, + capacity / 2)` so an unclean shutdown (where the last 1s of + decrements weren't snapshotted) can't lead to a post-restart + burst. Costs at most half a window's headroom in exchange for + a hard guarantee against accidental 429-storms. + + Capacity 0 (paused account) round-trips intact. + """ + capacity = max(0, int(snap.get("capacity", 1))) + b = cls(capacity=capacity) + saved_tokens = float(snap.get("tokens", capacity)) + b.tokens = max(0.0, min(saved_tokens, capacity / 2.0)) + b.locked_until = float(snap.get("locked_until", 0.0)) + b.last_refill = time.monotonic() + return b + + +# RateGate result categories. The route layer dispatches on this rather +# than catching exceptions for control flow. +GateOutcome = Literal["ok", "rate_limited", "locked", "needs_relogin", "suspended", "no_account", "error"] + + +@dataclass +class GateResult: + outcome: GateOutcome + value: object = None # ok: serialized payload; rate_limited: {retry_after_s}; error: error message + account_id: str | None = None + + +class RateGate: + """Per-request orchestration: cache -> pool -> bucket -> twikit. + + The class is deliberately tiny — it owns no state of its own beyond + references to the pool and the cache, so a route handler can do: + + result = await gate.execute( + endpoint="search_tweet", + cache_key=("search", q, product, count, cursor), + cache_ttl=60, + op=lambda client: client.search_tweet(q, product, count, cursor), + serializer=serialize_tweet_result, + ) + + and never has to think about buckets, semaphores, or 429s. Errors + from twikit are mapped to `GateResult.outcome` so the route returns + structured JSON instead of raising HTTPException out of business + code. + + `block_ceiling_s` is the longest the gate will wait inside a + `bucket.acquire()` before returning a 429 to the caller. The MCP + shim translates that 429 into a polite "retry in N seconds" + response so the LLM backs off — don't make this larger than ~10s + or the agent will think the tool is hung and start spawning + parallel calls. + """ + + def __init__(self, pool, cache, *, block_ceiling_s: float = 10.0) -> None: + self.pool = pool + self.cache = cache + self.block_ceiling_s = block_ceiling_s + + async def execute( + self, + *, + endpoint: str, + op: Callable[[object], Awaitable[T]], + serializer: Callable[[T], object], + cache_key: tuple | None = None, + cache_ttl: int = 0, + skip_cache: bool = False, + ) -> GateResult: + # 1. Cache check. + if cache_key is not None and cache_ttl > 0 and not skip_cache: + cached = self.cache.get(cache_key) + if cached is not None: + return GateResult(outcome="ok", value=cached) + + # 2. Pick an account. + account = await self.pool.pick(endpoint) + if account is None: + return GateResult(outcome="no_account", value={"error": "No active Twitter account available"}) + + bucket = account.bucket(endpoint) + + # 3. Decide whether to wait or to bounce. The wait check uses + # the *current* bucket state; the queue ahead of us on the + # semaphore can make the real wait longer. That's fine — we'll + # re-check inside `bucket.acquire()` and bail out (via the loop + # in Bucket.acquire) if the bucket ends up locked while we + # waited. + wait = bucket.time_until_available() + if wait > self.block_ceiling_s: + return GateResult( + outcome="rate_limited", + value={"retry_after_s": math.ceil(wait), "endpoint": endpoint}, + account_id=account.id, + ) + + # 4. Serialize twikit calls on this account so cookies don't race. + # Import lazily to avoid hard dependency at module-import time + # (handy for unit-testing Bucket/RateGate against fake clients). + from twikit.errors import ( + AccountLocked, + AccountSuspended, + TooManyRequests, + Unauthorized, + ) + + async with account.concurrency: + # Acquire the bucket *inside* the semaphore. If we did it + # outside, two concurrent requests to the same account + # could each find tokens>=1, both decrement, and both + # serialize on the semaphore — over-consuming by one token + # per cycle. Holding the semaphore makes the acquire + # atomic per account. + await bucket.acquire() + try: + raw = await op(account.client) + except TooManyRequests as e: + bucket.mark_rate_limited(getattr(e, "rate_limit_reset", None)) + # Record this — `account_state` audit lets the operator + # tell whether the trust_multiplier needs to come down. + self.pool.record_429(account.id, endpoint) + return GateResult( + outcome="rate_limited", + value={ + "retry_after_s": math.ceil(bucket.time_until_available()), + "endpoint": endpoint, + }, + account_id=account.id, + ) + except AccountLocked as e: + self.pool.mark_locked(account.id, str(e)) + return GateResult( + outcome="locked", + value={"error": "Account is locked (Arkose challenge). Solve in browser and re-login."}, + account_id=account.id, + ) + except AccountSuspended as e: + self.pool.mark_suspended(account.id, str(e)) + return GateResult( + outcome="suspended", + value={"error": "Account has been suspended by X."}, + account_id=account.id, + ) + except Unauthorized as e: + self.pool.mark_needs_relogin(account.id, str(e)) + return GateResult( + outcome="needs_relogin", + value={"error": "Session expired. Re-login required."}, + account_id=account.id, + ) + except asyncio.CancelledError: + # Don't swallow cancellation — the route was aborted + # (client disconnected) or the account is being deleted + # (pool.remove() is mid-flight). Propagate so FastAPI + # cleans up. + raise + except Exception as e: + logger.exception("Twikit call failed: endpoint=%s account=%s", endpoint, account.id) + return GateResult(outcome="error", value={"error": f"{type(e).__name__}: {e}"}, account_id=account.id) + + # 5. Serialize + cache. + try: + serialized = serializer(raw) + except Exception as e: + logger.exception("Serializer failed: endpoint=%s", endpoint) + return GateResult(outcome="error", value={"error": f"serialize failed: {e}"}, account_id=account.id) + + # Never cache errors; we got here on success. + if cache_key is not None and cache_ttl > 0: + try: + self.cache.set(cache_key, serialized, ttl=cache_ttl) + except Exception as e: + logger.warning("Cache set failed (non-fatal): %s", e) + + return GateResult(outcome="ok", value=serialized, account_id=account.id) diff --git a/backend/apps/twitter/serializers.py b/backend/apps/twitter/serializers.py new file mode 100644 index 00000000..f338e2e0 --- /dev/null +++ b/backend/apps/twitter/serializers.py @@ -0,0 +1,182 @@ +"""twikit object -> plain dict serializers. + +Strict whitelist: we only pass through fields we know the LLM benefits +from seeing. Two reasons to keep this tight: + +1. The LLM context is precious. Dumping every Twitter internal field + (gibberish like `core.user_results.result.legacy.advertiser_account_type`) + blows context for no gain. +2. Cache values are persisted to sqlite as JSON. We need every + serialized value to be JSON-safe, and twikit objects aren't (they + hold a back-reference to the Client). + +Missing-field policy: `_safe()` swallows AttributeError and any other +exception twikit's property descriptors raise (they dereference into +`_data['legacy'][...]` and 404 on half-populated tweets), logging at +DEBUG and returning a default. We trade silent-degradation for not +surfacing drift, on the bet that the lifespan's smoke probe + the +gate's per-call exception handlers are the right place to catch real +twikit/X wire changes — the serializer just has to keep the cache JSON +viable. + +`media_to_dict` handles the polymorphic Photo/Video/AnimatedGif subclass +case by reading the `type` attribute and only pulling the URLs each +subclass actually exposes. Streams (subclass of Video) reuse Video's +serializer. +""" + +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + + +def _safe(obj: object, attr: str, default: Any = None) -> Any: + """Pull `attr` off `obj`, swallowing AttributeError + property errors. + + twikit's properties dereference into `_data['legacy'][...]` which + sometimes 404s on freshly-fetched-but-half-populated tweets. Rather + than mark the whole serialization as failed, log once and continue + with the default. + """ + try: + return getattr(obj, attr, default) + except Exception as e: # noqa: BLE001 — twikit properties can raise *anything* + logger.debug("serializer: %r raised on attr %s: %s", obj, attr, e) + return default + + +def media_to_dict(media: object) -> dict: + """Photo / Video / AnimatedGif / Stream → plain dict.""" + return { + "id": _safe(media, "id"), + "type": _safe(media, "type"), + "url": _safe(media, "url") or _safe(media, "media_url"), + # Photos expose `media_url`; videos expose `streams` (list of + # variants). Serialize both possibilities, swallow whichever + # doesn't apply. + "alt_text": _safe(media, "alt_text"), + } + + +def user_to_dict(user: object) -> dict: + """twikit.User → JSON-safe dict, whitelisted fields only.""" + return { + "id": _safe(user, "id"), + "handle": _safe(user, "screen_name"), + "name": _safe(user, "name"), + "description": _safe(user, "description"), + "location": _safe(user, "location"), + "url": _safe(user, "url"), + "profile_image_url": _safe(user, "profile_image_url"), + "profile_banner_url": _safe(user, "profile_banner_url"), + "created_at": _safe(user, "created_at"), + "is_blue_verified": _safe(user, "is_blue_verified"), + "verified": _safe(user, "verified"), + "followers_count": _safe(user, "followers_count"), + "following_count": _safe(user, "following_count"), + "statuses_count": _safe(user, "statuses_count"), + "media_count": _safe(user, "media_count"), + "listed_count": _safe(user, "listed_count"), + "favourites_count": _safe(user, "favourites_count"), + "pinned_tweet_ids": _safe(user, "pinned_tweet_ids") or [], + } + + +def tweet_to_dict(tweet: object, *, include_replies: bool = False) -> dict: + """twikit.Tweet → JSON-safe dict, recursively flattening quotes. + + `include_replies` is opt-in because the recursive .replies attribute + can be a `Result[Tweet]` of arbitrary length. We materialize the + current page only (no auto-pagination). + """ + user = _safe(tweet, "user") + quote = _safe(tweet, "quote") + retweeted = _safe(tweet, "retweeted_tweet") + media_list = _safe(tweet, "media") or [] + + out: dict = { + "id": _safe(tweet, "id"), + "created_at": _safe(tweet, "created_at"), + "text": _safe(tweet, "text"), + "lang": _safe(tweet, "lang"), + "in_reply_to": _safe(tweet, "in_reply_to"), + "is_quote_status": _safe(tweet, "is_quote_status"), + "possibly_sensitive": _safe(tweet, "possibly_sensitive"), + "view_count": _safe(tweet, "view_count"), + "reply_count": _safe(tweet, "reply_count"), + "favorite_count": _safe(tweet, "favorite_count"), + "retweet_count": _safe(tweet, "retweet_count"), + "quote_count": _safe(tweet, "quote_count"), + "bookmark_count": _safe(tweet, "bookmark_count"), + "hashtags": _safe(tweet, "hashtags") or [], + "urls": _safe(tweet, "urls") or [], + "media": [media_to_dict(m) for m in media_list], + "user": user_to_dict(user) if user is not None else None, + # Nested tweets recurse but don't drill into THEIR replies/ + # quote/retweet to keep payload sizes bounded. + "quote": _shallow_tweet_to_dict(quote) if quote is not None else None, + "retweeted_tweet": _shallow_tweet_to_dict(retweeted) if retweeted is not None else None, + } + + if include_replies: + replies = _safe(tweet, "replies") + if replies is not None: + out["replies"] = result_to_dict(replies, tweet_to_dict) + + return out + + +def _shallow_tweet_to_dict(tweet: object) -> dict: + """Tweet → dict but without recursing further into quote/retweet chains.""" + user = _safe(tweet, "user") + return { + "id": _safe(tweet, "id"), + "created_at": _safe(tweet, "created_at"), + "text": _safe(tweet, "text"), + "user": user_to_dict(user) if user is not None else None, + "reply_count": _safe(tweet, "reply_count"), + "favorite_count": _safe(tweet, "favorite_count"), + "retweet_count": _safe(tweet, "retweet_count"), + } + + +def result_to_dict(result: object, item_serializer) -> dict: + """twikit.utils.Result → {items, next_cursor, previous_cursor}. + + Result is iterable (yields items in this page). It also carries + cursor strings the caller threads back through subsequent calls + for pagination. We don't auto-fetch the next page — that's the + caller's job, and counts against the rate-limit budget separately. + + Tolerates `result is None` (the previous shape — `for item in None` + then `for item in result or []` — handled it accidentally, but the + TypeError fallback was a no-op since the second loop also raised + on None). + """ + if result is None: + return {"items": [], "next_cursor": None, "previous_cursor": None} + + items = [] + try: + for item in result: + items.append(item_serializer(item)) + except TypeError: + # Some twikit endpoints return bare lists rather than Result + # instances; iter() on a list works fine, so a TypeError here + # really means "this object isn't iterable at all" — log and + # treat as empty. + logger.debug("serializer: result_to_dict got non-iterable %r", result) + return { + "items": [], + "next_cursor": _safe(result, "next_cursor"), + "previous_cursor": _safe(result, "previous_cursor"), + } + + return { + "items": items, + "next_cursor": _safe(result, "next_cursor"), + "previous_cursor": _safe(result, "previous_cursor"), + } diff --git a/backend/apps/twitter/twitter.py b/backend/apps/twitter/twitter.py new file mode 100644 index 00000000..e24c5a4c --- /dev/null +++ b/backend/apps/twitter/twitter.py @@ -0,0 +1,752 @@ +"""Twitter SubApp: FastAPI router + lifespan + module-level pool/gate. + +The router exposes two flavors of route, both under `/api/twitter`: + +- Admin: `/accounts/login`, `/accounts`, `/accounts/{id}` (PATCH/DELETE), + `/accounts/{id}/verify`, `/accounts/{id}/health`. These manage the + pool itself — adding accounts, checking session health, tuning the + trust_multiplier. +- Tool reads: `/search`, `/user`, `/user/{id}/tweets`, `/tweet/{id}`, + `/tweet/{id}/replies`. These are what the MCP shim calls on behalf + of the LLM. Every one goes through `RateGate.execute` so cache hits + short-circuit twikit, and 429s are translated into structured + responses (`HTTP 429 + {"retry_after_s": N}`) the shim relays to the + agent. + +Module globals (`_pool`, `_gate`, etc.) are initialized inside the +lifespan context. Routes that touch them check for `None` and return +503 if the SubApp isn't ready yet (shouldn't happen in normal startup, +but defensive — tests sometimes import the router without running the +lifespan). +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import time +from contextlib import asynccontextmanager +from typing import Optional +from uuid import uuid4 + +from fastapi import HTTPException, Query, Response + +from backend.apps.twitter import persistence, serializers +from backend.apps.twitter.cache import TTLCache +from backend.apps.twitter.models import ( + AccountHealth, + BucketSnapshot, + LoginRequest, + TrustUpdateRequest, + TwitterAccount, +) +from backend.apps.twitter.pool import AccountPool, ManagedAccount +from backend.apps.twitter.ratelimit import GateResult, RateGate +from backend.config.Apps import SubApp + +logger = logging.getLogger(__name__) + + +# Surface twikit drift at import time. If the wheel didn't ship (broken +# packaged build, missing dependency in dev), the routes will still 503 +# correctly — but without this line the operator would have no idea +# whether twikit was the problem. +# +# We also apply our x_client_transaction regex patch here, *before* any +# Client/GuestClient is constructed (the patch hits a module-level class so +# instance order doesn't matter, but doing it at import keeps the dependency +# graph obvious). The patch is a workaround for twikit#408 and can be removed +# once a fixed twikit release is on PyPI. +try: + import twikit as _twikit + logger.info("twitter: twikit version %s", getattr(_twikit, "__version__", "unknown")) + from backend.apps.twitter import _twikit_patches as _twikit_patches_mod + _twikit_patches_mod.apply() +except ImportError as _e: + logger.error("twitter: twikit not importable (%s); SubApp will 503 on every call", _e) + + +# --------------------------------------------------------------------------- +# Module-level singletons (initialized inside the lifespan) +# --------------------------------------------------------------------------- + +_pool: AccountPool | None = None +_gate: RateGate | None = None +_cache: TTLCache | None = None +_snapshot_task: asyncio.Task | None = None +_probe_task: asyncio.Task | None = None +SNAPSHOT_INTERVAL_S = 1.0 + + +def _require_started() -> AccountPool: + """Raise 503 if the SubApp wasn't initialized. + + The token + middleware already returned by this point, so a 503 + here is unambiguous: the backend is up but the Twitter SubApp's + lifespan didn't run (likely an import-time crash in twikit). + """ + if _pool is None: + raise HTTPException(503, "Twitter SubApp not ready") + return _pool + + +# --------------------------------------------------------------------------- +# Lifespan: load state, start snapshot loop, smoke-probe twikit +# --------------------------------------------------------------------------- + +async def _periodic_snapshot(pool: AccountPool) -> None: + """Persist bucket state to sqlite at SNAPSHOT_INTERVAL_S. + + Sized at 1 second so a crash loses at most one second of decrements; + combined with `Bucket.restore`'s `min(saved, capacity/2)` clamp the + worst-case post-crash state is "we burned half a window's headroom", + not "we 429-storm Twitter." + """ + try: + while True: + try: + pool.snapshot_all() + except Exception: + logger.exception("twitter: periodic snapshot failed") + await asyncio.sleep(SNAPSHOT_INTERVAL_S) + except asyncio.CancelledError: + # Final snapshot on shutdown so the on-disk state matches + # whatever decrement just happened. + try: + pool.snapshot_all() + except Exception: + logger.exception("twitter: final snapshot on shutdown failed") + raise + + +async def _smoke_probe(pool: AccountPool) -> None: + """One-shot `client.user()` on the first active account. + + This is our canary for twikit-wire-shape drift. If X rotated their + GraphQL query IDs since the pinned twikit version was tested, + we'll see an unexpected exception class here and we can both log + loudly and write an audit row so /health surfaces the problem. + + Runs once at startup, never repeats. Uses the `_self_user` bucket + so it can't accidentally drain user-visible budget. + + Skipped entirely in dev (`OPENSWARM_TWITTER_SKIP_PROBE=1`) so that + hot-reload cycles don't burn the verify window on every restart. + """ + if os.environ.get("OPENSWARM_TWITTER_SKIP_PROBE") == "1": + logger.info("twitter: smoke probe skipped (OPENSWARM_TWITTER_SKIP_PROBE=1)") + return + + for managed in pool.accounts: + if managed.state != "active": + continue + try: + bucket = managed.bucket("_self_user") + if bucket.time_until_available() > 1.0: + # Don't burn budget at startup; skip and rely on the + # first /verify call to surface drift instead. + pool.audit_lifecycle(managed.id, "smoke_probe_skip", "no budget") + return + async with managed.concurrency: + await bucket.acquire() + await managed.client.user() + logger.info("twitter: startup smoke probe ok (%s)", managed.id) + pool.mark_active(managed.id) + pool.audit_lifecycle(managed.id, "smoke_probe_ok") + pool.commit() + return + except asyncio.CancelledError: + raise + except Exception as e: + # Don't crash startup — log so the operator notices, mark + # the account so the UI shows the failure, and write an + # audit row so /health.recent_429_count's siblings can + # surface the drift event to the operator without trawling + # logs. + logger.error( + "twitter: startup smoke probe FAILED for %s (%s): %s — " + "may indicate twikit/X wire-shape drift", + managed.id, + type(e).__name__, + e, + ) + pool.mark_needs_relogin(managed.id, f"smoke probe failed: {e}") + pool.audit_lifecycle(managed.id, "smoke_probe_fail", f"{type(e).__name__}: {e}") + pool.commit() + return + + +async def _hydrate_pool(pool: AccountPool) -> None: + """Reconstruct ManagedAccount instances from accounts.json + cookies. + + twikit's Client.set_cookies takes a dict directly. We read the + saved cookies file (twikit-format) and feed it back in. If a + cookies file is missing, we still register the account but mark + it needs_relogin so the UI prompts a re-login. + """ + import json as _json + from twikit import Client + + for raw in persistence.load_accounts(): + try: + record = TwitterAccount(**raw) + except Exception as e: + logger.warning("twitter: skip malformed account %s: %s", raw, e) + continue + + client = Client(language="en-US", proxy=record.proxy or None) + cookie_path = persistence.cookies_path(record.id) + if os.path.isfile(cookie_path): + try: + with open(cookie_path) as f: + cookies = _json.load(f) + client.set_cookies(cookies) + except Exception as e: + logger.warning( + "twitter: %s cookies unreadable (%s); marking needs_relogin", + record.id, + e, + ) + record.state = "needs_relogin" + record.last_error = "cookies unreadable" + else: + record.state = "needs_relogin" + record.last_error = "no cookies on disk" + + await pool.add(record, client) + + +def _persist_accounts(pool: AccountPool) -> None: + """Flush the public account state to accounts.json. + + Called after every mutation (login/delete/state change). Cheap — + accounts.json is small and rewritten atomically. + """ + persistence.save_accounts([a.record.model_dump() for a in pool.accounts]) + + +@asynccontextmanager +async def twitter_lifespan(): + """Initialize SubApp state, start background tasks, clean up on exit.""" + global _pool, _gate, _cache, _snapshot_task, _probe_task + + persistence.ensure_dirs() + conn = persistence.open_state_db() + persistence.trim_audit(conn) + + _cache = TTLCache(conn) + _pool = AccountPool(conn) + _gate = RateGate(_pool, _cache, block_ceiling_s=10.0) + + await _hydrate_pool(_pool) + # Hold the probe task so we can cancel it on shutdown — a bare + # `asyncio.create_task` here was leaking on fast restarts and + # surfacing as `Task was destroyed but it is pending!` warnings + # under pytest. + _probe_task = asyncio.create_task(_smoke_probe(_pool)) + _snapshot_task = asyncio.create_task(_periodic_snapshot(_pool)) + + try: + yield + finally: + for task in (_probe_task, _snapshot_task): + if task is None: + continue + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + except Exception: + logger.exception("twitter: background task raised during shutdown") + # Save cookies for every *active* account one last time. Cookies + # can rotate during a session (twikit may refresh ct0 silently), + # so this final flush is the difference between "next restart + # works" and "next restart needs_relogin." We deliberately skip + # non-active accounts: their `client` was constructed from a + # broken cookies file (or none at all) and calling save_cookies() + # would clobber whatever the user had on disk with garbage. + if _pool is not None: + for managed in _pool.accounts: + if managed.state != "active": + continue + try: + managed.client.save_cookies(persistence.cookies_path(managed.id)) + persistence.chmod_cookies(persistence.cookies_path(managed.id)) + except Exception as e: + logger.warning("twitter: final cookie save failed for %s: %s", managed.id, e) + # Only persist if we actually have accounts in memory. An empty + # in-memory pool overwriting accounts.json is the + # `import_cookies.py`-while-backend-shuts-down clobber: the + # script's write to disk gets nuked by our shutdown writing + # `[]` back. Skipping the persist when there's nothing to + # persist can only ever destroy information, never add it, so + # this guard is strictly safer. Mutating routes (`accounts_login`, + # `accounts_delete`, etc.) already call `_persist_accounts` on + # their own paths, so we're not relying on shutdown to flush + # legitimate state changes. + if _pool.accounts: + _persist_accounts(_pool) + try: + conn.close() + except Exception: + pass + + +twitter = SubApp("twitter", twitter_lifespan) + + +# --------------------------------------------------------------------------- +# Per-endpoint cache TTLs +# --------------------------------------------------------------------------- + +_CACHE_TTLS = { + "search_tweet": 60, + "get_user_by_screen_name": 300, + "get_user_by_id": 300, + "get_user_tweets": 120, + "get_tweet_by_id": 30, +} + + +def _gate_result_to_response(result: GateResult, response: Response) -> object: + """Translate `GateResult.outcome` into HTTP status + body. + + Keeps every tool route's tail identical: result -> response. + """ + if result.outcome == "ok": + return result.value + if result.outcome == "rate_limited": + response.status_code = 429 + # Return the structured retry-after the shim translates to MCP. + return result.value + if result.outcome == "no_account": + response.status_code = 503 + return result.value + if result.outcome in ("locked", "needs_relogin", "suspended"): + response.status_code = 409 # account is in a bad state, not a 500 + return result.value + # outcome == "error" + response.status_code = 502 + return result.value + + +# --------------------------------------------------------------------------- +# Admin routes +# --------------------------------------------------------------------------- + +@twitter.router.post("/accounts/login") +async def accounts_login(body: LoginRequest): + """Log in and persist cookies. Handles both fresh accounts and re-login. + + The password is used once to call `client.login()` and then dropped + on the floor — never stored, never logged. + + Re-login semantics: if any existing account has the same handle as + the one we just logged into, we mutate that ManagedAccount in place + so bucket state survives. Otherwise we create a fresh record. + """ + pool = _require_started() + persistence.ensure_dirs() + + from twikit import Client + + # We don't know the handle until after login (the user may have + # logged in by email/phone). Create a Client, log in, *then* match. + client = Client(language="en-US") + # Random tmp filename per request — a fixed `_tmp_login` path used + # to race when two browser tabs (or one double-submit) hit + # /accounts/login concurrently and the second login could clobber + # the first's cookies before os.replace ran. + cookies_temp = persistence.cookies_path(f"_tmp_login_{uuid4().hex}") + try: + await client.login( + auth_info_1=body.auth_info_1, + auth_info_2=body.auth_info_2, + password=body.password, + totp_secret=body.totp_secret, + cookies_file=cookies_temp, + ) + except Exception as e: + logger.warning("twitter: login failed for auth_info_1=%s: %s", body.auth_info_1, type(e).__name__) + # Make sure we don't leave a partial cookies file behind. + try: + os.remove(cookies_temp) + except OSError: + pass + raise HTTPException(401, f"Login failed: {type(e).__name__}: {e}") + + # Look up the just-authenticated handle so we can match against an + # existing record (re-login path). + try: + me = await client.user() + handle = getattr(me, "screen_name", None) + except Exception as e: + logger.warning("twitter: post-login client.user() failed: %s", e) + handle = None + + existing = pool.by_handle(handle) if handle else None + if existing is not None: + # Re-login: keep the existing record's id, label, role, trust. + existing.record.handle = handle or existing.record.handle + existing.record.state = "active" + existing.record.last_error = None + existing.record.last_verified_at = time.time() + if body.label: + existing.record.label = body.label + target_id = existing.id + else: + record = TwitterAccount( + label=body.label or (handle or body.auth_info_1), + handle=handle, + role=body.role, + ) + record.state = "active" + record.last_verified_at = time.time() + target_id = record.id + + # Move the temp cookies file to its permanent home (chmod 0600). + final_cookies_path = persistence.cookies_path(target_id) + try: + os.replace(cookies_temp, final_cookies_path) + except OSError as e: + logger.warning("twitter: cookies rename failed: %s; falling back to save_cookies()", e) + client.save_cookies(final_cookies_path) + persistence.chmod_cookies(final_cookies_path) + + if existing is not None: + await pool.add(existing.record, client) + managed = existing + else: + managed = await pool.add(record, client) + + pool.audit_lifecycle(managed.id, "login_ok", handle or "") + pool.commit() + _persist_accounts(pool) + + return {"account": managed.record.model_dump()} + + +@twitter.router.get("/accounts") +async def accounts_list(): + pool = _require_started() + return {"accounts": [a.record.model_dump() for a in pool.accounts]} + + +@twitter.router.patch("/accounts/{account_id}") +async def accounts_patch(account_id: str, body: TrustUpdateRequest): + pool = _require_started() + managed = pool.get(account_id) + if managed is None: + raise HTTPException(404, "Account not found") + managed.record.trust_multiplier = body.trust_multiplier + managed.rescale_buckets() + _persist_accounts(pool) + return {"account": managed.record.model_dump()} + + +@twitter.router.delete("/accounts/{account_id}") +async def accounts_delete(account_id: str): + pool = _require_started() + managed = pool.get(account_id) + if managed is None: + # Idempotent: deleting an absent account isn't an error. + return {"removed": True} + + clean = await pool.remove(account_id) + persistence.delete_cookies(account_id) + pool.audit_lifecycle(account_id, "delete") + pool.commit() + _persist_accounts(pool) + return {"removed": True, "clean": clean} + + +async def _verify_account(pool: AccountPool, managed: ManagedAccount) -> bool: + """Internal: call client.user() under the `_self_user` bucket. + + Asymmetric audit policy historically only logged verify_ok; we now + also log verify_fail with the twikit exception class so /health can + surface "why is this account stuck" without trawling the python log. + """ + from twikit.errors import ( + AccountLocked, + AccountSuspended, + TooManyRequests, + Unauthorized, + ) + + bucket = managed.bucket("_self_user") + if bucket.time_until_available() > 1.0: + # Don't wait — verify isn't critical. Return the current + # state's truthiness so callers can fall back to "still active + # from last successful verify." + pool.audit_lifecycle(managed.id, "verify_skip", "no budget") + pool.commit() + return managed.state == "active" + async with managed.concurrency: + # Acquire the bucket inside the semaphore — same atomicity + # rule as RateGate.execute, otherwise a concurrent verify + + # tool call could each grab a token even though only one was + # available. + await bucket.acquire() + try: + me = await managed.client.user() + handle = getattr(me, "screen_name", None) or managed.record.handle + managed.record.handle = handle + pool.mark_active(managed.id) + # Cookies sometimes refresh in the response; persist any + # diff. Cheapest correct thing: always save on verify_ok. + try: + managed.client.save_cookies(persistence.cookies_path(managed.id)) + persistence.chmod_cookies(persistence.cookies_path(managed.id)) + except Exception as e: + logger.warning("twitter: cookie save after verify failed: %s", e) + pool.audit_lifecycle(managed.id, "verify_ok") + pool.commit() + return True + except TooManyRequests as e: + bucket.mark_rate_limited(getattr(e, "rate_limit_reset", None)) + pool.audit_lifecycle(managed.id, "verify_fail", f"TooManyRequests: {e}") + pool.commit() + return False + except AccountLocked as e: + pool.mark_locked(managed.id, str(e)) + pool.audit_lifecycle(managed.id, "verify_fail", f"AccountLocked: {e}") + pool.commit() + return False + except AccountSuspended as e: + pool.mark_suspended(managed.id, str(e)) + pool.audit_lifecycle(managed.id, "verify_fail", f"AccountSuspended: {e}") + pool.commit() + return False + except Unauthorized as e: + pool.mark_needs_relogin(managed.id, str(e)) + pool.audit_lifecycle(managed.id, "verify_fail", f"Unauthorized: {e}") + pool.commit() + return False + except asyncio.CancelledError: + raise + except Exception as e: + logger.exception("twitter: verify failed for %s", managed.id) + pool.audit_lifecycle(managed.id, "verify_fail", f"{type(e).__name__}: {e}") + pool.commit() + return False + + +@twitter.router.post("/accounts/{account_id}/verify") +async def accounts_verify(account_id: str): + pool = _require_started() + managed = pool.get(account_id) + if managed is None: + raise HTTPException(404, "Account not found") + ok = await _verify_account(pool, managed) + _persist_accounts(pool) + return {"ok": ok, "account": managed.record.model_dump()} + + +@twitter.router.get("/accounts/{account_id}/health") +async def accounts_health(account_id: str) -> AccountHealth: + """Snapshot of an account's runtime state. Does NOT call twikit. + + Served entirely from in-memory state + an audit-log query for the + recent_429_count. Safe to poll aggressively from the frontend. + """ + pool = _require_started() + managed = pool.get(account_id) + if managed is None: + raise HTTPException(404, "Account not found") + + snaps: list[BucketSnapshot] = [] + for endpoint, b in managed._buckets.items(): + snaps.append(BucketSnapshot( + endpoint=endpoint, + capacity=b.capacity, + tokens=round(b.tokens, 2), + locked_until=b.locked_until, + seconds_until_available=round(b.time_until_available(), 2), + )) + + recent_429 = pool.recent_429s(account_id, since_s=24 * 3600) + + return AccountHealth( + id=managed.id, + label=managed.record.label, + handle=managed.record.handle, + state=managed.state, + role=managed.role, + trust_multiplier=managed.trust_multiplier, + last_verified_at=managed.record.last_verified_at, + last_error=managed.record.last_error, + recent_429_count=recent_429, + buckets=snaps, + ) + + +# --------------------------------------------------------------------------- +# Tool read routes — these are what the MCP shim hits +# --------------------------------------------------------------------------- + +@twitter.router.get("/search") +async def tool_search( + response: Response, + q: str = Query(..., min_length=1), + product: str = Query("Latest", pattern="^(Top|Latest|Media)$"), + count: int = Query(20, ge=1, le=50), + cursor: Optional[str] = None, +): + _require_started() + if _gate is None: + raise HTTPException(503, "Twitter SubApp not ready") + + async def op(client): + return await client.search_tweet(q, product, count, cursor) + + def serialize(result): + return serializers.result_to_dict(result, serializers.tweet_to_dict) + + res = await _gate.execute( + endpoint="search_tweet", + op=op, + serializer=serialize, + cache_key=("search", q, product, count, cursor or ""), + cache_ttl=_CACHE_TTLS["search_tweet"], + ) + return _gate_result_to_response(res, response) + + +@twitter.router.get("/user") +async def tool_get_user( + response: Response, + handle: Optional[str] = None, + user_id: Optional[str] = Query(None, alias="id"), +): + """Lookup by handle or by id; exactly one required.""" + _require_started() + if _gate is None: + raise HTTPException(503, "Twitter SubApp not ready") + if bool(handle) == bool(user_id): + raise HTTPException(400, "specify exactly one of: handle, id") + + if handle: + endpoint = "get_user_by_screen_name" + norm_handle = handle.lstrip("@") + + async def op(client): + return await client.get_user_by_screen_name(norm_handle) + + cache_key = ("user_by_handle", norm_handle.lower()) + else: + endpoint = "get_user_by_id" + + async def op(client): + return await client.get_user_by_id(user_id) + + cache_key = ("user_by_id", user_id) + + res = await _gate.execute( + endpoint=endpoint, + op=op, + serializer=serializers.user_to_dict, + cache_key=cache_key, + cache_ttl=_CACHE_TTLS[endpoint], + ) + return _gate_result_to_response(res, response) + + +@twitter.router.get("/user/{user_id}/tweets") +async def tool_get_user_tweets( + user_id: str, + response: Response, + type: str = Query("Tweets", pattern="^(Tweets|Replies|Media|Likes)$"), + count: int = Query(20, ge=1, le=50), + cursor: Optional[str] = None, +): + _require_started() + if _gate is None: + raise HTTPException(503, "Twitter SubApp not ready") + + async def op(client): + return await client.get_user_tweets(user_id, type, count, cursor) + + def serialize(result): + return serializers.result_to_dict(result, serializers.tweet_to_dict) + + res = await _gate.execute( + endpoint="get_user_tweets", + op=op, + serializer=serialize, + cache_key=("user_tweets", user_id, type, count, cursor or ""), + cache_ttl=_CACHE_TTLS["get_user_tweets"], + ) + return _gate_result_to_response(res, response) + + +@twitter.router.get("/tweet/{tweet_id}") +async def tool_get_tweet(tweet_id: str, response: Response): + _require_started() + if _gate is None: + raise HTTPException(503, "Twitter SubApp not ready") + + async def op(client): + return await client.get_tweet_by_id(tweet_id) + + res = await _gate.execute( + endpoint="get_tweet_by_id", + op=op, + serializer=lambda t: serializers.tweet_to_dict(t, include_replies=False), + cache_key=("tweet", tweet_id), + cache_ttl=_CACHE_TTLS["get_tweet_by_id"], + ) + return _gate_result_to_response(res, response) + + +@twitter.router.get("/tweet/{tweet_id}/replies") +async def tool_get_tweet_replies( + tweet_id: str, + response: Response, + cursor: Optional[str] = None, +): + """Replies share the get_tweet_by_id endpoint — same bucket. + + twikit's `get_tweet_by_id(id, cursor=...)` returns a Tweet whose + `.replies` is a Result of replies for that cursor page. We + serialize only the replies (not the parent tweet) so the agent + isn't fed duplicate context. + """ + _require_started() + if _gate is None: + raise HTTPException(503, "Twitter SubApp not ready") + + async def op(client): + return await client.get_tweet_by_id(tweet_id, cursor=cursor) + + def serialize(tweet): + replies = getattr(tweet, "replies", None) + if replies is None: + return {"items": [], "next_cursor": None, "previous_cursor": None} + return serializers.result_to_dict(replies, serializers.tweet_to_dict) + + res = await _gate.execute( + endpoint="get_tweet_by_id", # shares the parent endpoint's bucket + op=op, + serializer=serialize, + cache_key=("tweet_replies", tweet_id, cursor or ""), + cache_ttl=_CACHE_TTLS["get_tweet_by_id"], + ) + return _gate_result_to_response(res, response) + + +# --------------------------------------------------------------------------- +# Stats — small operator endpoint (not consumed by MCP) +# --------------------------------------------------------------------------- + +@twitter.router.get("/stats") +async def stats(): + pool = _require_started() + cache_stats = _cache.stats() if _cache else {} + return { + "accounts": len(pool.accounts), + "active_accounts": sum(1 for a in pool.accounts if a.state == "active"), + "cache": cache_stats, + } diff --git a/backend/apps/twitter_mcp_shim/__init__.py b/backend/apps/twitter_mcp_shim/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/twitter_mcp_shim/__main__.py b/backend/apps/twitter_mcp_shim/__main__.py new file mode 100644 index 00000000..fdb0fd58 --- /dev/null +++ b/backend/apps/twitter_mcp_shim/__main__.py @@ -0,0 +1,5 @@ +"""Module entrypoint so `python -m backend.apps.twitter_mcp_shim` works.""" +from backend.apps.twitter_mcp_shim.server import main + +if __name__ == "__main__": + main() diff --git a/backend/apps/twitter_mcp_shim/server.py b/backend/apps/twitter_mcp_shim/server.py new file mode 100644 index 00000000..90adc6da --- /dev/null +++ b/backend/apps/twitter_mcp_shim/server.py @@ -0,0 +1,438 @@ +"""Stdio MCP shim for the Twitter SubApp. + +Stdlib-only on purpose so the subprocess starts fast. Carries no +credentials: each MCP tool call is forwarded as a small HTTP request +to the local OpenSwarm backend (`OPENSWARM_BASE_URL`), authenticated +with our per-install token. The SubApp owns twikit, cookies, and rate- +limit state — this shim is pure protocol translation. + +Authentication. The bearer token comes from either: + +1. `OPENSWARM_AUTH_TOKEN_FILE` — preferred. The shim re-reads this file + on every `_call()` so a backend restart that rotates the token can't + strand a long-lived shim subprocess in 401-land. The token is also + cached in memory between calls; we only re-read when a 401 happens + or after a small interval (5s), so the steady-state cost is one + shared stat() per call, not a file open. +2. `OPENSWARM_AUTH_TOKEN` — fallback for callers that prefer the + env-value pattern (matches `agent_manager`'s convention for other + internal MCP servers). + +Rate-limit response handling. The SubApp returns HTTP 429 with a JSON +body `{"retry_after_s": N}` when a tool call would exceed the gate's +block ceiling. We translate that into a structured MCP error content +("Rate limited; retry in N seconds") so the LLM backs off cleanly +instead of timing out and spawning parallel calls. +""" + +from __future__ import annotations + +import json +import os +import sys +import time +import urllib.error +import urllib.parse +import urllib.request + +# --------------------------------------------------------------------------- +# Configuration (env vars set by tools_lib at spawn time) +# --------------------------------------------------------------------------- + +BACKEND_BASE = ( + os.environ.get("OPENSWARM_BASE_URL") + or f"http://127.0.0.1:{os.environ.get('OPENSWARM_PORT', '8324')}" +).rstrip("/") + +AUTH_TOKEN_FILE = os.environ.get("OPENSWARM_AUTH_TOKEN_FILE", "") +_ENV_TOKEN = os.environ.get("OPENSWARM_AUTH_TOKEN", "") + +# Cache the last-read token + read time so steady-state HTTP is fast. +# We re-read the file whenever the cached entry is older than this +# (cheap insurance against a rotation we missed). +_TOKEN_CACHE_TTL_S = 5.0 +_token_cache: tuple[float, str] = (0.0, "") + + +def _read_token() -> str: + """Return the current bearer token, preferring the file over env. + + Re-reading the file on every call would mean a syscall per request, + but the in-memory cache trims that to one read per 5s. If the file + isn't set, fall through to the env var (matches the rest of the + codebase's convention). + """ + global _token_cache + now = time.time() + if AUTH_TOKEN_FILE: + cached_at, cached = _token_cache + if cached and (now - cached_at) < _TOKEN_CACHE_TTL_S: + return cached + try: + with open(AUTH_TOKEN_FILE, "r", encoding="utf-8") as f: + tok = f.read().strip() + if tok: + _token_cache = (now, tok) + return tok + except OSError: + # File missing or unreadable — fall back to env var so the + # shim doesn't hard-fail just because the file rotated mid- + # read. + pass + return _ENV_TOKEN + + +# --------------------------------------------------------------------------- +# Tool definitions (input schemas mirror the SubApp's route signatures) +# --------------------------------------------------------------------------- + +TOOLS = [ + { + "name": "twitter_search", + "description": ( + "Search recent tweets matching a query. Returns a page of tweets and a " + "`next_cursor` you can pass back to paginate. `product` selects ranking: " + "'Latest' (recent), 'Top' (engagement-weighted), 'Media' (only tweets " + "with images/video)." + ), + "inputSchema": { + "type": "object", + "properties": { + "q": {"type": "string", "description": "Search query"}, + "product": { + "type": "string", + "enum": ["Top", "Latest", "Media"], + "default": "Latest", + }, + "count": {"type": "integer", "minimum": 1, "maximum": 50, "default": 20}, + "cursor": {"type": "string"}, + }, + "required": ["q"], + }, + }, + { + "name": "twitter_get_user", + "description": ( + "Look up a Twitter/X user by handle OR by numeric id. Provide exactly one " + "of `handle` (without the @) or `user_id`." + ), + "inputSchema": { + "type": "object", + "properties": { + "handle": {"type": "string", "description": "Screen name without @"}, + "user_id": {"type": "string", "description": "Numeric user ID"}, + }, + }, + }, + { + "name": "twitter_get_user_tweets", + "description": ( + "Page through a user's tweets. `type` selects which timeline: 'Tweets' " + "(originals + retweets, default), 'Replies' (only replies), 'Media' " + "(only tweets with media), 'Likes' (the user's likes)." + ), + "inputSchema": { + "type": "object", + "properties": { + "user_id": {"type": "string", "description": "Numeric user ID"}, + "type": { + "type": "string", + "enum": ["Tweets", "Replies", "Media", "Likes"], + "default": "Tweets", + }, + "count": {"type": "integer", "minimum": 1, "maximum": 50, "default": 20}, + "cursor": {"type": "string"}, + }, + "required": ["user_id"], + }, + }, + { + "name": "twitter_get_tweet", + "description": "Fetch a single tweet by id, including its author + media + counts.", + "inputSchema": { + "type": "object", + "properties": {"tweet_id": {"type": "string"}}, + "required": ["tweet_id"], + }, + }, + { + "name": "twitter_get_tweet_replies", + "description": ( + "Fetch a page of replies to a tweet. `cursor` paginates through deeper " + "reply pages — pass back the `next_cursor` from a prior call." + ), + "inputSchema": { + "type": "object", + "properties": { + "tweet_id": {"type": "string"}, + "cursor": {"type": "string"}, + }, + "required": ["tweet_id"], + }, + }, +] + + +# --------------------------------------------------------------------------- +# HTTP plumbing +# --------------------------------------------------------------------------- + +def _call( + method: str, + path: str, + *, + query: dict | None = None, + body: dict | None = None, + timeout: float = 30.0, +) -> tuple[int, dict | str]: + """Single HTTP hop to the Twitter SubApp. + + Returns (status_code, parsed_body_or_text). The shim's tool handlers + interpret status codes directly: + - 200: success + - 429: rate-limited -> surface retry_after_s to the agent + - 409: account is in a bad state (locked/needs_relogin/suspended) + - 503: backend unreachable or no active account + - others: surfaced verbatim as MCP errors + + On 401 we invalidate the token cache and retry exactly once. The + cache TTL is 5s, so without this the shim would 401-loop for up to + five seconds after a backend restart that rotated the bearer + token. One retry covers the common case (rotation) without + looping forever on a real auth failure. + """ + return _call_with_retry(method, path, query=query, body=body, timeout=timeout, _retry_on_401=True) + + +def _call_with_retry( + method: str, + path: str, + *, + query: dict | None, + body: dict | None, + timeout: float, + _retry_on_401: bool, +) -> tuple[int, dict | str]: + global _token_cache + + token = _read_token() + if not token: + return 0, "OPENSWARM_AUTH_TOKEN missing — shim can't authenticate to backend" + + url = f"{BACKEND_BASE}{path}" + if query: + cleaned = {k: v for k, v in query.items() if v is not None and v != ""} + if cleaned: + url += "?" + urllib.parse.urlencode(cleaned) + + headers = { + "Authorization": f"Bearer {token}", + "Accept": "application/json", + } + data: bytes | None = None + if body is not None: + headers["Content-Type"] = "application/json" + data = json.dumps(body).encode("utf-8") + + req = urllib.request.Request(url, data=data, headers=headers, method=method) + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + text = resp.read().decode("utf-8", errors="replace") + try: + return resp.status, json.loads(text) if text else {} + except json.JSONDecodeError: + return resp.status, text + except urllib.error.HTTPError as e: + if e.code == 401 and _retry_on_401: + # Bust the in-memory cache so _read_token re-reads the + # token file (which may have rotated). Then retry exactly + # once — flag is dropped so an actually-bad token surfaces + # as the 401 rather than looping. + _token_cache = (0.0, "") + return _call_with_retry( + method, path, + query=query, body=body, timeout=timeout, + _retry_on_401=False, + ) + text = "" + try: + text = e.read().decode("utf-8", errors="replace") if e.fp else "" + except Exception: + pass + try: + return e.code, json.loads(text) if text else {} + except json.JSONDecodeError: + return e.code, text or str(e) + except urllib.error.URLError as e: + return 0, f"Backend unreachable: {e.reason}" + except Exception as e: + return 0, f"Request failed: {e!r}" + + +# --------------------------------------------------------------------------- +# MCP response helpers +# --------------------------------------------------------------------------- + +def _ok(payload) -> dict: + if isinstance(payload, str): + return {"content": [{"type": "text", "text": payload}]} + return {"content": [{"type": "text", "text": json.dumps(payload, indent=2, default=str)}]} + + +def _err(text: str) -> dict: + return {"content": [{"type": "text", "text": f"Error: {text}"}], "isError": True} + + +def _rate_limited(body: dict) -> dict: + """Surface a 429 in MCP form so the LLM can back off cleanly. + + There's no structured retry-after in MCP, so we put the integer in + the text. Most LLMs honor "Rate limited; retry in N seconds" by + pausing or returning to the user. + """ + retry = body.get("retry_after_s") if isinstance(body, dict) else None + endpoint = body.get("endpoint") if isinstance(body, dict) else None + if retry is None: + return _err("Rate limited (no retry hint).") + msg = f"Rate limited on {endpoint or 'twitter'}; retry in {int(retry)} seconds." + return {"content": [{"type": "text", "text": msg}], "isError": True} + + +def _handle_response(status: int, body) -> dict: + """Translate a backend HTTP response into an MCP tool result.""" + if status == 200: + return _ok(body) + if status == 429: + return _rate_limited(body if isinstance(body, dict) else {}) + if status == 409: + # Account is in a bad state (locked/needs_relogin/suspended). + msg = (body.get("error") if isinstance(body, dict) else None) or "Account is unavailable" + return _err(msg) + if status == 503: + msg = (body.get("error") if isinstance(body, dict) else None) or "Twitter backend not ready" + return _err(msg) + if status == 0: + return _err(str(body)) + # Everything else: surface verbatim so we don't paper over real bugs. + return _err(f"HTTP {status}: {body}") + + +# --------------------------------------------------------------------------- +# Tool dispatch +# --------------------------------------------------------------------------- + +def handle_tool_call(name: str, args: dict) -> dict: + if name == "twitter_search": + q = str(args.get("q", "")).strip() + if not q: + return _err("Missing required argument: q") + status, body = _call( + "GET", + "/api/twitter/search", + query={ + "q": q, + "product": args.get("product", "Latest"), + "count": args.get("count", 20), + "cursor": args.get("cursor"), + }, + ) + return _handle_response(status, body) + + if name == "twitter_get_user": + handle = str(args.get("handle", "")).strip().lstrip("@") + user_id = str(args.get("user_id", "")).strip() + if bool(handle) == bool(user_id): + return _err("Specify exactly one of: handle, user_id") + query = {"handle": handle} if handle else {"id": user_id} + status, body = _call("GET", "/api/twitter/user", query=query) + return _handle_response(status, body) + + if name == "twitter_get_user_tweets": + user_id = str(args.get("user_id", "")).strip() + if not user_id: + return _err("Missing required argument: user_id") + status, body = _call( + "GET", + f"/api/twitter/user/{urllib.parse.quote(user_id, safe='')}/tweets", + query={ + "type": args.get("type", "Tweets"), + "count": args.get("count", 20), + "cursor": args.get("cursor"), + }, + ) + return _handle_response(status, body) + + if name == "twitter_get_tweet": + tweet_id = str(args.get("tweet_id", "")).strip() + if not tweet_id: + return _err("Missing required argument: tweet_id") + status, body = _call("GET", f"/api/twitter/tweet/{urllib.parse.quote(tweet_id, safe='')}") + return _handle_response(status, body) + + if name == "twitter_get_tweet_replies": + tweet_id = str(args.get("tweet_id", "")).strip() + if not tweet_id: + return _err("Missing required argument: tweet_id") + status, body = _call( + "GET", + f"/api/twitter/tweet/{urllib.parse.quote(tweet_id, safe='')}/replies", + query={"cursor": args.get("cursor")}, + ) + return _handle_response(status, body) + + return _err(f"Unknown tool: {name}") + + +# --------------------------------------------------------------------------- +# JSON-RPC stdio loop (mirrors backend/apps/discord_mcp_shim/server.py) +# --------------------------------------------------------------------------- + +def _send(id_, result=None, error=None): + msg = {"jsonrpc": "2.0", "id": id_} + if error is not None: + msg["error"] = error + else: + msg["result"] = result + sys.stdout.write(json.dumps(msg) + "\n") + sys.stdout.flush() + + +def main(): + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + msg = json.loads(line) + except json.JSONDecodeError: + continue + + method = msg.get("method") + id_ = msg.get("id") + params = msg.get("params", {}) or {} + + if method == "initialize": + _send(id_, { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "openswarm-twitter", "version": "1.0.0"}, + }) + elif method == "notifications/initialized": + pass + elif method == "tools/list": + _send(id_, {"tools": TOOLS}) + elif method == "tools/call": + name = params.get("name", "") + args = params.get("arguments", {}) or {} + try: + _send(id_, handle_tool_call(name, args)) + except Exception as e: + _send(id_, _err(f"shim crashed: {e!r}")) + elif method == "ping": + _send(id_, {}) + elif id_ is not None: + _send(id_, error={"code": -32601, "message": f"Method not found: {method}"}) + + +if __name__ == "__main__": + main() diff --git a/backend/config/paths.py b/backend/config/paths.py index 63024e14..7fa1b215 100644 --- a/backend/config/paths.py +++ b/backend/config/paths.py @@ -34,6 +34,7 @@ OUTPUTS_WORKSPACE_DIR = os.path.join(DATA_ROOT, "outputs_workspace") SKILLS_WORKSPACE_DIR = os.path.join(DATA_ROOT, "skills_workspace") DASHBOARD_LAYOUT_DIR = os.path.join(DATA_ROOT, "dashboard_layout") BUILTIN_PERMISSIONS_PATH = os.path.join(DATA_ROOT, "builtin_permissions.json") +TWITTER_DIR = os.path.join(DATA_ROOT, "twitter") # Per-install auth token for the localhost WS + HTTP API. Regenerated # every backend start. Only code running as the current OS user (Electron diff --git a/backend/main.py b/backend/main.py index bc536a8f..24756af3 100644 --- a/backend/main.py +++ b/backend/main.py @@ -43,11 +43,12 @@ from backend.apps.subscription.router import subscription from backend.apps.auth.router import auth from backend.apps.web.web import web from backend.apps.agents.anthropic_proxy import anthropic_proxy +from backend.apps.twitter.twitter import twitter from fastapi.middleware.cors import CORSMiddleware from fastapi import WebSocket, WebSocketDisconnect import json -main_app = MainApp([health, agents, skills, tools_lib, modes, settings, mcp_registry, skill_registry, outputs, dashboards, service, subscription, auth, web, anthropic_proxy]) +main_app = MainApp([health, agents, skills, tools_lib, modes, settings, mcp_registry, skill_registry, outputs, dashboards, service, subscription, auth, web, anthropic_proxy, twitter]) app = main_app.app # Generate per-install auth token BEFORE we bind the HTTP port. By the diff --git a/backend/requirements.txt b/backend/requirements.txt index 2f166bf0..756b33fd 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -16,6 +16,25 @@ python-dotenv==1.1.1 Pillow httpx>=0.27.0 trafilatura +# Twitter MCP SubApp (backend.apps.twitter). twikit drives X's internal +# GraphQL API on behalf of a user's logged-in account. Pinned to an exact +# version because twikit breaks every few weeks as X rotates query IDs; +# the SubApp's lifespan does a smoke probe on startup and logs loudly if +# the wire shape drifted. +twikit==2.3.3 +# Cloudflare TLS-fingerprint bypass for twikit. Stock httpx → OpenSSL +# binding → recognizable Python-library JA3/JA4 → 403 on every request +# to x.com since late 2025. `httpx-curl-cffi` is a BaseTransport that +# wraps curl-impersonate (BoringSSL for Chrome targets), producing a +# byte-identical TLS ClientHello + HTTP/2 SETTINGS frame to a real +# browser. `_twikit_patches._apply_tls_transport_patch` wires it into +# twikit's namespace at import time and picks the best chrome target +# the installed curl-cffi exposes (see `_IMPERSONATE_PREFERENCE`), so +# we don't need a tight version pin. `curl-cffi>=0.7.0` is a generous +# floor that covers every chrome target in our preference list back to +# chrome120. See twikit#396. +curl-cffi>=0.7.0 +httpx-curl-cffi>=0.1.5 # tzlocal: dev-mode fallback for resolving the user's IANA timezone when # Electron's OPENSWARM_TIMEZONE env var isn't set (i.e. `bash run.sh`). # Packaged builds get the env var directly so this is a safety net. diff --git a/backend/tests/test_twitter_persistence.py b/backend/tests/test_twitter_persistence.py new file mode 100644 index 00000000..5003c4c8 --- /dev/null +++ b/backend/tests/test_twitter_persistence.py @@ -0,0 +1,190 @@ +"""Persistence-layer tests for the Twitter SubApp. + +These exercise the small sqlite schema, accounts.json round-trip, and +the bucket-snapshot/restore path that's load-bearing for crash safety +(see `Bucket.restore` semantics). +""" + +from __future__ import annotations + +import json +import os +import tempfile + +import pytest + + +@pytest.fixture +def tmp_twitter_dir(monkeypatch): + """Re-point `TWITTER_DIR` at a tempdir so writes don't touch DATA_ROOT.""" + with tempfile.TemporaryDirectory() as d: + monkeypatch.setattr("backend.apps.twitter.persistence.TWITTER_DIR", d) + monkeypatch.setattr( + "backend.apps.twitter.persistence.ACCOUNTS_PATH", + os.path.join(d, "accounts.json"), + ) + monkeypatch.setattr( + "backend.apps.twitter.persistence.COOKIES_DIR", + os.path.join(d, "cookies"), + ) + monkeypatch.setattr( + "backend.apps.twitter.persistence.STATE_DB_PATH", + os.path.join(d, "state.sqlite"), + ) + yield d + + +def test_ensure_dirs_creates_with_strict_modes(tmp_twitter_dir): + from backend.apps.twitter.persistence import COOKIES_DIR, ensure_dirs + + ensure_dirs() + assert os.path.isdir(tmp_twitter_dir) + assert os.path.isdir(COOKIES_DIR) + # On macOS / Linux, mode 0700 means user-only access. + mode = os.stat(COOKIES_DIR).st_mode & 0o777 + assert mode == 0o700, f"expected 0700, got {oct(mode)}" + + +def test_accounts_round_trip(tmp_twitter_dir): + from backend.apps.twitter.persistence import load_accounts, save_accounts + + assert load_accounts() == [] + save_accounts([{"id": "a1", "label": "main"}]) + loaded = load_accounts() + assert loaded == [{"id": "a1", "label": "main"}] + + +def test_accounts_missing_file_returns_empty(tmp_twitter_dir): + from backend.apps.twitter.persistence import load_accounts + + assert load_accounts() == [] + + +def test_accounts_atomic_write_replaces_tmp(tmp_twitter_dir): + """save_accounts writes via a tmp file + os.replace; tmp shouldn't linger.""" + from backend.apps.twitter.persistence import ACCOUNTS_PATH, save_accounts + + save_accounts([{"id": "x"}]) + tmp = ACCOUNTS_PATH + ".tmp" + assert not os.path.exists(tmp) + assert os.path.isfile(ACCOUNTS_PATH) + + +def test_open_state_db_creates_schema(tmp_twitter_dir): + from backend.apps.twitter.persistence import open_state_db + + conn = open_state_db() + tables = {r[0] for r in conn.execute( + "SELECT name FROM sqlite_master WHERE type='table'" + ).fetchall()} + assert {"twitter_buckets", "twitter_cache", "twitter_audit"} <= tables + + +def test_save_load_bucket_round_trip(tmp_twitter_dir): + from backend.apps.twitter.persistence import load_buckets, open_state_db, save_bucket + + conn = open_state_db() + save_bucket( + conn, + account_id="a1", + endpoint="search_tweet", + snapshot={"capacity": 50, "tokens": 17.5, "locked_until": 12345.0}, + ) + conn.commit() + rows = load_buckets(conn, "a1") + assert rows == { + "search_tweet": {"capacity": 50, "tokens": 17.5, "locked_until": 12345.0} + } + + +def test_save_bucket_upserts_on_conflict(tmp_twitter_dir): + """Re-saving the same (account, endpoint) updates in place.""" + from backend.apps.twitter.persistence import load_buckets, open_state_db, save_bucket + + conn = open_state_db() + save_bucket(conn, "a1", "search_tweet", {"capacity": 50, "tokens": 50.0, "locked_until": 0}) + save_bucket(conn, "a1", "search_tweet", {"capacity": 50, "tokens": 3.0, "locked_until": 999}) + conn.commit() + rows = load_buckets(conn, "a1") + assert rows["search_tweet"]["tokens"] == pytest.approx(3.0) + assert rows["search_tweet"]["locked_until"] == pytest.approx(999.0) + # And only one row, not two. + (n,) = conn.execute("SELECT COUNT(*) FROM twitter_buckets WHERE account_id='a1'").fetchone() + assert n == 1 + + +def test_delete_buckets_for(tmp_twitter_dir): + from backend.apps.twitter.persistence import delete_buckets_for, open_state_db, save_bucket + + conn = open_state_db() + save_bucket(conn, "a1", "search_tweet", {"capacity": 50, "tokens": 50.0, "locked_until": 0}) + save_bucket(conn, "a2", "search_tweet", {"capacity": 50, "tokens": 50.0, "locked_until": 0}) + delete_buckets_for(conn, "a1") + (n_a1,) = conn.execute("SELECT COUNT(*) FROM twitter_buckets WHERE account_id='a1'").fetchone() + (n_a2,) = conn.execute("SELECT COUNT(*) FROM twitter_buckets WHERE account_id='a2'").fetchone() + assert n_a1 == 0 + assert n_a2 == 1 + + +def test_audit_and_recent_429s(tmp_twitter_dir): + from backend.apps.twitter.persistence import audit, open_state_db, recent_429s + + conn = open_state_db() + audit(conn, "a1", "search_tweet", "429", "rate limit") + audit(conn, "a1", "search_tweet", "429", "rate limit") + audit(conn, "a1", "search_tweet", "login_ok", None) + conn.commit() + assert recent_429s(conn, "a1", since_s=3600) == 2 + assert recent_429s(conn, "a2", since_s=3600) == 0 + + +def test_trim_audit_drops_old(tmp_twitter_dir, monkeypatch): + import time + from backend.apps.twitter.persistence import audit, open_state_db, trim_audit + + conn = open_state_db() + # Insert a row 60 days in the past (default keep_days=30 should drop). + old_ts = time.time() - 60 * 86400 + conn.execute( + "INSERT INTO twitter_audit (ts, account_id, endpoint, event, detail) " + "VALUES (?, 'a1', 'x', 'login_ok', NULL)", + (old_ts,), + ) + audit(conn, "a1", "x", "login_ok") # fresh + conn.commit() + trim_audit(conn, keep_days=30) + (n,) = conn.execute("SELECT COUNT(*) FROM twitter_audit").fetchone() + assert n == 1 + + +def test_cookies_path_and_chmod(tmp_twitter_dir): + from backend.apps.twitter.persistence import ( + chmod_cookies, + cookies_path, + ensure_dirs, + ) + + ensure_dirs() + p = cookies_path("acct-uuid-1") + # Simulate twikit having written cookies with default mode. + with open(p, "w") as f: + json.dump({"auth_token": "x", "ct0": "y"}, f) + os.chmod(p, 0o644) + chmod_cookies(p) + mode = os.stat(p).st_mode & 0o777 + assert mode == 0o600 + + +def test_delete_cookies_is_idempotent(tmp_twitter_dir): + from backend.apps.twitter.persistence import delete_cookies, ensure_dirs, cookies_path + + ensure_dirs() + # No-op when file doesn't exist. + delete_cookies("never-existed") + # Create then delete. + p = cookies_path("acct1") + with open(p, "w") as f: + f.write("{}") + assert os.path.isfile(p) + delete_cookies("acct1") + assert not os.path.exists(p) diff --git a/backend/tests/test_twitter_pool.py b/backend/tests/test_twitter_pool.py new file mode 100644 index 00000000..8d1eef7a --- /dev/null +++ b/backend/tests/test_twitter_pool.py @@ -0,0 +1,445 @@ +"""AccountPool tests: pick ordering, add/remove lifecycle, error hooks. + +No twikit needed — we feed in a sentinel object as the "client" since +the pool itself never calls twikit (RateGate does). +""" + +from __future__ import annotations + +import asyncio +import os +import tempfile + +import pytest + + +@pytest.fixture +def tmp_state_db(monkeypatch): + """Isolated sqlite + accounts dir per test.""" + with tempfile.TemporaryDirectory() as d: + monkeypatch.setattr("backend.apps.twitter.persistence.TWITTER_DIR", d) + monkeypatch.setattr( + "backend.apps.twitter.persistence.ACCOUNTS_PATH", + os.path.join(d, "accounts.json"), + ) + monkeypatch.setattr( + "backend.apps.twitter.persistence.COOKIES_DIR", + os.path.join(d, "cookies"), + ) + monkeypatch.setattr( + "backend.apps.twitter.persistence.STATE_DB_PATH", + os.path.join(d, "state.sqlite"), + ) + from backend.apps.twitter.persistence import open_state_db + conn = open_state_db() + yield conn + conn.close() + + +def _make_pool(conn): + from backend.apps.twitter.pool import AccountPool + return AccountPool(conn) + + +def _make_record(id_="a1", state="active", role="primary", trust=0.4, handle=None): + from backend.apps.twitter.models import TwitterAccount + return TwitterAccount(id=id_, state=state, role=role, trust_multiplier=trust, handle=handle) + + +# --------------------------------------------------------------------------- +# pick() +# --------------------------------------------------------------------------- + +def test_pick_empty_returns_none(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + return await pool.pick("search_tweet") + + assert asyncio.run(_run()) is None + + +def test_pick_skips_non_active_accounts(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1", state="locked"), object()) + await pool.add(_make_record("a2", state="suspended"), object()) + return await pool.pick("search_tweet") + + assert asyncio.run(_run()) is None + + +def test_pick_returns_active(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1", state="locked"), object()) + await pool.add(_make_record("a2", state="active"), object()) + chosen = await pool.pick("search_tweet") + return chosen.id + + assert asyncio.run(_run()) == "a2" + + +def test_pick_prefers_account_with_more_budget(tmp_state_db, monkeypatch): + """Two active accounts: the one whose bucket is unlocked wins.""" + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("hot"), object()) + await pool.add(_make_record("cold"), object()) + + # Drain the "hot" account's bucket to force a long wait. + hot = pool.get("hot") + b = hot.bucket("search_tweet") + b.tokens = 0.0 + import time as t + b.locked_until = t.time() + 999 # blocked for a long time + + chosen = await pool.pick("search_tweet") + return chosen.id + + assert asyncio.run(_run()) == "cold" + + +# --------------------------------------------------------------------------- +# add() — runtime registration +# --------------------------------------------------------------------------- + +def test_add_inserts_with_buckets_lazy(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + managed = await pool.add(_make_record("a1"), object()) + # No bucket created until requested. + assert managed._buckets == {} + b = managed.bucket("search_tweet") + # Trust 0.4 * 50 default = 20. + assert b.capacity == 20 + + asyncio.run(_run()) + + +def test_add_idempotent_for_relogin(tmp_state_db): + """Re-login: same account.id, same handle — must reuse the existing + ManagedAccount (so we don't lose bucket state on a cookie refresh).""" + async def _run(): + pool = _make_pool(tmp_state_db) + first = await pool.add(_make_record("a1"), object()) + # Consume some budget so we can verify it survives the re-add. + b = first.bucket("search_tweet") + b.tokens = 3.0 + new_client = object() + second = await pool.add(_make_record("a1"), new_client) + assert second is first, "must mutate in place, not replace" + assert second.client is new_client + assert second.bucket("search_tweet").tokens == pytest.approx(3.0) + + asyncio.run(_run()) + + +def test_add_restores_buckets_from_disk(tmp_state_db): + """Snapshot persisted from a previous run gets restored on add().""" + from backend.apps.twitter.persistence import save_bucket + + save_bucket( + tmp_state_db, + account_id="a1", + endpoint="search_tweet", + snapshot={"capacity": 20, "tokens": 4.0, "locked_until": 1234.0}, + ) + tmp_state_db.commit() + + async def _run(): + pool = _make_pool(tmp_state_db) + managed = await pool.add(_make_record("a1"), object()) + b = managed.bucket("search_tweet") + # restore() clamps to capacity / 2 = 10, and 4 < 10 so kept at 4. + assert b.tokens == pytest.approx(4.0) + assert b.locked_until == pytest.approx(1234.0) + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- +# remove() — cancellation safety +# --------------------------------------------------------------------------- + +def test_remove_clean_path(tmp_state_db): + """No in-flight call: remove() returns True immediately.""" + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1"), object()) + # Touch a bucket so we have something to clean up. + pool.get("a1").bucket("search_tweet") + pool.snapshot_all() + result = await pool.remove("a1") + return result, pool.get("a1") + + clean, after = asyncio.run(_run()) + assert clean is True + assert after is None + + +def test_remove_missing_account_is_noop(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + return await pool.remove("never-existed") + + assert asyncio.run(_run()) is True + + +def test_remove_waits_for_in_flight_call(tmp_state_db, monkeypatch): + """If a call is mid-flight under the semaphore, remove() should wait.""" + + async def _run(): + from backend.apps.twitter.pool import AccountPool + pool = AccountPool(tmp_state_db) + await pool.add(_make_record("a1"), object()) + managed = pool.get("a1") + + # Manually take the semaphore as if a twikit call were running. + await managed.concurrency.acquire() + + async def _release_later(): + await asyncio.sleep(0.05) + managed.concurrency.release() + + rel = asyncio.create_task(_release_later()) + clean = await pool.remove("a1") + await rel + return clean + + assert asyncio.run(_run()) is True + + +def test_remove_timeout_returns_false(tmp_state_db, monkeypatch): + """If the in-flight call doesn't yield within REMOVE_TIMEOUT_S, we + bail and let the caller route a 503.""" + monkeypatch.setattr("backend.apps.twitter.pool.REMOVE_TIMEOUT_S", 0.05) + + async def _run(): + from backend.apps.twitter.pool import AccountPool + pool = AccountPool(tmp_state_db) + await pool.add(_make_record("a1"), object()) + managed = pool.get("a1") + # Hold the semaphore for longer than the timeout. + await managed.concurrency.acquire() + clean = await pool.remove("a1") + managed.concurrency.release() + return clean + + assert asyncio.run(_run()) is False + + +def test_remove_wipes_buckets_from_disk(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1"), object()) + pool.get("a1").bucket("search_tweet") + pool.snapshot_all() + # Sanity: bucket row exists. + (n_before,) = tmp_state_db.execute( + "SELECT COUNT(*) FROM twitter_buckets WHERE account_id='a1'" + ).fetchone() + assert n_before == 1 + + await pool.remove("a1") + (n_after,) = tmp_state_db.execute( + "SELECT COUNT(*) FROM twitter_buckets WHERE account_id='a1'" + ).fetchone() + assert n_after == 0 + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- +# lifecycle hooks +# --------------------------------------------------------------------------- + +def test_mark_locked_flips_state_and_audits(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1"), object()) + pool.mark_locked("a1", "arkose") + tmp_state_db.commit() + acct = pool.get("a1") + return acct.state, acct.record.last_error + + state, err = asyncio.run(_run()) + assert state == "locked" + assert "arkose" in (err or "") + # And an audit row was written. + (n,) = tmp_state_db.execute( + "SELECT COUNT(*) FROM twitter_audit WHERE account_id='a1' AND event='locked'" + ).fetchone() + assert n == 1 + + +def test_mark_suspended_routes_correctly(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1"), object()) + pool.mark_suspended("a1", "policy violation") + return pool.get("a1").state + + assert asyncio.run(_run()) == "suspended" + + +def test_mark_needs_relogin_routes_correctly(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1"), object()) + pool.mark_needs_relogin("a1", "401") + return pool.get("a1").state + + assert asyncio.run(_run()) == "needs_relogin" + + +def test_mark_active_after_recovery(tmp_state_db): + """A locked account that re-logins successfully goes back to active.""" + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1", state="locked"), object()) + pool.mark_active("a1") + acct = pool.get("a1") + return acct.state, acct.record.last_error, acct.record.last_verified_at > 0 + + state, err, verified = asyncio.run(_run()) + assert state == "active" + assert err is None + assert verified is True + + +# --------------------------------------------------------------------------- +# rescale_buckets (PATCH /accounts/{id} support) +# --------------------------------------------------------------------------- + +def test_rescale_buckets_lowers_cap_and_clamps_tokens(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1", trust=0.4), object()) + managed = pool.get("a1") + b = managed.bucket("search_tweet") + assert b.capacity == 20 + b.tokens = 18.0 + + # Halve the multiplier — new cap should be 10, and tokens + # should be clamped down with it. + managed.record.trust_multiplier = 0.2 + managed.rescale_buckets() + b2 = managed.bucket("search_tweet") + assert b2.capacity == 10 + assert b2.tokens <= 10 + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- +# by_handle lookup (re-login path) +# --------------------------------------------------------------------------- + +def test_by_handle_case_insensitive(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1", handle="OpenSwarm"), object()) + return pool.by_handle("@openswarm").id + + assert asyncio.run(_run()) == "a1" + + +def test_by_handle_missing_returns_none(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + await pool.add(_make_record("a1", handle="X"), object()) + return pool.by_handle("Y") + + assert asyncio.run(_run()) is None + + +# --------------------------------------------------------------------------- +# Re-login client-swap: must drain in-flight calls +# --------------------------------------------------------------------------- + +def test_add_waits_for_in_flight_call_on_relogin(tmp_state_db): + """A second `add` on the same id swaps the Client. If a twikit call + is mid-flight under the semaphore, the swap must wait — otherwise + the in-flight call sees its cookies replaced halfway through. + """ + + async def _run(): + from backend.apps.twitter.pool import AccountPool + pool = AccountPool(tmp_state_db) + first_client = object() + await pool.add(_make_record("a1"), first_client) + managed = pool.get("a1") + + # Pretend a twikit call is in flight by taking the semaphore. + await managed.concurrency.acquire() + + new_client = object() + # Schedule a release shortly after we kick off `add`. + async def _release_later(): + await asyncio.sleep(0.05) + managed.concurrency.release() + + rel = asyncio.create_task(_release_later()) + result = await pool.add(_make_record("a1"), new_client) + await rel + + # We got back the SAME ManagedAccount with the NEW client. + assert result is managed + assert result.client is new_client + + asyncio.run(_run()) + + +def test_add_relogin_times_out_on_stuck_call(tmp_state_db, monkeypatch): + """If the in-flight call never releases, the swap proceeds anyway. + + The alternative is stranding a re-login indefinitely on a stuck + twikit call — much worse than a one-call cookie-drift surprise. + """ + monkeypatch.setattr("backend.apps.twitter.pool.REPLACE_CLIENT_TIMEOUT_S", 0.05) + + async def _run(): + from backend.apps.twitter.pool import AccountPool + pool = AccountPool(tmp_state_db) + await pool.add(_make_record("a1"), object()) + managed = pool.get("a1") + # Hold the semaphore for longer than the timeout. + await managed.concurrency.acquire() + + new_client = object() + result = await pool.add(_make_record("a1"), new_client) + managed.concurrency.release() + return result.client is new_client + + assert asyncio.run(_run()) is True + + +# --------------------------------------------------------------------------- +# audit/commit/recent_429s helpers (keep routes out of pool._conn) +# --------------------------------------------------------------------------- + +def test_audit_lifecycle_writes_underscore_lifecycle_endpoint(tmp_state_db): + async def _run(): + pool = _make_pool(tmp_state_db) + pool.audit_lifecycle("a1", "delete", "user-requested") + pool.commit() + + asyncio.run(_run()) + rows = tmp_state_db.execute( + "SELECT endpoint, event, detail FROM twitter_audit WHERE account_id='a1'" + ).fetchall() + assert rows == [("_lifecycle", "delete", "user-requested")] + + +def test_recent_429s_helper_returns_audit_count(tmp_state_db): + """pool.recent_429s wraps persistence.recent_429s so /health doesn't + need to know about the sqlite connection.""" + + async def _run(): + pool = _make_pool(tmp_state_db) + pool.record_429("a1", "search_tweet") + pool.record_429("a1", "search_tweet") + pool.commit() + return pool.recent_429s("a1", since_s=3600) + + assert asyncio.run(_run()) == 2 diff --git a/backend/tests/test_twitter_ratelimit.py b/backend/tests/test_twitter_ratelimit.py new file mode 100644 index 00000000..e32bc906 --- /dev/null +++ b/backend/tests/test_twitter_ratelimit.py @@ -0,0 +1,497 @@ +"""Unit tests for the Twitter SubApp's rate-limit and cache primitives. + +No live twikit calls — `RateGate` is exercised via fake `pool` / `cache` +stand-ins so we can deterministically drive each error branch and assert +the resulting `GateResult.outcome`. + +Why these exist: +- Bucket math is load-bearing for the whole system; a refill bug ships + as "user gets locked out of X." +- The `min(saved, capacity/2)` restore rule is a correctness invariant + for crash safety; future refactors must not regress it. +- Errors-not-cached is a security/correctness rule (don't memoize a + rate-limit-storm as the canonical answer). +""" + +from __future__ import annotations + +import asyncio +import sqlite3 +import time +from dataclasses import dataclass, field +from typing import Any + +import pytest + + +# --------------------------------------------------------------------------- +# Bucket +# --------------------------------------------------------------------------- + +def test_bucket_starts_full(): + from backend.apps.twitter.ratelimit import Bucket + + b = Bucket(capacity=10) + assert b.tokens == pytest.approx(10.0) + assert b.time_until_available() == 0.0 + + +def test_bucket_refills_proportionally(monkeypatch): + """1 minute elapsed -> 1/15 of capacity refilled.""" + from backend.apps.twitter.ratelimit import Bucket + + b = Bucket(capacity=150) + b.tokens = 0.0 + fake_now = 1000.0 + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.monotonic", lambda: fake_now) + b.last_refill = fake_now + fake_now = 1000.0 + 60.0 # 1 minute later + # Expected refill: 150 * (60 / 900) = 10 tokens. + b._refill() + assert b.tokens == pytest.approx(10.0, abs=0.01) + + +def test_bucket_acquire_returns_when_token_available(): + from backend.apps.twitter.ratelimit import Bucket + + async def _run(): + b = Bucket(capacity=5) + before = b.tokens + await b.acquire() + assert b.tokens == pytest.approx(before - 1.0) + + asyncio.run(_run()) + + +def test_bucket_time_until_available_respects_locked_until(monkeypatch): + """Even with tokens available, locked_until in the future should win.""" + from backend.apps.twitter.ratelimit import Bucket + + b = Bucket(capacity=10) + assert b.tokens >= 1.0 + fake_wall = 5000.0 + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.time", lambda: fake_wall) + b.locked_until = fake_wall + 30.0 + wait = b.time_until_available() + assert wait == pytest.approx(30.0, abs=0.5) + + +def test_bucket_locked_until_in_past_is_noop(monkeypatch): + from backend.apps.twitter.ratelimit import Bucket + + b = Bucket(capacity=10) + fake_wall = 5000.0 + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.time", lambda: fake_wall) + b.locked_until = fake_wall - 10.0 # past + assert b.time_until_available() == 0.0 + + +def test_bucket_mark_rate_limited_zeros_tokens_and_locks(monkeypatch): + from backend.apps.twitter.ratelimit import Bucket + + b = Bucket(capacity=10) + fake_wall = 5000.0 + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.time", lambda: fake_wall) + b.mark_rate_limited(reset_at=fake_wall + 60.0) + assert b.tokens == 0.0 + # locked_until = reset + 1-5s jitter, so 61.0 <= locked_until - now <= 65.0. + assert fake_wall + 61.0 <= b.locked_until <= fake_wall + 65.0 + + +def test_bucket_mark_rate_limited_no_reset_falls_back_to_full_window(monkeypatch): + """If twikit didn't surface a reset header, cool down for one window.""" + from backend.apps.twitter.ratelimit import Bucket, WINDOW_S + + b = Bucket(capacity=10) + fake_wall = 5000.0 + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.time", lambda: fake_wall) + b.mark_rate_limited(reset_at=None) + # locked_until ~ now + WINDOW_S + jitter + assert fake_wall + WINDOW_S + 1.0 <= b.locked_until <= fake_wall + WINDOW_S + 5.0 + + +def test_bucket_snapshot_restore_round_trip(monkeypatch): + """Snapshot then restore preserves locked_until and caps tokens at capacity/2.""" + from backend.apps.twitter.ratelimit import Bucket + + fake_wall = 5000.0 + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.time", lambda: fake_wall) + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.monotonic", lambda: 7000.0) + + b = Bucket(capacity=100) + b.tokens = 80.0 + b.locked_until = fake_wall + 42.0 + + snap = b.snapshot() + assert snap["locked_until"] == pytest.approx(fake_wall + 42.0) + + restored = Bucket.restore(snap) + # tokens were 80 but capacity/2 is 50, so restore should clamp to 50. + assert restored.tokens == pytest.approx(50.0) + assert restored.locked_until == pytest.approx(fake_wall + 42.0) + assert restored.capacity == 100 + + +def test_bucket_restore_with_low_saved_tokens(monkeypatch): + """If saved tokens < capacity/2, keep the saved value (don't inflate).""" + from backend.apps.twitter.ratelimit import Bucket + + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.monotonic", lambda: 1.0) + restored = Bucket.restore({"capacity": 100, "tokens": 5.0, "locked_until": 0.0}) + assert restored.tokens == pytest.approx(5.0) + + +def test_bucket_capacity_zero_does_not_divide_by_zero(): + """Paused account (trust_multiplier=0) → capacity=0 in the bucket math. + + Used to crash with ZeroDivisionError in time_until_available because + the refill rate is capacity/WINDOW_S. We now treat capacity=0 as + "infinite wait" with a 24h sentinel so pick() naturally deprioritizes + paused accounts. + """ + from backend.apps.twitter.ratelimit import Bucket + + b = Bucket(capacity=0) + # Should not raise. + wait = b.time_until_available() + # Sentinel is large enough that any other account beats it in pick(). + assert wait >= 3600.0 + + +def test_bucket_capacity_zero_round_trips_through_snapshot(monkeypatch): + """Restoring a snapshot with capacity=0 keeps it at 0 (not floored to 1).""" + from backend.apps.twitter.ratelimit import Bucket + + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.monotonic", lambda: 1.0) + restored = Bucket.restore({"capacity": 0, "tokens": 0.0, "locked_until": 0.0}) + assert restored.capacity == 0 + + +# --------------------------------------------------------------------------- +# TTLCache +# --------------------------------------------------------------------------- + +@pytest.fixture +def cache_conn(): + """In-memory sqlite with the cache schema; isolated per test.""" + conn = sqlite3.connect(":memory:") + conn.execute( + """ + CREATE TABLE twitter_cache ( + key TEXT PRIMARY KEY, + value_json TEXT NOT NULL, + expires_at REAL NOT NULL + ) + """ + ) + conn.commit() + return conn + + +def test_cache_get_miss_returns_none(cache_conn): + from backend.apps.twitter.cache import TTLCache + + c = TTLCache(cache_conn) + assert c.get(("x",)) is None + + +def test_cache_set_then_get(cache_conn): + from backend.apps.twitter.cache import TTLCache + + c = TTLCache(cache_conn) + c.set(("user", "openai"), {"id": "123"}, ttl=60) + assert c.get(("user", "openai")) == {"id": "123"} + + +def test_cache_expiry(cache_conn): + """An entry with a past expiry should miss and be evicted from memory.""" + from backend.apps.twitter.cache import TTLCache + + c = TTLCache(cache_conn) + c.set(("x",), 1, ttl=60) + # Force expiry by rewriting the in-memory entry. + c._mem[c._mem and list(c._mem.keys())[0]] = (time.time() - 1, 1) + assert c.get(("x",)) is None + + +def test_cache_warm_from_disk_drops_expired(cache_conn): + from backend.apps.twitter.cache import TTLCache, _normalize_key + + # Pre-populate with one fresh, one stale entry. + cache_conn.execute( + "INSERT INTO twitter_cache (key, value_json, expires_at) VALUES (?, ?, ?)", + (_normalize_key(("fresh",)), '{"a": 1}', time.time() + 60), + ) + cache_conn.execute( + "INSERT INTO twitter_cache (key, value_json, expires_at) VALUES (?, ?, ?)", + (_normalize_key(("stale",)), '{"a": 2}', time.time() - 60), + ) + cache_conn.commit() + + c = TTLCache(cache_conn) + assert c.get(("fresh",)) == {"a": 1} + assert c.get(("stale",)) is None + # The stale row should have been deleted from disk on warm-up. + n = cache_conn.execute("SELECT COUNT(*) FROM twitter_cache").fetchone()[0] + assert n == 1 + + +def test_cache_invalidate(cache_conn): + from backend.apps.twitter.cache import TTLCache + + c = TTLCache(cache_conn) + c.set(("x",), 42, ttl=60) + c.invalidate(("x",)) + assert c.get(("x",)) is None + + +# --------------------------------------------------------------------------- +# RateGate +# --------------------------------------------------------------------------- + +@dataclass +class _FakeAccount: + id: str = "acct1" + client: object = field(default_factory=lambda: object()) + concurrency: asyncio.Semaphore = field(default_factory=lambda: asyncio.Semaphore(1)) + _buckets: dict = field(default_factory=dict) + + def bucket(self, endpoint: str): + from backend.apps.twitter.ratelimit import Bucket, DEFAULT_BUDGETS + if endpoint not in self._buckets: + self._buckets[endpoint] = Bucket(capacity=DEFAULT_BUDGETS.get(endpoint, 10)) + return self._buckets[endpoint] + + +@dataclass +class _FakePool: + """Minimal AccountPool stand-in for gate tests. + + Tracks side-effects (mark_locked / mark_suspended / etc.) as flags + so tests can assert the gate routed each twikit error correctly. + """ + + account: _FakeAccount | None = field(default_factory=_FakeAccount) + locked: list[str] = field(default_factory=list) + suspended: list[str] = field(default_factory=list) + relogin: list[str] = field(default_factory=list) + rate_limited_log: list[tuple[str, str]] = field(default_factory=list) + + async def pick(self, endpoint: str): + return self.account + + def record_429(self, acct_id: str, endpoint: str) -> None: + self.rate_limited_log.append((acct_id, endpoint)) + + def mark_locked(self, acct_id: str, reason: str) -> None: + self.locked.append(acct_id) + + def mark_suspended(self, acct_id: str, reason: str) -> None: + self.suspended.append(acct_id) + + def mark_needs_relogin(self, acct_id: str, reason: str) -> None: + self.relogin.append(acct_id) + + +class _FakeCache: + """In-memory cache with .get/.set shape compatible with TTLCache.""" + def __init__(self) -> None: + self.store: dict[tuple, Any] = {} + + def get(self, key): + return self.store.get(key) + + def set(self, key, value, *, ttl): + self.store[key] = value + + +def test_gate_cache_hit_short_circuits_bucket(): + """A cache hit should never touch pool.pick or twikit.""" + from backend.apps.twitter.ratelimit import RateGate + + cache = _FakeCache() + cache.store[("k",)] = {"hit": True} + pool = _FakePool(account=None) # would explode if asked + gate = RateGate(pool, cache) + + async def op(_client): + raise AssertionError("op must not be called on cache hit") + + async def _run(): + return await gate.execute( + endpoint="search_tweet", + op=op, + serializer=lambda x: x, + cache_key=("k",), + cache_ttl=60, + ) + + result = asyncio.run(_run()) + assert result.outcome == "ok" + assert result.value == {"hit": True} + + +def test_gate_no_account_returns_no_account(): + from backend.apps.twitter.ratelimit import RateGate + + pool = _FakePool(account=None) + gate = RateGate(pool, _FakeCache()) + + async def op(_client): + return None + + async def _run(): + return await gate.execute( + endpoint="search_tweet", + op=op, + serializer=lambda x: x, + ) + + result = asyncio.run(_run()) + assert result.outcome == "no_account" + + +def test_gate_returns_rate_limited_when_wait_exceeds_ceiling(monkeypatch): + """If the bucket says we'd wait > block_ceiling_s, return 429 instead of blocking.""" + from backend.apps.twitter.ratelimit import RateGate + + pool = _FakePool() + gate = RateGate(pool, _FakeCache(), block_ceiling_s=5.0) + # Force the search bucket to be locked far in the future. + bucket = pool.account.bucket("search_tweet") + fake_wall = 1000.0 + monkeypatch.setattr("backend.apps.twitter.ratelimit.time.time", lambda: fake_wall) + bucket.locked_until = fake_wall + 120.0 + bucket.tokens = 0.0 + + async def op(_client): + raise AssertionError("op must not be called when ceiling exceeded") + + async def _run(): + return await gate.execute( + endpoint="search_tweet", + op=op, + serializer=lambda x: x, + ) + + result = asyncio.run(_run()) + assert result.outcome == "rate_limited" + assert result.value["retry_after_s"] >= 120 + + +def test_gate_ok_path_calls_serializer_and_caches(): + from backend.apps.twitter.ratelimit import RateGate + + pool = _FakePool() + cache = _FakeCache() + gate = RateGate(pool, cache) + + async def op(_client): + return "raw" + + async def _run(): + return await gate.execute( + endpoint="search_tweet", + op=op, + serializer=lambda x: {"serialized": x}, + cache_key=("k",), + cache_ttl=60, + ) + + result = asyncio.run(_run()) + assert result.outcome == "ok" + assert result.value == {"serialized": "raw"} + assert cache.store == {("k",): {"serialized": "raw"}} + + +def test_gate_does_not_cache_errors(): + """Error outcomes must not leave anything in the cache.""" + from backend.apps.twitter.ratelimit import RateGate + from twikit.errors import TooManyRequests + + pool = _FakePool() + cache = _FakeCache() + gate = RateGate(pool, cache) + + err = TooManyRequests("rate limited", headers={"x-rate-limit-reset": str(int(time.time() + 30))}) + + async def op(_client): + raise err + + async def _run(): + return await gate.execute( + endpoint="search_tweet", + op=op, + serializer=lambda x: x, + cache_key=("k",), + cache_ttl=60, + ) + + result = asyncio.run(_run()) + assert result.outcome == "rate_limited" + assert cache.store == {}, "errors must never end up in the cache" + assert pool.rate_limited_log == [("acct1", "search_tweet")] + + +def test_gate_account_locked_routes_to_locked_outcome(): + from backend.apps.twitter.ratelimit import RateGate + from twikit.errors import AccountLocked + + pool = _FakePool() + gate = RateGate(pool, _FakeCache()) + + async def op(_client): + raise AccountLocked("arkose") + + async def _run(): + return await gate.execute( + endpoint="search_tweet", + op=op, + serializer=lambda x: x, + ) + + result = asyncio.run(_run()) + assert result.outcome == "locked" + assert pool.locked == ["acct1"] + + +def test_gate_account_suspended_routes_to_suspended_outcome(): + from backend.apps.twitter.ratelimit import RateGate + from twikit.errors import AccountSuspended + + pool = _FakePool() + gate = RateGate(pool, _FakeCache()) + + async def op(_client): + raise AccountSuspended("suspended") + + async def _run(): + return await gate.execute( + endpoint="search_tweet", + op=op, + serializer=lambda x: x, + ) + + assert asyncio.run(_run()).outcome == "suspended" + assert pool.suspended == ["acct1"] + + +def test_gate_unauthorized_routes_to_needs_relogin(): + from backend.apps.twitter.ratelimit import RateGate + from twikit.errors import Unauthorized + + pool = _FakePool() + gate = RateGate(pool, _FakeCache()) + + async def op(_client): + raise Unauthorized("expired cookies") + + async def _run(): + return await gate.execute( + endpoint="search_tweet", + op=op, + serializer=lambda x: x, + ) + + assert asyncio.run(_run()).outcome == "needs_relogin" + assert pool.relogin == ["acct1"] diff --git a/backend/tests/test_twitter_routes.py b/backend/tests/test_twitter_routes.py new file mode 100644 index 00000000..d6ea44f7 --- /dev/null +++ b/backend/tests/test_twitter_routes.py @@ -0,0 +1,414 @@ +"""Integration tests for the Twitter SubApp's HTTP routes. + +We don't need (and won't get) a live twikit session in CI; instead we +build the SubApp's state by hand and exercise the routes via +FastAPI's TestClient. The pool's `client` field is a Mock with +AsyncMock methods, so we can drive `TooManyRequests` / +`AccountLocked` / `Unauthorized` paths deterministically. + +These tests cover the load-bearing route-layer behaviors: + +- The full GateResult -> HTTP status mapping (429 with retry_after_s, + 409 for locked/suspended/needs_relogin, 503 for no_account). +- DELETE wipes cookies + audit log + accounts.json. +- /health is pure-memory (no twikit call) so the test never has to + mock anything to call it. +- Errors are not cached (a 429 followed by a success must still hit + twikit on the second call). +""" + +from __future__ import annotations + +import json +import os +import tempfile +import time +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + + +@pytest.fixture +def isolated_subapp(monkeypatch): + """Stand up the twitter SubApp's state against a temp DATA_ROOT. + + We avoid `backend.main` entirely so we don't pull in the full + OpenSwarm app (and its 15-second startup time). Instead we mount + the SubApp's router directly on a fresh FastAPI app. + """ + with tempfile.TemporaryDirectory() as d: + monkeypatch.setattr("backend.apps.twitter.persistence.TWITTER_DIR", d) + monkeypatch.setattr( + "backend.apps.twitter.persistence.ACCOUNTS_PATH", + os.path.join(d, "accounts.json"), + ) + monkeypatch.setattr( + "backend.apps.twitter.persistence.COOKIES_DIR", + os.path.join(d, "cookies"), + ) + monkeypatch.setattr( + "backend.apps.twitter.persistence.STATE_DB_PATH", + os.path.join(d, "state.sqlite"), + ) + + # Build the SubApp's singletons inline (skip the lifespan path + # which hydrates from disk + spawns the snapshot task). + from backend.apps.twitter import persistence, twitter as tw + from backend.apps.twitter.cache import TTLCache + from backend.apps.twitter.pool import AccountPool + from backend.apps.twitter.ratelimit import RateGate + + persistence.ensure_dirs() + conn = persistence.open_state_db() + tw._cache = TTLCache(conn) + tw._pool = AccountPool(conn) + tw._gate = RateGate(tw._pool, tw._cache, block_ceiling_s=10.0) + + app = FastAPI() + app.include_router(tw.twitter.router, prefix="/api/twitter") + client = TestClient(app) + + yield { + "client": client, + "pool": tw._pool, + "cache": tw._cache, + "gate": tw._gate, + "conn": conn, + "tmp": d, + } + + # Reset module-level singletons so tests don't bleed state. + tw._pool = None + tw._gate = None + tw._cache = None + conn.close() + + +def _add_fake_account(pool, account_id="a1", state="active", handle="me"): + """Register a Mock-backed account on the pool. + + The twikit.Client surface used by routes is async (`search_tweet`, + `get_user_by_screen_name`, etc.), so AsyncMock is the right shape. + `save_cookies` and `set_cookies` are sync on the real client; use + plain MagicMock for those. + """ + from backend.apps.twitter.models import TwitterAccount + + record = TwitterAccount(id=account_id, state=state, handle=handle, label="Test") + + client = MagicMock(name="twikit.Client") + client.search_tweet = AsyncMock() + client.get_user_by_screen_name = AsyncMock() + client.get_user_by_id = AsyncMock() + client.get_user_tweets = AsyncMock() + client.get_tweet_by_id = AsyncMock() + client.user = AsyncMock() + client.save_cookies = MagicMock() + + import asyncio + # `pool.add` is async because it acquires asyncio.Lock; spin a + # fresh loop just for this call so we don't lean on the deprecated + # get_event_loop() implicit-loop behavior. + loop = asyncio.new_event_loop() + try: + loop.run_until_complete(pool.add(record, client)) + finally: + loop.close() + return record, client + + +# --------------------------------------------------------------------------- +# /accounts/{id}/health — pure-memory; no twikit calls needed +# --------------------------------------------------------------------------- + +def test_health_returns_in_memory_state(isolated_subapp): + """Health is the keep-it-cheap endpoint; should not call twikit.""" + _add_fake_account(isolated_subapp["pool"], account_id="a1") + r = isolated_subapp["client"].get("/api/twitter/accounts/a1/health") + assert r.status_code == 200 + body = r.json() + assert body["id"] == "a1" + assert body["state"] == "active" + assert body["recent_429_count"] == 0 + + +def test_health_missing_returns_404(isolated_subapp): + r = isolated_subapp["client"].get("/api/twitter/accounts/missing/health") + assert r.status_code == 404 + + +def test_accounts_list_excludes_cookies(isolated_subapp): + _add_fake_account(isolated_subapp["pool"]) + r = isolated_subapp["client"].get("/api/twitter/accounts") + assert r.status_code == 200 + body = r.json() + assert len(body["accounts"]) == 1 + # No password / cookies / sensitive fields should ever appear. + flat = json.dumps(body) + assert "password" not in flat.lower() + assert "auth_token" not in flat.lower() + assert "ct0" not in flat.lower() + + +# --------------------------------------------------------------------------- +# Tool reads — happy path +# --------------------------------------------------------------------------- + +def _fake_tweet(tweet_id="1"): + """Build a duck-typed tweet that survives the serializer's _safe().""" + m = MagicMock() + m.id = tweet_id + m.created_at = "2024-01-01" + m.text = "hi" + m.lang = "en" + m.in_reply_to = None + m.is_quote_status = False + m.possibly_sensitive = False + m.view_count = 100 + m.reply_count = 0 + m.favorite_count = 0 + m.retweet_count = 0 + m.quote_count = 0 + m.bookmark_count = 0 + m.hashtags = [] + m.urls = [] + m.media = [] + m.user = None + m.quote = None + m.retweeted_tweet = None + m.replies = None + return m + + +def _fake_result(items): + """Mimic twikit.utils.Result enough for `result_to_dict`.""" + r = MagicMock() + r.__iter__ = lambda self: iter(items) + r.next_cursor = "next-x" + r.previous_cursor = None + return r + + +def test_search_happy_path(isolated_subapp): + _, client = _add_fake_account(isolated_subapp["pool"]) + client.search_tweet.return_value = _fake_result([_fake_tweet("1"), _fake_tweet("2")]) + + r = isolated_subapp["client"].get("/api/twitter/search?q=hello&count=2") + assert r.status_code == 200 + body = r.json() + assert [t["id"] for t in body["items"]] == ["1", "2"] + assert body["next_cursor"] == "next-x" + + +def test_search_uses_cache_on_second_call(isolated_subapp): + _, client = _add_fake_account(isolated_subapp["pool"]) + client.search_tweet.return_value = _fake_result([_fake_tweet("1")]) + + api = isolated_subapp["client"] + api.get("/api/twitter/search?q=hello&count=2") + api.get("/api/twitter/search?q=hello&count=2") + assert client.search_tweet.await_count == 1, "cache hit must skip twikit" + + +# --------------------------------------------------------------------------- +# Error -> HTTP mapping +# --------------------------------------------------------------------------- + +def test_too_many_requests_returns_429_with_retry_after_s(isolated_subapp): + """TooManyRequests from twikit -> HTTP 429 + retry_after_s body.""" + from twikit.errors import TooManyRequests + + _, client = _add_fake_account(isolated_subapp["pool"]) + reset_at = int(time.time() + 90) + client.search_tweet.side_effect = TooManyRequests( + "rate limited", + headers={"x-rate-limit-reset": str(reset_at)}, + ) + + r = isolated_subapp["client"].get("/api/twitter/search?q=hi") + assert r.status_code == 429 + body = r.json() + assert "retry_after_s" in body + # Should be in the ballpark of 90 (server told us ~90s). + assert 80 <= body["retry_after_s"] <= 100 + + +def test_429_is_not_cached(isolated_subapp): + """Errors must never end up memoized. + + Sequence: first call -> 429; bucket clears; second call -> success. + If the cache held the 429, the second call would also 429. + """ + from twikit.errors import TooManyRequests + + _, client = _add_fake_account(isolated_subapp["pool"]) + client.search_tweet.side_effect = [ + TooManyRequests("rate limited", headers={"x-rate-limit-reset": str(int(time.time() + 1))}), + _fake_result([_fake_tweet("1")]), + ] + + api = isolated_subapp["client"] + r1 = api.get("/api/twitter/search?q=cached-test") + assert r1.status_code == 429 + + # Clear the bucket lock so the second call can go through. + pool = isolated_subapp["pool"] + pool.get("a1").bucket("search_tweet").locked_until = 0.0 + pool.get("a1").bucket("search_tweet").tokens = 5.0 + + r2 = api.get("/api/twitter/search?q=cached-test") + assert r2.status_code == 200 + # Both calls must hit twikit — the 429 didn't poison the cache. + assert client.search_tweet.await_count == 2 + + +def test_account_locked_returns_409(isolated_subapp): + from twikit.errors import AccountLocked + + _, client = _add_fake_account(isolated_subapp["pool"]) + client.search_tweet.side_effect = AccountLocked("arkose") + r = isolated_subapp["client"].get("/api/twitter/search?q=x") + assert r.status_code == 409 + # Pool state should reflect. + assert isolated_subapp["pool"].get("a1").state == "locked" + + +def test_unauthorized_returns_409_and_marks_needs_relogin(isolated_subapp): + from twikit.errors import Unauthorized + + _, client = _add_fake_account(isolated_subapp["pool"]) + client.search_tweet.side_effect = Unauthorized("expired") + r = isolated_subapp["client"].get("/api/twitter/search?q=x") + assert r.status_code == 409 + assert isolated_subapp["pool"].get("a1").state == "needs_relogin" + + +def test_no_account_returns_503(isolated_subapp): + """No accounts in the pool -> 503 (we don't have anyone to ask).""" + r = isolated_subapp["client"].get("/api/twitter/search?q=x") + assert r.status_code == 503 + + +def test_inactive_account_treated_as_no_account(isolated_subapp): + _, _client = _add_fake_account(isolated_subapp["pool"], state="locked") + r = isolated_subapp["client"].get("/api/twitter/search?q=x") + # pick() returns None because no active accounts -> 503. + assert r.status_code == 503 + + +# --------------------------------------------------------------------------- +# /user — handle vs id +# --------------------------------------------------------------------------- + +def test_get_user_by_handle(isolated_subapp): + _, client = _add_fake_account(isolated_subapp["pool"]) + fake_user = MagicMock() + fake_user.id = "999" + fake_user.screen_name = "openai" + fake_user.name = "OpenAI" + fake_user.description = "" + fake_user.location = "" + fake_user.url = "" + fake_user.profile_image_url = "" + fake_user.profile_banner_url = "" + fake_user.created_at = "" + fake_user.is_blue_verified = False + fake_user.verified = False + fake_user.followers_count = 1 + fake_user.following_count = 2 + fake_user.statuses_count = 3 + fake_user.media_count = 4 + fake_user.listed_count = 5 + fake_user.favourites_count = 6 + fake_user.pinned_tweet_ids = [] + client.get_user_by_screen_name.return_value = fake_user + + r = isolated_subapp["client"].get("/api/twitter/user?handle=openai") + assert r.status_code == 200 + assert r.json()["handle"] == "openai" + + +def test_get_user_requires_exactly_one_arg(isolated_subapp): + api = isolated_subapp["client"] + assert api.get("/api/twitter/user").status_code == 400 + assert api.get("/api/twitter/user?handle=a&id=b").status_code == 400 + + +# --------------------------------------------------------------------------- +# PATCH /accounts/{id} — trust_multiplier +# --------------------------------------------------------------------------- + +def test_patch_trust_multiplier_rescales_buckets(isolated_subapp): + _add_fake_account(isolated_subapp["pool"]) + # Force-create the search_tweet bucket so the rescale has something to do. + isolated_subapp["pool"].get("a1").bucket("search_tweet") + + api = isolated_subapp["client"] + r = api.patch("/api/twitter/accounts/a1", json={"trust_multiplier": 0.2}) + assert r.status_code == 200 + pool = isolated_subapp["pool"] + b = pool.get("a1").bucket("search_tweet") + # 50 * 0.2 = 10. + assert b.capacity == 10 + + +def test_patch_trust_multiplier_zero_pauses_account(isolated_subapp): + """trust_multiplier=0 is allowed and zeros the bucket capacity. + + Operators reach for this when an account starts misbehaving and they + want to leave it in the pool (cookies, audit history) without it + being picked. Bucket.time_until_available() returns a long sentinel + in that state so pick() naturally deprioritizes it. + """ + _add_fake_account(isolated_subapp["pool"]) + isolated_subapp["pool"].get("a1").bucket("search_tweet") + + api = isolated_subapp["client"] + r = api.patch("/api/twitter/accounts/a1", json={"trust_multiplier": 0.0}) + assert r.status_code == 200 + b = isolated_subapp["pool"].get("a1").bucket("search_tweet") + assert b.capacity == 0 + + +def test_patch_trust_multiplier_validates_range(isolated_subapp): + _add_fake_account(isolated_subapp["pool"]) + api = isolated_subapp["client"] + assert api.patch("/api/twitter/accounts/a1", json={"trust_multiplier": 1.5}).status_code == 422 + assert api.patch("/api/twitter/accounts/a1", json={"trust_multiplier": -0.1}).status_code == 422 + + +# --------------------------------------------------------------------------- +# DELETE — wipes everything +# --------------------------------------------------------------------------- + +def test_delete_wipes_cookies_and_buckets(isolated_subapp): + from backend.apps.twitter import persistence as pers + + _add_fake_account(isolated_subapp["pool"]) + # Drop a fake cookies file so DELETE has something to wipe. + cpath = pers.cookies_path("a1") + with open(cpath, "w") as f: + f.write("{}") + + # Touch a bucket + snapshot so we have a row to clean up. + pool = isolated_subapp["pool"] + pool.get("a1").bucket("search_tweet") + pool.snapshot_all() + + api = isolated_subapp["client"] + r = api.delete("/api/twitter/accounts/a1") + assert r.status_code == 200 + assert pool.get("a1") is None + assert not os.path.exists(cpath) + (n,) = isolated_subapp["conn"].execute( + "SELECT COUNT(*) FROM twitter_buckets WHERE account_id='a1'" + ).fetchone() + assert n == 0 + + +def test_delete_missing_account_is_noop(isolated_subapp): + r = isolated_subapp["client"].delete("/api/twitter/accounts/never-existed") + assert r.status_code == 200 + assert r.json()["removed"] is True diff --git a/backend/tests/test_twitter_shim.py b/backend/tests/test_twitter_shim.py new file mode 100644 index 00000000..39d72ea2 --- /dev/null +++ b/backend/tests/test_twitter_shim.py @@ -0,0 +1,329 @@ +"""Shim tests: JSON-RPC framing + HTTP-to-backend translation. + +We don't spawn a subprocess; instead we patch `urllib.request.urlopen` +to fake the backend's HTTP responses and call `handle_tool_call` / +`main()` directly. That keeps the test fast and lets us assert the +shim's response shape (MCP content/isError) deterministically. + +Critical behaviors covered: + +- Auth-token reading: shim picks up a token rotation between calls + (re-reads the file rather than caching forever). +- HTTP 429 -> MCP error content carrying retry_after_s. +- HTTP 409 (account locked/needs_relogin) -> MCP error. +- 200 with payload -> non-error MCP content. +- Missing required arg -> early MCP error without calling the backend. +- The JSON-RPC framing for initialize/tools/list/tools/call matches + the shape mcp clients expect. +""" + +from __future__ import annotations + +import io +import json +import os +import sys +import tempfile +from unittest.mock import patch + +import pytest + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def fresh_shim(monkeypatch): + """Re-import the shim with a chosen env so each test gets clean state. + + `OPENSWARM_AUTH_TOKEN_FILE` etc. are read at module-import time; + re-importing forces the env to take effect. + """ + with tempfile.TemporaryDirectory() as tmp: + token_file = os.path.join(tmp, "auth.token") + with open(token_file, "w") as f: + f.write("token-v1") + monkeypatch.setenv("OPENSWARM_AUTH_TOKEN_FILE", token_file) + monkeypatch.setenv("OPENSWARM_AUTH_TOKEN", "") # avoid env fallback masking the file test + monkeypatch.setenv("OPENSWARM_BASE_URL", "http://127.0.0.1:8324") + + if "backend.apps.twitter_mcp_shim.server" in sys.modules: + del sys.modules["backend.apps.twitter_mcp_shim.server"] + import backend.apps.twitter_mcp_shim.server as server # noqa: E402 + + yield {"server": server, "token_file": token_file} + + +def _fake_urlopen(status: int, body): + """Build a context-manager-style mock matching urlopen()'s return shape.""" + class _Resp: + def __init__(self): + self.status = status + payload = json.dumps(body).encode() if not isinstance(body, (bytes, str)) else ( + body.encode() if isinstance(body, str) else body + ) + self._payload = payload + + def read(self): + return self._payload + + def __enter__(self): + return self + + def __exit__(self, *_a): + return False + + return _Resp() + + +def _fake_http_error(code: int, body): + """Build a `urllib.error.HTTPError` mimicking a non-2xx response.""" + import urllib.error + + payload = json.dumps(body).encode() if not isinstance(body, str) else body.encode() + fp = io.BytesIO(payload) + return urllib.error.HTTPError( + url="http://x/", code=code, msg="error", hdrs=None, fp=fp, + ) + + +# --------------------------------------------------------------------------- +# Token plumbing +# --------------------------------------------------------------------------- + +def test_token_read_from_file(fresh_shim): + server = fresh_shim["server"] + assert server._read_token() == "token-v1" + + +def test_token_rotation_picked_up_after_cache_ttl(fresh_shim, monkeypatch): + """After the cache TTL, the shim should re-read the file.""" + server = fresh_shim["server"] + assert server._read_token() == "token-v1" + + with open(fresh_shim["token_file"], "w") as f: + f.write("token-v2") + + # Fast-forward past the cache TTL. + import time as t + monkeypatch.setattr(server, "_token_cache", (t.time() - 100.0, "token-v1")) + assert server._read_token() == "token-v2" + + +def test_token_falls_back_to_env_when_file_missing(monkeypatch): + monkeypatch.setenv("OPENSWARM_AUTH_TOKEN_FILE", "/nope/never/exists.token") + monkeypatch.setenv("OPENSWARM_AUTH_TOKEN", "env-only-token") + if "backend.apps.twitter_mcp_shim.server" in sys.modules: + del sys.modules["backend.apps.twitter_mcp_shim.server"] + import backend.apps.twitter_mcp_shim.server as server # noqa: E402 + assert server._read_token() == "env-only-token" + + +# --------------------------------------------------------------------------- +# Tool dispatch -> HTTP +# --------------------------------------------------------------------------- + +def test_search_happy_path(fresh_shim): + server = fresh_shim["server"] + with patch.object(server.urllib.request, "urlopen") as urlopen_mock: + urlopen_mock.return_value = _fake_urlopen(200, {"items": [{"id": "1"}], "next_cursor": None}) + result = server.handle_tool_call("twitter_search", {"q": "hello"}) + assert "isError" not in result, result + # Payload arrives as JSON string in `content[0].text`. + text = result["content"][0]["text"] + parsed = json.loads(text) + assert parsed["items"][0]["id"] == "1" + + +def test_search_missing_q_is_local_error(fresh_shim): + server = fresh_shim["server"] + with patch.object(server.urllib.request, "urlopen") as urlopen_mock: + result = server.handle_tool_call("twitter_search", {"q": " "}) + assert result.get("isError") is True + # Backend must not be called when required args are missing. + urlopen_mock.assert_not_called() + + +def test_get_user_requires_exactly_one_arg(fresh_shim): + server = fresh_shim["server"] + with patch.object(server.urllib.request, "urlopen") as urlopen_mock: + r1 = server.handle_tool_call("twitter_get_user", {}) + r2 = server.handle_tool_call("twitter_get_user", {"handle": "x", "user_id": "1"}) + assert r1.get("isError") is True + assert r2.get("isError") is True + urlopen_mock.assert_not_called() + + +def test_get_user_strips_at_prefix(fresh_shim): + server = fresh_shim["server"] + with patch.object(server.urllib.request, "urlopen") as urlopen_mock: + urlopen_mock.return_value = _fake_urlopen(200, {"handle": "openai"}) + server.handle_tool_call("twitter_get_user", {"handle": "@openai"}) + # Assert we sent `?handle=openai`, not `?handle=@openai`. + sent_req = urlopen_mock.call_args[0][0] + assert "handle=openai" in sent_req.full_url + assert "%40" not in sent_req.full_url + + +# --------------------------------------------------------------------------- +# HTTP status -> MCP response mapping +# --------------------------------------------------------------------------- + +def test_429_surfaces_retry_after(fresh_shim): + """The shim must translate {retry_after_s} into agent-readable text.""" + server = fresh_shim["server"] + with patch.object(server.urllib.request, "urlopen") as urlopen_mock: + urlopen_mock.side_effect = _fake_http_error(429, { + "retry_after_s": 47, + "endpoint": "search_tweet", + }) + result = server.handle_tool_call("twitter_search", {"q": "hi"}) + assert result.get("isError") is True + text = result["content"][0]["text"] + assert "47 seconds" in text + assert "search_tweet" in text + + +def test_409_surfaces_account_state_error(fresh_shim): + server = fresh_shim["server"] + with patch.object(server.urllib.request, "urlopen") as urlopen_mock: + urlopen_mock.side_effect = _fake_http_error(409, {"error": "Account is locked"}) + result = server.handle_tool_call("twitter_search", {"q": "hi"}) + assert result.get("isError") is True + assert "locked" in result["content"][0]["text"].lower() + + +def test_503_surfaces_no_account_error(fresh_shim): + server = fresh_shim["server"] + with patch.object(server.urllib.request, "urlopen") as urlopen_mock: + urlopen_mock.side_effect = _fake_http_error(503, {"error": "No active Twitter account"}) + result = server.handle_tool_call("twitter_search", {"q": "hi"}) + assert result.get("isError") is True + assert "no active" in result["content"][0]["text"].lower() + + +def test_backend_unreachable_surfaces_clean_error(fresh_shim): + """If urlopen raises URLError, we should return an MCP error, not crash.""" + import urllib.error + server = fresh_shim["server"] + with patch.object(server.urllib.request, "urlopen") as urlopen_mock: + urlopen_mock.side_effect = urllib.error.URLError("connection refused") + result = server.handle_tool_call("twitter_search", {"q": "hi"}) + assert result.get("isError") is True + + +def test_401_triggers_single_token_refresh_retry(fresh_shim): + """A 401 should bust the token cache and retry exactly once. + + The cache TTL is 5s; without retry the shim would 401-loop for + that long after a backend restart rotated the token. The test + rotates the file mid-call and checks that the second urlopen got + the new token. + """ + server = fresh_shim["server"] + + # Sequence: first call -> 401, second call -> 200. Capture + # Authorization headers from each Request to confirm we re-read. + seen_tokens = [] + + def fake_urlopen(req, *_a, **_kw): + seen_tokens.append(req.headers.get("Authorization")) + if len(seen_tokens) == 1: + raise _fake_http_error(401, {"error": "unauthorized"}) + return _fake_urlopen(200, {"ok": True}) + + # Rotate the on-disk token between the two calls. The first + # `_read_token()` already populated the cache with "token-v1". + # When we 401, the cache is busted, so the second `_read_token()` + # reads the rotated file. + with patch.object(server.urllib.request, "urlopen", side_effect=fake_urlopen): + # Pre-warm the cache so the first call carries token-v1. + assert server._read_token() == "token-v1" + with open(fresh_shim["token_file"], "w") as f: + f.write("token-v2") + result = server.handle_tool_call("twitter_search", {"q": "hi"}) + + assert "isError" not in result, result + assert seen_tokens == ["Bearer token-v1", "Bearer token-v2"] + + +def test_401_does_not_retry_more_than_once(fresh_shim): + """If the token rotation didn't fix the 401, surface it (don't loop).""" + server = fresh_shim["server"] + n_calls = [] + + def fake_urlopen(req, *_a, **_kw): + n_calls.append(1) + raise _fake_http_error(401, {"error": "unauthorized"}) + + with patch.object(server.urllib.request, "urlopen", side_effect=fake_urlopen): + result = server.handle_tool_call("twitter_search", {"q": "hi"}) + + assert result.get("isError") is True + # Two attempts: original + one retry. Not three. + assert len(n_calls) == 2 + + +# --------------------------------------------------------------------------- +# JSON-RPC framing +# --------------------------------------------------------------------------- + +def _drive_stdio(server, requests: list[dict]) -> list[dict]: + """Feed a sequence of JSON-RPC frames through `main()` and capture replies.""" + stdin = io.StringIO("\n".join(json.dumps(r) for r in requests) + "\n") + stdout = io.StringIO() + with patch.object(server.sys, "stdin", stdin), patch.object(server.sys, "stdout", stdout): + server.main() + out = stdout.getvalue().splitlines() + return [json.loads(line) for line in out if line.strip()] + + +def test_stdio_initialize_returns_protocol_metadata(fresh_shim): + server = fresh_shim["server"] + replies = _drive_stdio(server, [{"jsonrpc": "2.0", "id": 1, "method": "initialize"}]) + assert len(replies) == 1 + assert replies[0]["id"] == 1 + assert replies[0]["result"]["serverInfo"]["name"] == "openswarm-twitter" + assert "tools" in replies[0]["result"]["capabilities"] + + +def test_stdio_tools_list_returns_five_tools(fresh_shim): + server = fresh_shim["server"] + replies = _drive_stdio(server, [{"jsonrpc": "2.0", "id": 2, "method": "tools/list"}]) + names = [t["name"] for t in replies[0]["result"]["tools"]] + assert names == [ + "twitter_search", + "twitter_get_user", + "twitter_get_user_tweets", + "twitter_get_tweet", + "twitter_get_tweet_replies", + ] + + +def test_stdio_tools_call_routes_through_handle(fresh_shim): + server = fresh_shim["server"] + with patch.object(server.urllib.request, "urlopen") as urlopen_mock: + urlopen_mock.return_value = _fake_urlopen(200, {"items": []}) + replies = _drive_stdio(server, [{ + "jsonrpc": "2.0", + "id": 5, + "method": "tools/call", + "params": {"name": "twitter_search", "arguments": {"q": "hi"}}, + }]) + assert replies[0]["id"] == 5 + assert "content" in replies[0]["result"] + assert "isError" not in replies[0]["result"] + + +def test_stdio_unknown_method_returns_jsonrpc_error(fresh_shim): + server = fresh_shim["server"] + replies = _drive_stdio(server, [{"jsonrpc": "2.0", "id": 9, "method": "bogus"}]) + assert replies[0]["error"]["code"] == -32601 + + +def test_stdio_ping_works(fresh_shim): + """MCP clients ping to keep the connection alive.""" + server = fresh_shim["server"] + replies = _drive_stdio(server, [{"jsonrpc": "2.0", "id": 10, "method": "ping"}]) + assert replies[0] == {"jsonrpc": "2.0", "id": 10, "result": {}} diff --git a/backend/tests/test_twitter_twikit_patches.py b/backend/tests/test_twitter_twikit_patches.py new file mode 100644 index 00000000..319b94f8 --- /dev/null +++ b/backend/tests/test_twitter_twikit_patches.py @@ -0,0 +1,275 @@ +"""Tests for the runtime twikit patches. + +These guard the regex + header changes we apply at import time as a workaround +for twikit#408. They do *not* hit the network; everything is exercised against +synthetic minified-JS snippets crafted to mimic the pre- and post-rotation +shapes that X has shipped. + +If twikit ships a real fix and we delete `_twikit_patches.py`, these tests +should be deleted with it. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +from backend.apps.twitter import _twikit_patches + +# Capture pristine twikit state *before* any test triggers the patch. We use +# these to fully unwind monkey-patched classes between tests so each test +# starts with a clean twikit, exercises `apply()`, and leaves no residue for +# the next one. Without this the autouse fixture's re-apply on teardown would +# leak Client.__init__ wrapping into the next test, breaking the +# "patch is disabled" assertions. +from twikit.client.client import Client as _Client # noqa: E402 +from twikit.x_client_transaction import transaction as _tx_mod # noqa: E402 + +_PRISTINE_CLIENT_INIT = _Client.__init__ +_PRISTINE_CLIENT_BASE_HEADERS = _Client._base_headers +_PRISTINE_TX_GET_INDICES = _tx_mod.ClientTransaction.get_indices +_PRISTINE_INDICES_REGEX = _tx_mod.INDICES_REGEX +_PRISTINE_ON_DEMAND_FILE_REGEX = _tx_mod.ON_DEMAND_FILE_REGEX + + +def _restore_pristine_twikit() -> None: + _Client.__init__ = _PRISTINE_CLIENT_INIT + _Client._base_headers = _PRISTINE_CLIENT_BASE_HEADERS + _tx_mod.ClientTransaction.get_indices = _PRISTINE_TX_GET_INDICES + _tx_mod.INDICES_REGEX = _PRISTINE_INDICES_REGEX + _tx_mod.ON_DEMAND_FILE_REGEX = _PRISTINE_ON_DEMAND_FILE_REGEX + _twikit_patches._APPLIED_TX = False + _twikit_patches._APPLIED_HEADERS = False + + +@pytest.fixture(autouse=True) +def _reapply_patch(): + """Reset twikit to its pristine state for each test, then re-apply the + real patch on teardown so subsequent suites (and the running process) + see the same patched state production does.""" + _restore_pristine_twikit() + yield + _restore_pristine_twikit() + _twikit_patches.apply() + + +def test_apply_is_idempotent(): + """Calling apply twice should not raise and should report success both times.""" + assert _twikit_patches.apply() is True + assert _twikit_patches.apply() is True + + +def test_apply_skipped_when_both_env_disabled(monkeypatch): + monkeypatch.setenv("OPENSWARM_TWITTER_DISABLE_TWIKIT_PATCH", "1") + monkeypatch.setenv("OPENSWARM_TWITTER_DISABLE_TWIKIT_HEADER_PATCH", "1") + assert _twikit_patches.apply() is False + + +def test_apply_returns_true_when_only_transaction_disabled(monkeypatch): + """Disabling just one patch should still let the other apply.""" + monkeypatch.setenv("OPENSWARM_TWITTER_DISABLE_TWIKIT_PATCH", "1") + assert _twikit_patches.apply() is True + assert _twikit_patches._APPLIED_TX is False + assert _twikit_patches._APPLIED_HEADERS is True + + +def test_apply_returns_true_when_only_header_disabled(monkeypatch): + monkeypatch.setenv("OPENSWARM_TWITTER_DISABLE_TWIKIT_HEADER_PATCH", "1") + assert _twikit_patches.apply() is True + assert _twikit_patches._APPLIED_TX is True + assert _twikit_patches._APPLIED_HEADERS is False + + +def test_indices_regex_matches_post_rotation_two_char_var(): + """The new minified output uses 1-2 char variable names like `xx[NN]`. + Original regex (`\\w{1}`) missed this; our replacement (`\\w{1,2}`) must + capture both the byte index and the modulus marker.""" + _twikit_patches.apply() + from twikit.x_client_transaction import transaction as _tx + + sample = "...computeKey:function(xx){return(xx[13],16),(xx[14],16),(xx[7],16)}..." + matches = [m.group(2) for m in _tx.INDICES_REGEX.finditer(sample)] + assert matches == ["13", "14", "7"] + + +def test_indices_regex_still_matches_pre_rotation_single_char_var(): + """Backwards-compat: if X ever rolls back to one-char names, we should + keep working without another patch.""" + _twikit_patches.apply() + from twikit.x_client_transaction import transaction as _tx + + sample = "...(x[2],16),(x[42],16),(x[45],16)..." + matches = [m.group(2) for m in _tx.INDICES_REGEX.finditer(sample)] + assert matches == ["2", "42", "45"] + + +def test_on_demand_file_regex_extracts_chunk_index(): + """The new home-page chunk map embeds `,NN:"ondemand.s"`; our regex + pulls out NN so the follow-up hash lookup can find the right chunk.""" + _twikit_patches.apply() + from twikit.x_client_transaction import transaction as _tx + + sample = '...,99:"prev",964:"ondemand.s",100:"next",964:"deadbeef1234"...' + match = _tx.ON_DEMAND_FILE_REGEX.search(sample) + assert match is not None + assert match.group(1) == "964" + + +def test_on_demand_hash_pattern_resolves_to_chunk_hash(): + """The second-pass lookup must find the hash for the chunk index + returned by ON_DEMAND_FILE_REGEX.""" + import re + + _twikit_patches.apply() + from twikit.x_client_transaction import transaction as _tx + + sample = '...,99:"prev",964:"ondemand.s",100:"next",964:"deadbeef1234"...' + hash_re = re.compile(_tx.ON_DEMAND_HASH_PATTERN.format("964")) + match = hash_re.search(sample) + assert match is not None + assert match.group(1) == "deadbeef1234" + + +def test_get_indices_extracts_full_index_list_from_synthetic_payload(): + """End-to-end exercise of the patched `get_indices` against a fake bs4 + response + fake httpx session that returns a synthetic minified JS body + matching the post-rotation INDICES_REGEX shape.""" + + import bs4 + + _twikit_patches.apply() + from twikit.x_client_transaction import transaction as _tx + + home_html = ( + '' + '' + '' + ) + home_soup = bs4.BeautifulSoup(home_html, "html.parser") + + on_demand_js = "function k(xx){return(xx[2],16),(xx[12],16),(xx[7],16)}" + + class _FakeResponse: + def __init__(self, text: str) -> None: + self.text = text + + class _FakeSession: + def __init__(self, body: str) -> None: + self._body = body + self.calls: list[tuple[str, str]] = [] + + async def request(self, method: str, url: str, headers: Any) -> _FakeResponse: + self.calls.append((method, url)) + return _FakeResponse(self._body) + + session = _FakeSession(on_demand_js) + + ct = _tx.ClientTransaction() + ct.home_page_response = home_soup + + row_index, byte_indices = asyncio.run( + ct.get_indices(home_soup, session, headers={}) + ) + + assert row_index == 2 + assert byte_indices == [12, 7] + assert session.calls == [ + ( + "GET", + "https://abs.twimg.com/responsive-web/client-web/ondemand.s.abcdef0123a.js", + ) + ] + + +def test_get_indices_raises_legacy_exception_when_payload_empty(): + """If neither pass finds anything, surface the exact same exception + string the unpatched library raises. This keeps the SubApp's lifecycle + audit log ('Couldn't get KEY_BYTE indices') and the smoke probe's drift + detection both pointing at the right symptom.""" + + import bs4 + + _twikit_patches.apply() + from twikit.x_client_transaction import transaction as _tx + + home_soup = bs4.BeautifulSoup( + "", + "html.parser", + ) + + class _NeverCalledSession: + async def request(self, *args: Any, **kwargs: Any): # pragma: no cover + raise AssertionError("session should not be hit when home page has no chunk map") + + ct = _tx.ClientTransaction() + ct.home_page_response = home_soup + + with pytest.raises(Exception, match="Couldn't get KEY_BYTE indices"): + asyncio.run(ct.get_indices(home_soup, _NeverCalledSession(), headers={})) + + +# ---- header patch tests -------------------------------------------------- + + +def test_header_patch_swaps_default_user_agent(): + """A fresh Client with no UA kwarg should end up on the Chrome 133 string.""" + _twikit_patches.apply() + from twikit.client.client import Client + + client = Client() + assert "Chrome/133" in client._user_agent + assert "Safari/537.36" in client._user_agent + + +def test_header_patch_respects_user_supplied_user_agent(): + """If the caller passes user_agent= we must not clobber it.""" + _twikit_patches.apply() + from twikit.client.client import Client + + custom = "MyCustomAgent/1.0" + client = Client(user_agent=custom) + assert client._user_agent == custom + + +def test_header_patch_adds_sec_ch_ua_headers_to_base_headers(): + """The patched _base_headers property must include the modern Sec-* set + so Cloudflare doesn't trivially flag the request as a non-browser.""" + _twikit_patches.apply() + from twikit.client.client import Client + + client = Client() + headers = client._base_headers + assert "Chrome" in headers["sec-ch-ua"] + assert headers["sec-ch-ua-mobile"] == "?0" + assert headers["sec-ch-ua-platform"] == '"macOS"' + assert headers["sec-fetch-dest"] == "empty" + assert headers["sec-fetch-mode"] == "cors" + assert headers["sec-fetch-site"] == "same-origin" + + +def test_header_patch_preserves_original_required_headers(): + """The merge must not strip the auth/CSRF headers twikit relies on.""" + _twikit_patches.apply() + from twikit.client.client import Client + + client = Client() + headers = client._base_headers + assert headers["content-type"] == "application/json" + assert headers["X-Twitter-Auth-Type"] == "OAuth2Session" + assert headers["X-Twitter-Active-User"] == "yes" + assert headers["authorization"].startswith("Bearer ") + + +def test_header_patch_disabled_via_env(monkeypatch): + """When disabled, the default UA stays on the original Safari string.""" + monkeypatch.setenv("OPENSWARM_TWITTER_DISABLE_TWIKIT_HEADER_PATCH", "1") + _twikit_patches.apply() + from twikit.client.client import Client + + client = Client() + assert "Version/17.5 Safari" in client._user_agent + assert "sec-ch-ua" not in client._base_headers