codemaster-cli 2.2.0__py3-none-any.whl

vibe/core/agents/__init__.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from vibe.core.agents.manager import AgentManager
+from vibe.core.agents.models import (
+    ACCEPT_EDITS,
+    AUTO_APPROVE,
+    BUILTIN_AGENTS,
+    DEFAULT,
+    EXPLORE,
+    PLAN,
+    PLAN_AGENT_TOOLS,
+    AgentProfile,
+    AgentSafety,
+    AgentType,
+    BuiltinAgentName,
+)
+
+__all__ = [
+    "ACCEPT_EDITS",
+    "AUTO_APPROVE",
+    "BUILTIN_AGENTS",
+    "DEFAULT",
+    "EXPLORE",
+    "PLAN",
+    "PLAN_AGENT_TOOLS",
+    "AgentManager",
+    "AgentProfile",
+    "AgentSafety",
+    "AgentType",
+    "BuiltinAgentName",
+]

vibe/core/agents/manager.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from logging import getLogger
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from vibe.core.agents.models import (
+    BUILTIN_AGENTS,
+    AgentProfile,
+    AgentType,
+    BuiltinAgentName,
+)
+from vibe.core.paths.config_paths import resolve_local_agents_dir
+from vibe.core.paths.global_paths import GLOBAL_AGENTS_DIR
+from vibe.core.utils import name_matches
+
+if TYPE_CHECKING:
+    from vibe.core.config import VibeConfig
+
+logger = getLogger("vibe")
+
+
+class AgentManager:
+    def __init__(
+        self,
+        config_getter: Callable[[], VibeConfig],
+        initial_agent: str = BuiltinAgentName.DEFAULT,
+    ) -> None:
+        self._config_getter = config_getter
+        self._search_paths = self._compute_search_paths(self._config)
+        self._available: dict[str, AgentProfile] = self._discover_agents()
+
+        custom_count = len(self._available) - len(BUILTIN_AGENTS)
+        if custom_count > 0:
+            custom_names = [
+                name for name in self._available if name not in BUILTIN_AGENTS
+            ]
+            logger.info(
+                "Discovered custom agents %s in %s",
+                " ".join(custom_names),
+                " ".join(str(p) for p in self._search_paths),
+            )
+
+        self.active_profile = self._available.get(
+            initial_agent, self._available[BuiltinAgentName.DEFAULT]
+        )
+        self._cached_config: VibeConfig | None = None
+
+    @property
+    def _config(self) -> VibeConfig:
+        return self._config_getter()
+
+    @property
+    def available_agents(self) -> dict[str, AgentProfile]:
+        if self._config.enabled_agents:
+            return {
+                name: profile
+                for name, profile in self._available.items()
+                if name_matches(name, self._config.enabled_agents)
+            }
+        if self._config.disabled_agents:
+            return {
+                name: profile
+                for name, profile in self._available.items()
+                if not name_matches(name, self._config.disabled_agents)
+            }
+        return dict(self._available)
+
+    @property
+    def config(self) -> VibeConfig:
+        if self._cached_config is None:
+            self._cached_config = self.active_profile.apply_to_config(self._config)
+        return self._cached_config
+
+    def switch_profile(self, name: str) -> None:
+        self.active_profile = self.get_agent(name)
+        self._cached_config = None
+
+    def invalidate_config(self) -> None:
+        self._cached_config = None
+
+    @staticmethod
+    def _compute_search_paths(config: VibeConfig) -> list[Path]:
+        paths: list[Path] = []
+        for path in config.agent_paths:
+            if path.is_dir():
+                paths.append(path)
+        if (agents_dir := resolve_local_agents_dir(Path.cwd())) is not None:
+            paths.append(agents_dir)
+        if GLOBAL_AGENTS_DIR.path.is_dir():
+            paths.append(GLOBAL_AGENTS_DIR.path)
+        unique: list[Path] = []
+        for p in paths:
+            rp = p.resolve()
+            if rp not in unique:
+                unique.append(rp)
+        return unique
+
+    def _discover_agents(self) -> dict[str, AgentProfile]:
+        agents: dict[str, AgentProfile] = dict(BUILTIN_AGENTS)
+
+        for base in self._search_paths:
+            if not base.is_dir():
+                continue
+            for agent_file in base.glob("*.toml"):
+                if not agent_file.is_file():
+                    continue
+                if (agent := self._try_load_agent(agent_file)) is not None:
+                    if agent.name in BUILTIN_AGENTS:
+                        logger.info(
+                            "Custom agent '%s' overrides builtin agent", agent.name
+                        )
+                    elif agent.name in agents:
+                        logger.debug(
+                            "Skipping duplicate agent '%s' at %s",
+                            agent.name,
+                            agent_file,
+                        )
+                        continue
+                    agents[agent.name] = agent
+
+        return agents
+
+    def _try_load_agent(self, agent_file: Path) -> AgentProfile | None:
+        try:
+            agent = AgentProfile.from_toml(agent_file)
+            agent.apply_to_config(self._config)
+            return agent
+        except Exception as e:
+            logger.warning("Failed to load agent at %s: %s", agent_file, e)
+            return None
+
+    def get_agent(self, name: str) -> AgentProfile:
+        if agent := self.available_agents.get(name):
+            return agent
+        raise ValueError(f"Agent '{name}' not found")
+
+    def get_subagents(self) -> list[AgentProfile]:
+        return [
+            a
+            for a in self.available_agents.values()
+            if a.agent_type == AgentType.SUBAGENT
+        ]
+
+    def get_agent_order(self) -> list[str]:
+        builtin_order: list[str] = [
+            BuiltinAgentName.DEFAULT,
+            BuiltinAgentName.PLAN,
+            BuiltinAgentName.ACCEPT_EDITS,
+            BuiltinAgentName.AUTO_APPROVE,
+        ]
+        primary_agents = [
+            name
+            for name, agent in self.available_agents.items()
+            if agent.agent_type == AgentType.AGENT
+        ]
+        order = [name for name in builtin_order if name in primary_agents]
+        custom = sorted(name for name in primary_agents if name not in builtin_order)
+        return order + custom
+
+    def next_agent(self, current: AgentProfile) -> AgentProfile:
+        order = self.get_agent_order()
+        idx = order.index(current.name) if current.name in order else -1
+        return self.available_agents[order[(idx + 1) % len(order)]]

