agentpool 2.1.9__py3-none-any.whl → 2.2.3__py3-none-any.whl

agentpool/utils/file_watcher.py

@@ -0,0 +1,269 @@
+"""File watcher utilities using watchfiles."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, Set as AbstractSet
+import contextlib
+from dataclasses import dataclass, field
+import logging
+from pathlib import Path
+from typing import Self
+
+from watchfiles import Change, awatch
+
+
+logger = logging.getLogger(__name__)
+
+
+# Callback type for file change notifications
+FileChangeCallback = Callable[[AbstractSet[tuple[Change, str]]], Awaitable[None]]
+
+
+@dataclass
+class FileWatcher:
+    """Async file watcher using watchfiles.
+
+    Watches specified paths for changes and calls a callback when changes occur.
+
+    Example:
+        ```python
+        async def on_change(changes):
+            for change_type, path in changes:
+                print(f"{change_type}: {path}")
+
+        watcher = FileWatcher(
+            paths=["/path/to/.git/HEAD"],
+            callback=on_change,
+        )
+
+        async with watcher:
+            # Watcher is running
+            await asyncio.sleep(60)
+        # Watcher stopped
+        ```
+    """
+
+    paths: list[str | Path]
+    """Paths to watch (files or directories)."""
+
+    callback: FileChangeCallback
+    """Async callback invoked when changes are detected."""
+
+    debounce: int = 100
+    """Debounce time in milliseconds."""
+
+    _task: asyncio.Task[None] | None = field(default=None, repr=False)
+    """Background watch task."""
+
+    _stop_event: asyncio.Event = field(default_factory=asyncio.Event, repr=False)
+    """Event to signal stop."""
+
+    async def start(self) -> None:
+        """Start watching for file changes."""
+        if self._task is not None:
+            return  # Already running
+
+        self._stop_event.clear()
+        self._task = asyncio.create_task(self._watch_loop())
+
+    async def stop(self) -> None:
+        """Stop watching for file changes."""
+        if self._task is None:
+            return
+
+        self._stop_event.set()
+        self._task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self._task
+        self._task = None
+
+    async def _watch_loop(self) -> None:
+        """Internal watch loop."""
+        str_paths = [str(p) for p in self.paths]
+
+        # Filter to only existing paths
+        existing_paths = [p for p in str_paths if Path(p).exists()]
+        if not existing_paths:
+            logger.warning("FileWatcher: no existing paths to watch from %s", str_paths)
+            return
+
+        logger.info("FileWatcher: starting watch loop for %s", existing_paths)
+        try:
+            async for changes in awatch(
+                *existing_paths,
+                debounce=self.debounce,
+                stop_event=self._stop_event,
+            ):
+                logger.info("FileWatcher detected changes: %s", changes)
+                # Don't let callback errors kill the watcher
+                try:
+                    await self.callback(changes)
+                except Exception:
+                    logger.exception("Error in file watcher callback")
+        except Exception:
+            logger.exception("FileWatcher watch loop failed")
+
+    async def __aenter__(self) -> Self:
+        """Start watcher on context enter."""
+        await self.start()
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        """Stop watcher on context exit."""
+        await self.stop()
+
+
+async def get_git_branch(repo_path: str | Path) -> str | None:
+    """Get the current git branch name.
+
+    Args:
+        repo_path: Path to the git repository
+
+    Returns:
+        Branch name or None if not a git repo or on detached HEAD
+
+    TODO: For remote/ACP support, this should accept an optional ExecutionEnvironment
+    and use env.execute_command() instead of subprocess. This would allow git commands
+    to run on the client side where the repository lives.
+    """
+    try:
+        proc = await asyncio.create_subprocess_exec(
+            "git",
+            "rev-parse",
+            "--abbrev-ref",
+            "HEAD",
+            cwd=str(repo_path),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, _ = await proc.communicate()
+        if proc.returncode != 0:
+            return None
+    except OSError:
+        return None
+    else:
+        branch = stdout.decode().strip()
+        return branch if branch != "HEAD" else None
+
+
+@dataclass
+class GitBranchWatcher:
+    """Watches for git branch changes using polling.
+
+    Polls the current git branch periodically and calls a callback when it changes.
+    Uses polling instead of file watching because git uses atomic renames which
+    are not reliably detected by inotify/watchfiles.
+
+    Example:
+        ```python
+        async def on_branch_change(branch: str | None):
+            print(f"Branch changed to: {branch}")
+
+        watcher = GitBranchWatcher(
+            repo_path="/path/to/repo",
+            callback=on_branch_change,
+        )
+
+        async with watcher:
+            await asyncio.sleep(60)
+        ```
+
+    TODO: For remote/ACP support, this should accept an ExecutionEnvironment
+    and run git commands through env.execute_command(). The polling would still
+    happen server-side, but the git commands would execute on the client.
+    """
+
+    repo_path: str | Path
+    """Path to the git repository."""
+
+    callback: Callable[[str | None], Awaitable[None]]
+    """Async callback invoked with new branch name when branch changes."""
+
+    poll_interval: float = 1.0
+    """Polling interval in seconds."""
+
+    _current_branch: str | None = field(default=None, repr=False)
+    """Cached current branch."""
+
+    _task: asyncio.Task[None] | None = field(default=None, repr=False)
+    """Background polling task."""
+
+    _stop_event: asyncio.Event = field(default_factory=asyncio.Event, repr=False)
+    """Event to signal stop."""
+
+    async def start(self) -> None:
+        """Start watching for branch changes."""
+        if self._task is not None:
+            return  # Already running
+
+        repo = Path(self.repo_path)
+        git_dir = repo / ".git"
+
+        # Handle git worktrees - .git might be a file pointing to the real git dir
+        if git_dir.is_file():
+            content = git_dir.read_text().strip()
+            if content.startswith("gitdir:"):
+                git_dir = Path(content[7:].strip())
+
+        if not git_dir.exists():
+            logger.warning("Git directory not found: %s", git_dir)
+            return
+
+        # Get initial branch
+        self._current_branch = await get_git_branch(self.repo_path)
+        logger.info(
+            "GitBranchWatcher started (polling), repo: %s, initial branch: %s",
+            self.repo_path,
+            self._current_branch,
+        )
+
+        self._stop_event.clear()
+        self._task = asyncio.create_task(self._poll_loop())
+
+    async def _poll_loop(self) -> None:
+        """Internal polling loop."""
+        while not self._stop_event.is_set():
+            try:
+                await asyncio.wait_for(
+                    self._stop_event.wait(),
+                    timeout=self.poll_interval,
+                )
+                break  # Stop event was set
+            except TimeoutError:
+                # Poll interval elapsed, check for changes
+                pass
+
+            try:
+                new_branch = await get_git_branch(self.repo_path)
+                if new_branch != self._current_branch:
+                    logger.info("Branch changed: %s -> %s", self._current_branch, new_branch)
+                    self._current_branch = new_branch
+                    await self.callback(new_branch)
+            except Exception:
+                logger.exception("Error polling git branch")
+
+    async def stop(self) -> None:
+        """Stop watching for branch changes."""
+        if self._task is None:
+            return
+
+        self._stop_event.set()
+        self._task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self._task
+        self._task = None
+
+    @property
+    def current_branch(self) -> str | None:
+        """Get the current cached branch name."""
+        return self._current_branch
+
+    async def __aenter__(self) -> Self:
+        """Start watcher on context enter."""
+        await self.start()
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        """Stop watcher on context exit."""
+        await self.stop()

agentpool/utils/identifiers.py

@@ -0,0 +1,121 @@
+"""Identifier generation utilities.
+
+Generates IDs that are lexicographically sortable by creation time.
+Format: {prefix}_{hex_timestamp}{random_base62}
+
+Compatible with OpenCode's identifier format.
+"""
+
+from __future__ import annotations
+
+import secrets
+import time
+from typing import Literal
+
+
+PrefixType = Literal["session", "message", "permission", "user", "part", "pty"]
+
+PREFIXES: dict[PrefixType, str] = {
+    "session": "ses",
+    "message": "msg",
+    "permission": "per",
+    "user": "usr",
+    "part": "prt",
+    "pty": "pty",
+}
+
+BASE62_CHARS = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+ID_LENGTH = 26  # Characters after prefix (12 hex + 14 base62)
+
+# State for monotonic ID generation
+_last_timestamp = 0
+_counter = 0
+
+
+def _random_base62(length: int) -> str:
+    """Generate random base62 string."""
+    return "".join(secrets.choice(BASE62_CHARS) for _ in range(length))
+
+
+def ascending(prefix: PrefixType, given: str | None = None) -> str:
+    """Generate an ascending (chronologically sortable) ID.
+
+    Args:
+        prefix: The type prefix for the ID
+        given: If provided, validate and return this ID instead of generating
+
+    Returns:
+        A sortable ID with the format {prefix}_{hex_timestamp}{random}
+
+    Raises:
+        ValueError: If given ID doesn't start with expected prefix
+    """
+    if given is not None:
+        expected_prefix = PREFIXES[prefix]
+        if not given.startswith(expected_prefix):
+            msg = f"ID {given} does not start with {expected_prefix}"
+            raise ValueError(msg)
+        return given
+
+    return _create(prefix, descending=False)
+
+
+def descending(prefix: PrefixType) -> str:
+    """Generate a descending (reverse chronologically sortable) ID.
+
+    Args:
+        prefix: The type prefix for the ID
+
+    Returns:
+        A reverse-sortable ID
+    """
+    return _create(prefix, descending=True)
+
+
+def _create(prefix: PrefixType, *, descending: bool = False) -> str:
+    """Create a new ID with timestamp encoding.
+
+    Args:
+        prefix: The type prefix
+        descending: If True, invert the timestamp for reverse sorting
+
+    Returns:
+        A new ID string
+    """
+    global _last_timestamp, _counter  # noqa: PLW0603
+
+    current_timestamp = int(time.time() * 1000)  # milliseconds
+
+    if current_timestamp != _last_timestamp:
+        _last_timestamp = current_timestamp
+        _counter = 0
+    _counter += 1
+
+    # Combine timestamp and counter
+    now = current_timestamp * 0x1000 + _counter
+
+    if descending:
+        now = ~now & 0xFFFFFFFFFFFF  # Invert for descending order (48 bits)
+
+    # Encode as 6 bytes (48 bits), big-endian
+    time_bytes = bytearray(6)
+    for i in range(6):
+        time_bytes[i] = (now >> (40 - 8 * i)) & 0xFF
+
+    time_hex = time_bytes.hex()
+
+    # Add random suffix (14 chars for 26 total after prefix)
+    random_suffix = _random_base62(ID_LENGTH - 12)
+
+    return f"{PREFIXES[prefix]}_{time_hex}{random_suffix}"
+
+
+def generate_session_id() -> str:
+    """Generate a unique, chronologically sortable session ID.
+
+    Convenience function for the common case.
+
+    Returns:
+        A session ID like 'ses_b71310fdf001ZHcn6VSpkaBcHi'
+    """
+    return ascending("session")

agentpool/utils/pydantic_ai_helpers.py

@@ -0,0 +1,46 @@
+"""Helper utilities for working with pydantic-ai message types."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai.messages import BaseToolCallPart
+
+
+if TYPE_CHECKING:
+    from pydantic_ai.messages import ToolCallPartDelta
+
+
+def safe_args_as_dict(
+    part: BaseToolCallPart | ToolCallPartDelta,
+    *,
+    default: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Safely extract args as dict from a tool call part.
+
+    Models can return malformed JSON for tool arguments, especially during
+    streaming when args are still being assembled. This helper catches parse
+    errors and returns a fallback value.
+
+    Args:
+        part: A tool call part (complete or delta) with args to extract
+        default: Value to return on parse failure. If None, returns {"_raw_args": ...}
+                 with the original unparsed args.
+
+    Returns:
+        The parsed arguments dict, or a fallback on parse failure.
+    """
+    if not isinstance(part, BaseToolCallPart):
+        # ToolCallPartDelta doesn't have args_as_dict
+        if default is not None:
+            return default
+        raw = getattr(part, "args", None)
+        return {"_raw_args": raw} if raw else {}
+    try:
+        return part.args_as_dict()
+    except ValueError:
+        # Model returned malformed JSON for tool args
+        if default is not None:
+            return default
+        # Preserve raw args for debugging/inspection
+        return {"_raw_args": part.args} if part.args else {}