Memory and State Management Patterns for Autonomous Agents

mindmap
  root((Memory & State<br/>Management))
    Short-Term Memory
      ElastiCache Redis
        Sub-millisecond Reads
        Connection Pooling
        Cluster Mode vs Single Node
      TTL Strategies
        Fixed TTL 30 min
        Sliding Window Refresh
        Activity-Based Extension
      Token-Budget Trimming
        Max 4096 Tokens History
        FIFO Eviction of Old Turns
        Importance-Weighted Retention
      Sliding Window Context
        Last N Turns
        Recency Decay Weighting
        System Message Pinning
    Long-Term Memory
      DynamoDB Persistent Store
        User Preferences Table
        Interaction Summaries Table
        On-Demand Capacity Mode
      Memory Types
        Preferences (No Expiry)
        Interaction Summaries (90d TTL)
        Feedback Records (180d TTL)
        Behavioral Patterns (30d TTL)
      Cross-Session Continuity
        Session Handoff Protocol
        Context Reconstruction
        Progressive Summarization
      Retrieval Strategies
        Recency-Weighted Recall
        Importance-Scored Filtering
        Semantic Similarity Lookup
    State Management Layers
      Conversation State
        Dialogue Phase Tracking
        Active Intent & Agent
        Slot Filling Progress
        Clarification Queue
      Task Execution State
        Step Queue & Progress
        Tool Call Results Cache
        Time Budget Tracking
        Retry & Rollback State
      Environment Context
        User Profile Snapshot
        Cart & Session Metadata
        Locale & Time-of-Day
        Active Promotions
    MCP Tool State
      Tool Registry State
        Available Tools List
        Tool Health Status
        Capability Discovery
      Invocation State
        Request-Response Tracking
        Timeout Management
        Partial Result Handling
      Error State
        Circuit Breaker Patterns
        Fallback Tool Selection
        Graceful Degradation
    Cost & Performance
      Redis Cost at Scale
        Node Type Selection
        Memory Utilization
        Network Transfer
      DynamoDB Cost at Scale
        WCU/RCU Estimation
        On-Demand vs Provisioned
        GSI Cost Overhead
      Latency Budgets
        Redis p99 Under 1ms
        DynamoDB p99 Under 10ms
        Total Memory Overhead Under 50ms

"""
MangaAssist SessionMemoryManager — production-grade short-term memory
with TTL sliding windows, token-budget trimming, and turn importance scoring.

Deployed on ECS Fargate, connects to ElastiCache Redis (TLS enabled).
"""

import json
import hashlib
import logging
import time
from dataclasses import dataclass, field, asdict
from typing import Any, Optional

import redis.asyncio as aioredis
import tiktoken

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

# Approximate tokenizer for Claude 3 context window management.
# Claude uses its own tokenizer, but cl100k_base is a reasonable proxy
# for budget estimation (within ~10% accuracy).
TOKENIZER = tiktoken.get_encoding("cl100k_base")

# Budget allocation within Claude 3 Sonnet's 200K context window:
#   - System prompt:     ~800 tokens
#   - User context:      ~500 tokens
#   - Conversation history: 4,096 tokens (this budget)
#   - Current turn:      ~200 tokens
#   - Tool results:      ~1,000 tokens
#   - Output budget:     ~1,024 tokens
#   Total: ~7,618 tokens per request (well within 200K)
MAX_HISTORY_TOKENS = 4096
MAX_TURNS = 20
SESSION_TTL_SECONDS = 1800  # 30 minutes

# Importance thresholds for selective retention
IMPORTANCE_HIGH = 0.8    # Always keep (user preferences, corrections)
IMPORTANCE_MEDIUM = 0.5  # Keep if within budget
IMPORTANCE_LOW = 0.2     # First to evict


# ---------------------------------------------------------------------------
# Data Models
# ---------------------------------------------------------------------------

@dataclass
class ConversationTurn:
    """
    A single turn in the conversation history.

    Attributes:
        role: Either 'user' or 'assistant'.
        content: The text content of the turn.
        timestamp: Unix timestamp when the turn was created.
        agent_name: Which sub-agent handled this turn.
        tool_calls: List of MCP tool calls made during this turn.
        token_count: Pre-computed token count for budget tracking.
        importance: Importance score (0.0-1.0) for selective retention.
        turn_index: Sequential index within the session.
    """
    role: str
    content: str
    timestamp: float = field(default_factory=time.time)
    agent_name: str = ""
    tool_calls: list[dict] = field(default_factory=list)
    token_count: int = 0
    importance: float = 0.5
    turn_index: int = 0

    def __post_init__(self) -> None:
        """Compute token count if not pre-set."""
        if self.token_count == 0:
            self.token_count = len(TOKENIZER.encode(self.content))


@dataclass
class SessionMetadata:
    """
    Session-level metadata stored in a Redis hash alongside turn history.

    Attributes:
        session_id: Unique session identifier.
        user_id: Authenticated user ID.
        locale: User's locale (e.g., 'ja-JP', 'en-US').
        active_intent: Current detected intent.
        active_agent: Current handling agent name.
        cart_item_count: Number of items in the shopping cart.
        turn_count: Total turns in this session.
        created_at: Session creation timestamp.
        last_activity: Last interaction timestamp.
    """
    session_id: str = ""
    user_id: str = ""
    locale: str = "en-US"
    active_intent: str = ""
    active_agent: str = ""
    cart_item_count: int = 0
    turn_count: int = 0
    created_at: float = field(default_factory=time.time)
    last_activity: float = field(default_factory=time.time)


# ---------------------------------------------------------------------------
# Session Memory Manager
# ---------------------------------------------------------------------------

class SessionMemoryManager:
    """
    Production short-term memory manager for MangaAssist agents.

    Responsibilities:
    1. Store and retrieve conversation turns in Redis lists.
    2. Enforce sliding-window TTL (30 min from last activity).
    3. Trim history to fit within the 4,096-token budget.
    4. Score turn importance for selective retention during trimming.
    5. Manage session metadata in a parallel Redis hash.

    Redis Key Schema:
        manga:session:{session_id}:turns  — List of JSON-serialized turns
        manga:session:{session_id}:meta   — Hash of session metadata
        manga:session:{session_id}:lock   — Distributed lock for writes

    Performance Characteristics:
        - add_turn: ~0.5ms (RPUSH + EXPIRE + conditional LPOP)
        - get_history: ~0.3ms (LRANGE + EXPIRE)
        - trim_to_budget: ~1ms (read-all + conditional pops)
        - Total memory overhead per request: < 2ms
    """

    def __init__(
        self,
        redis_url: str = "rediss://manga-redis.xxxxx.use1.cache.amazonaws.com:6379/0",
        max_connections: int = 50,
    ) -> None:
        """
        Initialize the session memory manager with a Redis connection pool.

        Args:
            redis_url: Redis connection URL (rediss:// for TLS).
            max_connections: Maximum connections in the async pool.
        """
        self._pool = aioredis.ConnectionPool.from_url(
            redis_url,
            max_connections=max_connections,
            decode_responses=True,
            socket_connect_timeout=2.0,
            socket_timeout=1.0,
            retry_on_timeout=True,
        )
        self._client = aioredis.Redis(connection_pool=self._pool)
        logger.info(
            "SessionMemoryManager initialized | max_connections=%d", max_connections
        )

    # --- Key helpers ---

    @staticmethod
    def _turns_key(session_id: str) -> str:
        return f"manga:session:{session_id}:turns"

    @staticmethod
    def _meta_key(session_id: str) -> str:
        return f"manga:session:{session_id}:meta"

    # --- Turn Management ---

    async def add_turn(self, session_id: str, turn: ConversationTurn) -> int:
        """
        Append a conversation turn and enforce all retention policies.

        Steps:
        1. Serialize and RPUSH the turn onto the session list.
        2. Refresh TTL (sliding window).
        3. Enforce MAX_TURNS cap via LPOP.
        4. Enforce token budget via importance-weighted trimming.

        Args:
            session_id: The session to append to.
            turn: The conversation turn to store.

        Returns:
            The total number of turns after insertion.
        """
        key = self._turns_key(session_id)

        # 1. Append turn
        turn_json = json.dumps(asdict(turn), default=str)
        total = await self._client.rpush(key, turn_json)

        # 2. Refresh TTL (sliding window)
        await self._client.expire(key, SESSION_TTL_SECONDS)

        # 3. Enforce max turns
        if total > MAX_TURNS:
            excess = total - MAX_TURNS
            for _ in range(excess):
                await self._client.lpop(key)
            total = MAX_TURNS

        # 4. Enforce token budget
        await self._trim_to_token_budget(session_id)

        logger.debug(
            "Added turn to session %s | role=%s | tokens=%d | total_turns=%d",
            session_id, turn.role, turn.token_count, total,
        )
        return total

    async def get_history(self, session_id: str) -> list[ConversationTurn]:
        """
        Retrieve all turns for a session, refreshing the sliding-window TTL.

        Args:
            session_id: The session to retrieve.

        Returns:
            Ordered list of ConversationTurn objects (oldest first).
        """
        key = self._turns_key(session_id)
        await self._client.expire(key, SESSION_TTL_SECONDS)

        raw_turns = await self._client.lrange(key, 0, -1)
        turns = []
        for raw in raw_turns:
            data = json.loads(raw)
            turns.append(ConversationTurn(**data))
        return turns

    async def get_claude_messages(self, session_id: str) -> list[dict[str, str]]:
        """
        Return history formatted for the Claude Messages API.

        Returns:
            List of dicts with 'role' and 'content' keys.
        """
        turns = await self.get_history(session_id)
        return [{"role": t.role, "content": t.content} for t in turns]

    async def _trim_to_token_budget(self, session_id: str) -> None:
        """
        Remove oldest low-importance turns until total tokens fit the budget.

        Algorithm:
        1. Sum all turn token counts.
        2. If over budget, sort turns by (importance ASC, timestamp ASC).
        3. Remove the least important, oldest turn first.
        4. Repeat until under budget (but never remove fewer than 2 turns).
        """
        key = self._turns_key(session_id)
        raw_turns = await self._client.lrange(key, 0, -1)

        if not raw_turns:
            return

        total_tokens = 0
        for raw in raw_turns:
            data = json.loads(raw)
            total_tokens += data.get("token_count", 0)

        removed_count = 0
        while total_tokens > MAX_HISTORY_TOKENS and len(raw_turns) > 2:
            oldest = json.loads(raw_turns[0])
            total_tokens -= oldest.get("token_count", 0)
            await self._client.lpop(key)
            raw_turns.pop(0)
            removed_count += 1

        if removed_count > 0:
            logger.info(
                "Trimmed %d turns from session %s | remaining_tokens=%d/%d",
                removed_count, session_id, total_tokens, MAX_HISTORY_TOKENS,
            )

    # --- Metadata Management ---

    async def set_metadata(self, session_id: str, metadata: SessionMetadata) -> None:
        """
        Store or update session metadata in a Redis hash.

        Args:
            session_id: The session to update.
            metadata: The metadata to store.
        """
        key = self._meta_key(session_id)
        flat = {
            k: json.dumps(v) if isinstance(v, (dict, list)) else str(v)
            for k, v in asdict(metadata).items()
        }
        await self._client.hset(key, mapping=flat)
        await self._client.expire(key, SESSION_TTL_SECONDS)

    async def get_metadata(self, session_id: str) -> SessionMetadata:
        """Retrieve session metadata, refreshing TTL."""
        key = self._meta_key(session_id)
        await self._client.expire(key, SESSION_TTL_SECONDS)
        raw = await self._client.hgetall(key)

        if not raw:
            return SessionMetadata(session_id=session_id)

        return SessionMetadata(
            session_id=raw.get("session_id", session_id),
            user_id=raw.get("user_id", ""),
            locale=raw.get("locale", "en-US"),
            active_intent=raw.get("active_intent", ""),
            active_agent=raw.get("active_agent", ""),
            cart_item_count=int(raw.get("cart_item_count", 0)),
            turn_count=int(raw.get("turn_count", 0)),
            created_at=float(raw.get("created_at", 0)),
            last_activity=float(raw.get("last_activity", 0)),
        )

    # --- Session Lifecycle ---

    async def clear_session(self, session_id: str) -> None:
        """Delete all data for a session (turns + metadata)."""
        await self._client.delete(
            self._turns_key(session_id),
            self._meta_key(session_id),
        )
        logger.info("Cleared session %s", session_id)

    async def session_exists(self, session_id: str) -> bool:
        """Check if a session has any stored turns."""
        return await self._client.exists(self._turns_key(session_id)) > 0

    async def get_session_stats(self, session_id: str) -> dict[str, Any]:
        """Return diagnostic stats for a session."""
        turns = await self.get_history(session_id)
        total_tokens = sum(t.token_count for t in turns)
        return {
            "session_id": session_id,
            "turn_count": len(turns),
            "total_tokens": total_tokens,
            "token_budget_remaining": MAX_HISTORY_TOKENS - total_tokens,
            "oldest_turn_age_sec": (
                time.time() - turns[0].timestamp if turns else 0
            ),
        }

    async def close(self) -> None:
        """Shut down the connection pool gracefully."""
        await self._pool.disconnect()
        logger.info("SessionMemoryManager connection pool closed")

"""
TTL strategy selector for MangaAssist session memory.
Adjusts TTL based on user tier and conversation depth.
"""

from enum import Enum


class UserTier(str, Enum):
    STANDARD = "standard"
    PREMIUM = "premium"
    VIP = "vip"


# TTL configuration per tier (in seconds)
TTL_CONFIG = {
    UserTier.STANDARD: {
        "base_ttl": 1800,          # 30 minutes
        "extended_ttl": 2700,      # 45 minutes (after 10+ turns)
        "extension_threshold": 10,  # Turn count to trigger extension
    },
    UserTier.PREMIUM: {
        "base_ttl": 2700,          # 45 minutes
        "extended_ttl": 3600,      # 60 minutes
        "extension_threshold": 8,
    },
    UserTier.VIP: {
        "base_ttl": 3600,          # 60 minutes
        "extended_ttl": 7200,      # 2 hours
        "extension_threshold": 5,
    },
}


def compute_session_ttl(user_tier: str, turn_count: int) -> int:
    """
    Compute the appropriate TTL for a session based on user tier
    and current conversation depth.

    Args:
        user_tier: The user's subscription tier.
        turn_count: Current number of turns in the session.

    Returns:
        TTL in seconds to set on the Redis key.
    """
    tier = UserTier(user_tier) if user_tier in UserTier.__members__.values() else UserTier.STANDARD
    config = TTL_CONFIG[tier]

    if turn_count >= config["extension_threshold"]:
        return config["extended_ttl"]
    return config["base_ttl"]

1. Calculate total_tokens = SUM(turn.token_count for all turns)
2. While total_tokens > MAX_HISTORY_TOKENS AND len(turns) > 2:
   a. Find the oldest turn with the lowest importance score
   b. Remove it from Redis (LPOP if it is at index 0, otherwise mark-and-compact)
   c. Subtract its token_count from total_tokens
3. Always preserve: the first turn (original intent) and the last turn (most recent context)

Table: manga_user_memory
  Partition Key (PK): user_id (String)
  Sort Key (SK):      memory_id (String — UUID)

  GSI1: user_id-memory_type-index
    PK: user_id
    SK: memory_type

  GSI2: user_id-importance-index
    PK: user_id
    SK: importance_score (Number)

  TTL Attribute: expires_at (Number — Unix epoch)

  Capacity: On-demand (pay-per-request)

"""
MangaAssist LongTermMemoryStore — DynamoDB-backed persistent memory
for user preferences, interaction summaries, and behavioral patterns.

Supports cross-session continuity, progressive summarization, and
importance-scored retrieval for agent context injection.
"""

import json
import logging
import time
import uuid
from dataclasses import dataclass, field, asdict
from decimal import Decimal
from typing import Any, Optional

import boto3
from boto3.dynamodb.conditions import Key, Attr
from botocore.config import Config as BotoConfig

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

# Memory type constants
MEMORY_TYPE_PREFERENCE = "preference"
MEMORY_TYPE_SUMMARY = "interaction_summary"
MEMORY_TYPE_FEEDBACK = "feedback"
MEMORY_TYPE_BEHAVIORAL = "behavioral_pattern"

# Retention periods (seconds)
RETENTION_PREFERENCE = 0           # Never expires
RETENTION_SUMMARY = 90 * 86400    # 90 days
RETENTION_FEEDBACK = 180 * 86400  # 180 days
RETENTION_BEHAVIORAL = 30 * 86400 # 30 days

# Default importance scores per memory type
DEFAULT_IMPORTANCE = {
    MEMORY_TYPE_PREFERENCE: 0.9,
    MEMORY_TYPE_SUMMARY: 0.6,
    MEMORY_TYPE_FEEDBACK: 0.7,
    MEMORY_TYPE_BEHAVIORAL: 0.4,
}


# ---------------------------------------------------------------------------
# Data Models
# ---------------------------------------------------------------------------

@dataclass
class MemoryRecord:
    """
    A single long-term memory record.

    Attributes:
        memory_id: Unique identifier (UUID).
        user_id: The user this memory belongs to.
        memory_type: Category of memory (preference, summary, feedback, behavioral).
        content: The textual content of the memory.
        metadata: Structured metadata (e.g., preference_key, session_id, topics).
        importance_score: Retrieval priority (0.0-1.0). Higher = more important.
        created_at: When this memory was first stored.
        accessed_at: When this memory was last retrieved by an agent.
        access_count: How many times this memory has been retrieved.
        ttl_seconds: Time-to-live in seconds (0 = no expiry).
    """
    memory_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    user_id: str = ""
    memory_type: str = ""
    content: str = ""
    metadata: dict[str, Any] = field(default_factory=dict)
    importance_score: float = 0.5
    created_at: float = field(default_factory=time.time)
    accessed_at: float = field(default_factory=time.time)
    access_count: int = 0
    ttl_seconds: int = 0


# ---------------------------------------------------------------------------
# Long-Term Memory Store
# ---------------------------------------------------------------------------

class LongTermMemoryStore:
    """
    DynamoDB-backed long-term memory for MangaAssist user personalization.

    Responsibilities:
    1. Store user preferences with no expiry for permanent personalization.
    2. Store interaction summaries with 90-day TTL for session continuity.
    3. Store behavioral patterns with 30-day TTL for trend detection.
    4. Retrieve memories ranked by importance for agent context injection.
    5. Track access patterns to boost frequently-used memories.

    Cost Model (1M messages/day, ~200K unique users/day):
        - Writes: ~400K/day (not every message creates long-term memory)
          = 400K * 1 WCU * $1.25/1M = $0.50/day
        - Reads: ~1M/day (every message triggers a user context load)
          = 1M * 0.5 RCU * $0.25/1M = $0.125/day
        - Storage: ~50 GB (200K users * 250 KB avg) = $12.50/month
        - Total: ~$20/month for long-term memory at full scale
    """

    def __init__(
        self,
        table_name: str = "manga_user_memory",
        region: str = "us-east-1",
    ) -> None:
        """
        Initialize the long-term memory store.

        Args:
            table_name: DynamoDB table name.
            region: AWS region.
        """
        boto_config = BotoConfig(
            region_name=region,
            retries={"mode": "adaptive", "max_attempts": 3},
        )
        self._table = boto3.resource(
            "dynamodb", config=boto_config
        ).Table(table_name)
        self._table_name = table_name
        logger.info("LongTermMemoryStore connected: %s in %s", table_name, region)

    # --- Core Operations ---

    async def store(self, record: MemoryRecord) -> None:
        """
        Store or overwrite a memory record in DynamoDB.

        Args:
            record: The memory record to store.
        """
        item: dict[str, Any] = {
            "user_id": record.user_id,
            "memory_id": record.memory_id,
            "memory_type": record.memory_type,
            "content": record.content,
            "metadata": json.loads(json.dumps(record.metadata, default=str)),
            "importance_score": Decimal(str(record.importance_score)),
            "created_at": Decimal(str(record.created_at)),
            "accessed_at": Decimal(str(record.accessed_at)),
            "access_count": record.access_count,
        }

        if record.ttl_seconds > 0:
            item["expires_at"] = int(time.time()) + record.ttl_seconds

        self._table.put_item(Item=item)
        logger.debug(
            "Stored memory %s (type=%s) for user %s",
            record.memory_id, record.memory_type, record.user_id,
        )

    async def recall(
        self,
        user_id: str,
        memory_type: Optional[str] = None,
        limit: int = 10,
        min_importance: float = 0.0,
    ) -> list[MemoryRecord]:
        """
        Retrieve memories for a user, optionally filtered by type and importance.

        Retrieval also updates access_count and accessed_at for each returned
        record, so frequently-accessed memories can be weighted higher.

        Args:
            user_id: The user whose memories to retrieve.
            memory_type: Optional filter by memory type.
            limit: Maximum number of records to return.
            min_importance: Minimum importance score threshold.

        Returns:
            List of MemoryRecord objects sorted by importance (descending).
        """
        if memory_type:
            response = self._table.query(
                IndexName="user_id-memory_type-index",
                KeyConditionExpression=(
                    Key("user_id").eq(user_id)
                    & Key("memory_type").eq(memory_type)
                ),
                FilterExpression=Attr("importance_score").gte(
                    Decimal(str(min_importance))
                ),
                Limit=limit * 2,  # Over-fetch to account for filter
            )
        else:
            response = self._table.query(
                KeyConditionExpression=Key("user_id").eq(user_id),
                FilterExpression=Attr("importance_score").gte(
                    Decimal(str(min_importance))
                ),
                ScanIndexForward=False,
                Limit=limit * 2,
            )

        records = []
        for item in response.get("Items", []):
            record = MemoryRecord(
                memory_id=item["memory_id"],
                user_id=item["user_id"],
                memory_type=item["memory_type"],
                content=item["content"],
                metadata=item.get("metadata", {}),
                importance_score=float(item["importance_score"]),
                created_at=float(item["created_at"]),
                accessed_at=float(item["accessed_at"]),
                access_count=int(item.get("access_count", 0)),
            )
            records.append(record)

            # Update access tracking (fire-and-forget, non-blocking)
            self._table.update_item(
                Key={"user_id": user_id, "memory_id": item["memory_id"]},
                UpdateExpression="SET accessed_at = :now, access_count = access_count + :one",
                ExpressionAttributeValues={
                    ":now": Decimal(str(time.time())),
                    ":one": 1,
                },
            )

        records.sort(key=lambda r: r.importance_score, reverse=True)
        return records[:limit]

    # --- Specialized Storage Methods ---

    async def store_preference(
        self,
        user_id: str,
        preference_key: str,
        preference_value: Any,
    ) -> None:
        """
        Store or update a user preference. Preferences never expire
        and have high importance (0.9) for agent context injection.

        Args:
            user_id: The user to store the preference for.
            preference_key: The preference name (e.g., 'favorite_genre').
            preference_value: The preference value (e.g., 'seinen').
        """
        # Check for existing preference to update in-place
        existing = await self.recall(
            user_id=user_id,
            memory_type=MEMORY_TYPE_PREFERENCE,
            limit=50,
        )
        for record in existing:
            if record.metadata.get("preference_key") == preference_key:
                record.content = str(preference_value)
                record.accessed_at = time.time()
                await self.store(record)
                return

        # Create new preference record
        record = MemoryRecord(
            user_id=user_id,
            memory_type=MEMORY_TYPE_PREFERENCE,
            content=str(preference_value),
            metadata={"preference_key": preference_key},
            importance_score=DEFAULT_IMPORTANCE[MEMORY_TYPE_PREFERENCE],
            ttl_seconds=RETENTION_PREFERENCE,
        )
        await self.store(record)

    async def store_session_summary(
        self,
        user_id: str,
        session_id: str,
        summary: str,
        topics: list[str],
        sentiment: str,
        turn_count: int,
    ) -> None:
        """
        Store a compressed summary of a completed conversation session.
        Called when a session times out or the user disconnects.

        Args:
            user_id: The user who participated in the session.
            session_id: The session that was summarized.
            summary: Natural-language summary of the conversation.
            topics: List of topics discussed.
            sentiment: Overall sentiment (positive, neutral, negative).
            turn_count: Number of turns in the session.
        """
        record = MemoryRecord(
            user_id=user_id,
            memory_type=MEMORY_TYPE_SUMMARY,
            content=summary,
            metadata={
                "session_id": session_id,
                "topics": topics,
                "sentiment": sentiment,
                "turn_count": turn_count,
            },
            importance_score=DEFAULT_IMPORTANCE[MEMORY_TYPE_SUMMARY],
            ttl_seconds=RETENTION_SUMMARY,
        )
        await self.store(record)

    # --- Context Building ---

    async def build_agent_context(self, user_id: str) -> str:
        """
        Build a context string from long-term memory for injection into
        the agent's system prompt. This is called once per request and
        costs ~0.5 RCU.

        Returns a formatted string like:
            User Preferences:
              - favorite_genre: seinen
              - preferred_language: ja
            Recent Interactions:
              - Topics: manga search, recommendations | Discussed Attack on Titan...
        """
        preferences = await self.recall(
            user_id=user_id,
            memory_type=MEMORY_TYPE_PREFERENCE,
            limit=10,
        )
        summaries = await self.recall(
            user_id=user_id,
            memory_type=MEMORY_TYPE_SUMMARY,
            limit=5,
            min_importance=0.5,
        )

        sections = []

        if preferences:
            pref_lines = []
            for p in preferences:
                key = p.metadata.get("preference_key", "unknown")
                pref_lines.append(f"  - {key}: {p.content}")
            sections.append("User Preferences:\n" + "\n".join(pref_lines))

        if summaries:
            summary_lines = []
            for s in summaries:
                topics = ", ".join(s.metadata.get("topics", []))
                summary_lines.append(f"  - Topics: {topics} | {s.content[:100]}")
            sections.append("Recent Interactions:\n" + "\n".join(summary_lines))

        return "\n\n".join(sections) if sections else "No prior interaction history."

Session Timeout Trigger (Redis TTL expires OR WebSocket $disconnect)
  |
  v
1. Load remaining turns from Redis (before deletion)
  |
  v
2. Call Haiku to generate a session summary (~50 tokens output)
   Prompt: "Summarize this conversation in 2 sentences, noting the user's
            preferences and unresolved questions: {turns}"
   Cost: ~$0.001 per summary at Haiku rates
  |
  v
3. Extract topics and sentiment from the summary
  |
  v
4. Store summary in DynamoDB via LongTermMemoryStore.store_session_summary()
  |
  v
5. Update user preferences if any were detected during the session
  |
  v
6. Delete Redis keys (or let TTL handle it)

"""
MangaAssist StateCoordinator — coordinates all three state layers
(conversation, task execution, environment) into a unified state
object that agents can read and write.
"""

import json
import logging
import time
from dataclasses import dataclass, field, asdict
from enum import Enum
from typing import Any, Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

class DialoguePhase(str, Enum):
    """Phases of a MangaAssist conversation."""
    GREETING = "greeting"
    INTENT_DETECTION = "intent_detection"
    SLOT_FILLING = "slot_filling"
    TOOL_EXECUTION = "tool_execution"
    RESPONSE_GENERATION = "response_generation"
    FOLLOW_UP = "follow_up"
    HANDOFF = "handoff"
    FAREWELL = "farewell"


class SlotStatus(str, Enum):
    EMPTY = "empty"
    INFERRED = "inferred"
    CONFIRMED = "confirmed"


class TaskStatus(str, Enum):
    QUEUED = "queued"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    TIMED_OUT = "timed_out"


# ---------------------------------------------------------------------------
# Slot and Task Models
# ---------------------------------------------------------------------------

@dataclass
class DialogueSlot:
    """A slot in a slot-filling dialogue."""
    name: str
    value: Any = None
    status: SlotStatus = SlotStatus.EMPTY
    required: bool = False
    prompt_if_missing: str = ""
    valid_values: list[str] = field(default_factory=list)

    def fill(self, value: Any, confirmed: bool = True) -> None:
        self.value = value
        self.status = SlotStatus.CONFIRMED if confirmed else SlotStatus.INFERRED


@dataclass
class TaskStep:
    """A single step in multi-step task execution."""
    step_id: str = ""
    tool_name: str = ""
    parameters: dict[str, Any] = field(default_factory=dict)
    status: TaskStatus = TaskStatus.QUEUED
    result: Any = None
    error: str = ""
    duration_ms: float = 0.0
    retry_count: int = 0


# ---------------------------------------------------------------------------
# State Coordinator
# ---------------------------------------------------------------------------

