v2.7.4 (#579)
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:
parent
90763daf81
commit
e9a9217cc8
113 changed files with 7202 additions and 541 deletions
21
vibe/setup/auth/__init__.py
Normal file
21
vibe/setup/auth/__init__.py
Normal 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",
|
||||
]
|
||||
147
vibe/setup/auth/browser_sign_in.py
Normal file
147
vibe/setup/auth/browser_sign_in.py
Normal 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("=")
|
||||
55
vibe/setup/auth/browser_sign_in_gateway.py
Normal file
55
vibe/setup/auth/browser_sign_in_gateway.py
Normal 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: ...
|
||||
319
vibe/setup/auth/http_browser_sign_in_gateway.py
Normal file
319
vibe/setup/auth/http_browser_sign_in_gateway.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue