codemaster-cli 2.2.0__py3-none-any.whl

vibe/core/teleport/git.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from git import InvalidGitRepositoryError, Repo
+from git.exc import GitCommandError
+from giturlparse import parse as parse_git_url
+
+from vibe.core.teleport.errors import (
+    ServiceTeleportError,
+    ServiceTeleportNotSupportedError,
+)
+from vibe.core.utils import AsyncExecutor
+
+
+@dataclass
+class GitRepoInfo:
+    remote_url: str
+    owner: str
+    repo: str
+    branch: str | None
+    commit: str
+    diff: str
+
+
+class GitRepository:
+    def __init__(self, workdir: Path | None = None) -> None:
+        self._workdir = workdir or Path.cwd()
+        self._repo: Repo | None = None
+        # For network I/O (fetch, push) and potentially slow git commands (diff, rev-list)
+        self._executor = AsyncExecutor(max_workers=2, timeout=60.0, name="git")
+
+    async def __aenter__(self) -> GitRepository:
+        return self
+
+    async def __aexit__(self, *_: object) -> None:
+        self._executor.shutdown(wait=False)
+
+    async def is_supported(self) -> bool:
+        try:
+            repo = self._repo_or_raise()
+        except ServiceTeleportNotSupportedError:
+            return False
+        return self._find_github_remote(repo) is not None
+
+    async def get_info(self) -> GitRepoInfo:
+        repo = self._repo_or_raise()
+
+        parsed = self._find_github_remote(repo)
+        if not parsed:
+            raise ServiceTeleportNotSupportedError(
+                "No GitHub remote found. Teleport only supports GitHub repositories."
+            )
+
+        try:
+            commit = repo.head.commit.hexsha
+        except (ValueError, TypeError) as e:
+            raise ServiceTeleportNotSupportedError(
+                "Could not determine current commit"
+            ) from e
+
+        if not commit:
+            raise ServiceTeleportNotSupportedError("Could not determine current commit")
+
+        owner, repo_name = parsed
+        branch = None if repo.head.is_detached else repo.active_branch.name
+        diff = await self._get_diff(repo)
+
+        return GitRepoInfo(
+            remote_url=self._to_https_url(owner, repo_name),
+            owner=owner,
+            repo=repo_name,
+            branch=branch,
+            commit=commit,
+            diff=diff,
+        )
+
+    async def is_commit_pushed(self, commit: str, remote: str = "origin") -> bool:
+        repo = self._repo_or_raise()
+        await self._fetch(repo, remote)
+        return await self._branch_contains(repo, commit, remote)
+
+    async def get_unpushed_commit_count(self, remote: str = "origin") -> int:
+        repo = self._repo_or_raise()
+
+        if repo.head.is_detached:
+            raise ServiceTeleportError(
+                "Cannot count unpushed commits: no current branch"
+            )
+        branch = repo.active_branch.name
+
+        await self._fetch(repo, remote)
+
+        result = await self._rev_list_count(repo, f"{remote}/{branch}..HEAD")
+        if result is not None:
+            return result
+
+        # Fallback: branch not pushed yet, count commits from default branch
+        default_branch = await self._get_remote_default_branch(repo, remote)
+        if default_branch:
+            result = await self._rev_list_count(repo, f"{default_branch}..HEAD")
+            if result is not None:
+                return result
+
+        raise ServiceTeleportError(f"Failed to count unpushed commits for {branch}")
+
+    async def push_current_branch(self, remote: str = "origin") -> bool:
+        repo = self._repo_or_raise()
+        if repo.head.is_detached:
+            return False
+        return await self._push(repo, repo.active_branch.name, remote)
+
+    def _repo_or_raise(self) -> Repo:
+        if self._repo is None:
+            try:
+                self._repo = Repo(self._workdir, search_parent_directories=True)
+            except InvalidGitRepositoryError as e:
+                raise ServiceTeleportNotSupportedError("Not a git repository") from e
+        return self._repo
+
+    def _find_github_remote(self, repo: Repo) -> tuple[str, str] | None:
+        for remote in repo.remotes:
+            for url in remote.urls:
+                if parsed := self._parse_github_url(url):
+                    return parsed
+        return None
+
+    async def _fetch(self, repo: Repo, remote: str) -> None:
+        try:
+            await self._executor.run(lambda: repo.remote(remote).fetch())
+        except (TimeoutError, ValueError, GitCommandError):
+            pass
+
+    async def _get_diff(self, repo: Repo) -> str:
+        def get_full_diff() -> str:
+            # Mark untracked files as intent-to-add so they appear in diff
+            repo.git.add("-N", ".")
+            return repo.git.diff("HEAD", binary=True)
+
+        try:
+            return await self._executor.run(get_full_diff)
+        except (TimeoutError, GitCommandError):
+            return ""
+
+    async def _branch_contains(self, repo: Repo, commit: str, remote: str) -> bool:
+        try:
+            out = await self._executor.run(
+                lambda: repo.git.branch("-r", "--contains", commit)
+            )
+            return any(ln.strip().startswith(f"{remote}/") for ln in out.splitlines())
+        except (TimeoutError, GitCommandError):
+            return False
+
+    async def _rev_list_count(self, repo: Repo, ref_range: str) -> int | None:
+        try:
+            out = await self._executor.run(
+                lambda: repo.git.rev_list("--count", ref_range)
+            )
+            return int(out)
+        except (TimeoutError, GitCommandError, ValueError):
+            return None
+
+    async def _ref_exists(self, repo: Repo, ref: str) -> bool:
+        try:
+            await self._executor.run(lambda: repo.git.rev_parse("--verify", ref))
+            return True
+        except (TimeoutError, GitCommandError):
+            return False
+
+    async def _get_remote_default_branch(self, repo: Repo, remote: str) -> str | None:
+        try:
+            ref = repo.remotes[remote].refs.HEAD.reference.name
+            if await self._ref_exists(repo, ref):
+                return ref
+        except (KeyError, IndexError, TypeError, AttributeError):
+            pass
+        return None
+
+    async def _push(self, repo: Repo, branch: str, remote: str) -> bool:
+        try:
+            await self._executor.run(lambda: repo.remote(remote).push(branch))
+            return True
+        except (TimeoutError, ValueError, GitCommandError):
+            return False
+
+    @staticmethod
+    def _parse_github_url(url: str) -> tuple[str, str] | None:
+        p = parse_git_url(url)
+        if p.github and p.owner and p.repo:
+            return p.owner, p.repo
+        return None
+
+    @staticmethod
+    def _to_https_url(owner: str, repo: str) -> str:
+        return f"https://github.com/{owner}/{repo}.git"

