feat(auth): browser-based Stytch OAuth login + environment switching

assemblyai_cli/__main__.py → aai_cli/__main__.py

@@ -1,4 +1,4 @@
-from assemblyai_cli.main import run
+from aai_cli.main import run
 
 if __name__ == "__main__":
     run()

assemblyai_cli/agent/audio.py → aai_cli/agent/audio.py

@@ -6,8 +6,8 @@
 from collections.abc import Callable, Iterator
 from typing import Any
 
-from assemblyai_cli.errors import CLIError
-from assemblyai_cli.microphone import _default_rate, _resample, audio_missing_error
+from aai_cli.errors import CLIError
+from aai_cli.microphone import _default_rate, _resample, audio_missing_error
 
 SAMPLE_RATE = 24000  # Voice Agent native PCM16 mono rate
 
@@ -291,5 +291,5 @@ def __iter__(self) -> Iterator[bytes]:
         return self._duplex.capture_frames()
 
 
-# Microphone capture (MicrophoneSource) lives in assemblyai_cli.microphone and is
+# Microphone capture (MicrophoneSource) lives in aai_cli.microphone and is
 # shared with `aai stream`; the agent's live mic+speaker run through DuplexAudio.

assemblyai_cli/agent/render.py → aai_cli/agent/render.py

@@ -4,7 +4,7 @@
 
 from rich.text import Text
 
-from assemblyai_cli.render import BaseRenderer
+from aai_cli.render import BaseRenderer
 
 
 def _labeled(label: str, body: str, *, style: str = "aai.label") -> Text:

assemblyai_cli/agent/session.py → aai_cli/agent/session.py

@@ -6,7 +6,7 @@
 import threading
 from typing import Any
 
-from assemblyai_cli.errors import APIError, CLIError, auth_failure, is_auth_failure
+from aai_cli.errors import APIError, CLIError, auth_failure, is_auth_failure
 
 WS_URL = "wss://agents.assemblyai.com/v1/ws"
 

aai_cli/auth/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from aai_cli.auth.flow import run_login_flow
+
+__all__ = ["run_login_flow"]

aai_cli/auth/ams.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from typing import Any, cast
+
+import httpx
+
+from aai_cli.auth import endpoints
+from aai_cli.errors import APIError, NotAuthenticated
+
+_TIMEOUT = 30.0
+
+
+def _detail(resp: httpx.Response) -> str:
+    try:
+        body = resp.json()
+        if isinstance(body, dict) and "detail" in body:
+            return str(body["detail"])
+    except Exception:  # noqa: BLE001,S110 - non-JSON error body; pass is intentional
+        pass
+    return resp.text or f"HTTP {resp.status_code}"
+
+
+def _json_or_raise(resp: httpx.Response) -> Any:
+    if resp.status_code in (401, 403):
+        raise NotAuthenticated(f"AMS rejected the login ({resp.status_code}): {_detail(resp)}")
+    if resp.status_code >= 400:
+        raise APIError(f"AMS request failed ({resp.status_code}): {_detail(resp)}")
+    return resp.json()
+
+
+def _client(session_jwt: str | None = None) -> httpx.Client:
+    """An AMS HTTP client; pass a session JWT to send the authenticated cookie."""
+    cookies = {"stytch_session_jwt": session_jwt} if session_jwt else None
+    return httpx.Client(base_url=endpoints.ams_base(), timeout=_TIMEOUT, cookies=cookies)
+
+
+def discover(token: str) -> dict[str, Any]:
+    """POST /v2/auth/discover with a discovery_oauth token -> {orgs, email, IST}."""
+    with _client() as client:
+        resp = client.post(
+            "/v2/auth/discover",
+            json={"token": token, "token_type": "discovery_oauth"},
+        )
+    return cast(dict[str, Any], _json_or_raise(resp))
+
+
+def exchange(intermediate_session_token: str, organization_id: str) -> dict[str, Any]:
+    """POST /v2/auth/exchange -> SignedInResponse {account, session_jwt, session_token}."""
+    with _client() as client:
+        resp = client.post(
+            "/v2/auth/exchange",
+            json={
+                "intermediate_session_token": intermediate_session_token,
+                "organization_id": organization_id,
+            },
+        )
+    return cast(dict[str, Any], _json_or_raise(resp))
+
+
+def get_auth(session_jwt: str) -> dict[str, Any]:
+    """GET /v1/auth (session cookie) -> account incl. `id`."""
+    with _client(session_jwt) as client:
+        resp = client.get("/v1/auth")
+    return cast(dict[str, Any], _json_or_raise(resp))
+
+
+def list_projects(account_id: int, session_jwt: str) -> list[dict[str, Any]]:
+    """GET /v1/users/accounts/{id}/projects -> [{project, tokens[]}]."""
+    with _client(session_jwt) as client:
+        resp = client.get(f"/v1/users/accounts/{account_id}/projects")
+    return cast(list[dict[str, Any]], _json_or_raise(resp))
+
+
+def create_token(
+    account_id: int, project_id: int, token_name: str, session_jwt: str
+) -> dict[str, Any]:
+    """POST /v1/users/accounts/{id}/tokens -> TokenSchema incl. `api_key`."""
+    with _client(session_jwt) as client:
+        resp = client.post(
+            f"/v1/users/accounts/{account_id}/tokens",
+            json={"project_id": project_id, "token_name": token_name},
+        )
+    return cast(dict[str, Any], _json_or_raise(resp))