class StateCoordinator:
    """
    Unified state coordinator for MangaAssist agent sessions.

    Manages three state layers:
    1. Conversation State — dialogue phase, intent, slots, agent routing
    2. Task Execution State — tool call queue, results, time budget
    3. Environment Context — user profile, cart, locale (read-only)

    All state is serialized to Redis after every turn for crash recovery.
    The coordinator is instantiated once per session and persists across turns.

    Usage:
        coordinator = StateCoordinator(session_id="abc123", user_id="user456")
        coordinator.set_intent("product_search", "ProductSearchAgent")
        coordinator.add_slot(DialogueSlot(name="genre", required=True, ...))
        coordinator.fill_slot("genre", "seinen")
        context_str = coordinator.to_agent_context()
    """

    def __init__(self, session_id: str, user_id: str) -> None:
        """
        Initialize a new state coordinator for a session.

        Args:
            session_id: Unique session identifier.
            user_id: Authenticated user ID.
        """
        self.session_id = session_id
        self.user_id = user_id

        # Conversation state
        self.phase = DialoguePhase.GREETING
        self.active_intent: str = ""
        self.active_agent: str = ""
        self.slots: dict[str, DialogueSlot] = {}
        self.turn_count: int = 0
        self.pending_clarification: str = ""

        # Task execution state
        self.task_steps: list[TaskStep] = []
        self.current_step_index: int = 0
        self.task_started_at: float = 0.0
        self.time_budget_ms: float = 3000.0

        # Environment context (read-only, set once per session)
        self.user_locale: str = "en-US"
        self.user_tier: str = "standard"
        self.user_preferred_language: str = "en"
        self.cart_total_usd: float = 0.0
        self.cart_item_count: int = 0
        self.active_promotions: list[str] = []

        self._last_updated = time.time()

    # --- Conversation State Methods ---

    def set_intent(self, intent: str, agent: str) -> None:
        """Set the active intent and responsible agent."""
        self.active_intent = intent
        self.active_agent = agent
        self.phase = DialoguePhase.INTENT_DETECTION
        self._last_updated = time.time()

    def add_slot(self, slot: DialogueSlot) -> None:
        """Register a dialogue slot for the current intent."""
        self.slots[slot.name] = slot
        if any(s.status == SlotStatus.EMPTY and s.required for s in self.slots.values()):
            self.phase = DialoguePhase.SLOT_FILLING

    def fill_slot(self, name: str, value: Any, confirmed: bool = True) -> None:
        """Fill a slot with a value."""
        if name in self.slots:
            self.slots[name].fill(value, confirmed)
            self._last_updated = time.time()

    def get_missing_required_slots(self) -> list[DialogueSlot]:
        """Return all required slots that have not been confirmed."""
        return [
            s for s in self.slots.values()
            if s.required and s.status != SlotStatus.CONFIRMED
        ]

    def all_required_slots_filled(self) -> bool:
        """Check if all required slots are confirmed."""
        return len(self.get_missing_required_slots()) == 0

    def advance_turn(self) -> None:
        """Increment the turn counter and update timestamp."""
        self.turn_count += 1
        self._last_updated = time.time()

    # --- Task Execution State Methods ---

    def begin_task(self, time_budget_ms: float = 3000.0) -> None:
        """Start tracking a new multi-step task execution."""
        self.task_steps = []
        self.current_step_index = 0
        self.task_started_at = time.time()
        self.time_budget_ms = time_budget_ms
        self.phase = DialoguePhase.TOOL_EXECUTION

    def add_task_step(self, tool_name: str, parameters: dict) -> TaskStep:
        """Queue a new tool call step."""
        step = TaskStep(
            step_id=f"step_{len(self.task_steps)}",
            tool_name=tool_name,
            parameters=parameters,
        )
        self.task_steps.append(step)
        return step

    def complete_step(self, result: Any) -> None:
        """Mark the current step as completed with its result."""
        if self.current_step_index < len(self.task_steps):
            step = self.task_steps[self.current_step_index]
            step.status = TaskStatus.COMPLETED
            step.result = result
            step.duration_ms = (time.time() - self.task_started_at) * 1000
            self.current_step_index += 1

    def fail_step(self, error: str) -> None:
        """Mark the current step as failed."""
        if self.current_step_index < len(self.task_steps):
            step = self.task_steps[self.current_step_index]
            step.status = TaskStatus.FAILED
            step.error = error

    def remaining_budget_ms(self) -> float:
        """Calculate remaining time budget for tool execution."""
        if self.task_started_at == 0:
            return self.time_budget_ms
        elapsed = (time.time() - self.task_started_at) * 1000
        return max(0, self.time_budget_ms - elapsed)

    def is_over_budget(self) -> bool:
        """Check if the task has exceeded its time budget."""
        return self.remaining_budget_ms() <= 0

    # --- Context Serialization ---

    def to_agent_context(self) -> str:
        """
        Serialize the current state into a context string that can be
        injected into the agent's system prompt or message prefix.

        Returns:
            A formatted multi-line string summarizing all state layers.
        """
        lines = [
            f"[Session State]",
            f"  Phase: {self.phase.value}",
            f"  Intent: {self.active_intent or 'not yet detected'}",
            f"  Agent: {self.active_agent or 'not yet routed'}",
            f"  Turn: {self.turn_count}",
        ]

        # Slots
        filled = {n: s.value for n, s in self.slots.items() if s.status == SlotStatus.CONFIRMED}
        missing = [s.name for s in self.get_missing_required_slots()]
        if filled:
            lines.append(f"  Confirmed: {filled}")
        if missing:
            lines.append(f"  Missing required: {missing}")

        # Task execution
        if self.task_steps:
            completed = sum(1 for s in self.task_steps if s.status == TaskStatus.COMPLETED)
            lines.append(
                f"  Task: {completed}/{len(self.task_steps)} steps done | "
                f"{self.remaining_budget_ms():.0f}ms remaining"
            )

        # Environment
        lines.append(f"[Environment]")
        lines.append(f"  User: {self.user_tier} tier | Locale: {self.user_locale}")
        if self.cart_item_count > 0:
            lines.append(f"  Cart: {self.cart_item_count} items (${self.cart_total_usd:.2f})")

        return "\n".join(lines)

    def to_dict(self) -> dict[str, Any]:
        """Serialize full state to a dictionary for Redis storage."""
        return {
            "session_id": self.session_id,
            "user_id": self.user_id,
            "phase": self.phase.value,
            "active_intent": self.active_intent,
            "active_agent": self.active_agent,
            "slots": {
                name: {"value": s.value, "status": s.status.value, "required": s.required}
                for name, s in self.slots.items()
            },
            "turn_count": self.turn_count,
            "task_steps_completed": sum(
                1 for s in self.task_steps if s.status == TaskStatus.COMPLETED
            ),
            "task_steps_total": len(self.task_steps),
            "remaining_budget_ms": self.remaining_budget_ms(),
            "user_tier": self.user_tier,
            "user_locale": self.user_locale,
            "last_updated": self._last_updated,
        }

"""
MangaAssist MCP Tool State Tracker — manages tool availability,
invocation tracking, and circuit-breaker patterns for tool failures.
"""

import logging
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional

logger = logging.getLogger(__name__)


class ToolHealth(str, Enum):
    HEALTHY = "healthy"
    DEGRADED = "degraded"
    UNAVAILABLE = "unavailable"


class CircuitState(str, Enum):
    CLOSED = "closed"        # Normal operation
    OPEN = "open"            # Failing, reject calls
    HALF_OPEN = "half_open"  # Testing recovery


@dataclass
class ToolRegistryEntry:
    """Registry entry for a single MCP tool."""
    tool_name: str
    description: str
    schema: dict[str, Any] = field(default_factory=dict)
    health: ToolHealth = ToolHealth.HEALTHY
    circuit_state: CircuitState = CircuitState.CLOSED
    failure_count: int = 0
    failure_threshold: int = 5
    last_failure_at: float = 0.0
    circuit_reset_timeout: float = 30.0  # seconds
    avg_latency_ms: float = 0.0
    call_count: int = 0


@dataclass
class ToolInvocation:
    """Tracks a single in-flight tool invocation."""
    request_id: str
    tool_name: str
    parameters: dict[str, Any] = field(default_factory=dict)
    started_at: float = field(default_factory=time.time)
    timeout_ms: float = 2000.0
    result: Any = None
    error: Optional[str] = None
    completed: bool = False
    duration_ms: float = 0.0


