vibe/vibe/cli/cli.py

from __future__ import annotations

import argparse
from pathlib import Path
import sys

from rich import print as rprint
from rich.console import Console
import tomli_w

from vibe import __version__
from vibe.cli.textual_ui.app import StartupOptions, run_textual_ui
from vibe.core.agent_loop import AgentLoop, TeleportError
from vibe.core.agents.models import BuiltinAgentName
from vibe.core.config import (
    MissingAPIKeyError,
    MissingPromptFileError,
    VibeConfig,
    load_dotenv_values,
)
from vibe.core.config.harness_files import get_harness_files_manager
from vibe.core.hooks.config import load_hooks_from_fs
from vibe.core.logger import logger
from vibe.core.paths import HISTORY_FILE
from vibe.core.programmatic import run_programmatic
from vibe.core.session import last_session_pointer
from vibe.core.session.session_loader import SessionLoader
from vibe.core.telemetry.build_metadata import build_entrypoint_metadata
from vibe.core.telemetry.types import EntrypointMetadata
from vibe.core.tracing import setup_tracing
from vibe.core.trusted_folders import find_trustable_files, trusted_folders_manager
from vibe.core.types import LLMMessage, OutputFormat, Role
from vibe.core.utils import ConversationLimitException
from vibe.setup.onboarding import run_onboarding


def _build_cli_entrypoint_metadata() -> EntrypointMetadata:
    return build_entrypoint_metadata(
        agent_entrypoint="cli",
        agent_version=__version__,
        client_name="vibe_cli",
        client_version=__version__,
    )


def get_initial_agent_name(args: argparse.Namespace, config: VibeConfig) -> str:
    if args.prompt is not None and args.agent is None:
        return BuiltinAgentName.AUTO_APPROVE
    return args.agent or config.default_agent


def get_prompt_from_stdin() -> str | None:
    if sys.stdin.isatty():
        return None
    try:
        if content := sys.stdin.read().strip():
            sys.stdin = sys.__stdin__ = open("/dev/tty")
            return content
    except KeyboardInterrupt:
        pass
    except OSError:
        return None

    return None


def load_config_or_exit(*, interactive: bool) -> VibeConfig:
    try:
        return VibeConfig.load()
    except MissingAPIKeyError as e:
        if not interactive:
            print(
                f"Error: {e}. Set the environment variable (e.g. in ~/.vibe/.env "
                "or your shell), or run `vibe --setup` once interactively.",
                file=sys.stderr,
            )
            sys.exit(1)
        run_onboarding(entrypoint_metadata=_build_cli_entrypoint_metadata())
        return VibeConfig.load()
    except MissingPromptFileError as e:
        rprint(f"[yellow]Invalid system prompt id: {e}[/]")
        sys.exit(1)
    except ValueError as e:
        rprint(f"[yellow]{e}[/]")
        sys.exit(1)


def warn_if_workdir_trust_is_unset() -> None:
    try:
        cwd = Path.cwd()
    except FileNotFoundError:
        return
    if cwd.resolve() == Path.home().resolve():
        return
    if trusted_folders_manager.is_trusted(cwd) is not None:
        return
    detected = find_trustable_files(cwd)
    if not detected:
        return
    files_str = ", ".join(detected)
    Console(stderr=True).print(
        f"[yellow]Warning:[/] {cwd} is not trusted; "
        f"project configuration ({files_str}) will be ignored. "
        "Re-run with --trust to trust this folder temporarily."
    )


def bootstrap_config_files() -> None:
    mgr = get_harness_files_manager()
    config_file = mgr.user_config_file
    if not config_file.exists():
        try:
            config_file.parent.mkdir(parents=True, exist_ok=True)
            with config_file.open("wb") as f:
                tomli_w.dump(VibeConfig.create_default(), f)
        except Exception as e:
            rprint(f"[yellow]Could not create default config file: {e}[/]")

    history_file = HISTORY_FILE.path
    if not history_file.exists():
        try:
            history_file.parent.mkdir(parents=True, exist_ok=True)
            history_file.write_text("Hello Vibe!\n", "utf-8")
        except Exception as e:
            rprint(f"[yellow]Could not create history file: {e}[/]")


def load_session(
    args: argparse.Namespace, config: VibeConfig
) -> tuple[list[LLMMessage], Path] | None:
    if not args.continue_session and not args.resume:
        return None

    if not config.session_logging.enabled:
        rprint(
            "[red]Session logging is disabled. "
            "Enable it in config to use --continue or --resume[/]"
        )
        sys.exit(1)

    session_to_load = None
    if args.continue_session:
        cwd = Path.cwd().resolve()
        pointer_session_id = last_session_pointer.load(config.session_logging)
        if pointer_session_id:
            session_to_load = SessionLoader.find_session_by_id(
                pointer_session_id, config.session_logging, working_directory=cwd
            )
        if not session_to_load:
            session_to_load = SessionLoader.find_latest_session(
                config.session_logging, working_directory=cwd
            )
        if not session_to_load:
            rprint(
                f"[red]No previous sessions found in "
                f"{config.session_logging.save_dir} for {cwd=}[/]"
            )
            sys.exit(1)
    elif args.resume is True:
        return None
    else:
        session_to_load = SessionLoader.find_session_by_id(
            args.resume, config.session_logging
        )
        if not session_to_load:
            rprint(
                f"[red]Session '{args.resume}' not found in "
                f"{config.session_logging.save_dir}[/]"
            )
            sys.exit(1)

    try:
        loaded_messages, _ = SessionLoader.load_session(session_to_load)
        return loaded_messages, session_to_load
    except Exception as e:
        rprint(f"[red]Failed to load session: {e}[/]")
        sys.exit(1)


def _resume_previous_session(
    agent_loop: AgentLoop, loaded_messages: list[LLMMessage], session_path: Path
) -> None:
    non_system_messages = [msg for msg in loaded_messages if msg.role != Role.system]
    agent_loop.messages.extend(non_system_messages)

    _, metadata = SessionLoader.load_session(session_path)
    session_id = metadata.get("session_id", agent_loop.session_id)
    agent_loop.session_id = session_id
    agent_loop.parent_session_id = metadata.get("parent_session_id")
    agent_loop.session_logger.resume_existing_session(session_id, session_path)

    logger.info(
        "Resumed session %s with %d messages", session_id, len(non_system_messages)
    )


def run_cli(args: argparse.Namespace) -> None:
    load_dotenv_values()
    bootstrap_config_files()

    if args.setup:
        run_onboarding(entrypoint_metadata=_build_cli_entrypoint_metadata())
        sys.exit(0)

    try:
        is_interactive = args.prompt is None
        config = load_config_or_exit(interactive=is_interactive)
        initial_agent_name = get_initial_agent_name(args, config)
        hook_config_result = load_hooks_from_fs(config)
        setup_tracing(config)

        if args.enabled_tools:
            config.enabled_tools = args.enabled_tools

        loaded_session = load_session(args, config)

        stdin_prompt = get_prompt_from_stdin()
        if args.prompt is not None:
            warn_if_workdir_trust_is_unset()
            config.disabled_tools = [
                *config.disabled_tools,
                "ask_user_question",
                "exit_plan_mode",
            ]
            programmatic_prompt = args.prompt or stdin_prompt
            if not programmatic_prompt:
                print(
                    "Error: No prompt provided for programmatic mode", file=sys.stderr
                )
                sys.exit(1)
            output_format = OutputFormat(
                args.output if hasattr(args, "output") else "text"
            )

            try:
                final_response = run_programmatic(
                    config=config,
                    prompt=programmatic_prompt or "",
                    max_turns=args.max_turns,
                    max_price=args.max_price,
                    output_format=output_format,
                    previous_messages=loaded_session[0] if loaded_session else None,
                    agent_name=initial_agent_name,
                    teleport=args.teleport and config.vibe_code_enabled,
                    headless=True,
                    hook_config_result=hook_config_result,
                )
                if final_response:
                    print(final_response)
                sys.exit(0)
            except ConversationLimitException as e:
                print(e, file=sys.stderr)
                sys.exit(1)
            except TeleportError as e:
                print(f"Teleport error: {e}", file=sys.stderr)
                sys.exit(1)
            except (RuntimeError, ValueError) as e:
                print(f"Error: {e}", file=sys.stderr)
                sys.exit(1)
        else:
            try:
                agent_loop = AgentLoop(
                    config,
                    agent_name=initial_agent_name,
                    enable_streaming=True,
                    entrypoint_metadata=_build_cli_entrypoint_metadata(),
                    defer_heavy_init=True,
                    hook_config_result=hook_config_result,
                )
            except ValueError as e:
                rprint(f"[red]Error:[/] {e}")
                sys.exit(1)

            if loaded_session:
                _resume_previous_session(agent_loop, *loaded_session)

            run_textual_ui(
                agent_loop=agent_loop,
                startup=StartupOptions(
                    initial_prompt=args.initial_prompt or stdin_prompt,
                    teleport_on_start=args.teleport,
                    show_resume_picker=args.resume is True,
                    is_resuming_session=loaded_session is not None,
                ),
            )

    except (KeyboardInterrupt, EOFError):
        rprint("\n[dim]Bye![/]")
        sys.exit(0)