Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai>
Co-authored-by: Kim-Adeline Miguel <kimadeline.miguel@mistral.ai>
Co-authored-by: Lucas Marandat <31749711+lucasmrdt@users.noreply.github.com>
Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com>
Co-authored-by: Paul Cacheux <paul.cacheux@mistral.ai>
Co-authored-by: Peter Evers <pevers90@gmail.com>
Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai>
Co-authored-by: Pierre Rossinès <pierre.rossines@protonmail.com>
Co-authored-by: Quentin <quentin.torroba@mistral.ai>
Co-authored-by: Simon Van de Kerckhove <simon.vandekerckhove@mistral.ai>
Co-authored-by: Val <102326092+vdeva@users.noreply.github.com>
Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Mathias Gesbert 2026-04-09 18:40:46 +02:00 committed by GitHub
parent 90763daf81
commit e9a9217cc8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
113 changed files with 7202 additions and 541 deletions

View file

@ -0,0 +1,21 @@
from __future__ import annotations
from vibe.setup.auth.browser_sign_in import BrowserSignInService
from vibe.setup.auth.browser_sign_in_gateway import (
BrowserSignInError,
BrowserSignInErrorCode,
BrowserSignInGateway,
BrowserSignInPollResult,
BrowserSignInProcess,
)
from vibe.setup.auth.http_browser_sign_in_gateway import HttpBrowserSignInGateway
__all__ = [
"BrowserSignInError",
"BrowserSignInErrorCode",
"BrowserSignInGateway",
"BrowserSignInPollResult",
"BrowserSignInProcess",
"BrowserSignInService",
"HttpBrowserSignInGateway",
]

View file

@ -0,0 +1,147 @@
from __future__ import annotations
import asyncio
import base64
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
import hashlib
import secrets
import webbrowser
from vibe.setup.auth.browser_sign_in_gateway import (
BrowserSignInError,
BrowserSignInErrorCode,
BrowserSignInGateway,
BrowserSignInProcess,
)
StatusCallback = Callable[[str], None]
BrowserOpener = Callable[[str], bool]
SleepFn = Callable[[float], Awaitable[None]]
NowFn = Callable[[], datetime]
class BrowserSignInService:
_max_consecutive_poll_failures = 3
def __init__(
self,
gateway: BrowserSignInGateway,
*,
open_browser: BrowserOpener | None = None,
sleep: SleepFn = asyncio.sleep,
now: NowFn | None = None,
poll_interval: float = 3.0,
) -> None:
self._gateway = gateway
self._open_browser = open_browser or webbrowser.open
self._sleep = sleep
self._now = now or (lambda: datetime.now(UTC))
self._poll_interval = poll_interval
async def aclose(self) -> None:
await self._gateway.aclose()
async def authenticate(self, status_callback: StatusCallback | None = None) -> str:
verifier, challenge = _generate_pkce_pair()
process = await self._gateway.create_process(challenge)
self._emit(status_callback, "opening_browser")
self._open_browser_or_raise(process.sign_in_url)
self._emit(status_callback, "waiting_for_browser_sign_in")
exchange_token = await self._wait_for_completion(process)
self._emit(status_callback, "exchanging")
api_key = await self._gateway.exchange(
process.process_id, exchange_token, verifier
)
self._emit(status_callback, "completed")
return api_key
async def _wait_for_completion(self, process: BrowserSignInProcess) -> str:
consecutive_poll_failures = 0
while self._now() < process.expires_at:
try:
payload = await self._gateway.poll(process.poll_url)
except BrowserSignInError as err:
if err.code is not BrowserSignInErrorCode.POLL_FAILED:
raise
consecutive_poll_failures += 1
if consecutive_poll_failures >= self._max_consecutive_poll_failures:
raise
await self._sleep_until_next_poll_or_timeout(process.expires_at)
continue
consecutive_poll_failures = 0
match payload.status:
case "pending":
await self._sleep_until_next_poll_or_timeout(process.expires_at)
case "completed":
if payload.exchange_token:
return payload.exchange_token
raise BrowserSignInError(
"Sign-in worked, but setup couldn't finish.",
code=BrowserSignInErrorCode.MISSING_EXCHANGE_TOKEN,
)
case "expired":
raise BrowserSignInError(
"Browser sign-in expired.", code=BrowserSignInErrorCode.EXPIRED
)
case "denied":
raise BrowserSignInError(
"Browser sign-in was denied.",
code=BrowserSignInErrorCode.DENIED,
)
case "error":
raise BrowserSignInError(
payload.message or "Browser sign-in failed.",
code=BrowserSignInErrorCode.PROVIDER_ERROR,
)
case _:
raise BrowserSignInError(
"Browser sign-in returned an unknown state.",
code=BrowserSignInErrorCode.UNKNOWN_STATE,
)
raise BrowserSignInError(
"Browser sign-in timed out.", code=BrowserSignInErrorCode.TIMED_OUT
)
async def _sleep_until_next_poll_or_timeout(self, expires_at: datetime) -> None:
remaining_seconds = (expires_at - self._now()).total_seconds()
if remaining_seconds <= 0:
raise BrowserSignInError(
"Browser sign-in timed out.", code=BrowserSignInErrorCode.TIMED_OUT
)
await self._sleep(min(self._poll_interval, remaining_seconds))
def _emit(self, callback: StatusCallback | None, status: str) -> None:
if callback is not None:
callback(status)
def _open_browser_or_raise(self, sign_in_url: str) -> None:
try:
browser_opened = self._open_browser(sign_in_url)
except Exception as err:
raise BrowserSignInError(
"Failed to open browser for sign-in.",
code=BrowserSignInErrorCode.OPEN_BROWSER_FAILED,
) from err
if not browser_opened:
raise BrowserSignInError(
"Failed to open browser for sign-in.",
code=BrowserSignInErrorCode.OPEN_BROWSER_FAILED,
)
def _generate_code_verifier() -> str:
return secrets.token_urlsafe(64)
def _generate_pkce_pair() -> tuple[str, str]:
verifier = _generate_code_verifier()
return verifier, _generate_code_challenge(verifier)
def _generate_code_challenge(verifier: str) -> str:
digest = hashlib.sha256(verifier.encode("ascii")).digest()
return base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")