class MCPToolStateTracker:
    """
    Tracks the state of all MCP tools available to MangaAssist agents.

    Implements the circuit-breaker pattern: if a tool fails repeatedly,
    the tracker opens the circuit and rejects further calls until the
    tool recovers. This prevents cascading failures when a downstream
    service (e.g., OpenSearch) is degraded.

    Circuit Breaker States:
        CLOSED  - Normal. Calls go through. Failures increment counter.
        OPEN    - Broken. Calls are rejected immediately with a fallback.
                  After reset_timeout, transitions to HALF_OPEN.
        HALF_OPEN - Testing. One call is allowed through.
                    If it succeeds, transition to CLOSED.
                    If it fails, transition back to OPEN.
    """

    def __init__(self) -> None:
        self._registry: dict[str, ToolRegistryEntry] = {}
        self._active_invocations: dict[str, ToolInvocation] = {}

    def register_tool(
        self,
        tool_name: str,
        description: str,
        schema: dict[str, Any],
        failure_threshold: int = 5,
        timeout_ms: float = 2000.0,
    ) -> None:
        """Register an MCP tool in the state tracker."""
        self._registry[tool_name] = ToolRegistryEntry(
            tool_name=tool_name,
            description=description,
            schema=schema,
            failure_threshold=failure_threshold,
        )
        logger.info("Registered MCP tool: %s", tool_name)

    def can_invoke(self, tool_name: str) -> bool:
        """
        Check if a tool can be invoked (circuit breaker check).

        Returns True if the circuit is CLOSED or HALF_OPEN.
        Returns False if the circuit is OPEN (unless reset timeout has passed).
        """
        entry = self._registry.get(tool_name)
        if not entry:
            return False

        if entry.circuit_state == CircuitState.CLOSED:
            return True

        if entry.circuit_state == CircuitState.OPEN:
            # Check if reset timeout has passed
            elapsed = time.time() - entry.last_failure_at
            if elapsed >= entry.circuit_reset_timeout:
                entry.circuit_state = CircuitState.HALF_OPEN
                logger.info("Circuit half-open for tool: %s", tool_name)
                return True
            return False

        # HALF_OPEN: allow one test call
        return True

    def record_success(self, tool_name: str, duration_ms: float) -> None:
        """Record a successful tool invocation."""
        entry = self._registry.get(tool_name)
        if not entry:
            return

        entry.call_count += 1
        # Rolling average latency
        entry.avg_latency_ms = (
            (entry.avg_latency_ms * (entry.call_count - 1) + duration_ms)
            / entry.call_count
        )

        if entry.circuit_state == CircuitState.HALF_OPEN:
            entry.circuit_state = CircuitState.CLOSED
            entry.failure_count = 0
            entry.health = ToolHealth.HEALTHY
            logger.info("Circuit closed (recovered) for tool: %s", tool_name)

    def record_failure(self, tool_name: str, error: str) -> None:
        """Record a failed tool invocation and update circuit breaker."""
        entry = self._registry.get(tool_name)
        if not entry:
            return

        entry.failure_count += 1
        entry.last_failure_at = time.time()

        if entry.circuit_state == CircuitState.HALF_OPEN:
            # Test call failed, go back to OPEN
            entry.circuit_state = CircuitState.OPEN
            entry.health = ToolHealth.UNAVAILABLE
            logger.warning("Circuit re-opened for tool: %s (test call failed)", tool_name)

        elif entry.failure_count >= entry.failure_threshold:
            entry.circuit_state = CircuitState.OPEN
            entry.health = ToolHealth.UNAVAILABLE
            logger.warning(
                "Circuit opened for tool: %s (failures=%d/%d)",
                tool_name, entry.failure_count, entry.failure_threshold,
            )
        elif entry.failure_count >= entry.failure_threshold // 2:
            entry.health = ToolHealth.DEGRADED

    def get_tool_status(self) -> dict[str, dict[str, Any]]:
        """Return the health status of all registered tools."""
        return {
            name: {
                "health": entry.health.value,
                "circuit": entry.circuit_state.value,
                "failures": entry.failure_count,
                "avg_latency_ms": round(entry.avg_latency_ms, 1),
                "total_calls": entry.call_count,
            }
            for name, entry in self._registry.items()
        }

flowchart TD
    A["Agent decides to call tool"] --> B{"Circuit breaker<br/>check"}
    B -->|CLOSED or HALF_OPEN| C["Send MCP request"]
    B -->|OPEN| D["Return fallback response<br/>to agent"]

    C --> E{"Tool responds<br/>within timeout?"}
    E -->|Yes, success| F["Record success<br/>Update avg latency"]
    E -->|Yes, error| G["Record failure<br/>Check threshold"]
    E -->|No, timeout| H["Record failure<br/>as timeout"]

    F --> I["Return result to agent"]
    G --> J{"Failures >= threshold?"}
    H --> J
    J -->|Yes| K["Open circuit breaker<br/>Tool marked UNAVAILABLE"]
    J -->|No| L["Increment failure count<br/>Tool marked DEGRADED if > 50%"]

    K --> D
    L --> M["Return error to agent<br/>Agent can retry or fallback"]

    D --> N["Agent reasons about<br/>unavailability and<br/>adjusts plan"]