agenthandler 0.92.0__py3-none-any.whl

agenthandler/store.py

@@ -0,0 +1,341 @@
+"""Durable state persistence — checkpoint / restore for long-running agents.
+
+Pluggable storage backend via the StateStore protocol.  Ships with two
+implementations:
+
+- MemoryStore  — dict-backed, for testing
+- SqliteStore  — stdlib sqlite3, production-ready single-process default
+
+Security notes:
+- policy_dict is the ORIGINAL immutable policy set at start(). It is never
+  updated from agent-controlled data. The Supervisor always rebuilds from this.
+- payload is agent-writable and UNTRUSTED. Treat it like user input on resume.
+- circuit_breaker_states and budget counters are system-controlled (written only
+  by the Supervisor internals, never by agent code).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import secrets
+import sqlite3
+import stat
+import threading
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Dict, List, Optional, Protocol, Tuple, runtime_checkable
+
+# Maximum payload size in bytes (1 MB). Payloads exceeding this are rejected.
+MAX_PAYLOAD_BYTES = 1_048_576
+
+
+class SessionStatus(Enum):
+    """Lifecycle status of an agent session."""
+
+    RUNNING = "running"
+    PAUSED = "paused"
+    STOPPED = "stopped"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+@dataclass
+class Checkpoint:
+    """Snapshot of a session's full state — enough to rehydrate a Supervisor.
+
+    Trust levels:
+    - SYSTEM-CONTROLLED (written only by Supervisor/SessionManager internals):
+      session_id, agent_id, status, iterations, tokens_used, token_limit,
+      iteration_limit, circuit_breaker_states, original_policy_dict, timestamp,
+      created_at
+    - AGENT-WRITABLE (untrusted — treat like user input):
+      payload
+    """
+
+    session_id: str
+    agent_id: str
+    status: SessionStatus
+
+    # Supervisor state (system-controlled)
+    iterations: int = 0
+    tokens_used: int = 0
+    token_limit: Optional[int] = None
+    iteration_limit: int = 5
+    circuit_breaker_states: Dict[str, Any] = field(default_factory=dict)
+
+    # IMMUTABLE original policy set at start() — never updated from agent data
+    policy_dict: Dict[str, Any] = field(default_factory=dict)
+
+    # Agent-writable payload — UNTRUSTED, treat like user input on resume
+    payload: Dict[str, Any] = field(default_factory=dict)
+
+    timestamp: str = ""
+    created_at: str = ""  # set once at creation, used for expiry
+
+    # If True, session is in-memory only — no auto-checkpointing, no crash recovery
+    stateless: bool = False
+
+    # How many times this session has been resumed (system-controlled, crash loop protection)
+    resume_count: int = 0
+
+    # Human-readable reason the session ended (set on stop/fail)
+    failure_reason: str = ""
+
+    # HMAC checksum of the original policy — verified on resume
+    policy_checksum: str = ""
+
+    # Persisted audit entries (list of dicts) — survives crashes
+    audit_log: List[Dict[str, Any]] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        if not self.timestamp:
+            self.timestamp = datetime.now(timezone.utc).isoformat()
+        if not self.created_at:
+            self.created_at = self.timestamp
+        if isinstance(self.status, str):
+            self.status = SessionStatus(self.status)
+
+    def to_dict(self) -> Dict[str, Any]:
+        d = asdict(self)
+        d["status"] = self.status.value
+        return d
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "Checkpoint":
+        return cls(
+            session_id=d["session_id"],
+            agent_id=d["agent_id"],
+            status=SessionStatus(d["status"]) if isinstance(d["status"], str) else d["status"],
+            iterations=d.get("iterations", 0),
+            tokens_used=d.get("tokens_used", 0),
+            token_limit=d.get("token_limit"),
+            iteration_limit=d.get("iteration_limit", 5),
+            circuit_breaker_states=d.get("circuit_breaker_states", {}),
+            policy_dict=d.get("policy_dict", {}),
+            payload=d.get("payload", {}),
+            timestamp=d.get("timestamp", ""),
+            created_at=d.get("created_at", d.get("timestamp", "")),
+            stateless=d.get("stateless", False),
+            resume_count=d.get("resume_count", 0),
+            failure_reason=d.get("failure_reason", ""),
+            policy_checksum=d.get("policy_checksum", ""),
+            audit_log=d.get("audit_log", []),
+        )
+
+
+def new_session_id() -> str:
+    """Generate a cryptographically random session ID (128 bits of entropy)."""
+    return secrets.token_hex(16)
+
+
+def validate_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
+    """Validate and sanitize an agent payload before storage.
+
+    Enforces size limits and ensures the payload is safe to serialize.
+    Raises ValueError if the payload is invalid.
+    """
+    serialized = json.dumps(payload, default=str)
+    if len(serialized.encode("utf-8")) > MAX_PAYLOAD_BYTES:
+        raise ValueError(
+            f"Payload exceeds maximum size ({MAX_PAYLOAD_BYTES} bytes). "
+            f"Got {len(serialized.encode('utf-8'))} bytes."
+        )
+    # Re-parse to strip any non-JSON-serializable objects that default=str masked
+    result: Dict[str, Any] = json.loads(serialized)
+    return result
+
+
+# ---------------------------------------------------------------------------
+# StateStore protocol
+# ---------------------------------------------------------------------------
+
+
+@runtime_checkable
+class StateStore(Protocol):
+    """Pluggable storage backend for session checkpoints."""
+
+    def save_checkpoint(self, checkpoint: Checkpoint) -> None: ...
+
+    def load_checkpoint(self, session_id: str) -> Optional[Checkpoint]: ...
+
+    def list_sessions(self) -> List[Checkpoint]: ...
+
+    def delete_session(self, session_id: str) -> bool: ...
+
+
+# ---------------------------------------------------------------------------
+# MemoryStore — dict-backed, for testing
+# ---------------------------------------------------------------------------
+
+
+class MemoryStore:
+    """In-memory checkpoint store. Thread-safe."""
+
+    def __init__(self) -> None:
+        self._data: Dict[str, Checkpoint] = {}
+        self._lock = threading.Lock()
+
+    def save_checkpoint(self, checkpoint: Checkpoint) -> None:
+        with self._lock:
+            self._data[checkpoint.session_id] = checkpoint
+
+    def load_checkpoint(self, session_id: str) -> Optional[Checkpoint]:
+        with self._lock:
+            return self._data.get(session_id)
+
+    def list_sessions(self) -> List[Checkpoint]:
+        with self._lock:
+            return list(self._data.values())
+
+    def delete_session(self, session_id: str) -> bool:
+        with self._lock:
+            if session_id in self._data:
+                del self._data[session_id]
+                return True
+            return False
+
+
+# ---------------------------------------------------------------------------
+# SqliteStore — stdlib sqlite3, production default
+# ---------------------------------------------------------------------------
+
+_CREATE_TABLE = """\
+CREATE TABLE IF NOT EXISTS checkpoints (
+    session_id   TEXT PRIMARY KEY,
+    agent_id     TEXT NOT NULL,
+    status       TEXT NOT NULL,
+    iterations   INTEGER NOT NULL DEFAULT 0,
+    tokens_used  INTEGER NOT NULL DEFAULT 0,
+    token_limit  INTEGER,
+    iteration_limit INTEGER NOT NULL DEFAULT 5,
+    circuit_breaker_states TEXT NOT NULL DEFAULT '{}',
+    policy_dict  TEXT NOT NULL DEFAULT '{}',
+    payload      TEXT NOT NULL DEFAULT '{}',
+    timestamp    TEXT NOT NULL,
+    created_at   TEXT NOT NULL DEFAULT '',
+    audit_log    TEXT NOT NULL DEFAULT '[]'
+)"""
+
+# Migration: add columns if upgrading from Phase 1 schema
+_MIGRATIONS = [
+    "ALTER TABLE checkpoints ADD COLUMN created_at TEXT NOT NULL DEFAULT ''",
+    "ALTER TABLE checkpoints ADD COLUMN audit_log TEXT NOT NULL DEFAULT '[]'",
+]
+
+
+class SqliteStore:
+    """SQLite-backed checkpoint store. Thread-safe (serialized via lock).
+
+    Uses stdlib sqlite3 — no extra dependencies.
+    Creates the database file with owner-only permissions (0600).
+
+    Usage:
+        store = SqliteStore("checkpoints.db")
+        store.save_checkpoint(checkpoint)
+    """
+
+    def __init__(self, db_path: str = "agenthandler_sessions.db"):
+        self._db_path = db_path
+        self._lock = threading.Lock()
+
+        # Create file with restricted permissions if it doesn't exist
+        if not os.path.exists(db_path):
+            fd = os.open(db_path, os.O_CREAT | os.O_WRONLY, stat.S_IRUSR | stat.S_IWUSR)
+            os.close(fd)
+
+        with self._connect() as conn:
+            conn.execute(_CREATE_TABLE)
+            # Run migrations for existing databases
+            for migration in _MIGRATIONS:
+                try:
+                    conn.execute(migration)
+                except sqlite3.OperationalError:
+                    pass  # column already exists
+
+    def _connect(self) -> sqlite3.Connection:
+        return sqlite3.connect(self._db_path)
+
+    def save_checkpoint(self, checkpoint: Checkpoint) -> None:
+        with self._lock:
+            with self._connect() as conn:
+                conn.execute(
+                    """\
+                    INSERT OR REPLACE INTO checkpoints
+                        (session_id, agent_id, status, iterations, tokens_used,
+                         token_limit, iteration_limit, circuit_breaker_states,
+                         policy_dict, payload, timestamp, created_at, audit_log)
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+                    (
+                        checkpoint.session_id,
+                        checkpoint.agent_id,
+                        checkpoint.status.value,
+                        checkpoint.iterations,
+                        checkpoint.tokens_used,
+                        checkpoint.token_limit,
+                        checkpoint.iteration_limit,
+                        json.dumps(checkpoint.circuit_breaker_states),
+                        json.dumps(checkpoint.policy_dict),
+                        json.dumps(checkpoint.payload),
+                        checkpoint.timestamp,
+                        checkpoint.created_at,
+                        json.dumps(checkpoint.audit_log),
+                    ),
+                )
+
+    def load_checkpoint(self, session_id: str) -> Optional[Checkpoint]:
+        with self._lock:
+            with self._connect() as conn:
+                row = conn.execute(
+                    "SELECT * FROM checkpoints WHERE session_id = ?", (session_id,)
+                ).fetchone()
+                if row is None:
+                    return None
+                return self._row_to_checkpoint(row)
+
+    def list_sessions(self) -> List[Checkpoint]:
+        with self._lock:
+            with self._connect() as conn:
+                rows = conn.execute("SELECT * FROM checkpoints ORDER BY timestamp DESC").fetchall()
+                return [self._row_to_checkpoint(r) for r in rows]
+
+    def delete_session(self, session_id: str) -> bool:
+        with self._lock:
+            with self._connect() as conn:
+                cursor = conn.execute("DELETE FROM checkpoints WHERE session_id = ?", (session_id,))
+                return cursor.rowcount > 0
+
+    def delete_expired(self, max_age_seconds: float) -> int:
+        """Delete sessions older than max_age_seconds. Returns count deleted."""
+        cutoff = datetime.now(timezone.utc).timestamp() - max_age_seconds
+        with self._lock:
+            with self._connect() as conn:
+                # created_at is ISO format — compare as strings (works for ISO 8601)
+                cutoff_iso = datetime.fromtimestamp(cutoff, tz=timezone.utc).isoformat()
+                cursor = conn.execute(
+                    "DELETE FROM checkpoints WHERE created_at != '' AND created_at < ?",
+                    (cutoff_iso,),
+                )
+                return cursor.rowcount
+
+    @staticmethod
+    def _row_to_checkpoint(row: Tuple[Any, ...]) -> Checkpoint:
+        # Handle both old (11-column) and new (13-column) schemas
+        created_at = row[11] if len(row) > 11 else ""
+        audit_log = json.loads(row[12]) if len(row) > 12 else []
+        return Checkpoint(
+            session_id=row[0],
+            agent_id=row[1],
+            status=SessionStatus(row[2]),
+            iterations=row[3],
+            tokens_used=row[4],
+            token_limit=row[5],
+            iteration_limit=row[6],
+            circuit_breaker_states=json.loads(row[7]),
+            policy_dict=json.loads(row[8]),
+            payload=json.loads(row[9]),
+            timestamp=row[10],
+            created_at=created_at,
+            audit_log=audit_log,
+        )

