From 8ed1505ceaab68f2fb2edb899ece3f7b9f5dd993 Mon Sep 17 00:00:00 2001
From: cthomas <caren@letta.com>
Date: Sun, 20 Jul 2025 23:00:57 -0700
Subject: [PATCH] feat: add special message update helper for agent (#3442)

---
 letta/agents/letta_agent.py          |  4 ++--
 letta/agents/letta_agent_batch.py    |  2 +-
 letta/agents/voice_agent.py          |  2 +-
 letta/services/agent_file_manager.py |  2 +-
 letta/services/agent_manager.py      | 24 ++++++++++++++++++++++++
 5 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/letta/agents/letta_agent.py b/letta/agents/letta_agent.py
index 41543d44..3c7960bc 100644
--- a/letta/agents/letta_agent.py
+++ b/letta/agents/letta_agent.py
@@ -1115,7 +1115,7 @@ class LettaAgent(BaseAgent):
                 in_context_messages=in_context_messages,
                 new_letta_messages=new_letta_messages,
             )
-        await self.agent_manager.set_in_context_messages_async(
+        await self.agent_manager.update_message_ids_async(
             agent_id=self.agent_id,
             message_ids=[m.id for m in new_in_context_messages],
             actor=self.actor,
@@ -1132,7 +1132,7 @@ class LettaAgent(BaseAgent):
         new_in_context_messages, updated = await self.summarizer.summarize(
             in_context_messages=in_context_messages, new_letta_messages=[], force=True
         )
-        return await self.agent_manager.set_in_context_messages_async(
+        return await self.agent_manager.update_message_ids_async(
             agent_id=self.agent_id, message_ids=[m.id for m in new_in_context_messages], actor=self.actor
         )
 
diff --git a/letta/agents/letta_agent_batch.py b/letta/agents/letta_agent_batch.py
index 0e9afb07..3bca27f4 100644
--- a/letta/agents/letta_agent_batch.py
+++ b/letta/agents/letta_agent_batch.py
@@ -515,7 +515,7 @@ class LettaAgentBatch(BaseAgent):
         for agent_id, new_msgs in msg_map.items():
             ast = ctx.agent_state_map[agent_id]
             if not ast.message_buffer_autoclear:
-                await self.agent_manager.set_in_context_messages_async(
+                await self.agent_manager.update_message_ids_async(
                     agent_id=agent_id,
                     message_ids=ast.message_ids + [m.id for m in new_msgs],
                     actor=self.actor,
diff --git a/letta/agents/voice_agent.py b/letta/agents/voice_agent.py
index 93ced5cb..9e30cb54 100644
--- a/letta/agents/voice_agent.py
+++ b/letta/agents/voice_agent.py
@@ -298,7 +298,7 @@ class VoiceAgent(BaseAgent):
             in_context_messages=in_context_messages, new_letta_messages=new_letta_messages
         )
 
-        await self.agent_manager.set_in_context_messages_async(
+        await self.agent_manager.update_message_ids_async(
             agent_id=self.agent_id, message_ids=[m.id for m in new_in_context_messages], actor=self.actor
         )
 
diff --git a/letta/services/agent_file_manager.py b/letta/services/agent_file_manager.py
index 001cd964..da278912 100644
--- a/letta/services/agent_file_manager.py
+++ b/letta/services/agent_file_manager.py
@@ -434,7 +434,7 @@ class AgentFileManager:
                 in_context_db_ids = [message_file_to_db_ids[message_schema_id] for message_schema_id in agent_schema.in_context_message_ids]
 
                 # Update agent with the correct message_ids
-                await self.agent_manager.set_in_context_messages_async(agent_id=agent_db_id, message_ids=in_context_db_ids, actor=actor)
+                await self.agent_manager.update_message_ids_async(agent_id=agent_db_id, message_ids=in_context_db_ids, actor=actor)
 
             # 8. Create file-agent relationships (depends on agents and files)
             for agent_schema in schema.agents:
diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py
index e25d2203..0f42f2d0 100644
--- a/letta/services/agent_manager.py
+++ b/letta/services/agent_manager.py
@@ -907,6 +907,30 @@ class AgentManager:
 
             return await agent.to_pydantic_async()
 
+    @enforce_types
+    @trace_method
+    async def update_message_ids_async(
+        self,
+        agent_id: str,
+        message_ids: List[str],
+        actor: PydanticUser,
+    ) -> None:
+        async with db_registry.async_session() as session:
+            query = select(AgentModel)
+            query = AgentModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
+            query = query.where(AgentModel.id == agent_id)
+            query = _apply_relationship_filters(query, include_relationships=[])
+
+            result = await session.execute(query)
+            agent = result.scalar_one_or_none()
+
+            agent.updated_at = datetime.now(timezone.utc)
+            agent.last_updated_by_id = actor.id
+            agent.message_ids = message_ids
+
+            await agent.update_async(db_session=session, actor=actor, no_commit=True, no_refresh=True)
+            await session.commit()
+
     # TODO: Make this general and think about how to roll this into sqlalchemybase
     @trace_method
     def list_agents(