openai-agents 0.2.8__py3-none-any.whl → 0.6.8__py3-none-any.whl

agents/extensions/memory/encrypt_session.py

@@ -0,0 +1,185 @@
+"""Encrypted Session wrapper for secure conversation storage.
+
+This module provides transparent encryption for session storage with automatic
+expiration of old data. When TTL expires, expired items are silently skipped.
+
+Usage::
+
+    from agents.extensions.memory import EncryptedSession, SQLAlchemySession
+
+    # Create underlying session (e.g. SQLAlchemySession)
+    underlying_session = SQLAlchemySession.from_url(
+        session_id="user-123",
+        url="postgresql+asyncpg://app:secret@db.example.com/agents",
+        create_tables=True,
+    )
+
+    # Wrap with encryption and TTL-based expiration
+    session = EncryptedSession(
+        session_id="user-123",
+        underlying_session=underlying_session,
+        encryption_key="your-encryption-key",
+        ttl=600,  # 10 minutes
+    )
+
+    await Runner.run(agent, "Hello", session=session)
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+from typing import Any, cast
+
+from cryptography.fernet import Fernet, InvalidToken
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from typing_extensions import Literal, TypedDict, TypeGuard
+
+from ...items import TResponseInputItem
+from ...memory.session import SessionABC
+
+
+class EncryptedEnvelope(TypedDict):
+    """TypedDict for encrypted message envelopes stored in the underlying session."""
+
+    __enc__: Literal[1]
+    v: int
+    kid: str
+    payload: str
+
+
+def _ensure_fernet_key_bytes(master_key: str) -> bytes:
+    """
+    Accept either a Fernet key (urlsafe-b64, 32 bytes after decode) or a raw string.
+    Returns raw bytes suitable for HKDF input.
+    """
+    if not master_key:
+        raise ValueError("encryption_key not set; required for EncryptedSession.")
+    try:
+        key_bytes = base64.urlsafe_b64decode(master_key)
+        if len(key_bytes) == 32:
+            return key_bytes
+    except Exception:
+        pass
+    return master_key.encode("utf-8")
+
+
+def _derive_session_fernet_key(master_key_bytes: bytes, session_id: str) -> Fernet:
+    hkdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=session_id.encode("utf-8"),
+        info=b"agents.session-store.hkdf.v1",
+    )
+    derived = hkdf.derive(master_key_bytes)
+    return Fernet(base64.urlsafe_b64encode(derived))
+
+
+def _to_json_bytes(obj: Any) -> bytes:
+    return json.dumps(obj, ensure_ascii=False, separators=(",", ":"), default=str).encode("utf-8")
+
+
+def _from_json_bytes(data: bytes) -> Any:
+    return json.loads(data.decode("utf-8"))
+
+
+def _is_encrypted_envelope(item: object) -> TypeGuard[EncryptedEnvelope]:
+    """Type guard to check if an item is an encrypted envelope."""
+    return (
+        isinstance(item, dict)
+        and item.get("__enc__") == 1
+        and "payload" in item
+        and "kid" in item
+        and "v" in item
+    )
+
+
+class EncryptedSession(SessionABC):
+    """Encrypted wrapper for Session implementations with TTL-based expiration.
+
+    This class wraps any SessionABC implementation to provide transparent
+    encryption/decryption of stored items using Fernet encryption with
+    per-session key derivation and automatic expiration of old data.
+
+    When items expire (exceed TTL), they are silently skipped during retrieval.
+
+    Note: Expired tokens are rejected based on the system clock of the application server.
+    To avoid valid tokens being rejected due to clock drift, ensure all servers in
+    your environment are synchronized using NTP.
+    """
+
+    def __init__(
+        self,
+        session_id: str,
+        underlying_session: SessionABC,
+        encryption_key: str,
+        ttl: int = 600,
+    ):
+        """
+        Args:
+            session_id: ID for this session
+            underlying_session: The real session store (e.g. SQLiteSession, SQLAlchemySession)
+            encryption_key: Master key (Fernet key or raw secret)
+            ttl: Token time-to-live in seconds (default 10 min)
+        """
+        self.session_id = session_id
+        self.underlying_session = underlying_session
+        self.ttl = ttl
+
+        master = _ensure_fernet_key_bytes(encryption_key)
+        self.cipher = _derive_session_fernet_key(master, session_id)
+        self._kid = "hkdf-v1"
+        self._ver = 1
+
+    def __getattr__(self, name):
+        return getattr(self.underlying_session, name)
+
+    def _wrap(self, item: TResponseInputItem) -> EncryptedEnvelope:
+        if isinstance(item, dict):
+            payload = item
+        elif hasattr(item, "model_dump"):
+            payload = item.model_dump()
+        elif hasattr(item, "__dict__"):
+            payload = item.__dict__
+        else:
+            payload = dict(item)
+
+        token = self.cipher.encrypt(_to_json_bytes(payload)).decode("utf-8")
+        return {"__enc__": 1, "v": self._ver, "kid": self._kid, "payload": token}
+
+    def _unwrap(self, item: TResponseInputItem | EncryptedEnvelope) -> TResponseInputItem | None:
+        if not _is_encrypted_envelope(item):
+            return cast(TResponseInputItem, item)
+
+        try:
+            token = item["payload"].encode("utf-8")
+            plaintext = self.cipher.decrypt(token, ttl=self.ttl)
+            return cast(TResponseInputItem, _from_json_bytes(plaintext))
+        except (InvalidToken, KeyError):
+            return None
+
+    async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
+        encrypted_items = await self.underlying_session.get_items(limit)
+        valid_items: list[TResponseInputItem] = []
+        for enc in encrypted_items:
+            item = self._unwrap(enc)
+            if item is not None:
+                valid_items.append(item)
+        return valid_items
+
+    async def add_items(self, items: list[TResponseInputItem]) -> None:
+        wrapped: list[EncryptedEnvelope] = [self._wrap(it) for it in items]
+        await self.underlying_session.add_items(cast(list[TResponseInputItem], wrapped))
+
+    async def pop_item(self) -> TResponseInputItem | None:
+        while True:
+            enc = await self.underlying_session.pop_item()
+            if not enc:
+                return None
+            item = self._unwrap(enc)
+            if item is not None:
+                return item
+
+    async def clear_session(self) -> None:
+        await self.underlying_session.clear_session()

agents/extensions/memory/redis_session.py

@@ -0,0 +1,261 @@
+"""Redis-powered Session backend.
+
+Usage::
+
+    from agents.extensions.memory import RedisSession
+
+    # Create from Redis URL
+    session = RedisSession.from_url(
+        session_id="user-123",
+        url="redis://localhost:6379/0",
+    )
+
+    # Or pass an existing Redis client that your application already manages
+    session = RedisSession(
+        session_id="user-123",
+        redis_client=my_redis_client,
+    )
+
+    await Runner.run(agent, "Hello", session=session)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from typing import Any
+
+try:
+    import redis.asyncio as redis
+    from redis.asyncio import Redis
+except ImportError as e:
+    raise ImportError(
+        "RedisSession requires the 'redis' package. Install it with: pip install redis"
+    ) from e
+
+from ...items import TResponseInputItem
+from ...memory.session import SessionABC
+
+
+class RedisSession(SessionABC):
+    """Redis implementation of :pyclass:`agents.memory.session.Session`."""
+
+    def __init__(
+        self,
+        session_id: str,
+        *,
+        redis_client: Redis,
+        key_prefix: str = "agents:session",
+        ttl: int | None = None,
+    ):
+        """Initializes a new RedisSession.
+
+        Args:
+            session_id (str): Unique identifier for the conversation.
+            redis_client (Redis[bytes]): A pre-configured Redis async client.
+            key_prefix (str, optional): Prefix for Redis keys to avoid collisions.
+                Defaults to "agents:session".
+            ttl (int | None, optional): Time-to-live in seconds for session data.
+                If None, data persists indefinitely. Defaults to None.
+        """
+        self.session_id = session_id
+        self._redis = redis_client
+        self._key_prefix = key_prefix
+        self._ttl = ttl
+        self._lock = asyncio.Lock()
+        self._owns_client = False  # Track if we own the Redis client
+
+        # Redis key patterns
+        self._session_key = f"{self._key_prefix}:{self.session_id}"
+        self._messages_key = f"{self._session_key}:messages"
+        self._counter_key = f"{self._session_key}:counter"
+
+    @classmethod
+    def from_url(
+        cls,
+        session_id: str,
+        *,
+        url: str,
+        redis_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> RedisSession:
+        """Create a session from a Redis URL string.
+
+        Args:
+            session_id (str): Conversation ID.
+            url (str): Redis URL, e.g. "redis://localhost:6379/0" or "rediss://host:6380".
+            redis_kwargs (dict[str, Any] | None): Additional keyword arguments forwarded to
+                redis.asyncio.from_url.
+            **kwargs: Additional keyword arguments forwarded to the main constructor
+                (e.g., key_prefix, ttl, etc.).
+
+        Returns:
+            RedisSession: An instance of RedisSession connected to the specified Redis server.
+        """
+        redis_kwargs = redis_kwargs or {}
+
+        redis_client = redis.from_url(url, **redis_kwargs)
+        session = cls(session_id, redis_client=redis_client, **kwargs)
+        session._owns_client = True  # We created the client, so we own it
+        return session
+
+    async def _serialize_item(self, item: TResponseInputItem) -> str:
+        """Serialize an item to JSON string. Can be overridden by subclasses."""
+        return json.dumps(item, separators=(",", ":"))
+
+    async def _deserialize_item(self, item: str) -> TResponseInputItem:
+        """Deserialize a JSON string to an item. Can be overridden by subclasses."""
+        return json.loads(item)  # type: ignore[no-any-return]  # json.loads returns Any but we know the structure
+
+    async def _get_next_id(self) -> int:
+        """Get the next message ID using Redis INCR for atomic increment."""
+        result = await self._redis.incr(self._counter_key)
+        return int(result)
+
+    async def _set_ttl_if_configured(self, *keys: str) -> None:
+        """Set TTL on keys if configured."""
+        if self._ttl is not None:
+            pipe = self._redis.pipeline()
+            for key in keys:
+                pipe.expire(key, self._ttl)
+            await pipe.execute()
+
+    # ------------------------------------------------------------------
+    # Session protocol implementation
+    # ------------------------------------------------------------------
+
+    async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
+        """Retrieve the conversation history for this session.
+
+        Args:
+            limit: Maximum number of items to retrieve. If None, retrieves all items.
+                   When specified, returns the latest N items in chronological order.
+
+        Returns:
+            List of input items representing the conversation history
+        """
+        async with self._lock:
+            if limit is None:
+                # Get all messages in chronological order
+                raw_messages = await self._redis.lrange(self._messages_key, 0, -1)  # type: ignore[misc]  # Redis library returns Union[Awaitable[T], T] in async context
+            else:
+                if limit <= 0:
+                    return []
+                # Get the latest N messages (Redis list is ordered chronologically)
+                # Use negative indices to get from the end - Redis uses -N to -1 for last N items
+                raw_messages = await self._redis.lrange(self._messages_key, -limit, -1)  # type: ignore[misc]  # Redis library returns Union[Awaitable[T], T] in async context
+
+            items: list[TResponseInputItem] = []
+            for raw_msg in raw_messages:
+                try:
+                    # Handle both bytes (default) and str (decode_responses=True) Redis clients
+                    if isinstance(raw_msg, bytes):
+                        msg_str = raw_msg.decode("utf-8")
+                    else:
+                        msg_str = raw_msg  # Already a string
+                    item = await self._deserialize_item(msg_str)
+                    items.append(item)
+                except (json.JSONDecodeError, UnicodeDecodeError):
+                    # Skip corrupted messages
+                    continue
+
+            return items
+
+    async def add_items(self, items: list[TResponseInputItem]) -> None:
+        """Add new items to the conversation history.
+
+        Args:
+            items: List of input items to add to the history
+        """
+        if not items:
+            return
+
+        async with self._lock:
+            pipe = self._redis.pipeline()
+
+            # Set session metadata with current timestamp
+            pipe.hset(
+                self._session_key,
+                mapping={
+                    "session_id": self.session_id,
+                    "created_at": str(int(time.time())),
+                    "updated_at": str(int(time.time())),
+                },
+            )
+
+            # Add all items to the messages list
+            serialized_items = []
+            for item in items:
+                serialized = await self._serialize_item(item)
+                serialized_items.append(serialized)
+
+            if serialized_items:
+                pipe.rpush(self._messages_key, *serialized_items)
+
+            # Update the session timestamp
+            pipe.hset(self._session_key, "updated_at", str(int(time.time())))
+
+            # Execute all commands
+            await pipe.execute()
+
+            # Set TTL if configured
+            await self._set_ttl_if_configured(
+                self._session_key, self._messages_key, self._counter_key
+            )
+
+    async def pop_item(self) -> TResponseInputItem | None:
+        """Remove and return the most recent item from the session.
+
+        Returns:
+            The most recent item if it exists, None if the session is empty
+        """
+        async with self._lock:
+            # Use RPOP to atomically remove and return the rightmost (most recent) item
+            raw_msg = await self._redis.rpop(self._messages_key)  # type: ignore[misc]  # Redis library returns Union[Awaitable[T], T] in async context
+
+            if raw_msg is None:
+                return None
+
+            try:
+                # Handle both bytes (default) and str (decode_responses=True) Redis clients
+                if isinstance(raw_msg, bytes):
+                    msg_str = raw_msg.decode("utf-8")
+                else:
+                    msg_str = raw_msg  # Already a string
+                return await self._deserialize_item(msg_str)
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                # Return None for corrupted messages (already removed)
+                return None
+
+    async def clear_session(self) -> None:
+        """Clear all items for this session."""
+        async with self._lock:
+            # Delete all keys associated with this session
+            await self._redis.delete(
+                self._session_key,
+                self._messages_key,
+                self._counter_key,
+            )
+
+    async def close(self) -> None:
+        """Close the Redis connection.
+
+        Only closes the connection if this session owns the Redis client
+        (i.e., created via from_url). If the client was injected externally,
+        the caller is responsible for managing its lifecycle.
+        """
+        if self._owns_client:
+            await self._redis.aclose()
+
+    async def ping(self) -> bool:
+        """Test Redis connectivity.
+
+        Returns:
+            True if Redis is reachable, False otherwise.
+        """
+        try:
+            await self._redis.ping()  # type: ignore[misc]  # Redis library returns Union[Awaitable[T], T] in async context
+            return True
+        except Exception:
+            return False