aury-agent 0.0.4__py3-none-any.whl

aury/agents/backends/snapshot/git.py

@@ -0,0 +1,95 @@
+"""Git-based snapshot backend."""
+from __future__ import annotations
+
+import asyncio
+from datetime import datetime
+from pathlib import Path
+
+from .types import Patch
+
+
+class GitSnapshotBackend:
+    """Git-based snapshot backend."""
+    
+    def __init__(self, worktree: str | Path, snapshot_dir: str | Path | None = None):
+        self.worktree = Path(worktree).resolve()
+        self.snapshot_dir = Path(snapshot_dir).resolve() if snapshot_dir else self.worktree / ".aury_snapshot"
+        self._initialized = False
+    
+    async def _ensure_repo(self) -> None:
+        if self._initialized:
+            return
+        git_dir = self.snapshot_dir / ".git" if not str(self.snapshot_dir).endswith(".git") else self.snapshot_dir
+        if not git_dir.exists():
+            self.snapshot_dir.mkdir(parents=True, exist_ok=True)
+            proc = await asyncio.create_subprocess_exec(
+                "git", "init",
+                cwd=str(self.worktree),
+                env={"GIT_DIR": str(self.snapshot_dir)},
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            await proc.wait()
+            await self._git_config("user.email", "aury@agent.local")
+            await self._git_config("user.name", "Aury Agent")
+        self._initialized = True
+    
+    async def _git_config(self, key: str, value: str) -> None:
+        proc = await asyncio.create_subprocess_exec(
+            "git", "config", key, value,
+            cwd=str(self.worktree),
+            env={"GIT_DIR": str(self.snapshot_dir)},
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        await proc.wait()
+    
+    async def _run_git(self, *args: str, check: bool = True) -> tuple[str, str]:
+        await self._ensure_repo()
+        proc = await asyncio.create_subprocess_exec(
+            "git", *args,
+            cwd=str(self.worktree),
+            env={"GIT_DIR": str(self.snapshot_dir)},
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, stderr = await proc.communicate()
+        if check and proc.returncode != 0:
+            raise RuntimeError(f"Git command failed: {stderr.decode()}")
+        return stdout.decode(), stderr.decode()
+    
+    async def track(self) -> str:
+        await self._ensure_repo()
+        await self._run_git("add", "-A", check=False)
+        timestamp = datetime.now().isoformat()
+        await self._run_git("commit", "-m", f"snapshot_{timestamp}", "--allow-empty", check=False)
+        stdout, _ = await self._run_git("rev-parse", "HEAD")
+        return stdout.strip()
+    
+    async def restore(self, snapshot_id: str) -> None:
+        await self._run_git("checkout", snapshot_id, "--", ".")
+    
+    async def revert(self, patches: list[Patch]) -> None:
+        for patch in patches:
+            for file_path in patch.files:
+                await self._run_git("checkout", "HEAD~1", "--", file_path, check=False)
+    
+    async def diff(self, snapshot_id: str) -> str:
+        stdout, _ = await self._run_git("diff", snapshot_id, "--", check=False)
+        return stdout
+    
+    async def patch(self, snapshot_id: str) -> Patch:
+        stdout, _ = await self._run_git("diff", "--stat", snapshot_id, "--", check=False)
+        files = []
+        additions = deletions = 0
+        for line in stdout.strip().split("\n"):
+            if "|" in line:
+                file_path = line.split("|")[0].strip()
+                files.append(file_path)
+                stat_part = line.split("|")[1].strip() if "|" in line else ""
+                additions += stat_part.count("+")
+                deletions += stat_part.count("-")
+        return Patch(files=files, additions=additions, deletions=deletions, diff=await self.diff(snapshot_id))
+
+
+__all__ = ["GitSnapshotBackend"]

aury/agents/backends/snapshot/hybrid.py

@@ -0,0 +1,125 @@
+"""Git + S3 hybrid snapshot backend."""
+from __future__ import annotations
+
+import asyncio
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from .git import GitSnapshotBackend
+
+
+class GitS3HybridBackend(GitSnapshotBackend):
+    """Git + S3 hybrid snapshot backend.
+    
+    Combines local Git commits with S3 bundle uploads for remote persistence.
+    """
+    
+    def __init__(
+        self,
+        worktree: str | Path,
+        s3_bucket: str,
+        session_id: str,
+        snapshot_dir: str | Path | None = None,
+        s3_client: Any | None = None,
+    ):
+        super().__init__(worktree, snapshot_dir)
+        self.s3_bucket = s3_bucket
+        self.session_id = session_id
+        self._s3 = s3_client
+        self._last_uploaded_commit: str | None = None
+    
+    @property
+    def s3(self):
+        if self._s3 is None:
+            try:
+                import boto3
+                self._s3 = boto3.client('s3')
+            except ImportError:
+                raise ImportError("boto3 is required for GitS3HybridBackend")
+        return self._s3
+    
+    async def track(self) -> str:
+        import json
+        commit_id = await super().track()
+        bundle_path = self.snapshot_dir / f"{commit_id}.bundle"
+        
+        if self._last_uploaded_commit:
+            await self._create_bundle(bundle_path, f"{self._last_uploaded_commit}..{commit_id}")
+            bundle_type = "incremental"
+        else:
+            await self._create_bundle(bundle_path, commit_id)
+            bundle_type = "full"
+        
+        s3_key = f"agents/{self.session_id}/snapshots/{commit_id}.bundle"
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(None, lambda: self.s3.upload_file(str(bundle_path), self.s3_bucket, s3_key))
+        
+        meta = {
+            "commit_id": commit_id,
+            "parent_commit": self._last_uploaded_commit,
+            "timestamp": datetime.now().isoformat(),
+            "bundle_type": bundle_type,
+        }
+        meta_key = f"agents/{self.session_id}/snapshots/{commit_id}.meta.json"
+        await loop.run_in_executor(None, lambda: self.s3.put_object(Bucket=self.s3_bucket, Key=meta_key, Body=json.dumps(meta)))
+        
+        if bundle_path.exists():
+            bundle_path.unlink()
+        
+        self._last_uploaded_commit = commit_id
+        return commit_id
+    
+    async def _create_bundle(self, path: Path, ref: str) -> None:
+        proc = await asyncio.create_subprocess_exec(
+            "git", "bundle", "create", str(path), ref,
+            cwd=str(self.worktree),
+            env={"GIT_DIR": str(self.snapshot_dir)},
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        await proc.wait()
+    
+    @classmethod
+    async def restore_from_s3(
+        cls,
+        s3_bucket: str,
+        session_id: str,
+        local_dir: str | Path,
+        s3_client: Any | None = None,
+    ) -> "GitS3HybridBackend":
+        import subprocess
+        if s3_client is None:
+            import boto3
+            s3_client = boto3.client('s3')
+        
+        prefix = f"agents/{session_id}/snapshots/"
+        response = s3_client.list_objects_v2(Bucket=s3_bucket, Prefix=prefix)
+        bundles = sorted([obj['Key'] for obj in response.get('Contents', []) if obj['Key'].endswith('.bundle')])
+        
+        if not bundles:
+            raise ValueError(f"No snapshots found for session: {session_id}")
+        
+        local_path = Path(local_dir)
+        local_path.mkdir(parents=True, exist_ok=True)
+        repo_path = local_path / "repo"
+        
+        for i, bundle_key in enumerate(bundles):
+            bundle_file = local_path / "temp.bundle"
+            s3_client.download_file(s3_bucket, bundle_key, str(bundle_file))
+            if i == 0:
+                subprocess.run(["git", "clone", str(bundle_file), str(repo_path)], check=True, capture_output=True)
+            else:
+                subprocess.run(["git", "fetch", str(bundle_file)], cwd=str(repo_path), check=True, capture_output=True)
+            bundle_file.unlink()
+        
+        result = subprocess.run(["git", "rev-parse", "HEAD"], cwd=str(repo_path), capture_output=True, text=True)
+        last_commit = result.stdout.strip()
+        
+        backend = cls(worktree=repo_path, s3_bucket=s3_bucket, session_id=session_id, snapshot_dir=repo_path / ".git", s3_client=s3_client)
+        backend._last_uploaded_commit = last_commit
+        backend._initialized = True
+        return backend
+
+
+__all__ = ["GitS3HybridBackend"]

aury/agents/backends/snapshot/memory.py

@@ -0,0 +1,86 @@
+"""In-memory snapshot backend for testing."""
+from __future__ import annotations
+
+from .types import Patch
+
+
+class InMemorySnapshotBackend:
+    """In-memory snapshot backend for testing."""
+    
+    def __init__(self) -> None:
+        self._snapshots: dict[str, dict[str, str]] = {}
+        self._files: dict[str, str] = {}
+        self._counter = 0
+    
+    def set_file(self, path: str, content: str) -> None:
+        self._files[path] = content
+    
+    def get_file(self, path: str) -> str | None:
+        return self._files.get(path)
+    
+    def delete_file(self, path: str) -> None:
+        self._files.pop(path, None)
+    
+    async def track(self) -> str:
+        self._counter += 1
+        snapshot_id = f"snap_{self._counter:04d}"
+        self._snapshots[snapshot_id] = dict(self._files)
+        return snapshot_id
+    
+    async def restore(self, snapshot_id: str) -> None:
+        if snapshot_id not in self._snapshots:
+            raise ValueError(f"Unknown snapshot: {snapshot_id}")
+        self._files = dict(self._snapshots[snapshot_id])
+    
+    async def revert(self, patches: list[Patch]) -> None:
+        snapshots = list(self._snapshots.items())
+        if len(snapshots) < 2:
+            return
+        prev_snapshot = snapshots[-2][1]
+        for patch in patches:
+            for file_path in patch.files:
+                if file_path in prev_snapshot:
+                    self._files[file_path] = prev_snapshot[file_path]
+                else:
+                    self._files.pop(file_path, None)
+    
+    async def diff(self, snapshot_id: str) -> str:
+        if snapshot_id not in self._snapshots:
+            return ""
+        snapshot = self._snapshots[snapshot_id]
+        lines = []
+        all_files = set(snapshot.keys()) | set(self._files.keys())
+        for path in sorted(all_files):
+            old = snapshot.get(path, "")
+            new = self._files.get(path, "")
+            if old != new:
+                lines.append(f"--- a/{path}")
+                lines.append(f"+++ b/{path}")
+        return "\n".join(lines)
+    
+    async def patch(self, snapshot_id: str) -> Patch:
+        if snapshot_id not in self._snapshots:
+            return Patch(files=[])
+        snapshot = self._snapshots[snapshot_id]
+        files = []
+        additions = deletions = 0
+        for path in set(snapshot.keys()) | set(self._files.keys()):
+            old = snapshot.get(path, "")
+            new = self._files.get(path, "")
+            if old != new:
+                files.append(path)
+                old_lines = len(old.split("\n")) if old else 0
+                new_lines = len(new.split("\n")) if new else 0
+                if new_lines > old_lines:
+                    additions += new_lines - old_lines
+                else:
+                    deletions += old_lines - new_lines
+        return Patch(files=files, additions=additions, deletions=deletions, diff=await self.diff(snapshot_id))
+    
+    def clear(self) -> None:
+        self._snapshots.clear()
+        self._files.clear()
+        self._counter = 0
+
+
+__all__ = ["InMemorySnapshotBackend"]

aury/agents/backends/snapshot/types.py

@@ -0,0 +1,59 @@
+"""Snapshot backend types and protocols."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+
+@dataclass
+class Patch:
+    """File change record."""
+    files: list[str]
+    additions: int = 0
+    deletions: int = 0
+    diff: str = ""
+    
+    def to_dict(self) -> dict:
+        return {
+            "files": self.files,
+            "additions": self.additions,
+            "deletions": self.deletions,
+            "diff": self.diff,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict) -> Patch:
+        return cls(
+            files=data.get("files", []),
+            additions=data.get("additions", 0),
+            deletions=data.get("deletions", 0),
+            diff=data.get("diff", ""),
+        )
+
+
+@runtime_checkable
+class SnapshotBackend(Protocol):
+    """Protocol for file state tracking."""
+    
+    async def track(self) -> str:
+        """Record current state, return snapshot ID."""
+        ...
+    
+    async def restore(self, snapshot_id: str) -> None:
+        """Fully restore to snapshot state."""
+        ...
+    
+    async def revert(self, patches: list[Patch]) -> None:
+        """Revert specific files based on patches."""
+        ...
+    
+    async def diff(self, snapshot_id: str) -> str:
+        """Get diff between current state and snapshot."""
+        ...
+    
+    async def patch(self, snapshot_id: str) -> Patch:
+        """Get changed files since snapshot."""
+        ...
+
+
+__all__ = ["Patch", "SnapshotBackend"]

aury/agents/backends/state/__init__.py

@@ -0,0 +1,29 @@
+"""State backend for framework internal storage.
+
+Used for storing session state, plan data, invocation records, messages, etc.
+This is NOT for user file operations - use FileBackend for that.
+
+Architecture:
+- StateStore: Low-level protocol for API layer to implement
+- StateBackend: High-level protocol with namespace support
+- StoreBasedStateBackend: Wraps StateStore as StateBackend
+
+Default implementations: SQLiteStateBackend, MemoryStateBackend
+"""
+from .types import StateBackend, StateStore, StoreBasedStateBackend
+from .sqlite import SQLiteStateBackend
+from .memory import MemoryStateBackend
+from .file import FileStateBackend
+from .composite import CompositeStateBackend
+
+__all__ = [
+    # Protocols
+    "StateBackend",
+    "StateStore",
+    # Implementations
+    "StoreBasedStateBackend",  # Wraps StateStore
+    "SQLiteStateBackend",  # Default
+    "MemoryStateBackend",  # For testing
+    "FileStateBackend",
+    "CompositeStateBackend",
+]

aury/agents/backends/state/composite.py

@@ -0,0 +1,49 @@
+"""Composite state backend that routes by namespace."""
+from __future__ import annotations
+
+from typing import Any
+
+from .types import StateBackend
+
+
+class CompositeStateBackend:
+    """Routes operations to different backends by namespace.
+    
+    Example:
+        backend = CompositeStateBackend({
+            "session": FileStateBackend("./data/sessions"),
+            "usage": PostgresBackend(db),
+        }, default=MemoryStateBackend())
+    """
+    
+    def __init__(
+        self,
+        backends: dict[str, StateBackend],
+        default: StateBackend | None = None,
+    ):
+        self._backends = backends
+        self._default = default
+    
+    def _get_backend(self, namespace: str) -> StateBackend:
+        backend = self._backends.get(namespace, self._default)
+        if backend is None:
+            raise ValueError(f"No backend for namespace: {namespace}")
+        return backend
+    
+    async def get(self, namespace: str, key: str) -> Any | None:
+        return await self._get_backend(namespace).get(namespace, key)
+    
+    async def set(self, namespace: str, key: str, value: Any) -> None:
+        await self._get_backend(namespace).set(namespace, key, value)
+    
+    async def delete(self, namespace: str, key: str) -> bool:
+        return await self._get_backend(namespace).delete(namespace, key)
+    
+    async def list(self, namespace: str, prefix: str = "") -> list[str]:
+        return await self._get_backend(namespace).list(namespace, prefix)
+    
+    async def exists(self, namespace: str, key: str) -> bool:
+        return await self._get_backend(namespace).exists(namespace, key)
+
+
+__all__ = ["CompositeStateBackend"]

aury/agents/backends/state/file.py

@@ -0,0 +1,57 @@
+"""File-based state backend."""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+
+class FileStateBackend:
+    """File-based state backend.
+    
+    Stores data as JSON files organized by namespace.
+    """
+    
+    def __init__(self, base_path: str | Path):
+        self.base_path = Path(base_path)
+        self.base_path.mkdir(parents=True, exist_ok=True)
+    
+    def _get_path(self, namespace: str, key: str) -> Path:
+        ns_path = self.base_path / namespace
+        ns_path.mkdir(parents=True, exist_ok=True)
+        return ns_path / f"{key}.json"
+    
+    async def get(self, namespace: str, key: str) -> Any | None:
+        import json
+        import aiofiles
+        path = self._get_path(namespace, key)
+        if not path.exists():
+            return None
+        async with aiofiles.open(path, "r") as f:
+            content = await f.read()
+            return json.loads(content)
+    
+    async def set(self, namespace: str, key: str, value: Any) -> None:
+        import json
+        import aiofiles
+        path = self._get_path(namespace, key)
+        async with aiofiles.open(path, "w") as f:
+            await f.write(json.dumps(value, default=str, ensure_ascii=False, indent=2))
+    
+    async def delete(self, namespace: str, key: str) -> bool:
+        path = self._get_path(namespace, key)
+        if path.exists():
+            path.unlink()
+            return True
+        return False
+    
+    async def list(self, namespace: str, prefix: str = "") -> list[str]:
+        ns_path = self.base_path / namespace
+        if not ns_path.exists():
+            return []
+        return [f.stem for f in ns_path.glob("*.json") if f.stem.startswith(prefix)]
+    
+    async def exists(self, namespace: str, key: str) -> bool:
+        return self._get_path(namespace, key).exists()
+
+
+__all__ = ["FileStateBackend"]

aury/agents/backends/state/memory.py

@@ -0,0 +1,52 @@
+"""In-memory state backend for testing."""
+from __future__ import annotations
+
+from typing import Any
+
+
+class MemoryStateBackend:
+    """In-memory state backend for testing."""
+    
+    def __init__(self) -> None:
+        self._data: dict[str, dict[str, Any]] = {}
+        self._messages: dict[str, list[dict[str, Any]]] = {}  # session_id -> messages
+    
+    async def get(self, namespace: str, key: str) -> Any | None:
+        return self._data.get(namespace, {}).get(key)
+    
+    async def set(self, namespace: str, key: str, value: Any) -> None:
+        if namespace not in self._data:
+            self._data[namespace] = {}
+        self._data[namespace][key] = value
+    
+    async def delete(self, namespace: str, key: str) -> bool:
+        if namespace in self._data and key in self._data[namespace]:
+            del self._data[namespace][key]
+            return True
+        return False
+    
+    async def list(self, namespace: str, prefix: str = "") -> list[str]:
+        if namespace not in self._data:
+            return []
+        return [k for k in self._data[namespace] if k.startswith(prefix)]
+    
+    async def exists(self, namespace: str, key: str) -> bool:
+        return namespace in self._data and key in self._data[namespace]
+    
+    async def add_message(self, session_id: str, message: dict[str, Any]) -> None:
+        """Add a message to session history."""
+        if session_id not in self._messages:
+            self._messages[session_id] = []
+        self._messages[session_id].append(message)
+    
+    async def get_messages(self, session_id: str) -> list[dict[str, Any]]:
+        """Get all messages for a session."""
+        return self._messages.get(session_id, [])
+    
+    def clear(self) -> None:
+        """Clear all data (for testing)."""
+        self._data.clear()
+        self._messages.clear()
+
+
+__all__ = ["MemoryStateBackend"]