vibe/core/agents/models.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import StrEnum, auto
+from pathlib import Path
+import tomllib
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from vibe.core.config import VibeConfig
+
+
+def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
+    result = base.copy()
+    for key, value in override.items():
+        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+            result[key] = _deep_merge(result[key], value)
+        else:
+            result[key] = value
+    return result
+
+
+class AgentSafety(StrEnum):
+    SAFE = auto()
+    NEUTRAL = auto()
+    DESTRUCTIVE = auto()
+    YOLO = auto()
+
+
+class AgentType(StrEnum):
+    AGENT = auto()
+    SUBAGENT = auto()
+
+
+class BuiltinAgentName(StrEnum):
+    DEFAULT = "default"
+    PLAN = "plan"
+    ACCEPT_EDITS = "accept-edits"
+    AUTO_APPROVE = "auto-approve"
+    EXPLORE = "explore"
+
+
+@dataclass(frozen=True)
+class AgentProfile:
+    name: str
+    display_name: str
+    description: str
+    safety: AgentSafety
+    agent_type: AgentType = AgentType.AGENT
+    overrides: dict[str, Any] = field(default_factory=dict)
+
+    def apply_to_config(self, base: VibeConfig) -> VibeConfig:
+        from vibe.core.config import VibeConfig as VC
+
+        merged = _deep_merge(base.model_dump(), self.overrides)
+        return VC.model_validate(merged)
+
+    @classmethod
+    def from_toml(cls, path: Path) -> AgentProfile:
+        with path.open("rb") as f:
+            data = tomllib.load(f)
+        return cls(
+            name=path.stem,
+            display_name=data.pop("display_name", path.stem.replace("-", " ").title()),
+            description=data.pop("description", ""),
+            safety=AgentSafety(data.pop("safety", AgentSafety.NEUTRAL)),
+            agent_type=AgentType(data.pop("agent_type", AgentType.AGENT)),
+            overrides=data,
+        )
+
+
+PLAN_AGENT_TOOLS = ["grep", "read_file", "todo", "ask_user_question", "task"]
+
+DEFAULT = AgentProfile(
+    BuiltinAgentName.DEFAULT,
+    "Default",
+    "Requires approval for tool executions",
+    AgentSafety.NEUTRAL,
+)
+PLAN = AgentProfile(
+    BuiltinAgentName.PLAN,
+    "Plan",
+    "Read-only agent for exploration and planning",
+    AgentSafety.SAFE,
+    overrides={"auto_approve": True, "enabled_tools": PLAN_AGENT_TOOLS},
+)
+ACCEPT_EDITS = AgentProfile(
+    BuiltinAgentName.ACCEPT_EDITS,
+    "Accept Edits",
+    "Auto-approves file edits only",
+    AgentSafety.DESTRUCTIVE,
+    overrides={
+        "tools": {
+            "write_file": {"permission": "always"},
+            "search_replace": {"permission": "always"},
+        }
+    },
+)
+AUTO_APPROVE = AgentProfile(
+    BuiltinAgentName.AUTO_APPROVE,
+    "Auto Approve",
+    "Auto-approves all tool executions",
+    AgentSafety.YOLO,
+    overrides={"auto_approve": True},
+)
+
+EXPLORE = AgentProfile(
+    name=BuiltinAgentName.EXPLORE,
+    display_name="Explore",
+    description="Read-only subagent for codebase exploration",
+    safety=AgentSafety.SAFE,
+    agent_type=AgentType.SUBAGENT,
+    overrides={"enabled_tools": ["grep", "read_file"], "system_prompt_id": "explore"},
+)
+
+BUILTIN_AGENTS: dict[str, AgentProfile] = {
+    BuiltinAgentName.DEFAULT: DEFAULT,
+    BuiltinAgentName.PLAN: PLAN,
+    BuiltinAgentName.ACCEPT_EDITS: ACCEPT_EDITS,
+    BuiltinAgentName.AUTO_APPROVE: AUTO_APPROVE,
+    BuiltinAgentName.EXPLORE: EXPLORE,
+}