vibe/core/teleport/nuage.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+from dataclasses import asdict
+import types
+from typing import Any
+
+import httpx
+from pydantic import BaseModel, Field
+
+from vibe.core.auth import EncryptedPayload, encrypt
+from vibe.core.teleport.errors import ServiceTeleportError
+
+
+class GitRepoConfig(BaseModel):
+    url: str
+    branch: str | None = None
+    commit: str | None = None
+
+
+class VibeSandboxConfig(BaseModel):
+    git_repo: GitRepoConfig | None = None
+
+
+class VibeNewSandbox(BaseModel):
+    type: str = "new"
+    config: VibeSandboxConfig = Field(default_factory=VibeSandboxConfig)
+    teleported_diffs: bytes | None = None
+
+
+class TeleportSession(BaseModel):
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    messages: list[dict[str, Any]] = Field(default_factory=list)
+
+
+class WorkflowParams(BaseModel):
+    prompt: str
+    sandbox: VibeNewSandbox
+    session: TeleportSession | None = None
+
+
+class WorkflowExecuteResponse(BaseModel):
+    execution_id: str
+
+
+class PublicKeyResult(BaseModel):
+    public_key: str
+
+
+class QueryResponse(BaseModel):
+    result: PublicKeyResult
+
+
+class CreateLeChatThreadInput(BaseModel):
+    encrypted_api_key: dict[str, str]
+    user_message: str
+    project_name: str | None = None
+
+
+class CreateLeChatThreadOutput(BaseModel):
+    chat_url: str
+
+
+class UpdateResponse(BaseModel):
+    result: CreateLeChatThreadOutput
+
+
+class NuageClient:
+    def __init__(
+        self,
+        base_url: str,
+        api_key: str,
+        workflow_id: str,
+        *,
+        client: httpx.AsyncClient | None = None,
+        timeout: float = 60.0,
+    ) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._api_key = api_key
+        self._workflow_id = workflow_id
+        self._client = client
+        self._owns_client = client is None
+        self._timeout = timeout
+
+    async def __aenter__(self) -> NuageClient:
+        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
+
+    @property
+    def _http_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 _headers(self) -> dict[str, str]:
+        return {
+            "Authorization": f"Bearer {self._api_key}",
+            "Content-Type": "application/json",
+        }
+
+    async def start_workflow(self, params: WorkflowParams) -> str:
+        response = await self._http_client.post(
+            f"{self._base_url}/v1/workflows/{self._workflow_id}/execute",
+            headers=self._headers(),
+            json={"input": params.model_dump(mode="json")},
+        )
+        if not response.is_success:
+            error_msg = f"Nuage workflow trigger failed: {response.text}"
+            # TODO(vibe-nuage): remove this once prod has shared vibe-nuage workers
+            if "Unauthorized" in response.text or "unauthorized" in response.text:
+                error_msg += (
+                    "\n\nHint: This version uses Mistral staging environment. "
+                    "Set STAGING_MISTRAL_API_KEY from https://console.globalaegis.net/"
+                )
+            raise ServiceTeleportError(error_msg)
+        result = WorkflowExecuteResponse.model_validate(response.json())
+        return result.execution_id
+
+    async def send_github_token(self, execution_id: str, token: str) -> None:
+        public_key_pem = await self._query_public_key(execution_id)
+        encrypted = encrypt(token, public_key_pem)
+        await self._signal_encrypted_token(execution_id, encrypted)
+
+    async def _query_public_key(self, execution_id: str) -> bytes:
+        response = await self._http_client.post(
+            f"{self._base_url}/v1/workflows/executions/{execution_id}/queries",
+            headers=self._headers(),
+            json={"name": "get_public_key", "input": {}},
+        )
+        if not response.is_success:
+            raise ServiceTeleportError(f"Failed to get public key: {response.text}")
+
+        result = QueryResponse.model_validate(response.json())
+        return result.result.public_key.encode("utf-8")
+
+    async def _signal_encrypted_token(
+        self, execution_id: str, encrypted: EncryptedPayload
+    ) -> None:
+        response = await self._http_client.post(
+            f"{self._base_url}/v1/workflows/executions/{execution_id}/signals",
+            headers=self._headers(),
+            json={"name": "github_token", "input": {"payload": asdict(encrypted)}},
+        )
+        if not response.is_success:
+            raise ServiceTeleportError(f"Failed to send GitHub token: {response.text}")
+
+    async def create_le_chat_thread(
+        self, execution_id: str, user_message: str, project_name: str | None = None
+    ) -> str:
+        public_key_pem = await self._query_public_key(execution_id)
+        encrypted = encrypt(self._api_key, public_key_pem)
+        input_data = CreateLeChatThreadInput(
+            encrypted_api_key={
+                k: v for k, v in asdict(encrypted).items() if v is not None
+            },
+            user_message=user_message,
+            project_name=project_name,
+        )
+        response = await self._http_client.post(
+            f"{self._base_url}/v1/workflows/executions/{execution_id}/updates",
+            headers=self._headers(),
+            json={"name": "create_le_chat_thread", "input": input_data.model_dump()},
+        )
+        if not response.is_success:
+            raise ServiceTeleportError(
+                f"Failed to create Le Chat thread: {response.text}"
+            )
+        result = UpdateResponse.model_validate(response.json())
+        return result.result.chat_url

