feat: Auto-Dreamer v2 paper-based memory consolidation (NeurIPS 2026)

mirix/agent/__init__.py

@@ -33,9 +33,13 @@
     "ReflexionAgent",
     "ResourceMemoryAgent",
     "SemanticMemoryAgent",
+    "AutoDreamAgent",
+    "AutoDreamV2Agent",
 ]
 
 from mirix.agent.agent import Agent, AgentState, save_agent
+from mirix.agent.auto_dream_agent import AutoDreamAgent
+from mirix.agent.auto_dream_v2_agent import AutoDreamV2Agent
 from mirix.agent.background_agent import BackgroundAgent
 from mirix.agent.core_memory_agent import CoreMemoryAgent
 from mirix.agent.episodic_memory_agent import EpisodicMemoryAgent

mirix/agent/agent.py

@@ -970,11 +970,24 @@ async def _handle_ai_response(
                     assert tool_call.id is not None  # should be defined
 
             if (
+                self.agent_state.agent_type is not None
+                and response_message.tool_calls is not None
+                and len(response_message.tool_calls) > 1
+            ):
+                # Memory agents must execute at most one tool call per turn to avoid
+                # duplicate state mutations (e.g. double-inserting the same memory entry).
+                self.logger.warning(
+                    "Memory agent %s returned %d tool calls; truncating to first to prevent duplicate mutations",
+                    self.agent_state.agent_type,
+                    len(response_message.tool_calls),
+                )
+                response_message.tool_calls = response_message.tool_calls[:1]
+            elif (
                 response_message.tool_calls is not None
                 and len(response_message.tool_calls) > 1
             ):
                 self.logger.info(
-                    "Memory agent %s returned %d tool call(s); executing each sequentially",
+                    "Agent %s returned %d tool call(s); executing each sequentially",
                     self.agent_state.agent_type,
                     len(response_message.tool_calls),
                 )

mirix/agent/agent_configs.py

@@ -62,4 +62,16 @@
         "attr_name": "agent_state",
         "include_base_tools": True,
     },
+    {
+        "name": "auto_dream_agent",
+        "agent_type": AgentType.auto_dream_agent,
+        "attr_name": "auto_dream_agent_state",
+        "include_base_tools": False,
+    },
+    {
+        "name": "auto_dream_v2_agent",
+        "agent_type": AgentType.auto_dream_v2_agent,
+        "attr_name": "auto_dream_v2_agent_state",
+        "include_base_tools": False,
+    },
 ]

mirix/agent/auto_dream_agent.py

@@ -0,0 +1,6 @@
+from mirix.agent import Agent
+
+
+class AutoDreamAgent(Agent):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)

mirix/agent/auto_dream_v2_agent.py

@@ -0,0 +1 @@
+

mirix/agent/meta_agent.py

@@ -41,6 +41,8 @@ def __init__(self):
         self.resource_memory_agent_state: Optional[AgentState] = None
         self.reflexion_agent_state: Optional[AgentState] = None
         self.background_agent_state: Optional[AgentState] = None
+        self.auto_dream_agent_state: Optional[AgentState] = None
+        self.auto_dream_v2_agent_state: Optional[AgentState] = None
 
     def set_agent_state(self, name: str, state: AgentState):
         """Set an agent state by name."""
@@ -68,6 +70,8 @@ def get_all_states(self) -> Dict[str, Optional[AgentState]]:
             "resource_memory_agent_state": self.resource_memory_agent_state,
             "reflexion_agent_state": self.reflexion_agent_state,
             "background_agent_state": self.background_agent_state,
+            "auto_dream_agent_state": self.auto_dream_agent_state,
+            "auto_dream_v2_agent_state": self.auto_dream_v2_agent_state,
         }
 
     def get_all_agent_states_list(self) -> List[Optional[AgentState]]:
@@ -82,6 +86,8 @@ def get_all_agent_states_list(self) -> List[Optional[AgentState]]:
             self.resource_memory_agent_state,
             self.reflexion_agent_state,
             self.background_agent_state,
+            self.auto_dream_agent_state,
+            self.auto_dream_v2_agent_state,
         ]
 
 
@@ -141,6 +147,18 @@ def get_all_agent_states_list(self) -> List[Optional[AgentState]]:
         "attr_name": "background_agent_state",
         "include_base_tools": False,
     },
+    {
+        "name": "auto_dream_agent",
+        "agent_type": AgentType.auto_dream_agent,
+        "attr_name": "auto_dream_agent_state",
+        "include_base_tools": False,
+    },
+    {
+        "name": "auto_dream_v2_agent",
+        "agent_type": AgentType.auto_dream_v2_agent,
+        "attr_name": "auto_dream_v2_agent_state",
+        "include_base_tools": False,
+    },
 ]
 
 
@@ -293,6 +311,10 @@ def _load_existing_agents(self, existing_agents: List[AgentState]):
                 self.memory_agent_states.reflexion_agent_state = agent_state
             elif agent_state.name == "background_agent":
                 self.memory_agent_states.background_agent_state = agent_state
