vibe/vibe/setup/onboarding/screens/api_key.py

from __future__ import annotations

import os
from typing import ClassVar

from dotenv import set_key
from textual.app import ComposeResult
from textual.binding import Binding, BindingType
from textual.containers import Center, Horizontal, Vertical
from textual.events import MouseUp
from textual.validation import Length
from textual.widgets import Input, Link, Static

from vibe.cli.clipboard import copy_selection_to_clipboard
from vibe.cli.textual_ui.widgets.no_markup_static import NoMarkupStatic
from vibe.core.config import DEFAULT_PROVIDERS, ProviderConfig, VibeConfig
from vibe.core.paths import GLOBAL_ENV_FILE
from vibe.core.telemetry.send import TelemetryClient
from vibe.core.types import Backend
from vibe.setup.onboarding.base import OnboardingScreen
from vibe.setup.onboarding.context import OnboardingContext

PROVIDER_HELP = {
    "mistral": ("https://console.mistral.ai/codestral/cli", "Mistral AI Studio")
}
CONFIG_DOCS_URL = (
    "https://github.com/mistralai/mistral-vibe?tab=readme-ov-file#configuration"
)


def _save_api_key_to_env_file(env_key: str, api_key: str) -> None:
    GLOBAL_ENV_FILE.path.parent.mkdir(parents=True, exist_ok=True)
    set_key(GLOBAL_ENV_FILE.path, env_key, api_key)


def persist_api_key(provider: ProviderConfig, api_key: str) -> str:
    env_key = provider.api_key_env_var
    if not env_key:
        return "env_var_error:<empty>"
    try:
        os.environ[env_key] = api_key
    except ValueError:
        return f"env_var_error:{env_key}"
    try:
        _save_api_key_to_env_file(env_key, api_key)
    except (OSError, ValueError) as err:
        return f"save_error:{err}"
    if provider.backend == Backend.MISTRAL:
        try:
            telemetry = TelemetryClient(config_getter=VibeConfig)
            telemetry.send_onboarding_api_key_added()
        except Exception:
            pass
    return "completed"


def _get_mistral_provider() -> ProviderConfig:
    return next(
        provider for provider in DEFAULT_PROVIDERS if provider.name == "mistral"
    )


def _resolve_onboarding_provider(
    provider: ProviderConfig | None = None,
) -> ProviderConfig:
    resolved_provider = provider or OnboardingContext.load().provider
    if resolved_provider.api_key_env_var:
        return resolved_provider
    return _get_mistral_provider()


class ApiKeyScreen(OnboardingScreen):
    BINDINGS: ClassVar[list[BindingType]] = [
        Binding("ctrl+c", "cancel", "Cancel", show=False),
        Binding("escape", "cancel", "Cancel", show=False),
    ]

    NEXT_SCREEN = None

    def __init__(self, provider: ProviderConfig | None = None) -> None:
        super().__init__()
        self.provider = _resolve_onboarding_provider(provider)

    def _compose_provider_link(self, provider_name: str) -> ComposeResult:
        if self.provider.name not in PROVIDER_HELP:
            return

        help_url, help_name = PROVIDER_HELP[self.provider.name]
        yield NoMarkupStatic(f"Grab your {provider_name} API key from the {help_name}:")
        yield Center(
            Horizontal(
                NoMarkupStatic("→ ", classes="link-chevron"),
                Link(help_url, url=help_url),
                classes="link-row",
            )
        )

    def _compose_config_docs(self) -> ComposeResult:
        yield Static("[dim]Learn more about Vibe configuration:[/]")
        yield Horizontal(
            NoMarkupStatic("→ ", classes="link-chevron"),
            Link(CONFIG_DOCS_URL, url=CONFIG_DOCS_URL),
            classes="link-row",
        )

    def compose(self) -> ComposeResult:
        provider_name = self.provider.name.capitalize()

        self.input_widget = Input(
            password=True,
            id="key",
            placeholder="Paste your API key here",
            validators=[Length(minimum=1, failure_description="No API key provided.")],
        )

        with Vertical(id="api-key-outer"):
            yield NoMarkupStatic("", classes="spacer")
            yield Center(NoMarkupStatic("One last thing...", id="api-key-title"))
            with Center():
                with Vertical(id="api-key-content"):
                    yield from self._compose_provider_link(provider_name)
                    yield NoMarkupStatic(
                        "...and paste it below to finish the setup:", id="paste-hint"
                    )
                    yield Center(Horizontal(self.input_widget, id="input-box"))
                    yield NoMarkupStatic("", id="feedback")
            yield NoMarkupStatic("", classes="spacer")
            yield Vertical(
                Vertical(*self._compose_config_docs(), id="config-docs-group"),
                id="config-docs-section",
            )

    def on_mount(self) -> None:
        self.input_widget.focus()

    def on_input_changed(self, event: Input.Changed) -> None:
        feedback = self.query_one("#feedback", NoMarkupStatic)
        input_box = self.query_one("#input-box")

        if event.validation_result is None:
            return

        input_box.remove_class("valid", "invalid")
        feedback.remove_class("error", "success")

        if event.validation_result.is_valid:
            feedback.update("Press Enter to submit ↵")
            feedback.add_class("success")
            input_box.add_class("valid")
            return

        descriptions = event.validation_result.failure_descriptions
        feedback.update(descriptions[0])
        feedback.add_class("error")
        input_box.add_class("invalid")

    def on_input_submitted(self, event: Input.Submitted) -> None:
        if event.validation_result and event.validation_result.is_valid:
            self._save_and_finish(event.value)

    def _save_and_finish(self, api_key: str) -> None:
        self.app.exit(persist_api_key(self.provider, api_key))

    def on_mouse_up(self, event: MouseUp) -> None:
        copy_selection_to_clipboard(self.app)