vibe/core/teleport/teleport.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+import base64
+from collections.abc import AsyncGenerator
+from pathlib import Path
+import types
+
+import httpx
+import zstandard
+
+from vibe.core.auth.github import GitHubAuthProvider
+from vibe.core.session.session_logger import SessionLogger
+from vibe.core.teleport.errors import ServiceTeleportError
+from vibe.core.teleport.git import GitRepoInfo, GitRepository
+from vibe.core.teleport.nuage import (
+    GitRepoConfig,
+    NuageClient,
+    TeleportSession,
+    VibeNewSandbox,
+    VibeSandboxConfig,
+    WorkflowParams,
+)
+from vibe.core.teleport.types import (
+    TeleportAuthCompleteEvent,
+    TeleportAuthRequiredEvent,
+    TeleportCheckingGitEvent,
+    TeleportCompleteEvent,
+    TeleportPushingEvent,
+    TeleportPushRequiredEvent,
+    TeleportPushResponseEvent,
+    TeleportSendEvent,
+    TeleportSendingGithubTokenEvent,
+    TeleportStartingWorkflowEvent,
+    TeleportYieldEvent,
+)
+
+# TODO(vibe-nuage): update URL once prod has shared vibe-nuage workers
+_NUAGE_EXECUTION_URL_TEMPLATE = "https://console.globalaegis.net/build/workflows/{workflow_id}?tab=executions&executionId={execution_id}"
+_DEFAULT_TELEPORT_PROMPT = "please continue where you left off"
+
+
+class TeleportService:
+    def __init__(
+        self,
+        session_logger: SessionLogger,
+        nuage_base_url: str,
+        nuage_workflow_id: str,
+        nuage_api_key: str,
+        workdir: Path | None = None,
+        *,
+        client: httpx.AsyncClient | None = None,
+        timeout: float = 60.0,
+    ) -> None:
+        self._session_logger = session_logger
+        self._nuage_base_url = nuage_base_url
+        self._nuage_workflow_id = nuage_workflow_id
+        self._nuage_api_key = nuage_api_key
+        self._git = GitRepository(workdir)
+        self._client = client
+        self._owns_client = client is None
+        self._timeout = timeout
+        self._github_auth: GitHubAuthProvider | None = None
+        self._nuage: NuageClient | None = None
+
+    async def __aenter__(self) -> TeleportService:
+        if self._client is None:
+            self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
+        self._github_auth = GitHubAuthProvider(client=self._client)
+        self._nuage = NuageClient(
+            self._nuage_base_url,
+            self._nuage_api_key,
+            self._nuage_workflow_id,
+            client=self._client,
+        )
+        await self._git.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        await self._git.__aexit__(exc_type, exc_val, exc_tb)
+        if self._owns_client and self._client:
+            await self._client.aclose()
+            self._client = None
+
+    @property
+    def _http_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
+
+    @property
+    def _github_auth_provider(self) -> GitHubAuthProvider:
+        if self._github_auth is None:
+            self._github_auth = GitHubAuthProvider(client=self._http_client)
+        return self._github_auth
+
+    @property
+    def _nuage_client(self) -> NuageClient:
+        if self._nuage is None:
+            self._nuage = NuageClient(
+                self._nuage_base_url,
+                self._nuage_api_key,
+                self._nuage_workflow_id,
+                client=self._http_client,
+            )
+        return self._nuage
+
+    async def check_supported(self) -> None:
+        await self._git.get_info()
+
+    async def is_supported(self) -> bool:
+        return await self._git.is_supported()
+
+    async def execute(
+        self, prompt: str | None, session: TeleportSession
+    ) -> AsyncGenerator[TeleportYieldEvent, TeleportSendEvent]:
+        prompt = prompt or _DEFAULT_TELEPORT_PROMPT
+        self._validate_config()
+
+        git_info = await self._git.get_info()
+
+        yield TeleportCheckingGitEvent()
+        if not await self._git.is_commit_pushed(git_info.commit):
+            unpushed_count = await self._git.get_unpushed_commit_count()
+            response = yield TeleportPushRequiredEvent(
+                unpushed_count=max(1, unpushed_count)
+            )
+            if (
+                not isinstance(response, TeleportPushResponseEvent)
+                or not response.approved
+            ):
+                raise ServiceTeleportError("Teleport cancelled: commit not pushed.")
+
+            yield TeleportPushingEvent()
+            await self._push_or_fail()
+
+        github_token = await self._github_auth_provider.get_valid_token()
+
+        if not github_token:
+            handle = await self._github_auth_provider.start_device_flow(
+                open_browser=True
+            )
+            yield TeleportAuthRequiredEvent(
+                user_code=handle.info.user_code,
+                verification_uri=handle.info.verification_uri,
+            )
+            github_token = await self._github_auth_provider.wait_for_token(handle)
+            yield TeleportAuthCompleteEvent()
+
+        yield TeleportStartingWorkflowEvent()
+
+        execution_id = await self._nuage_client.start_workflow(
+            WorkflowParams(
+                prompt=prompt, sandbox=self._build_sandbox(git_info), session=session
+            )
+        )
+
+        yield TeleportSendingGithubTokenEvent()
+        await self._nuage_client.send_github_token(execution_id, github_token)
+
+        chat_url = _NUAGE_EXECUTION_URL_TEMPLATE.format(
+            workflow_id=self._nuage_workflow_id, execution_id=execution_id
+        )
+        # chat_url = await nuage.create_le_chat_thread(
+        #     execution_id=execution_id, user_message=prompt
+        # )
+
+        yield TeleportCompleteEvent(url=chat_url)
+
+    async def _push_or_fail(self) -> None:
+        if not await self._git.push_current_branch():
+            raise ServiceTeleportError("Failed to push current branch to remote.")
+
+    def _validate_config(self) -> None:
+        # TODO(vibe-nuage): update error message once prod has shared vibe-nuage workers
+        if not self._nuage_api_key:
+            raise ServiceTeleportError(
+                "STAGING_MISTRAL_API_KEY not set. "
+                "Set it from https://console.globalaegis.net/ to use teleport."
+            )
+
+    def _build_sandbox(self, git_info: GitRepoInfo) -> VibeNewSandbox:
+        return VibeNewSandbox(
+            config=VibeSandboxConfig(
+                git_repo=GitRepoConfig(
+                    url=git_info.remote_url,
+                    branch=git_info.branch,
+                    commit=git_info.commit,
+                )
+            ),
+            teleported_diffs=self._compress_diff(git_info.diff or ""),
+        )
+
+    def _compress_diff(self, diff: str, max_size: int = 1_000_000) -> bytes | None:
+        if not diff:
+            return None
+        compressed = zstandard.ZstdCompressor().compress(diff.encode("utf-8"))
+        encoded = base64.b64encode(compressed)
+        if len(encoded) > max_size:
+            raise ServiceTeleportError(
+                "Diff too large to teleport. Please commit and push your changes first."
+            )
+        return encoded