agenthandler/streams.py

@@ -0,0 +1,339 @@
+"""Streaming input connectors — Kafka, Kinesis, Redis Streams.
+
+Subscribe an agent to a real-time event stream. Each message fires the
+registered pipeline with the message payload as input.
+
+All three connectors implement the same StreamConnector protocol:
+- start() — begin consuming
+- stop() — stop consuming
+- Messages fire the registered callback asynchronously
+
+Optional dependencies:
+- Kafka: ``pip install aiokafka``
+- Kinesis: ``pip install aioboto3``
+- Redis Streams: ``pip install redis``
+
+Usage:
+    from agenthandler.streams import KafkaStreamConnector
+
+    async def on_message(payload):
+        # Run pipeline with the payload as input
+        return await pipeline.run()
+
+    connector = KafkaStreamConnector(
+        bootstrap_servers="localhost:9092",
+        topic="agent-events",
+        group_id="my-agent",
+        on_message=on_message,
+    )
+    await connector.start()
+    # ... later ...
+    await connector.stop()
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any, Awaitable, Callable, Dict, Optional, Protocol, runtime_checkable
+
+MessageHandler = Callable[[Dict[str, Any]], Awaitable[Any]]
+
+
+@runtime_checkable
+class StreamConnector(Protocol):
+    """Protocol for streaming input connectors."""
+
+    async def start(self) -> None: ...
+    async def stop(self) -> None: ...
+
+
+# ---------------------------------------------------------------------------
+# Kafka
+# ---------------------------------------------------------------------------
+
+
+class KafkaStreamConnector:
+    """Consume messages from a Kafka topic and fire an agent pipeline.
+
+    Uses aiokafka (optional dependency). Offsets are committed after
+    successful handler execution — at-least-once delivery.
+
+    Args:
+        bootstrap_servers: Comma-separated Kafka broker list.
+        topic: Topic to consume from.
+        group_id: Consumer group ID for offset tracking.
+        on_message: Async callback fired with each message payload.
+        auto_offset_reset: "earliest" or "latest".
+        value_deserializer: How to decode message bytes (default: JSON).
+    """
+
+    def __init__(
+        self,
+        bootstrap_servers: str,
+        topic: str,
+        group_id: str,
+        on_message: MessageHandler,
+        auto_offset_reset: str = "latest",
+        value_deserializer: Optional[Callable[[bytes], Any]] = None,
+    ):
+        self._bootstrap = bootstrap_servers
+        self._topic = topic
+        self._group_id = group_id
+        self._on_message = on_message
+        self._auto_offset_reset = auto_offset_reset
+        self._deserializer = value_deserializer or (lambda b: json.loads(b.decode()))
+        self._consumer: Any = None
+        self._running = False
+        self._task: Optional[asyncio.Task[None]] = None
+
+    async def start(self) -> None:
+        """Start consuming messages."""
+        try:
+            from aiokafka import AIOKafkaConsumer  # type: ignore[import-not-found]
+        except ImportError as e:
+            raise ImportError("aiokafka required: pip install aiokafka") from e
+
+        self._consumer = AIOKafkaConsumer(
+            self._topic,
+            bootstrap_servers=self._bootstrap,
+            group_id=self._group_id,
+            auto_offset_reset=self._auto_offset_reset,
+            value_deserializer=self._deserializer,
+            enable_auto_commit=False,
+        )
+        await self._consumer.start()
+        self._running = True
+        self._task = asyncio.create_task(self._consume_loop())
+
+    async def _consume_loop(self) -> None:
+        """Main consume loop — fires on_message for each record."""
+        try:
+            async for msg in self._consumer:
+                if not self._running:
+                    break
+                payload = {
+                    "key": msg.key.decode() if msg.key else None,
+                    "value": msg.value,
+                    "topic": msg.topic,
+                    "partition": msg.partition,
+                    "offset": msg.offset,
+                    "timestamp": msg.timestamp,
+                }
+                try:
+                    await self._on_message(payload)
+                    await self._consumer.commit()
+                except Exception:
+                    # Don't commit — message will be redelivered
+                    pass
+        except Exception:
+            pass
+
+    async def stop(self) -> None:
+        """Stop consuming."""
+        self._running = False
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except (asyncio.CancelledError, Exception):
+                pass
+        if self._consumer:
+            await self._consumer.stop()
+            self._consumer = None
+
+
+# ---------------------------------------------------------------------------
+# AWS Kinesis
+# ---------------------------------------------------------------------------
+
+
+class KinesisStreamConnector:
+    """Consume records from an AWS Kinesis Data Stream.
+
+    Uses aioboto3 (optional dependency). Polls shards in a loop and
+    fires on_message for each record.
+
+    Args:
+        stream_name: Kinesis stream name.
+        on_message: Async callback fired with each record payload.
+        region_name: AWS region.
+        shard_iterator_type: "LATEST" or "TRIM_HORIZON".
+        poll_interval: Seconds between shard polls.
+    """
+
+    def __init__(
+        self,
+        stream_name: str,
+        on_message: MessageHandler,
+        region_name: str = "us-east-1",
+        shard_iterator_type: str = "LATEST",
+        poll_interval: float = 1.0,
+    ):
+        self._stream = stream_name
+        self._on_message = on_message
+        self._region = region_name
+        self._iterator_type = shard_iterator_type
+        self._poll_interval = poll_interval
+        self._running = False
+        self._task: Optional[asyncio.Task[None]] = None
+        self._session: Any = None
+
+    async def start(self) -> None:
+        """Start consuming from all shards."""
+        try:
+            import aioboto3  # type: ignore[import-not-found]
+        except ImportError as e:
+            raise ImportError("aioboto3 required: pip install aioboto3") from e
+
+        self._session = aioboto3.Session()
+        self._running = True
+        self._task = asyncio.create_task(self._consume_loop())
+
+    async def _consume_loop(self) -> None:
+        """Poll Kinesis shards for records."""
+        async with self._session.client("kinesis", region_name=self._region) as client:
+            # Get shards
+            desc = await client.describe_stream(StreamName=self._stream)
+            shards = desc["StreamDescription"]["Shards"]
+
+            # Get shard iterators
+            iterators = {}
+            for shard in shards:
+                resp = await client.get_shard_iterator(
+                    StreamName=self._stream,
+                    ShardId=shard["ShardId"],
+                    ShardIteratorType=self._iterator_type,
+                )
+                iterators[shard["ShardId"]] = resp["ShardIterator"]
+
+            # Poll loop
+            while self._running:
+                for shard_id, iterator in list(iterators.items()):
+                    if iterator is None:
+                        continue
+                    try:
+                        resp = await client.get_records(ShardIterator=iterator, Limit=100)
+                        for record in resp.get("Records", []):
+                            payload = {
+                                "data": json.loads(record["Data"])
+                                if isinstance(record["Data"], bytes)
+                                else record["Data"],
+                                "partition_key": record.get("PartitionKey"),
+                                "sequence_number": record.get("SequenceNumber"),
+                                "shard_id": shard_id,
+                            }
+                            try:
+                                await self._on_message(payload)
+                            except Exception:
+                                pass
+                        iterators[shard_id] = resp.get("NextShardIterator")
+                    except Exception:
+                        pass
+                await asyncio.sleep(self._poll_interval)
+
+    async def stop(self) -> None:
+        """Stop consuming."""
+        self._running = False
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except (asyncio.CancelledError, Exception):
+                pass
+
+
+# ---------------------------------------------------------------------------
+# Redis Streams
+# ---------------------------------------------------------------------------
+
+
+class RedisStreamConnector:
+    """Consume messages from a Redis Stream.
+
+    Uses redis-py with consumer groups for at-least-once delivery.
+
+    Args:
+        url: Redis connection URL (e.g., "redis://localhost:6379").
+        stream: Stream key name.
+        group: Consumer group name.
+        consumer: Consumer name within the group.
+        on_message: Async callback fired with each message.
+        block_ms: How long to block waiting for messages.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        stream: str,
+        group: str,
+        consumer: str,
+        on_message: MessageHandler,
+        block_ms: int = 5000,
+    ):
+        self._url = url
+        self._stream = stream
+        self._group = group
+        self._consumer = consumer
+        self._on_message = on_message
+        self._block_ms = block_ms
+        self._redis: Any = None
+        self._running = False
+        self._task: Optional[asyncio.Task[None]] = None
+
+    async def start(self) -> None:
+        """Start consuming from the Redis stream."""
+        try:
+            import redis.asyncio as redis  # type: ignore[import-not-found]
+        except ImportError as e:
+            raise ImportError("redis required: pip install redis") from e
+
+        self._redis = redis.from_url(self._url, decode_responses=True)
+
+        # Create consumer group (idempotent)
+        try:
+            await self._redis.xgroup_create(self._stream, self._group, id="$", mkstream=True)
+        except Exception:
+            pass  # group already exists
+
+        self._running = True
+        self._task = asyncio.create_task(self._consume_loop())
+
+    async def _consume_loop(self) -> None:
+        """Main consume loop — uses XREADGROUP for consumer group semantics."""
+        while self._running:
+            try:
+                messages = await self._redis.xreadgroup(
+                    groupname=self._group,
+                    consumername=self._consumer,
+                    streams={self._stream: ">"},
+                    count=10,
+                    block=self._block_ms,
+                )
+                for stream_name, entries in messages:
+                    for msg_id, fields in entries:
+                        payload = {
+                            "id": msg_id,
+                            "stream": stream_name,
+                            "fields": fields,
+                        }
+                        try:
+                            await self._on_message(payload)
+                            await self._redis.xack(self._stream, self._group, msg_id)
+                        except Exception:
+                            pass
+            except Exception:
+                await asyncio.sleep(1)
+
+    async def stop(self) -> None:
+        """Stop consuming."""
+        self._running = False
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except (asyncio.CancelledError, Exception):
+                pass
+        if self._redis:
+            await self._redis.close()
+            self._redis = None