code-analysis-client-dev 1.6.341__py3-none-any.whl

code_analysis_client/worktree.py

@@ -0,0 +1,533 @@
+"""
+Worktree management on top of CodeAnalysisAsyncClient.
+
+Wraps ``worktree_*`` MCP commands with typed client methods, queued and synchronous,
+following the FileSessionClient pattern.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from code_analysis_client.exceptions import ClientValidationError
+from code_analysis_client.responses import unwrap_command_result, FacadeCaller
+
+if TYPE_CHECKING:
+    from code_analysis_client.client import CodeAnalysisAsyncClient
+
+
+class WorktreeConflictError(ClientValidationError):
+    """Raised when worktree_sync encounters WORKTREE_SYNC_CONFLICT."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        conflicts: Optional[List[Dict[str, Any]]] = None,
+        preserved_locked_files: Optional[List[str]] = None,
+    ) -> None:
+        """Initialize with conflict details."""
+        super().__init__(message, field="worktree_sync")
+        self.conflicts = conflicts or []
+        self.preserved_locked_files = preserved_locked_files or []
+
+
+def _unwrap(data: Dict[str, Any]) -> Dict[str, Any]:
+    """Unwrap command result payloads."""
+    return unwrap_command_result(data)
+
+
+def _require_non_empty(value: Any, *, field: str) -> str:
+    """Return stripped text or raise ClientValidationError for an empty field."""
+    text = str(value or "").strip()
+    if not text:
+        raise ClientValidationError(f"{field} is required", field=field)
+    return text
+
+
+class WorktreeClient:
+    """Worktree lifecycle and sync management commands."""
+
+    __slots__ = ("_client",)
+
+    def __init__(self, client: CodeAnalysisAsyncClient) -> None:
+        """Store the underlying async client used to call worktree commands."""
+        self._client = FacadeCaller(client)  # type: ignore[assignment]
+
+    # === Lifecycle commands ===
+
+    async def ensure(
+        self,
+        project_id: str,
+        name: str,
+        *,
+        branch: Optional[str] = None,
+        start_point: Optional[str] = None,
+        create_branch: bool = True,
+        owner_session_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Create a worktree if absent; idempotent if already present with same branch.
+
+        Calls ``worktree_ensure`` per candidate-WT-1.md section 5.1.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+        if branch is not None:
+            params["branch"] = branch
+        if start_point is not None:
+            params["start_point"] = start_point
+        if not create_branch:
+            params["create_branch"] = False
+        if owner_session_id is not None:
+            params["owner_session_id"] = owner_session_id
+
+        return _unwrap(
+            await self._client.call_validated("worktree_ensure", params)
+        )
+
+    async def list(
+        self,
+        project_id: str,
+        *,
+        owner_session_id: Optional[str] = None,
+        only_names: Optional[List[str]] = None,
+    ) -> Dict[str, Any]:
+        """List worktrees in a project.
+
+        Calls ``worktree_list`` per candidate-WT-1.md section 5.2.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+
+        params: Dict[str, Any] = {"project_id": project_id}
+        if owner_session_id is not None:
+            params["owner_session_id"] = owner_session_id
+        if only_names is not None:
+            params["only_names"] = only_names
+
+        return _unwrap(
+            await self._client.call_validated("worktree_list", params)
+        )
+
+    async def status(
+        self,
+        project_id: str,
+        name: str,
+    ) -> Dict[str, Any]:
+        """Get git status of one worktree (working tree + index).
+
+        Calls ``worktree_status`` per candidate-WT-1.md section 5.3.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+
+        return _unwrap(
+            await self._client.call_validated("worktree_status", params)
+        )
+
+    async def diff(
+        self,
+        project_id: str,
+        name: str,
+        *,
+        base: Optional[str] = None,
+        path: Optional[str] = None,
+        include_patch: bool = False,
+    ) -> Dict[str, Any]:
+        """Diff worktree against base branch.
+
+        Calls ``worktree_diff`` per candidate-WT-1.md section 5.4.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+        if base is not None:
+            params["base"] = base
+        if path is not None:
+            params["path"] = path
+        if include_patch:
+            params["include_patch"] = True
+
+        return _unwrap(
+            await self._client.call_validated("worktree_diff", params)
+        )
+
+    # === Sync and conflict resolution ===
+
+    async def sync(
+        self,
+        project_id: str,
+        name: str,
+        *,
+        strategy: Optional[str] = None,
+        base: Optional[str] = None,
+        auto_poll: bool = True,
+        timeout: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """Sync worktree from main branch (merge or rebase).
+
+        Calls ``worktree_sync`` per candidate-WT-1.md section 5.5.
+        Raises WorktreeConflictError on WORKTREE_SYNC_CONFLICT.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+        if strategy is not None:
+            params["strategy"] = strategy
+        if base is not None:
+            params["base"] = base
+
+        result = await self._client.call_validated(
+            "worktree_sync",
+            params,
+            auto_poll=auto_poll,
+            timeout=timeout,
+        )
+
+        # Check for WORKTREE_SYNC_CONFLICT error
+        if isinstance(result, dict) and not result.get("success", True):
+            error = result.get("error", {})
+            if error.get("code") == "WORKTREE_SYNC_CONFLICT":
+                details = error.get("details", {})
+                raise WorktreeConflictError(
+                    "Worktree sync conflict",
+                    conflicts=details.get("conflicts"),
+                    preserved_locked_files=details.get("preserved_locked_files"),
+                )
+
+        return _unwrap(result)
+
+    async def conflicts(
+        self,
+        project_id: str,
+        name: str,
+    ) -> Dict[str, Any]:
+        """List merge/rebase conflicts with per-file markers.
+
+        Calls ``worktree_conflicts`` per candidate-WT-1.md section 5.6.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+
+        return _unwrap(
+            await self._client.call_validated("worktree_conflicts", params)
+        )
+
+    async def continue_(
+        self,
+        project_id: str,
+        name: str,
+        resolved_paths: List[str],
+        *,
+        message: Optional[str] = None,
+        auto_poll: bool = True,
+        timeout: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """Finalize merge or rebase after resolving conflicts.
+
+        Calls ``worktree_continue`` per candidate-WT-1.md section 5.7.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+        if not resolved_paths:
+            raise ClientValidationError(
+                "resolved_paths must be non-empty",
+                field="resolved_paths",
+            )
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+            "resolved_paths": resolved_paths,
+        }
+        if message is not None:
+            params["message"] = message
+
+        return _unwrap(
+            await self._client.call_validated(
+                "worktree_continue",
+                params,
+                auto_poll=auto_poll,
+                timeout=timeout,
+            )
+        )
+
+    async def abort(
+        self,
+        project_id: str,
+        name: str,
+        *,
+        auto_poll: bool = True,
+        timeout: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """Abort in-progress merge or rebase.
+
+        Calls ``worktree_abort`` per candidate-WT-1.md section 5.8.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+
+        return _unwrap(
+            await self._client.call_validated(
+                "worktree_abort",
+                params,
+                auto_poll=auto_poll,
+                timeout=timeout,
+            )
+        )
+
+    # === Commit and removal ===
+
+    async def commit(
+        self,
+        project_id: str,
+        name: str,
+        message: str,
+        *,
+        paths: Optional[List[str]] = None,
+        all: bool = False,
+        auto_poll: bool = True,
+        timeout: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """Commit staged or all changes in worktree.
+
+        Calls ``worktree_commit`` per candidate-WT-1.md section 5.9.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+        message = _require_non_empty(message, field="message")
+
+        if not paths and not all:
+            raise ClientValidationError(
+                "At least one of paths or all=true is required",
+                field="paths",
+            )
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+            "message": message,
+        }
+        if paths is not None:
+            params["paths"] = paths
+        if all:
+            params["all"] = True
+
+        return _unwrap(
+            await self._client.call_validated(
+                "worktree_commit",
+                params,
+                auto_poll=auto_poll,
+                timeout=timeout,
+            )
+        )
+
+    async def remove(
+        self,
+        project_id: str,
+        name: str,
+        *,
+        force: bool = False,
+        delete_branch: bool = True,
+        auto_poll: bool = True,
+        timeout: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """Remove a worktree and optionally its branch.
+
+        Calls ``worktree_remove`` per candidate-WT-1.md section 5.10.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+        if force:
+            params["force"] = True
+        if not delete_branch:
+            params["delete_branch"] = False
+
+        return _unwrap(
+            await self._client.call_validated(
+                "worktree_remove",
+                params,
+                auto_poll=auto_poll,
+                timeout=timeout,
+            )
+        )
+
+    async def prune(
+        self,
+        project_id: str,
+        *,
+        name: Optional[str] = None,
+        ttl_seconds: Optional[int] = None,
+        only_names: Optional[List[str]] = None,
+        dry_run: bool = False,
+        auto_poll: bool = True,
+        timeout: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """TTL safety net: reap idle worktrees whose sessions are gone.
+
+        Calls ``worktree_prune`` per candidate-WT-1.md section 5.11.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+
+        params: Dict[str, Any] = {"project_id": project_id}
+        if name is not None:
+            params["name"] = name
+        if ttl_seconds is not None:
+            params["ttl_seconds"] = ttl_seconds
+        if only_names is not None:
+            params["only_names"] = only_names
+        if dry_run:
+            params["dry_run"] = True
+
+        return _unwrap(
+            await self._client.call_validated(
+                "worktree_prune",
+                params,
+                auto_poll=auto_poll,
+                timeout=timeout,
+            )
+        )
+
+    # === Lock and maintenance ===
+
+    async def lock(
+        self,
+        project_id: str,
+        name: str,
+        *,
+        reason: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Lock a worktree to protect from removal.
+
+        Calls ``worktree_lock`` per candidate-WT-1.md section 5.12.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+        if reason is not None:
+            params["reason"] = reason
+
+        return _unwrap(
+            await self._client.call_validated("worktree_lock", params)
+        )
+
+    async def unlock(
+        self,
+        project_id: str,
+        name: str,
+    ) -> Dict[str, Any]:
+        """Unlock a worktree.
+
+        Calls ``worktree_unlock`` per candidate-WT-1.md section 5.13.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+
+        return _unwrap(
+            await self._client.call_validated("worktree_unlock", params)
+        )
+
+    async def move(
+        self,
+        project_id: str,
+        name: str,
+        new_name: str,
+        *,
+        auto_poll: bool = True,
+        timeout: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """Move/rename a worktree.
+
+        Calls ``worktree_move`` per candidate-WT-1.md section 5.14.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+        new_name = _require_non_empty(new_name, field="new_name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+            "new_name": new_name,
+        }
+
+        return _unwrap(
+            await self._client.call_validated(
+                "worktree_move",
+                params,
+                auto_poll=auto_poll,
+                timeout=timeout,
+            )
+        )
+
+    async def repair(
+        self,
+        project_id: str,
+        name: str,
+        *,
+        auto_poll: bool = True,
+        timeout: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """Repair worktree after manual filesystem move/restore.
+
+        Calls ``worktree_repair`` per candidate-WT-1.md section 5.15.
+        """
+        project_id = _require_non_empty(project_id, field="project_id")
+        name = _require_non_empty(name, field="name")
+
+        params: Dict[str, Any] = {
+            "project_id": project_id,
+            "name": name,
+        }
+
+        return _unwrap(
+            await self._client.call_validated(
+                "worktree_repair",
+                params,
+                auto_poll=auto_poll,
+                timeout=timeout,
+            )
+        )