vibe/core/teleport/types.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from vibe.core.types import BaseEvent
+
+
+class TeleportAuthRequiredEvent(BaseEvent):
+    user_code: str
+    verification_uri: str
+
+
+class TeleportAuthCompleteEvent(BaseEvent):
+    pass
+
+
+class TeleportStartingWorkflowEvent(BaseEvent):
+    pass
+
+
+class TeleportCheckingGitEvent(BaseEvent):
+    pass
+
+
+class TeleportPushRequiredEvent(BaseEvent):
+    unpushed_count: int = 1
+
+
+class TeleportPushResponseEvent(BaseEvent):
+    approved: bool
+
+
+class TeleportPushingEvent(BaseEvent):
+    pass
+
+
+class TeleportSendingGithubTokenEvent(BaseEvent):
+    pass
+
+
+class TeleportCompleteEvent(BaseEvent):
+    url: str
+
+
+type TeleportYieldEvent = (
+    TeleportAuthRequiredEvent
+    | TeleportAuthCompleteEvent
+    | TeleportCheckingGitEvent
+    | TeleportPushRequiredEvent
+    | TeleportPushingEvent
+    | TeleportStartingWorkflowEvent
+    | TeleportSendingGithubTokenEvent
+    | TeleportCompleteEvent
+)
+
+type TeleportSendEvent = TeleportPushResponseEvent | None