+            elif agent_state.name == "auto_dream_agent":
+                self.memory_agent_states.auto_dream_agent_state = agent_state
+            elif agent_state.name == "auto_dream_v2_agent":
+                self.memory_agent_states.auto_dream_v2_agent_state = agent_state
 
         printv(f"[Mirix.Agent.{self.agent_state.name}] INFO: Loaded existing memory agent states")
 

mirix/client/remote_client.py

@@ -1849,6 +1849,73 @@ async def list_memory_components(
 
         return await self._request("GET", "/memory/components", params=params, headers=headers)
 
+    async def auto_dream(
+        self,
+        user_id: str,
+        memory_types: Optional[List[str]] = None,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Run the auto-dream self-reflection pipeline for a user.
+
+        Consolidates memories by removing redundancies and resolving conflicts
+        across the specified memory types.
+
+        Args:
+            user_id: End-user whose memories to consolidate
+            memory_types: Which memory types to process (default: all except core)
+
+        Returns:
+            AutoDreamResponse dict with stats on removed/merged/resolved items
+        """
+        await self._ensure_user_exists(user_id, headers=headers)
+
+        body: Dict[str, Any] = {}
+        if memory_types is not None:
+            body["memory_types"] = memory_types
+
+        return await self._request(
+            "POST",
+            "/memory/auto_dream",
+            params={"user_id": user_id},
+            json=body,
+            headers=headers,
+        )
+
+    async def auto_dream_v2(
+        self,
+        user_id: str,
+        memory_types: Optional[List[str]] = None,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Run the Auto-Dreamer v2 offline memory consolidation pipeline for a user.
+
+        Implements region rewriting inspired by the Auto-Dreamer paper: the agent
+        surveys the working region, synthesizes generalized replacements, and retires
+        the source entries.
+
+        Args:
+            user_id: End-user whose memories to consolidate
+            memory_types: Which memory types to include (default: episodic, semantic, procedural)
+
+        Returns:
+            AutoDreamV2Response dict with working-region stats
+        """
+        await self._ensure_user_exists(user_id, headers=headers)
+
+        body: Dict[str, Any] = {}
+        if memory_types is not None:
+            body["memory_types"] = memory_types
+
+        return await self._request(
+            "POST",
+            "/memory/auto_dream_v2",
+            params={"user_id": user_id},
+            json=body,
+            headers=headers,
+        )
+
     # ========================================================================
     # LangChain/Composio/CrewAI Integration (Not Supported)
     # ========================================================================

mirix/constants.py

@@ -132,6 +132,8 @@
     "finish_memory_update",
     "list_memory_within_timerange",
 ]
+AUTO_DREAM_V2_READ_TOOLS = ["search_in_memory", "list_memory_within_timerange", "get_source_messages", "check_memory"]
+AUTO_DREAM_V2_WRITE_TOOLS = ["synthesize_memory", "retire_memory_region", "finish_memory_update"]
 ALL_TOOLS = list(
     set(
         BASE_TOOLS
@@ -146,6 +148,8 @@
         + CHAT_AGENT_TOOLS
         + EXTRAS_TOOLS
         + MCP_TOOLS
+        + AUTO_DREAM_V2_READ_TOOLS
+        + AUTO_DREAM_V2_WRITE_TOOLS
     )
 )
 

mirix/functions/function_sets/base.py

@@ -1,8 +1,25 @@
-from typing import Optional
+import asyncio
+import logging
+from typing import List, Optional
 
 from mirix.agent import Agent, AgentState
 from mirix.utils import convert_timezone_to_utc
 
+logger = logging.getLogger(__name__)
+
+
+def _fire_access_log(results: List[dict], memory_type: str, user_id: str, org_id: Optional[str]) -> None:
+    """Fire-and-forget access log write so retrieval tracking doesn't block search."""
+    entries = [(r["id"], memory_type) for r in results if "id" in r]
+    if not entries:
+        return
+    from mirix.services.memory_access_log_manager import MemoryAccessLogManager
+    mgr = MemoryAccessLogManager()
+    try:
+        asyncio.create_task(mgr.log_access(entries, user_id, org_id))
+    except RuntimeError:
+        pass
+
 
 async def send_message(self: "Agent", agent_state: "AgentState", message: str) -> Optional[str]:
     """
@@ -159,6 +176,7 @@ async def search_in_memory(
             for x in episodic_memory
         ]
         if memory_type == "episodic":
+            _fire_access_log(formatted_results_from_episodic, "episodic", self.user.id, getattr(self.actor, "organization_id", None))
             return formatted_results_from_episodic, len(formatted_results_from_episodic)
 
     if memory_type == "resource" or memory_type == "all":
@@ -185,6 +203,7 @@ async def search_in_memory(
             for x in resource_memories
         ]
         if memory_type == "resource":
+            _fire_access_log(formatted_results_resource, "resource", self.user.id, getattr(self.actor, "organization_id", None))
             return formatted_results_resource, len(formatted_results_resource)
 
     if memory_type == "procedural" or memory_type == "all":
@@ -209,6 +228,7 @@ async def search_in_memory(
             for x in procedural_memories
         ]
         if memory_type == "procedural":
+            _fire_access_log(formatted_results_procedural, "procedural", self.user.id, getattr(self.actor, "organization_id", None))
             return formatted_results_procedural, len(formatted_results_procedural)
 
     if memory_type == "knowledge_vault" or memory_type == "all":
@@ -235,6 +255,7 @@ async def search_in_memory(
             for x in knowledge_vault_memories
         ]
         if memory_type == "knowledge_vault":
+            _fire_access_log(formatted_results_knowledge_vault, "knowledge_vault", self.user.id, getattr(self.actor, "organization_id", None))
             return formatted_results_knowledge_vault, len(formatted_results_knowledge_vault)
 
     if memory_type == "semantic" or memory_type == "all":
@@ -261,12 +282,19 @@ async def search_in_memory(
             for x in semantic_memories
         ]
         if memory_type == "semantic":
+            _fire_access_log(formatted_results_semantic, "semantic", self.user.id, getattr(self.actor, "organization_id", None))
             return formatted_results_semantic, len(formatted_results_semantic)
 
     else:
         raise ValueError(
             f"Memory type '{memory_type}' is not supported. Please choose from 'episodic', 'resource', 'procedural', 'knowledge_vault', 'semantic'."
         )
+    org_id = getattr(self.actor, "organization_id", None)
+    _fire_access_log(formatted_results_from_episodic, "episodic", self.user.id, org_id)
+    _fire_access_log(formatted_results_resource, "resource", self.user.id, org_id)
+    _fire_access_log(formatted_results_procedural, "procedural", self.user.id, org_id)
+    _fire_access_log(formatted_results_knowledge_vault, "knowledge_vault", self.user.id, org_id)
+    _fire_access_log(formatted_results_semantic, "semantic", self.user.id, org_id)
     return (
         formatted_results_from_episodic
         + formatted_results_resource

mirix/functions/function_sets/memory_tools.py

@@ -1205,3 +1205,62 @@ async def finish_memory_update(self: "Agent") -> str:
         str: Empty string (no response content).
     """
     return ""
+
+
+async def check_memory(self: "Agent", entries: List[dict]) -> str:
+    """
+    Inspect specific memory entries by their IDs. Use this to read the full content of entries
+    you already have IDs for, without performing a new search.
+
+    Mirrors the paper's ``check_memory(ids=[...])`` navigation tool for the working region.
+
+    Args:
+        entries: List of dicts, each with ``"id"`` (str) and ``"type"``
+                 (one of "episodic", "semantic", "procedural", "resource", "knowledge_vault").
+                 Up to 30 entries per call.
+
+    Returns:
+        str: JSON-formatted list of entry details, or error messages for entries not found.
+    """
+    import json
+
+    if self.user is None:
+        return "[error] user not set"
+
+    entries = entries[:30]
+    tz = "UTC"
+    results = []
+
+    for spec in entries:
+        eid = spec.get("id", "")
+        etype = spec.get("type", "")
+        if not eid or not etype:
+            results.append({"id": eid, "type": etype, "error": "malformed spec"})
+            continue
+        try:
+            if etype == "episodic":
+                entry = await self.episodic_memory_manager.get_episodic_memory_by_id(eid, user=self.user)
+                data = entry.model_dump(exclude_none=True)
+            elif etype == "semantic":
+                entry = await self.semantic_memory_manager.get_semantic_item_by_id(eid, user=self.user, timezone_str=tz)
+                data = entry.model_dump(exclude_none=True)
+            elif etype == "procedural":
+                entry = await self.procedural_memory_manager.get_item_by_id(eid, user=self.user, timezone_str=tz)
+                data = entry.model_dump(exclude_none=True)
+            elif etype == "resource":
+                entry = await self.resource_memory_manager.get_item_by_id(eid, user=self.user)
+                data = entry.model_dump(exclude_none=True)
+            elif etype == "knowledge_vault":
+                entry = await self.knowledge_vault_manager.get_item_by_id(eid, user=self.user, timezone_str=tz)
+                data = entry.model_dump(exclude_none=True)
+            else:
+                results.append({"id": eid, "type": etype, "error": f"unknown type: {etype}"})
+                continue
+            for key in list(data.keys()):
+                if key.endswith("_embedding"):
+                    del data[key]
+            results.append(data)
+        except Exception as e:
+            results.append({"id": eid, "type": etype, "error": str(e)})
+
+    return json.dumps(results, ensure_ascii=False, default=str, indent=2)

mirix/orm/__init__.py

@@ -7,6 +7,7 @@
 from mirix.orm.episodic_memory import EpisodicEvent
 from mirix.orm.file import FileMetadata
 from mirix.orm.knowledge_vault import KnowledgeVaultItem
+from mirix.orm.memory_access_log import MemoryAccessLog
 from mirix.orm.message import Message
 from mirix.orm.organization import Organization
 from mirix.orm.procedural_memory import ProceduralMemoryItem
@@ -28,6 +29,7 @@
     "EpisodicEvent",
     "FileMetadata",
     "KnowledgeVaultItem",
+    "MemoryAccessLog",
     "Message",
     "Organization",
     "ProceduralMemoryItem",