aury-agent 0.0.4__py3-none-any.whl

aury/agents/backends/session/memory.py

@@ -0,0 +1,93 @@
+"""In-memory session backend implementation."""
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+
+
+class InMemorySessionBackend:
+    """In-memory implementation of SessionBackend.
+    
+    Suitable for testing and simple single-process use cases.
+    Data is lost when the process exits.
+    """
+    
+    def __init__(self) -> None:
+        self._sessions: dict[str, dict[str, Any]] = {}
+        # Index: user_id -> list of session_ids
+        self._user_sessions: dict[str, list[str]] = {}
+    
+    async def create(
+        self,
+        id: str,
+        data: dict[str, Any],
+        user_id: str | None = None,
+    ) -> None:
+        """Create a new session."""
+        session_data = {
+            "id": id,
+            "user_id": user_id,
+            "created_at": datetime.now().isoformat(),
+            "updated_at": datetime.now().isoformat(),
+            **data,
+        }
+        self._sessions[id] = session_data
+        
+        # Update user index
+        if user_id:
+            if user_id not in self._user_sessions:
+                self._user_sessions[user_id] = []
+            self._user_sessions[user_id].append(id)
+    
+    async def get(self, id: str) -> dict[str, Any] | None:
+        """Get session by ID."""
+        return self._sessions.get(id)
+    
+    async def update(self, id: str, data: dict[str, Any]) -> None:
+        """Update session data."""
+        if id in self._sessions:
+            self._sessions[id].update(data)
+            self._sessions[id]["updated_at"] = datetime.now().isoformat()
+    
+    async def delete(self, id: str) -> bool:
+        """Delete a session."""
+        if id not in self._sessions:
+            return False
+        
+        session = self._sessions.pop(id)
+        
+        # Update user index
+        user_id = session.get("user_id")
+        if user_id and user_id in self._user_sessions:
+            self._user_sessions[user_id] = [
+                sid for sid in self._user_sessions[user_id] if sid != id
+            ]
+        
+        return True
+    
+    async def list(
+        self,
+        user_id: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[dict[str, Any]]:
+        """List sessions."""
+        if user_id:
+            # Filter by user
+            session_ids = self._user_sessions.get(user_id, [])
+            sessions = [
+                self._sessions[sid]
+                for sid in session_ids
+                if sid in self._sessions
+            ]
+        else:
+            sessions = list(self._sessions.values())
+        
+        # Sort by created_at descending (newest first)
+        sessions.sort(key=lambda s: s.get("created_at", ""), reverse=True)
+        
+        # Apply pagination
+        return sessions[offset:offset + limit]
+
+
+__all__ = ["InMemorySessionBackend"]

aury/agents/backends/session/types.py

@@ -0,0 +1,124 @@
+"""Session backend types and protocols."""
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ...core.context import InvocationContext
+
+
+@runtime_checkable
+class SessionBackend(Protocol):
+    """Protocol for session management.
+    
+    Sessions represent a conversation scope. Each session can contain
+    multiple invocations (user turns).
+    
+    All methods accept an optional `ctx` (InvocationContext) parameter.
+    When ctx is provided, session_id/user_id can be extracted from it.
+    
+    Example usage:
+        # With explicit params
+        await backend.create("sess_123", {"root_agent_id": "agent"}, user_id="user_1")
+        
+        # With ctx (auto-extract session_id)
+        await backend.get(ctx=ctx)  # uses ctx.session_id
+        
+        # List user's sessions
+        sessions = await backend.list(user_id="user_1")
+    """
+    
+    async def create(
+        self,
+        data: dict[str, Any],
+        *,
+        id: str | None = None,
+        user_id: str | None = None,
+        ctx: "InvocationContext | None" = None,
+    ) -> str:
+        """Create a new session.
+        
+        Args:
+            data: Session data (root_agent_id, metadata, etc.)
+            id: Session ID (auto-generated if None)
+            user_id: Optional user ID for multi-tenant isolation
+            ctx: Optional InvocationContext
+            
+        Returns:
+            Created session ID
+        """
+        ...
+    
+    async def get(
+        self,
+        id: str | None = None,
+        *,
+        ctx: "InvocationContext | None" = None,
+    ) -> dict[str, Any] | None:
+        """Get session by ID.
+        
+        Args:
+            id: Session ID (or extracted from ctx.session_id)
+            ctx: Optional InvocationContext
+            
+        Returns:
+            Session data dict or None if not found
+        """
+        ...
+    
+    async def update(
+        self,
+        data: dict[str, Any],
+        *,
+        id: str | None = None,
+        ctx: "InvocationContext | None" = None,
+    ) -> None:
+        """Update session data.
+        
+        Args:
+            data: Fields to update (partial update)
+            id: Session ID (or extracted from ctx.session_id)
+            ctx: Optional InvocationContext
+        """
+        ...
+    
+    async def delete(
+        self,
+        id: str | None = None,
+        *,
+        ctx: "InvocationContext | None" = None,
+    ) -> bool:
+        """Delete a session.
+        
+        Args:
+            id: Session ID (or extracted from ctx.session_id)
+            ctx: Optional InvocationContext
+            
+        Returns:
+            True if deleted, False if not found
+        """
+        ...
+    
+    async def list(
+        self,
+        *,
+        user_id: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+        ctx: "InvocationContext | None" = None,
+    ) -> list[dict[str, Any]]:
+        """List sessions.
+        
+        Args:
+            user_id: Optional filter by user
+            limit: Max sessions to return
+            offset: Offset for pagination
+            ctx: Optional InvocationContext
+            
+        Returns:
+            List of session data dicts
+        """
+        ...
+
+
+__all__ = ["SessionBackend"]

aury/agents/backends/shell/__init__.py

@@ -0,0 +1,11 @@
+"""Shell backend for command execution.
+
+Supports different execution environments:
+- Local shell
+- Docker/nsjail sandbox
+- E2B cloud sandbox
+"""
+from .types import ShellBackend, ShellResult
+from .local import LocalShellBackend
+
+__all__ = ["ShellBackend", "ShellResult", "LocalShellBackend"]

aury/agents/backends/shell/local.py

@@ -0,0 +1,110 @@
+"""Local shell execution backend."""
+from __future__ import annotations
+
+import asyncio
+import os
+import time
+from typing import AsyncIterator
+
+from .types import ShellResult
+
+
+class LocalShellBackend:
+    """Local shell execution backend."""
+    
+    def __init__(self, default_cwd: str | None = None, shell: str = "/bin/bash"):
+        self.default_cwd = default_cwd
+        self.shell = shell
+    
+    async def execute(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: int = 120,
+    ) -> ShellResult:
+        """Execute command locally."""
+        work_dir = cwd or self.default_cwd or os.getcwd()
+        start_time = time.time()
+        
+        full_env = os.environ.copy()
+        if env:
+            full_env.update(env)
+        
+        try:
+            proc = await asyncio.create_subprocess_shell(
+                command,
+                cwd=work_dir,
+                env=full_env,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            
+            stdout, stderr = await asyncio.wait_for(
+                proc.communicate(),
+                timeout=timeout,
+            )
+            
+            return ShellResult(
+                stdout=stdout.decode("utf-8", errors="replace"),
+                stderr=stderr.decode("utf-8", errors="replace"),
+                exit_code=proc.returncode or 0,
+                command=command,
+                cwd=work_dir,
+                duration_ms=int((time.time() - start_time) * 1000),
+            )
+            
+        except asyncio.TimeoutError:
+            return ShellResult(
+                stdout="",
+                stderr=f"Command timed out after {timeout}s",
+                exit_code=-1,
+                command=command,
+                cwd=work_dir,
+                duration_ms=timeout * 1000,
+            )
+        except Exception as e:
+            return ShellResult(
+                stdout="",
+                stderr=str(e),
+                exit_code=-1,
+                command=command,
+                cwd=work_dir,
+                duration_ms=int((time.time() - start_time) * 1000),
+            )
+    
+    async def execute_stream(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: int = 120,
+    ) -> AsyncIterator[str]:
+        """Execute with streaming output."""
+        work_dir = cwd or self.default_cwd or os.getcwd()
+        
+        full_env = os.environ.copy()
+        if env:
+            full_env.update(env)
+        
+        proc = await asyncio.create_subprocess_shell(
+            command,
+            cwd=work_dir,
+            env=full_env,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        
+        try:
+            async with asyncio.timeout(timeout):
+                while True:
+                    line = await proc.stdout.readline()
+                    if not line:
+                        break
+                    yield line.decode("utf-8", errors="replace")
+        except asyncio.TimeoutError:
+            proc.kill()
+            yield f"\n[Timeout after {timeout}s]\n"
+
+
+__all__ = ["LocalShellBackend"]

aury/agents/backends/shell/types.py

@@ -0,0 +1,55 @@
+"""Shell backend types and protocols."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import AsyncIterator, Protocol, runtime_checkable
+
+
+@dataclass
+class ShellResult:
+    """Result of shell command execution."""
+    stdout: str
+    stderr: str
+    exit_code: int
+    command: str = ""
+    cwd: str = ""
+    duration_ms: int = 0
+    
+    @property
+    def success(self) -> bool:
+        return self.exit_code == 0
+    
+    @property
+    def output(self) -> str:
+        """Combined output."""
+        if self.success:
+            return self.stdout
+        return f"{self.stdout}\n{self.stderr}".strip()
+
+
+@runtime_checkable
+class ShellBackend(Protocol):
+    """Protocol for shell command execution."""
+    
+    async def execute(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: int = 120,
+    ) -> ShellResult:
+        """Execute a shell command."""
+        ...
+    
+    async def execute_stream(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: int = 120,
+    ) -> AsyncIterator[str]:
+        """Execute command with streaming output."""
+        ...
+
+
+__all__ = ["ShellResult", "ShellBackend"]

aury/agents/backends/shell.py

@@ -0,0 +1,209 @@
+"""Shell backend protocol for command execution.
+
+Supports different execution environments:
+- Local shell
+- Docker/nsjail sandbox
+- E2B cloud sandbox
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import AsyncIterator, Literal, Protocol, runtime_checkable
+
+
+@dataclass
+class ShellResult:
+    """Result of shell command execution."""
+    stdout: str
+    stderr: str
+    exit_code: int
+    
+    # Execution metadata
+    command: str = ""
+    cwd: str = ""
+    duration_ms: int = 0
+    
+    @property
+    def success(self) -> bool:
+        return self.exit_code == 0
+    
+    @property
+    def output(self) -> str:
+        """Combined output (stdout + stderr if error)."""
+        if self.success:
+            return self.stdout
+        return f"{self.stdout}\n{self.stderr}".strip()
+    
+    def to_dict(self) -> dict:
+        return {
+            "stdout": self.stdout,
+            "stderr": self.stderr,
+            "exit_code": self.exit_code,
+            "command": self.command,
+            "cwd": self.cwd,
+            "duration_ms": self.duration_ms,
+        }
+
+
+@runtime_checkable
+class ShellBackend(Protocol):
+    """Protocol for shell command execution.
+    
+    Implementations:
+    - LocalShellBackend - Execute on local machine
+    - SandboxShellBackend - Execute in Docker/nsjail sandbox
+    - E2BShellBackend - Execute in E2B cloud sandbox
+    """
+    
+    async def execute(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: int = 120,
+    ) -> ShellResult:
+        """Execute a shell command.
+        
+        Args:
+            command: Command to execute
+            cwd: Working directory (optional)
+            env: Environment variables (optional)
+            timeout: Timeout in seconds (default: 120)
+            
+        Returns:
+            ShellResult with stdout, stderr, exit_code
+        """
+        ...
+    
+    async def execute_stream(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: int = 120,
+    ) -> AsyncIterator[str]:
+        """Execute command with streaming output.
+        
+        Args:
+            command: Command to execute
+            cwd: Working directory (optional)
+            env: Environment variables (optional)
+            timeout: Timeout in seconds (default: 120)
+            
+        Yields:
+            Output chunks as they arrive
+        """
+        ...
+
+
+class LocalShellBackend:
+    """Local shell execution backend."""
+    
+    def __init__(self, default_cwd: str | None = None, default_shell: str = "/bin/bash"):
+        self.default_cwd = default_cwd
+        self.default_shell = default_shell
+    
+    async def execute(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: int = 120,
+    ) -> ShellResult:
+        """Execute command locally."""
+        import asyncio
+        import os
+        import time
+        
+        work_dir = cwd or self.default_cwd or os.getcwd()
+        start_time = time.time()
+        
+        # Merge environment
+        full_env = os.environ.copy()
+        if env:
+            full_env.update(env)
+        
+        try:
+            proc = await asyncio.create_subprocess_shell(
+                command,
+                cwd=work_dir,
+                env=full_env,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            
+            stdout, stderr = await asyncio.wait_for(
+                proc.communicate(),
+                timeout=timeout,
+            )
+            
+            return ShellResult(
+                stdout=stdout.decode("utf-8", errors="replace"),
+                stderr=stderr.decode("utf-8", errors="replace"),
+                exit_code=proc.returncode or 0,
+                command=command,
+                cwd=work_dir,
+                duration_ms=int((time.time() - start_time) * 1000),
+            )
+            
+        except asyncio.TimeoutError:
+            return ShellResult(
+                stdout="",
+                stderr=f"Command timed out after {timeout} seconds",
+                exit_code=-1,
+                command=command,
+                cwd=work_dir,
+                duration_ms=timeout * 1000,
+            )
+        except Exception as e:
+            return ShellResult(
+                stdout="",
+                stderr=str(e),
+                exit_code=-1,
+                command=command,
+                cwd=work_dir,
+                duration_ms=int((time.time() - start_time) * 1000),
+            )
+    
+    async def execute_stream(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: int = 120,
+    ) -> AsyncIterator[str]:
+        """Execute with streaming output."""
+        import asyncio
+        import os
+        
+        work_dir = cwd or self.default_cwd or os.getcwd()
+        
+        full_env = os.environ.copy()
+        if env:
+            full_env.update(env)
+        
+        proc = await asyncio.create_subprocess_shell(
+            command,
+            cwd=work_dir,
+            env=full_env,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        
+        async def read_with_timeout():
+            try:
+                async with asyncio.timeout(timeout):
+                    while True:
+                        line = await proc.stdout.readline()
+                        if not line:
+                            break
+                        yield line.decode("utf-8", errors="replace")
+            except asyncio.TimeoutError:
+                proc.kill()
+                yield f"\n[Timeout after {timeout} seconds]\n"
+        
+        async for chunk in read_with_timeout():
+            yield chunk
+
+
+__all__ = ["ShellBackend", "ShellResult", "LocalShellBackend"]

aury/agents/backends/snapshot/__init__.py

@@ -0,0 +1,19 @@
+"""Snapshot backend for file state tracking and revert.
+
+Supports different snapshot strategies:
+- Git-based (local)
+- In-memory (testing)
+- Git + S3 hybrid (cloud persistence)
+"""
+from .types import Patch, SnapshotBackend
+from .memory import InMemorySnapshotBackend
+from .git import GitSnapshotBackend
+from .hybrid import GitS3HybridBackend
+
+__all__ = [
+    "SnapshotBackend",
+    "Patch",
+    "InMemorySnapshotBackend",
+    "GitSnapshotBackend",
+    "GitS3HybridBackend",
+]