aai_cli/auth/discovery.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from urllib.parse import urlencode
+
+from aai_cli.auth import endpoints
+
+
+def build_start_url() -> str:
+    """The Stytch B2B OAuth discovery *start* URL the browser opens.
+
+    Client-side endpoint authenticated by the project's public token. After the
+    user authenticates with the provider, Stytch redirects to our loopback
+    `discovery_redirect_url` with `?stytch_token_type=discovery_oauth&token=...`.
+    No custom state is appended: the redirect is exact-match validated and the
+    server is loopback-only, single-shot.
+    """
+    base = (
+        f"{endpoints.stytch_domain()}"
+        f"/v1/b2b/public/oauth/{endpoints.STYTCH_OAUTH_PROVIDER}/discovery/start"
+    )
+    params = {
+        "public_token": endpoints.stytch_public_token(),
+        "discovery_redirect_url": endpoints.redirect_uri(),
+    }
+    return f"{base}?{urlencode(params)}"

aai_cli/auth/endpoints.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import os
+
+from aai_cli import environments
+
+# Constant across environments.
+STYTCH_OAUTH_PROVIDER = "google"
+CLI_TOKEN_NAME = "AssemblyAI CLI"  # noqa: S105 - display name, not a credential
+
+# Fixed loopback (Stytch does exact-match redirect validation; 8585 is registered).
+LOOPBACK_HOST = "127.0.0.1"
+LOOPBACK_PORT = int(os.environ.get("AAI_AUTH_PORT", "8585"))
+LOOPBACK_PATH = "/callback"
+
+
+# Environment-specific values resolve from the active environment (see
+# aai_cli.environments). The B2B OAuth flow starts on Stytch's API domain (not a
+# vanity CNAME) so the CSRF state cookie is set on the same domain as the OAuth
+# callback — otherwise Stytch returns `oauth_invalid_state`.
+def stytch_domain() -> str:
+    return environments.active().stytch_domain
+
+
+def stytch_public_token() -> str:
+    return environments.active().stytch_public_token
+
+
+def ams_base() -> str:
+    return environments.active().ams_base
+
+
+def signup_url() -> str:
+    return environments.active().signup_url
+
+
+def redirect_uri() -> str:
+    """The exact loopback redirect URL registered in Stytch."""
+    return f"http://{LOOPBACK_HOST}:{LOOPBACK_PORT}{LOOPBACK_PATH}"