View file

@ -0,0 +1,55 @@
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from enum import StrEnum, auto
from typing import Literal, Protocol
class BrowserSignInErrorCode(StrEnum):
START_FAILED = auto()
POLL_FAILED = auto()
UNKNOWN_STATE = auto()
EXCHANGE_FAILED = auto()
MISSING_API_KEY = auto()
MISSING_EXCHANGE_TOKEN = auto()
EXPIRED = auto()
DENIED = auto()
PROVIDER_ERROR = auto()
TIMED_OUT = auto()
OPEN_BROWSER_FAILED = auto()
class BrowserSignInError(Exception):
def __init__(
self, message: str, *, code: BrowserSignInErrorCode | None = None
) -> None:
super().__init__(message)
self.code = code
@dataclass
class BrowserSignInProcess:
process_id: str
sign_in_url: str
poll_url: str
expires_at: datetime
@dataclass
class BrowserSignInPollResult:
status: Literal["pending", "completed", "expired", "denied", "error"]
exchange_token: str | None = None
message: str | None = None
class BrowserSignInGateway(Protocol):
async def create_process(self, code_challenge: str) -> BrowserSignInProcess: ...
async def poll(self, poll_url: str) -> BrowserSignInPollResult: ...
async def exchange(
self, process_id: str, exchange_token: str, code_verifier: str
) -> str: ...
async def aclose(self) -> None: ...

View file

