feat: add special message update helper for agent (#3442)

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
         )
 

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,

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
         )
 

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:

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(