vibe/core/auth/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from vibe.core.auth.crypto import EncryptedPayload, decrypt, encrypt
+from vibe.core.auth.github import GitHubAuthProvider
+
+__all__ = ["EncryptedPayload", "GitHubAuthProvider", "decrypt", "encrypt"]

vibe/core/auth/crypto.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import base64
+import binascii
+from dataclasses import dataclass
+import os
+
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+_AES_KEY_SIZE = 32
+_NONCE_SIZE = 12
+_MIN_RSA_KEY_SIZE = 2048
+_MAX_ENCRYPTED_KEY_SIZE = 1024
+_MAX_CIPHERTEXT_SIZE = 2 * 1024 * 1024  # Workflow transport limit: 2MB
+_PAYLOAD_VERSION = 1
+_ALG = "RSA-OAEP-SHA256"
+_ENC = "A256GCM"
+
+
+@dataclass(frozen=True)
+class EncryptedPayload:
+    encrypted_key: str
+    nonce: str
+    ciphertext: str
+    version: int | None = None
+    alg: str | None = None
+    enc: str | None = None
+    kid: str | None = None
+    purpose: str | None = None
+
+
+def _b64decode_strict(value: str, field_name: str) -> bytes:
+    try:
+        return base64.b64decode(value, validate=True)
+    except (binascii.Error, ValueError) as exc:
+        raise ValueError(f"Invalid base64 for {field_name}") from exc
+
+
+def _validate_payload_lengths(
+    encrypted_key: bytes, nonce: bytes, ciphertext: bytes
+) -> None:
+    if len(encrypted_key) > _MAX_ENCRYPTED_KEY_SIZE:
+        raise ValueError("Encrypted key too large")
+    if len(nonce) != _NONCE_SIZE:
+        raise ValueError("Invalid nonce size")
+    if not ciphertext:
+        raise ValueError("Ciphertext is empty")
+    if len(ciphertext) > _MAX_CIPHERTEXT_SIZE:
+        raise ValueError("Ciphertext exceeds maximum allowed size")
+
+
+def _build_aad(payload: EncryptedPayload) -> bytes | None:
+    if payload.version is None or payload.version <= 0:
+        return None
+    alg = payload.alg or _ALG
+    enc = payload.enc or _ENC
+    parts = [f"v={payload.version}", f"alg={alg}", f"enc={enc}"]
+    if payload.kid:
+        parts.append(f"kid={payload.kid}")
+    if payload.purpose:
+        parts.append(f"purpose={payload.purpose}")
+    return "|".join(parts).encode("utf-8")
+
+
+def encrypt(plaintext: str, public_key_pem: bytes) -> EncryptedPayload:
+    public_key = serialization.load_pem_public_key(public_key_pem)
+    if not isinstance(public_key, RSAPublicKey):
+        raise TypeError("Expected RSA public key")
+
+    if public_key.key_size < _MIN_RSA_KEY_SIZE:
+        raise ValueError(f"RSA key size must be at least {_MIN_RSA_KEY_SIZE} bits")
+
+    aes_key = os.urandom(_AES_KEY_SIZE)
+    nonce = os.urandom(_NONCE_SIZE)
+
+    payload = EncryptedPayload(
+        encrypted_key="",
+        nonce="",
+        ciphertext="",
+        version=_PAYLOAD_VERSION,
+        alg=_ALG,
+        enc=_ENC,
+    )
+    aad = _build_aad(payload)
+    aesgcm = AESGCM(aes_key)
+    ciphertext = aesgcm.encrypt(nonce, plaintext.encode("utf-8"), aad)
+
+    encrypted_key = public_key.encrypt(
+        aes_key,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA256()),
+            algorithm=hashes.SHA256(),
+            label=None,
+        ),
+    )
+
+    return EncryptedPayload(
+        encrypted_key=base64.b64encode(encrypted_key).decode("ascii"),
+        nonce=base64.b64encode(nonce).decode("ascii"),
+        ciphertext=base64.b64encode(ciphertext).decode("ascii"),
+        version=payload.version,
+        alg=payload.alg,
+        enc=payload.enc,
+    )
+
+
+def decrypt(payload: EncryptedPayload, private_key_pem: bytes) -> str:
+    private_key = serialization.load_pem_private_key(private_key_pem, password=None)
+    if not isinstance(private_key, RSAPrivateKey):
+        raise TypeError("Expected RSA private key")
+
+    if private_key.key_size < _MIN_RSA_KEY_SIZE:
+        raise ValueError(f"RSA key size must be at least {_MIN_RSA_KEY_SIZE} bits")
+
+    encrypted_key = _b64decode_strict(payload.encrypted_key, "encrypted_key")
+    nonce = _b64decode_strict(payload.nonce, "nonce")
+    ciphertext = _b64decode_strict(payload.ciphertext, "ciphertext")
+    _validate_payload_lengths(encrypted_key, nonce, ciphertext)
+
+    aes_key = private_key.decrypt(
+        encrypted_key,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA256()),
+            algorithm=hashes.SHA256(),
+            label=None,
+        ),
+    )
+
+    if len(aes_key) != _AES_KEY_SIZE:
+        raise ValueError("Invalid AES key size after decryption")
+
+    aesgcm = AESGCM(aes_key)
+    aad = _build_aad(payload)
+    return aesgcm.decrypt(nonce, ciphertext, aad).decode("utf-8")

