src/agent/memory.py

"""
Per-session memory store for the Claude Managed Agent.

Stores arbitrary key→value pairs scoped to one session.
Used by the CMA tools store_memory and get_memory.

Designed as a thin wrapper so the backing store can be swapped to Redis or
SQLite without changing the tool interface — just replace the _store dict.
"""

from __future__ import annotations


class SessionMemory:
    """In-memory key-value store scoped to one CMA session."""

    def __init__(self) -> None:
        self._store: dict = {}

    def set(self, key: str, value) -> None:
        """Store *value* under *key*, overwriting any previous value."""
        self._store[key] = value

    def get(self, key: str, default=None):
        """Retrieve the value for *key*, returning *default* if absent."""
        return self._store.get(key, default)

    def all(self) -> dict:
        """Return a shallow copy of all stored key-value pairs."""
        return dict(self._store)

    def delete(self, key: str) -> None:
        """Remove *key* from the store (no-op if absent)."""
        self._store.pop(key, None)

    def __repr__(self) -> str:
        return f"SessionMemory({list(self._store.keys())})"