world-model-optimizer 0.2.0__py3-none-any.whl

wmo/connect/connector.py

@@ -0,0 +1,79 @@
+"""ContextConnector protocol + a small registry (mirrors `wmo.ingest.adapter`).
+
+A connector owns one service end to end: the interactive auth flow, a cheap credential check,
+and pulling content normalized into `ContextItem`s. Connectors register themselves on import and
+are looked up by name (`get_connector`) or listed (`list_connectors`).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+from wmo.connect.types import ConnectorAuth, ContextItem, PullQuery
+
+
+@dataclass
+class ConnectUI:
+    """Presentation callbacks a connector uses during interactive auth.
+
+    Connectors never print: the CLI layer builds one of these over its rich Console; tests pass
+    recording lambdas.
+
+    Attributes:
+        open_url: Show and/or open the browser authorization URL.
+        present_code: Show a device-flow verification URI and user code.
+        prompt_secret: Ask the user for a secret (e.g. a pasted API token); returns the value.
+        info: Show a short status message.
+    """
+
+    open_url: Callable[[str], None]
+    present_code: Callable[[str, str], None]
+    prompt_secret: Callable[[str], str]
+    info: Callable[[str], None]
+
+
+@runtime_checkable
+class ContextConnector(Protocol):
+    """One service's auth + pull logic, producing normalized `ContextItem`s.
+
+    Attributes:
+        name: Registry key and `ContextItem.source` value (e.g. "github").
+        label: Human name shown in pickers (e.g. "GitHub").
+    """
+
+    name: str
+    label: str
+
+    def connect(self, ui: ConnectUI) -> ConnectorAuth:
+        """Run the interactive auth flow (OAuth or token prompt) and return the credential."""
+        ...
+
+    def verify(self, auth: ConnectorAuth) -> str:
+        """Cheap API call returning a human identity string; raises ConnectError on bad auth."""
+        ...
+
+    def pull(self, auth: ConnectorAuth, query: PullQuery) -> list[ContextItem]:
+        """Fetch content matching `query`, normalized and capped at `query.limit` items."""
+        ...
+
+
+_CONNECTORS: dict[str, ContextConnector] = {}
+
+
+def register_connector(connector: ContextConnector) -> None:
+    """Register a connector under its `name` (typically at module import time)."""
+    _CONNECTORS[connector.name] = connector
+
+
+def get_connector(name: str) -> ContextConnector:
+    """Look up a registered connector by name."""
+    if name not in _CONNECTORS:
+        raise ValueError(f"no context connector registered for {name!r}; have {list(_CONNECTORS)}")
+    return _CONNECTORS[name]
+
+
+def list_connectors() -> list[str]:
+    """Names of all registered connectors, sorted (what the CLI picker shows)."""
+    return sorted(_CONNECTORS)

wmo/connect/credentials.py

@@ -0,0 +1,164 @@
+"""Connector credentials, stored once per user.
+
+Like the platform credential (`wmo/platform/credentials.py`), connector tokens are user-global:
+`~/.wmo/connectors.toml` (directory overridable via `$WMO_HOME`, exact file overridable via
+`$WMO_CONNECTORS_PATH`), one TOML table per connector name. A `WMO_<NAME>_TOKEN` environment
+variable injects a token-kind credential without a file, so CI and headless runs never write one
+to disk; a few connectors additionally honor a service-native variable (`ENV_TOKEN_ALIASES`)
+that deployments already carry.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+import tomllib
+from pathlib import Path
+
+import tomli_w
+from pydantic import ValidationError
+
+from wmo.connect.types import ConnectError, ConnectorAuth
+from wmo.core.types import JsonObject
+from wmo.platform.credentials import wmo_home
+
+ENV_CONNECTORS_PATH = "WMO_CONNECTORS_PATH"
+
+CONNECTORS_FILENAME = "connectors.toml"
+
+
+def connectors_path() -> Path:
+    """Where the connector credential file lives (`$WMO_CONNECTORS_PATH` wins over `$WMO_HOME`)."""
+    override = os.environ.get(ENV_CONNECTORS_PATH)
+    return Path(override) if override else wmo_home() / CONNECTORS_FILENAME
+
+
+def token_env_var(name: str) -> str:
+    """The env var that injects a plain token for connector `name`: WMO_<NAME>_TOKEN."""
+    return f"WMO_{name.upper().replace('-', '_')}_TOKEN"
+
+
+# Service-native env vars accepted as token sources per connector, consulted after the generic
+# WMO_<NAME>_TOKEN override. BRAVE_SEARCH_API_KEY is the key the grounding engine already uses
+# (wmo/engine/grounding.py) and deployments already carry, so the brave connector honors it too.
+ENV_TOKEN_ALIASES: dict[str, tuple[str, ...]] = {"brave": ("BRAVE_SEARCH_API_KEY",)}
+
+
+def token_env_vars(name: str) -> list[str]:
+    """Every env var that can inject a token for connector `name`, in precedence order."""
+    return [token_env_var(name), *ENV_TOKEN_ALIASES.get(name, ())]
+
+
+def resolve_env_token(name: str) -> tuple[str, str] | None:
+    """The first set env token for connector `name` as (var, token), or None when none is set.
+
+    Set-but-empty vars are treated as unset.
+    """
+    for var in token_env_vars(name):
+        token = os.environ.get(var)
+        if token:
+            return var, token
+    return None
+
+
+def load_connector_auth(name: str) -> ConnectorAuth | None:
+    """Load the stored credential for one connector, or None when not connected.
+
+    A non-empty env token (`WMO_<NAME>_TOKEN`, then any `ENV_TOKEN_ALIASES` entry) takes
+    precedence over the file and yields a token-kind auth (a set-but-empty var is treated as
+    unset).
+    """
+    resolved = resolve_env_token(name)
+    if resolved is not None:
+        return ConnectorAuth(kind="token", access_token=resolved[1])
+    section = _read_sections().get(name)
+    if section is None:
+        return None
+    return _validate_section(name, section)
+
+
+def save_connector_auth(name: str, auth: ConnectorAuth) -> Path:
+    """Persist one connector's credential, keeping every other connector's table intact.
+
+    Mirrors the platform credential writer: refuses symlinks, writes through a 0600 mkstemp,
+    and swaps into place atomically.
+
+    Raises:
+        ValueError: If the target path is a symlink.
+    """
+    sections = _read_sections()
+    sections[name] = {
+        key: value for key, value in auth.model_dump(mode="json").items() if value is not None
+    }
+    return _write_sections(sections)
+
+
+def delete_connector_auth(name: str) -> bool:
+    """Remove one connector's credential; returns whether one existed.
+
+    Deleting the last entry removes the file entirely.
+    """
+    sections = _read_sections()
+    if name not in sections:
+        return False
+    del sections[name]
+    if sections:
+        _write_sections(sections)
+    else:
+        connectors_path().unlink()
+    return True
+
+
+def list_connected() -> dict[str, ConnectorAuth]:
+    """Every connector with a stored credential, by name (sorted).
+
+    Only reads the file: connections injected purely via `WMO_<NAME>_TOKEN` env vars are not
+    enumerable and do not appear here (per-name lookups still see them via
+    `load_connector_auth`).
+    """
+    return {
+        name: _validate_section(name, section) for name, section in sorted(_read_sections().items())
+    }
+
+
+def _read_sections() -> dict[str, JsonObject]:
+    """All connector tables in the credential file (empty when the file is absent)."""
+    path = connectors_path()
+    if not path.exists():
+        return {}
+    document = tomllib.loads(path.read_text(encoding="utf-8"))
+    return {name: section for name, section in document.items() if isinstance(section, dict)}
+
+
+def _validate_section(name: str, section: JsonObject) -> ConnectorAuth:
+    """Parse one connector table, turning pydantic errors into an actionable ConnectError."""
+    try:
+        return ConnectorAuth.model_validate(section)
+    except ValidationError as exc:
+        raise ConnectError(
+            f"invalid [{name}] entry in {connectors_path()}: {exc}; "
+            f"the stored credential for {name} is malformed; provide a valid token or "
+            "reauthorize the connection"
+        ) from exc
+
+
+def _write_sections(sections: dict[str, JsonObject]) -> Path:
+    """Atomically rewrite the credential file with owner-only permissions."""
+    path = connectors_path()
+    if path.is_symlink():
+        msg = (
+            f"refusing to write connector credentials through the symlink {path}; "
+            f"remove the link or point ${ENV_CONNECTORS_PATH} elsewhere"
+        )
+        raise ValueError(msg)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    # mkstemp creates the file 0600: that IS the owner-only mechanism.
+    fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=f"{path.name}.")
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8") as fh:
+            fh.write(tomli_w.dumps(sections))
+        os.replace(tmp_name, path)
+    except BaseException:
+        os.unlink(tmp_name)
+        raise
+    return path