aai_cli/auth/flow.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import webbrowser
+from typing import Any
+
+from aai_cli import output
+from aai_cli.auth import ams, discovery, endpoints, loopback
+from aai_cli.errors import APIError
+
+
+def _require(mapping: Any, key: str, what: str) -> Any:
+    """Pull a required field out of an AMS response, or raise a clean APIError.
+
+    AMS only returns HTTP errors for outright failures; a 200 with an unexpected
+    shape would otherwise KeyError into an ugly traceback, so map that to the same
+    "run login again" message the rest of the flow uses.
+    """
+    value = mapping.get(key) if isinstance(mapping, dict) else None
+    if value is None:
+        raise APIError(
+            f"Login failed: the server response was missing {what}. Run 'aai login' again."
+        )
+    return value
+
+
+def _open_browser(url: str) -> None:
+    """Open the system browser, falling back to printing the URL."""
+    output.console.print(f"Opening your browser to sign in:\n  {url}")
+    try:
+        webbrowser.open(url)
+    except Exception:  # noqa: BLE001 - opening a browser is best-effort
+        output.console.print(
+            "[aai.muted]Could not open a browser; open the URL above manually.[/aai.muted]"
+        )
+
+
+def _capture() -> loopback.CallbackResult:
+    return loopback.capture_callback()
+
+
+def _is_reusable_cli_token(token: dict[str, Any]) -> bool:
+    """A live 'AssemblyAI CLI' token whose key the list actually exposes."""
+    # List endpoints may key the display name as either "name" or "token_name"
+    # (the latter matches the create payload); accept either so we don't mint a
+    # duplicate every login. A token whose api_key the list omits can't be reused.
+    name = token.get("name") or token.get("token_name")
+    return (
+        name == endpoints.CLI_TOKEN_NAME
+        and not token.get("is_disabled")
+        and bool(token.get("api_key"))
+    )
+
+
+def find_or_create_cli_key(account_id: int, session_jwt: str) -> str:
+    """Return the existing 'AssemblyAI CLI' key, or create one in the first project."""
+    projects = ams.list_projects(account_id, session_jwt)
+    if not projects:
+        raise APIError("Your account has no project to create an API key in.")
+    for entry in projects:
+        for token in entry.get("tokens", []):
+            if _is_reusable_cli_token(token):
+                return str(token["api_key"])
+    project_id = _require(_require(projects[0], "project", "a project"), "id", "a project id")
+    created = ams.create_token(account_id, project_id, endpoints.CLI_TOKEN_NAME, session_jwt)
+    return str(_require(created, "api_key", "an API key"))
+
+
+def run_login_flow() -> str:
+    """Drive the full browser + AMS login and return a usable AssemblyAI API key."""
+    _open_browser(discovery.build_start_url())
+    result = _capture()
+
+    if result.error == "timeout":
+        raise APIError("Login timed out waiting for the browser. Run 'aai login' again.")
+    if result.token_type != "discovery_oauth" or not result.token:  # noqa: S105
+        raise APIError("Login did not return a valid OAuth token. Run 'aai login' again.")
+
+    disc = ams.discover(result.token)
+    organizations = disc.get("organizations") or []
+    if not organizations:
+        raise APIError(
+            "Signed in, but this identity has no AssemblyAI account yet. "
+            f"Create one at {endpoints.signup_url()}, then run 'aai login' again."
+        )
+    if len(organizations) > 1:
+        chosen = organizations[0].get("organization_name") or organizations[0].get(
+            "organization_id", "the first"
+        )
+        output.console.print(
+            f"[aai.muted]Found {len(organizations)} organizations; signing in to "
+            f"'{chosen}'.[/aai.muted]"
+        )
+    organization_id = _require(organizations[0], "organization_id", "an organization id")
+
+    intermediate_session_token = _require(disc, "intermediate_session_token", "a session token")
+    signed_in = ams.exchange(intermediate_session_token, organization_id)
+    session_jwt = _require(signed_in, "session_jwt", "a session token")
+    # `exchange` already returns the signed-in account, so read the id from it
+    # rather than making a second GET /v1/auth round-trip.
+    account_id = _require(_require(signed_in, "account", "an account"), "id", "an account id")
+    return find_or_create_cli_key(int(account_id), session_jwt)