@ -0,0 +1,319 @@
from __future__ import annotations
from collections.abc import Mapping
from datetime import UTC, datetime
import posixpath
import types
from typing import Any, Literal, TypedDict, cast
from urllib.parse import SplitResult, unquote, urlsplit
import httpx
from vibe.core.logger import logger
from vibe.setup.auth.browser_sign_in_gateway import (
BrowserSignInError,
BrowserSignInErrorCode,
BrowserSignInGateway,
BrowserSignInPollResult,
BrowserSignInProcess,
)
class CreateProcessPayload(TypedDict):
process_id: str
sign_in_url: str
poll_url: str
expires_at: str
class PollPayload(TypedDict, total=False):
status: Literal["pending", "completed", "expired", "denied", "error"]
exchange_token: str
message: str
class ExchangePayload(TypedDict, total=False):
api_key: str
HTTP_GONE = 410
_DEFAULT_PORTS = {"http": 80, "https": 443}
class HttpBrowserSignInGateway(BrowserSignInGateway):
def __init__(
self,
browser_base_url: str,
api_base_url: str,
*,
client: httpx.AsyncClient | None = None,
) -> None:
self._browser_base_url = browser_base_url.rstrip("/")
self._api_base_url = api_base_url.rstrip("/")
self._client = client
self._should_manage_client = client is None
async def __aenter__(self) -> HttpBrowserSignInGateway:
self._ensure_client()
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: types.TracebackType | None,
) -> None:
await self.aclose()
async def aclose(self) -> None:
if self._should_manage_client and self._client is not None:
await self._client.aclose()
self._client = None
async def create_process(self, code_challenge: str) -> BrowserSignInProcess:
message = "Failed to start browser sign-in."
code = BrowserSignInErrorCode.START_FAILED
try:
response = await self._ensure_client().post(
f"{self._api_base_url}/vibe/sign-in",
json={
"code_challenge": code_challenge,
"code_challenge_method": "S256",
},
)
except httpx.HTTPError as err:
logger.warning(
"Browser sign-in start request failed for api_base_url=%s: %s",
self._api_base_url,
err,
exc_info=True,
)
raise BrowserSignInError(message, code=code) from err
if not response.is_success:
logger.warning(
"Browser sign-in start request returned status_code=%s for api_base_url=%s",
response.status_code,
self._api_base_url,
)
raise BrowserSignInError(message, code=code)
payload = _response_json_or_raise(response, message=message, code=code)
data = cast(CreateProcessPayload, payload)
try:
return BrowserSignInProcess(
process_id=data["process_id"],
sign_in_url=_validate_url_against_base_url(
data["sign_in_url"],
base_url=self._browser_base_url,
message=message,
code=code,
),
poll_url=_validate_url_against_base_url(
data["poll_url"],
base_url=self._api_base_url,
message=message,
code=code,
),
expires_at=_parse_expires_at(data["expires_at"]),
)
except BrowserSignInError:
sign_in_url_details, poll_url_details = _build_start_payload_log_details(
data
)
logger.warning(
"Browser sign-in start payload validation failed for browser_base_url=%s api_base_url=%s sign_in_url=%s poll_url=%s",
self._browser_base_url,
self._api_base_url,
sign_in_url_details,
poll_url_details,
exc_info=True,
)
raise
except (KeyError, TypeError, ValueError) as err:
sign_in_url_details, poll_url_details = _build_start_payload_log_details(
payload
)
logger.warning(
"Browser sign-in start payload parsing failed for api_base_url=%s payload_keys=%s sign_in_url=%s poll_url=%s",
self._api_base_url,
sorted(str(key) for key in payload),
sign_in_url_details,
poll_url_details,
exc_info=True,
)
raise BrowserSignInError(message, code=code) from err
async def poll(self, poll_url: str) -> BrowserSignInPollResult:
message = "Browser sign-in status could not be retrieved."
code = BrowserSignInErrorCode.POLL_FAILED
validated_poll_url = _validate_url_against_base_url(
poll_url, base_url=self._api_base_url, message=message, code=code
)
try:
response = await self._ensure_client().get(validated_poll_url)
except httpx.HTTPError as err:
raise BrowserSignInError(message, code=code) from err
if response.status_code == HTTP_GONE:
return BrowserSignInPollResult(status="expired")
if not response.is_success:
raise BrowserSignInError(message, code=code)
raw_payload = _response_json_or_raise(response, message=message, code=code)
payload = cast(PollPayload, raw_payload)
status = payload.get("status")
if status not in {"pending", "completed", "expired", "denied", "error"}:
raise BrowserSignInError(
"Browser sign-in returned an unknown state.",
code=BrowserSignInErrorCode.UNKNOWN_STATE,
)
return BrowserSignInPollResult(
status=status,
exchange_token=payload.get("exchange_token"),
message=payload.get("message"),
)
async def exchange(
self, process_id: str, exchange_token: str, code_verifier: str
) -> str:
message = "Failed to exchange browser sign-in for an API key."
code = BrowserSignInErrorCode.EXCHANGE_FAILED
try:
response = await self._ensure_client().post(
f"{self._api_base_url}/vibe/sign-in/{process_id}/exchange",
json={"exchange_token": exchange_token, "code_verifier": code_verifier},
)
except httpx.HTTPError as err:
raise BrowserSignInError(message, code=code) from err
if not response.is_success:
raise BrowserSignInError(message, code=code)
raw_payload = _response_json_or_raise(response, message=message, code=code)
payload = cast(ExchangePayload, raw_payload)
if api_key := payload.get("api_key"):
return api_key
raise BrowserSignInError(
"Browser sign-in exchange did not return an API key.",
code=BrowserSignInErrorCode.MISSING_API_KEY,
)
def _ensure_client(self) -> httpx.AsyncClient:
if self._client is None:
self._client = httpx.AsyncClient()
self._should_manage_client = True
return self._client
def _parse_expires_at(value: str) -> datetime:
return datetime.fromisoformat(value.replace("Z", "+00:00")).astimezone(UTC)
def _validate_url_against_base_url(
value: str, *, base_url: str, message: str, code: BrowserSignInErrorCode
) -> str:
current_url = urlsplit(value)
base = urlsplit(base_url)
safe_url_details = _build_safe_url_log_details(value)
try:
current_origin = _normalized_origin(current_url)
base_origin = _normalized_origin(base)
except ValueError as err:
logger.warning(
"Browser sign-in URL origin validation failed for returned_url=%s expected_base_url=%s",
safe_url_details,
base_url,
exc_info=True,
)
raise BrowserSignInError(message, code=code) from err
if current_origin != base_origin:
logger.warning(
"Browser sign-in URL host validation failed for returned_url=%s expected_base_url=%s",
safe_url_details,
base_url,
)
raise BrowserSignInError(message, code=code)
if not _is_path_under_base_path(current_url.path, base.path):
logger.warning(
"Browser sign-in URL path validation failed for returned_url=%s expected_base_url=%s",
safe_url_details,
base_url,
)
raise BrowserSignInError(message, code=code)
return value
def _is_path_under_base_path(path: str, base_path: str) -> bool:
normalized_path = _normalize_url_path(path)
normalized_base_path = _normalize_url_path(base_path).rstrip("/")
if not normalized_base_path:
return True
if normalized_path == normalized_base_path:
return True
return normalized_path.startswith(f"{normalized_base_path}/")
def _normalized_origin(parsed: SplitResult) -> tuple[str, str | None, int | None]:
scheme = parsed.scheme.lower()
port = parsed.port
return scheme, parsed.hostname, _effective_port(scheme, port)
def _effective_port(scheme: str, port: int | None) -> int | None:
if port is not None:
return port
return _DEFAULT_PORTS.get(scheme)
def _normalize_url_path(path: str) -> str:
if not path:
return "/"
decoded_path = unquote(path)
return posixpath.normpath(decoded_path)
def _build_start_payload_log_details(payload: Mapping[str, object]) -> tuple[str, str]:
return (
_build_safe_url_log_details(payload.get("sign_in_url")),
_build_safe_url_log_details(payload.get("poll_url")),
)
def _response_json_or_raise(
response: httpx.Response, *, message: str, code: BrowserSignInErrorCode
) -> dict[str, Any]:
try:
payload = response.json()
except ValueError as err:
raise BrowserSignInError(message, code=code) from err
if not isinstance(payload, Mapping):
raise BrowserSignInError(message, code=code)
return dict(payload)
def _build_safe_url_log_details(value: object) -> str:
if not isinstance(value, str):
return "unavailable"
parsed = urlsplit(value)
return (
f"scheme={parsed.scheme or None!r} "
f"hostname={parsed.hostname!r} "
f"port={_safe_port(parsed)!r} "
f"has_path={bool(parsed.path)} "
f"has_query={bool(parsed.query)} "
f"has_fragment={bool(parsed.fragment)}"
)
def _safe_port(parsed: SplitResult) -> int | None:
try:
return parsed.port
except ValueError:
return None