wmo/connect/github.py

@@ -0,0 +1,321 @@
+"""GitHub context connector: device-flow OAuth, then repo issues, pull requests, and the README.
+
+`connect` runs the RFC 8628 device flow against the shared "github" OAuth app (client id only,
+no secret needed), `verify` resolves the token to its login via `GET /user`, and `pull` fetches
+one repository's issues and pull requests (the combined listing, newest first) plus its README,
+normalized into `ContextItem`s. A `PullQuery.query` switches the listing to the GitHub search
+API; `PullQuery.since` maps to the listing's `since` parameter (or an `updated:>=` qualifier in
+search).
+"""
+
+from __future__ import annotations
+
+import base64
+import logging
+import re
+from datetime import UTC, datetime
+from typing import cast
+
+import httpx
+
+from wmo.connect.apps import get_app
+from wmo.connect.connector import ConnectUI, register_connector
+from wmo.connect.oauth import run_device_flow
+from wmo.connect.types import (
+    ConnectError,
+    ConnectorAuth,
+    ContextItem,
+    ItemKind,
+    PullQuery,
+    opt_str,
+    transport_errors,
+)
+from wmo.core.types import JsonObject, JsonValue
+
+logger = logging.getLogger(__name__)
+
+_API_BASE = "https://api.github.com"
+_API_HOST = "api.github.com"
+_API_VERSION = "2022-11-28"
+_PER_PAGE = 100
+_TIMEOUT_SECONDS = 30.0
+_SOURCE = "github"
+
+_TARGET_RE = re.compile(r"^([^/\s]+)/([^/\s]+)$")
+
+# GitHub App tokens only reach repositories where the app is installed; users pick those
+# repos on their app's installation page after connecting (a PAT needs no installation).
+_INSTALL_HINT = (
+    "GitHub App tokens only reach repositories the app is installed on: install yours on the "
+    "target repos from https://github.com/settings/apps (your app > Install App)"
+)
+
+
+def _parse_target(target: str | None) -> tuple[str, str]:
+    """Split an "owner/repo" pull target, raising an actionable error on anything else."""
+    match = _TARGET_RE.match(target.strip()) if target else None
+    if match is None:
+        raise ConnectError(
+            f"github pull needs a repository target like 'owner/repo' (got {target!r}); "
+            "pass --target <owner>/<repo>"
+        )
+    return match.group(1), match.group(2)
+
+
+def _format_reset(reset: str) -> str:
+    """A rate-limit reset epoch as ISO-8601 UTC (falls back to the raw header value)."""
+    try:
+        moment = datetime.fromtimestamp(int(reset), tz=UTC)
+    except (OSError, OverflowError, ValueError):
+        return reset
+    return moment.isoformat(timespec="seconds")
+
+
+def _raise_for_response(response: httpx.Response, *, doing: str, repo: str | None = None) -> None:
+    """Turn GitHub error statuses into ConnectErrors that say what to do next."""
+    status = response.status_code
+    if status < 400:
+        return
+    if status == 401:
+        raise ConnectError(
+            f"github rejected the stored credential during {doing} (HTTP 401); "
+            "the token is invalid or expired and the connection must be reauthorized"
+        )
+    if status == 403 and response.headers.get("X-RateLimit-Remaining") == "0":
+        reset = response.headers.get("X-RateLimit-Reset", "unknown")
+        raise ConnectError(
+            f"github rate limit exceeded during {doing}; it resets at {_format_reset(reset)} "
+            f"(X-RateLimit-Reset {reset}); wait for the reset or connect a higher-limit account"
+        )
+    if status == 404 and repo is not None:
+        raise ConnectError(
+            f"github repository {repo!r} was not found during {doing} (HTTP 404): it does not "
+            "exist, the connected account cannot see it, or (GitHub App auth) the app is not "
+            "installed on it; check the owner/repo spelling, install your app on the repo "
+            "(https://github.com/settings/apps > your app > Install App), or reauthorize the "
+            "connection with an account that has access"
+        )
+    raise ConnectError(
+        f"github {doing} failed (HTTP {status}): {response.text[:200]}; "
+        "check the request (token scopes, target, query) and retry"
+    )
+
+
+def _json_object(response: httpx.Response) -> JsonObject:
+    """The response body as a JSON object ({} when it is anything else)."""
+    try:
+        raw = response.json()
+    except ValueError:
+        return {}
+    return cast(JsonObject, raw) if isinstance(raw, dict) else {}
+
+
+def _page_rows(response: httpx.Response, *, nested: bool) -> list[JsonObject]:
+    """One page of issue rows: the raw JSON list, or the search wrapper's `items` list."""
+    try:
+        raw = response.json()
+    except ValueError:
+        return []
+    rows = raw.get("items") if nested and isinstance(raw, dict) else raw
+    if not isinstance(rows, list):
+        return []
+    return [cast(JsonObject, row) for row in rows if isinstance(row, dict)]
+
+
+def _paginate(
+    client: httpx.Client,
+    url: str,
+    params: dict[str, str],
+    *,
+    limit: int,
+    repo: str,
+    doing: str,
+    nested: bool,
+) -> list[JsonObject]:
+    """Follow `Link: rel="next"` pagination, stopping as soon as `limit` rows are in hand."""
+    rows: list[JsonObject] = []
+    next_url: str | None = url
+    next_params: dict[str, str] | None = params
+    while next_url is not None and len(rows) < limit:
+        response = client.get(next_url, params=next_params)
+        _raise_for_response(response, doing=doing, repo=repo)
+        page = _page_rows(response, nested=nested)
+        rows.extend(page)
+        next_url = response.links.get("next", {}).get("url")
+        next_params = None  # the next URL already carries its query string
+        if not page and next_url:
+            break  # an empty page that still advertises a next link would loop forever
+    return rows[:limit]
+
+
+def _issue_item(row: JsonObject, *, owner: str, repo: str) -> ContextItem:
+    """Normalize one combined-listing (or search) row into an issue or pull-request item."""
+    number = row.get("number")
+    suffix = str(number) if isinstance(number, int) else str(row.get("id") or "unknown")
+    title = opt_str(row.get("title")) or "(untitled)"
+    body = row.get("body")
+    user = row.get("user")
+    author = opt_str(user.get("login")) if isinstance(user, dict) else None
+    labels_value = row.get("labels")
+    labels: list[JsonValue] = []
+    if isinstance(labels_value, list):
+        labels = [
+            label["name"]
+            for label in labels_value
+            if isinstance(label, dict) and isinstance(label.get("name"), str)
+        ]
+    comments = row.get("comments")
+    return ContextItem(
+        id=f"{owner}/{repo}#{suffix}",
+        source=_SOURCE,
+        kind=ItemKind.PULL_REQUEST if "pull_request" in row else ItemKind.ISSUE,
+        title=f"#{suffix} {title}",
+        body=body if isinstance(body, str) else "",
+        url=opt_str(row.get("html_url")),
+        created_at=opt_str(row.get("created_at")),
+        updated_at=opt_str(row.get("updated_at")),
+        metadata={
+            "state": opt_str(row.get("state")),
+            "labels": labels,
+            "author": author,
+            "comments": comments if isinstance(comments, int) else 0,
+        },
+    )
+
+
+class GitHubConnector:
+    """GitHub connector: device-flow OAuth plus repository issue/PR/README pulls.
+
+    Args:
+        transport: Injected httpx transport threaded into every HTTP call (tests pass
+            `httpx.MockTransport`); None means the real network.
+    """
+
+    name = _SOURCE
+    label = "GitHub"
+
+    def __init__(self, *, transport: httpx.BaseTransport | None = None) -> None:
+        self._transport = transport
+
+    def connect(self, ui: ConnectUI) -> ConnectorAuth:
+        """Authorize via the RFC 8628 device flow and stamp the credential with its identity."""
+        app = get_app(self.name)
+        auth = run_device_flow(
+            app, scopes=app.scopes, present=ui.present_code, transport=self._transport
+        )
+        account = self.verify(auth)
+        ui.info(f"connected to GitHub as {account}")
+        ui.info(_INSTALL_HINT)
+        return auth.model_copy(update={"account": account})
+
+    def verify(self, auth: ConnectorAuth) -> str:
+        """Resolve the credential to its GitHub identity via the cheapest call, `GET /user`."""
+        with self._client(auth) as client, transport_errors(_API_HOST):
+            response = client.get("/user")
+        _raise_for_response(response, doing="the GitHub identity check")
+        payload = _json_object(response)
+        login = opt_str(payload.get("login"))
+        if login is None:
+            raise ConnectError(
+                "github's /user endpoint returned no login for the stored credential; "
+                "the token is invalid or expired and the connection must be reauthorized"
+            )
+        name = opt_str(payload.get("name"))
+        return f"{login} ({name})" if name else login
+
+    def pull(self, auth: ConnectorAuth, query: PullQuery) -> list[ContextItem]:
+        """Fetch a repository's issues, pull requests, and README, newest first, up to the limit.
+
+        `query.target` must be "owner/repo". Issues and pull requests come from the combined
+        `GET /repos/{owner}/{repo}/issues` listing (sorted by update time, descending, honoring
+        `query.since`); when `query.query` is set, the GitHub search API replaces that listing.
+        The README is appended as one DOCUMENT item when the limit leaves room for it.
+        """
+        owner, repo = _parse_target(query.target)
+        with self._client(auth) as client, transport_errors(_API_HOST):
+            rows = self._issue_rows(client, owner, repo, query)
+            items = [_issue_item(row, owner=owner, repo=repo) for row in rows]
+            if len(items) < query.limit:
+                readme = self._readme_item(client, owner, repo)
+                if readme is not None:
+                    items.append(readme)
+        logger.debug("pulled %d github items from %s/%s", len(items), owner, repo)
+        return items
+
+    def _client(self, auth: ConnectorAuth) -> httpx.Client:
+        """An API client carrying the credential and GitHub's versioned-media headers."""
+        return httpx.Client(
+            base_url=_API_BASE,
+            headers={
+                "Authorization": f"Bearer {auth.access_token}",
+                "Accept": "application/vnd.github+json",
+                "X-GitHub-Api-Version": _API_VERSION,
+            },
+            timeout=_TIMEOUT_SECONDS,
+            transport=self._transport,
+        )
+
+    def _issue_rows(
+        self, client: httpx.Client, owner: str, repo: str, query: PullQuery
+    ) -> list[JsonObject]:
+        """Issue/PR rows, newest first: the combined listing, or search when a query is set."""
+        target = f"{owner}/{repo}"
+        per_page = str(min(query.limit, _PER_PAGE))
+        if query.query:
+            q = f"repo:{target} {query.query}"
+            if query.since:
+                q += f" updated:>={query.since}"
+            if query.until:
+                q += f" updated:<={query.until}"
+            params = {"q": q, "sort": "updated", "order": "desc", "per_page": per_page}
+            return _paginate(
+                client,
+                "/search/issues",
+                params,
+                limit=query.limit,
+                repo=target,
+                doing="the GitHub issue search",
+                nested=True,
+            )
+        params = {"state": "all", "sort": "updated", "direction": "desc", "per_page": per_page}
+        if query.since:
+            params["since"] = query.since
+        return _paginate(
+            client,
+            f"/repos/{owner}/{repo}/issues",
+            params,
+            limit=query.limit,
+            repo=target,
+            doing="the GitHub issue listing",
+            nested=False,
+        )
+
+    def _readme_item(self, client: httpx.Client, owner: str, repo: str) -> ContextItem | None:
+        """The repo README as one DOCUMENT item, or None when the repo has no README (404)."""
+        response = client.get(f"/repos/{owner}/{repo}/readme")
+        if response.status_code == 404:
+            logger.debug("no README in %s/%s; skipping", owner, repo)
+            return None
+        _raise_for_response(response, doing="the GitHub README fetch", repo=f"{owner}/{repo}")
+        payload = _json_object(response)
+        content = payload.get("content")
+        body = ""
+        if isinstance(content, str) and content:
+            try:
+                body = base64.b64decode(content).decode("utf-8", errors="replace")
+            except ValueError:
+                logger.debug("undecodable README content in %s/%s; empty body", owner, repo)
+        name = opt_str(payload.get("name")) or "README"
+        path = opt_str(payload.get("path")) or name
+        return ContextItem(
+            id=f"{owner}/{repo}:{path}",
+            source=_SOURCE,
+            kind=ItemKind.DOCUMENT,
+            title=name,
+            body=body,
+            url=opt_str(payload.get("html_url")),
+            metadata={"path": path},
+        )
+
+
+register_connector(GitHubConnector())