monorepoly 0.2.0__py3-none-any.whl

monorepoly/hosted.py

@@ -0,0 +1,155 @@
+"""Hosted MCP server over HTTP: every request carries its user's GitLab OAuth token."""
+
+import asyncio
+import collections
+import hashlib
+import logging
+import threading
+import time
+from collections.abc import Callable
+
+import requests
+from mcp.server.auth.middleware.auth_context import get_access_token
+from mcp.server.auth.provider import AccessToken
+from mcp.server.auth.settings import AuthSettings
+from requests.exceptions import RequestException
+from starlette.applications import Starlette
+
+from monorepoly.access import Access, load_access
+from monorepoly.config import ServerSettings
+from monorepoly.embed import Embedder, make_embedder
+from monorepoly.exceptions import MonorepolyAccessError
+from monorepoly.gitlab import make_oauth_gitlab
+from monorepoly.judge import Judge, make_judge
+from monorepoly.mcp_server import build_server
+from monorepoly.search import QueryContext
+from monorepoly.store import Store, connect
+
+_LOGGER = logging.getLogger(__name__)
+
+_ACCESS_TTL_SECONDS = 300
+# Behind Cloud Run the Host header is the service's own name, so listen on every interface;
+# this also leaves the SDK's localhost-only DNS rebinding protection off.
+_HOST = '0.0.0.0'
+_MAX_CACHED_USERS = 1000
+_SCOPE = 'read_api'  # Lists the group's projects for `load_access`.
+_TIMEOUT_SECONDS = 10
+
+
+class GitLabTokenVerifier:
+    """Verifies bearer tokens with the GitLab instance that issued them."""
+
+    def __init__(self, *, gitlab_url: str) -> None:
+        self._gitlab_url = gitlab_url.rstrip('/')
+
+    async def verify_token(self, token: str) -> AccessToken | None:
+        """Returns the token's grant, or None unless GitLab confirms it."""
+        return await asyncio.to_thread(self._verify, token)
+
+    def _verify(self, token: str) -> AccessToken | None:
+        # Fail closed: GitLab tokens are opaque, so only GitLab can vouch for one. Asked on every
+        # request, so a revoked token stops working at once.
+        try:
+            response = requests.get(
+                f'{self._gitlab_url}/oauth/token/info',
+                headers={'Authorization': f'Bearer {token}'},
+                timeout=_TIMEOUT_SECONDS,
+            )
+            if response.status_code == 401:
+                return None
+            response.raise_for_status()
+            info = response.json()
+        except (RequestException, ValueError) as err:
+            _LOGGER.warning('GitLab did not verify a token: %s', err)
+            return None
+        expires_in = info.get('expires_in')
+        return AccessToken(
+            client_id=(info.get('application') or {}).get('uid') or '',
+            expires_at=int(time.time()) + expires_in if expires_in else None,
+            scopes=info.get('scope') or [],
+            subject=str(info.get('resource_owner_id')),
+            token=token,
+        )
+
+
+class AccessCache:
+    """Remembers each user's access briefly: loading it lists the group's projects."""
+
+    def __init__(self, *, gitlab_url: str, group_id: str) -> None:
+        self._entries: collections.OrderedDict[str, tuple[float, Access]] = (
+            collections.OrderedDict()
+        )
+        self._gitlab_url = gitlab_url
+        self._group_id = group_id
+        self._lock = threading.Lock()
+
+    def get(self, token: str) -> Access:
+        """Returns the token user's access, loading it from GitLab when not cached."""
+        # Keyed by hash, so a memory dump does not hand out working tokens.
+        key = hashlib.sha256(token.encode()).hexdigest()
+        with self._lock:
+            entry = self._entries.get(key)
+            if entry and entry[0] > time.monotonic():
+                return entry[1]
+        access = load_access(
+            gitlab=make_oauth_gitlab(token=token, url=self._gitlab_url), group_id=self._group_id
+        )
+        with self._lock:
+            self._entries[key] = (time.monotonic() + _ACCESS_TTL_SECONDS, access)
+            self._entries.move_to_end(key)
+            while len(self._entries) > _MAX_CACHED_USERS:
+                self._entries.popitem(last=False)
+        return access
+
+
+def token_context(
+    *,
+    cache: AccessCache,
+    embedder: Embedder,
+    judge: Judge,
+    store: Store,
+) -> Callable[[], QueryContext]:
+    """Returns a function giving the current request's user their query context."""
+
+    def context() -> QueryContext:
+        token = get_access_token()
+        # Fail closed: the auth middleware rejects such requests before any tool runs.
+        if token is None:
+            raise MonorepolyAccessError('The request has no GitLab token!')
+        return QueryContext(
+            access=cache.get(token.token), embedder=embedder, judge=judge, store=store
+        )
+
+    return context
+
+
+def build_app(*, context: Callable[[], QueryContext], settings: ServerSettings) -> Starlette:
+    """Returns the ASGI app serving the tools at `/mcp` to GitLab-authenticated users."""
+    auth = AuthSettings(
+        issuer_url=settings.gitlab_url.rstrip('/'),
+        required_scopes=[_SCOPE],
+        resource_server_url=f'{settings.public_url.rstrip("/")}/mcp',
+        # GitLab tokens name no audience. Any `read_api` token of the user is accepted: it already
+        # reads everything the index would show them.
+        validate_token_resource=False,
+    )
+    server = build_server(
+        auth=auth,
+        context=context,
+        token_verifier=GitLabTokenVerifier(gitlab_url=settings.gitlab_url),
+    )
+    # Stateless, so any Cloud Run instance can answer any request.
+    return server.streamable_http_app(json_response=True, host=_HOST, stateless_http=True)
+
+
+def serve_http(*, port: int, settings: ServerSettings) -> None:
+    """Runs the hosted server until stopped."""
+    import uvicorn
+
+    context = token_context(
+        cache=AccessCache(gitlab_url=settings.gitlab_url, group_id=settings.gitlab_group_id),
+        embedder=make_embedder(project=settings.google_cloud_project),
+        judge=make_judge(),
+        store=Store(db=connect(db=settings.mongodb_db, uri=settings.mongodb_uri)),
+    )
+    uvicorn.run(build_app(context=context, settings=settings), host=_HOST, port=port)