vibe/core/auth/github.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+import types
+import webbrowser
+
+import httpx
+import keyring
+import keyring.errors
+
+GITHUB_CLIENT_ID = "Ov23liJ7sk5kFDMEyvDT"
+
+_SERVICE_NAME = "vibe"
+_KEYRING_USERNAME = "github_token"
+_DEVICE_CODE_URL = "https://github.com/login/device/code"
+_TOKEN_URL = "https://github.com/login/oauth/access_token"
+_VALIDATE_URL = "https://api.github.com/user"
+_SCOPES = "repo read:org write:org workflow read:user user:email"
+
+
+class GitHubAuthError(Exception):
+    pass
+
+
+@dataclass
+class DeviceFlowInfo:
+    user_code: str
+    verification_uri: str
+
+
+@dataclass
+class DeviceFlowHandle:
+    device_code: str
+    expires_in: int
+    info: DeviceFlowInfo
+
+
+class GitHubAuthProvider:
+    def __init__(
+        self,
+        client_id: str = GITHUB_CLIENT_ID,
+        *,
+        client: httpx.AsyncClient | None = None,
+        timeout: float = 60.0,
+    ) -> None:
+        self._client_id = client_id
+        self._client = client
+        self._owns_client = client is None
+        self._timeout = timeout
+
+    async def __aenter__(self) -> GitHubAuthProvider:
+        if self._client is None:
+            self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        if self._owns_client and self._client:
+            await self._client.aclose()
+            self._client = None
+
+    def _get_client(self) -> httpx.AsyncClient:
+        if self._client is None:
+            self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
+            self._owns_client = True
+        return self._client
+
+    def get_token(self) -> str | None:
+        try:
+            return keyring.get_password(_SERVICE_NAME, _KEYRING_USERNAME)
+        except keyring.errors.KeyringError:
+            return None
+
+    def has_token(self) -> bool:
+        return bool(self.get_token())
+
+    def delete_token(self) -> None:
+        try:
+            keyring.delete_password(_SERVICE_NAME, _KEYRING_USERNAME)
+        except keyring.errors.KeyringError:
+            pass
+
+    async def get_valid_token(self) -> str | None:
+        token = self.get_token()
+        if not token:
+            return None
+        if await self._is_token_valid(token):
+            return token
+        self.delete_token()
+        return None
+
+    async def _is_token_valid(self, token: str) -> bool:
+        client = self._get_client()
+        try:
+            response = await client.get(
+                _VALIDATE_URL,
+                headers={
+                    "Authorization": f"Bearer {token}",
+                    "Accept": "application/vnd.github+json",
+                },
+            )
+            return response.is_success
+        except httpx.HTTPError:
+            return False
+
+    async def start_device_flow(self, open_browser: bool = True) -> DeviceFlowHandle:
+        client = self._get_client()
+        response = await client.post(
+            _DEVICE_CODE_URL,
+            data={"client_id": self._client_id, "scope": _SCOPES},
+            headers={"Accept": "application/json"},
+        )
+        if not response.is_success:
+            raise GitHubAuthError(f"Failed to initiate device flow: {response.text}")
+
+        data = response.json()
+
+        if open_browser:
+            webbrowser.open(data["verification_uri"])
+
+        return DeviceFlowHandle(
+            device_code=data["device_code"],
+            expires_in=data["expires_in"],
+            info=DeviceFlowInfo(data["user_code"], data["verification_uri"]),
+        )
+
+    async def wait_for_token(self, handle: DeviceFlowHandle) -> str:
+        client = self._get_client()
+        token = await self._poll_for_token(
+            client, handle.device_code, handle.expires_in, interval=1
+        )
+        self._save_token(token)
+        return token
+
+    def _save_token(self, token: str) -> None:
+        try:
+            keyring.set_password(_SERVICE_NAME, _KEYRING_USERNAME, token)
+        except keyring.errors.KeyringError as e:
+            raise GitHubAuthError(f"Failed to save token to keyring: {e}") from e
+
+    async def _poll_for_token(
+        self,
+        client: httpx.AsyncClient,
+        device_code: str,
+        expires_in: int,
+        interval: int,
+    ) -> str:
+        elapsed = 0.0
+        while elapsed < expires_in:
+            await asyncio.sleep(interval)
+            elapsed += interval
+
+            response = await client.post(
+                _TOKEN_URL,
+                data={
+                    "client_id": self._client_id,
+                    "device_code": device_code,
+                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                },
+                headers={"Accept": "application/json"},
+            )
+            result = response.json()
+
+            if "access_token" in result:
+                return result["access_token"]
+
+            error = result.get("error")
+            if error == "slow_down":
+                interval = result.get("interval", interval + 5)
+            elif error in {"expired_token", "access_denied"}:
+                raise GitHubAuthError(f"Authentication failed: {error}")
+
+        raise GitHubAuthError("Authentication timed out")