aai_cli/auth/loopback.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from urllib.parse import parse_qs, urlparse
+
+from aai_cli.auth import endpoints
+from aai_cli.errors import APIError
+
+_SUCCESS_HTML = (
+    b"<html><body style='font-family:sans-serif'>"
+    b"<h2>Signed in.</h2><p>You can close this tab and return to the terminal.</p>"
+    b"</body></html>"
+)
+
+
+@dataclass
+class CallbackResult:
+    token: str | None = None
+    token_type: str | None = None
+    error: str | None = None
+
+
+def capture_callback(timeout: float = 120.0) -> CallbackResult:
+    """Bind the fixed loopback port, capture one OAuth callback, return its token.
+
+    Returns a CallbackResult; `error="timeout"` if no callback arrives in time.
+    """
+    result = CallbackResult()
+    done = threading.Event()
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_GET(self) -> None:  # stdlib API name
+            parsed = urlparse(self.path)
+            if parsed.path != endpoints.LOOPBACK_PATH:
+                self.send_response(404)
+                self.end_headers()
+                return
+            qs = parse_qs(parsed.query)
+            result.token = next(iter(qs.get("token", [])), None)
+            result.token_type = next(iter(qs.get("stytch_token_type", [])), None)
+            self.send_response(200)
+            self.send_header("Content-Type", "text/html")
+            self.end_headers()
+            self.wfile.write(_SUCCESS_HTML)
+            done.set()
+
+        def log_message(self, *args: object) -> None:  # silence stderr logging
+            pass
+
+    try:
+        server = HTTPServer((endpoints.LOOPBACK_HOST, endpoints.LOOPBACK_PORT), Handler)
+    except OSError as exc:
+        raise APIError(
+            f"Could not start the login callback server on "
+            f"{endpoints.LOOPBACK_HOST}:{endpoints.LOOPBACK_PORT} ({exc}). "
+            "Close whatever is using that port and run 'aai login' again."
+        ) from exc
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    try:
+        if not done.wait(timeout):
+            result.error = "timeout"
+    finally:
+        server.shutdown()
+        thread.join(timeout=5)
+    return result

assemblyai_cli/client.py → aai_cli/client.py

@@ -12,8 +12,8 @@
     StreamingParameters,
 )
 
-from assemblyai_cli import stdio
-from assemblyai_cli.errors import APIError, CLIError, UsageError, auth_failure, is_auth_failure
+from aai_cli import environments, stdio
+from aai_cli.errors import APIError, CLIError, UsageError, auth_failure, is_auth_failure
 
 SAMPLE_AUDIO_URL = "https://assembly.ai/wildfires.mp3"
 
@@ -32,6 +32,7 @@ def resolve_audio_source(source: str | None, *, sample: bool) -> str:
 
 def _configure(api_key: str) -> None:
     aai.settings.api_key = api_key
+    aai.settings.base_url = environments.active().api_base
 
 
 def validate_key(api_key: str) -> bool:
@@ -156,7 +157,7 @@ def stream_audio(
     `params` is a fully-built StreamingParameters (sample_rate/speech_model/etc).
     """
     sc = StreamingClient(
-        StreamingClientOptions(api_key=api_key, api_host="streaming.assemblyai.com")
+        StreamingClientOptions(api_key=api_key, api_host=environments.active().streaming_host)
     )
 
     def _guard(cb: Callable[[Any], Any]) -> Callable[[Any, Any], None]: