letta-nightly 0.11.3.dev20250819104229__py3-none-any.whl → 0.11.4.dev20250820213507__py3-none-any.whl

letta/server/rest_api/streaming_response.py

@@ -6,6 +6,7 @@ import asyncio
 import json
 from collections.abc import AsyncIterator
 
+import anyio
 from fastapi.responses import StreamingResponse
 from starlette.types import Send
 
@@ -15,6 +16,7 @@ from letta.schemas.letta_ping import LettaPing
 from letta.schemas.user import User
 from letta.server.rest_api.utils import capture_sentry_exception
 from letta.services.job_manager import JobManager
+from letta.settings import settings
 
 logger = get_logger(__name__)
 
@@ -175,8 +177,24 @@ class StreamingResponseWithStatusCode(StreamingResponse):
 
     body_iterator: AsyncIterator[str | bytes]
     response_started: bool = False
+    _client_connected: bool = True
 
     async def stream_response(self, send: Send) -> None:
+        if settings.use_asyncio_shield:
+            try:
+                await asyncio.shield(self._protected_stream_response(send))
+            except asyncio.CancelledError:
+                logger.info(f"Stream response was cancelled, but shielded task should continue")
+            except anyio.ClosedResourceError:
+                logger.info("Client disconnected, but shielded task should continue")
+                self._client_connected = False
+            except Exception as e:
+                logger.error(f"Error in protected stream response: {e}")
+                raise
+        else:
+            await self._protected_stream_response(send)
+
+    async def _protected_stream_response(self, send: Send) -> None:
         more_body = True
         try:
             first_chunk = await self.body_iterator.__anext__()
@@ -188,21 +206,25 @@ class StreamingResponseWithStatusCode(StreamingResponse):
             if isinstance(first_chunk_content, str):
                 first_chunk_content = first_chunk_content.encode(self.charset)
 
-            await send(
-                {
-                    "type": "http.response.start",
-                    "status": self.status_code,
-                    "headers": self.raw_headers,
-                }
-            )
-            self.response_started = True
-            await send(
-                {
-                    "type": "http.response.body",
-                    "body": first_chunk_content,
-                    "more_body": more_body,
-                }
-            )
+            try:
+                await send(
+                    {
+                        "type": "http.response.start",
+                        "status": self.status_code,
+                        "headers": self.raw_headers,
+                    }
+                )
+                self.response_started = True
+                await send(
+                    {
+                        "type": "http.response.body",
+                        "body": first_chunk_content,
+                        "more_body": more_body,
+                    }
+                )
+            except anyio.ClosedResourceError:
+                logger.info("Client disconnected during initial response, continuing processing without sending more chunks")
+                self._client_connected = False
 
             async for chunk in self.body_iterator:
                 if isinstance(chunk, tuple):
@@ -219,13 +241,21 @@ class StreamingResponseWithStatusCode(StreamingResponse):
                 if isinstance(content, str):
                     content = content.encode(self.charset)
                 more_body = True
-                await send(
-                    {
-                        "type": "http.response.body",
-                        "body": content,
-                        "more_body": more_body,
-                    }
-                )
+
+                # Only attempt to send if client is still connected
+                if self._client_connected:
+                    try:
+                        await send(
+                            {
+                                "type": "http.response.body",
+                                "body": content,
+                                "more_body": more_body,
+                            }
+                        )
+                    except anyio.ClosedResourceError:
+                        logger.info("Client disconnected, continuing processing without sending more data")
+                        self._client_connected = False
+                        # Continue processing but don't try to send more data
 
         # Handle explicit job cancellations (should not throw error)
         except JobCancelledException as exc:
@@ -243,13 +273,17 @@ class StreamingResponseWithStatusCode(StreamingResponse):
                     }
                 )
                 raise
-            await send(
-                {
-                    "type": "http.response.body",
-                    "body": cancellation_event,
-                    "more_body": more_body,
-                }
-            )
+            if self._client_connected:
+                try:
+                    await send(
+                        {
+                            "type": "http.response.body",
+                            "body": cancellation_event,
+                            "more_body": more_body,
+                        }
+                    )
+                except anyio.ClosedResourceError:
+                    self._client_connected = False
             return
 
         # Handle client timeouts (should throw error to inform user)
@@ -268,13 +302,17 @@ class StreamingResponseWithStatusCode(StreamingResponse):
                     }
                 )
                 raise
-            await send(
-                {
-                    "type": "http.response.body",
-                    "body": error_event,
-                    "more_body": more_body,
-                }
-            )
+            if self._client_connected:
+                try:
+                    await send(
+                        {
+                            "type": "http.response.body",
+                            "body": error_event,
+                            "more_body": more_body,
+                        }
+                    )
+                except anyio.ClosedResourceError:
+                    self._client_connected = False
             capture_sentry_exception(exc)
             return
 
@@ -293,14 +331,22 @@ class StreamingResponseWithStatusCode(StreamingResponse):
                     }
                 )
                 raise
-            await send(
-                {
-                    "type": "http.response.body",
-                    "body": error_event,
-                    "more_body": more_body,
-                }
-            )
+            if self._client_connected:
+                try:
+                    await send(
+                        {
+                            "type": "http.response.body",
+                            "body": error_event,
+                            "more_body": more_body,
+                        }
+                    )
+                except anyio.ClosedResourceError:
+                    self._client_connected = False
+
             capture_sentry_exception(exc)
             return
-        if more_body:
-            await send({"type": "http.response.body", "body": b"", "more_body": False})
+        if more_body and self._client_connected:
+            try:
+                await send({"type": "http.response.body", "body": b"", "more_body": False})
+            except anyio.ClosedResourceError:
+                self._client_connected = False

letta/server/server.py

@@ -23,7 +23,8 @@ from letta.config import LettaConfig
 from letta.constants import LETTA_TOOL_EXECUTION_DIR
 from letta.data_sources.connectors import DataConnector, load_data
 from letta.errors import HandleNotFoundError
-from letta.functions.mcp_client.types import MCPServerType, MCPTool, SSEServerConfig, StdioServerConfig
+from letta.functions.mcp_client.types import MCPServerType, MCPTool, MCPToolHealth, SSEServerConfig, StdioServerConfig
+from letta.functions.schema_validator import validate_complete_json_schema
 from letta.groups.helpers import load_multi_agent
 from letta.helpers.datetime_helpers import get_utc_time
 from letta.helpers.json_helpers import json_dumps, json_loads
@@ -40,7 +41,7 @@ from letta.schemas.block import Block, BlockUpdate, CreateBlock
 from letta.schemas.embedding_config import EmbeddingConfig
 
 # openai schemas
-from letta.schemas.enums import JobStatus, MessageStreamStatus, ProviderCategory, ProviderType, SandboxType
+from letta.schemas.enums import JobStatus, MessageStreamStatus, ProviderCategory, ProviderType, SandboxType, ToolSourceType
 from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate
 from letta.schemas.group import GroupCreate, ManagerType, SleeptimeManager, VoiceSleeptimeManager
 from letta.schemas.job import Job, JobUpdate
@@ -856,6 +857,9 @@ class SyncServer(Server):
             request.llm_config = await self.get_cached_llm_config_async(actor=actor, **config_params)
             log_event(name="end get_cached_llm_config", attributes=config_params)
 
+        if request.reasoning is None:
+            request.reasoning = request.llm_config.enable_reasoner or request.llm_config.put_inner_thoughts_in_kwargs
+
         if request.embedding_config is None:
             if request.embedding is None:
                 if settings.default_embedding_handle is None:
@@ -1099,33 +1103,6 @@ class SyncServer(Server):
     def get_recall_memory_summary(self, agent_id: str, actor: User) -> RecallMemorySummary:
         return RecallMemorySummary(size=self.message_manager.size(actor=actor, agent_id=agent_id))
 
-    def get_agent_archival(
-        self,
-        user_id: str,
-        agent_id: str,
-        after: Optional[str] = None,
-        before: Optional[str] = None,
-        limit: Optional[int] = 100,
-        order_by: Optional[str] = "created_at",
-        reverse: Optional[bool] = False,
-        query_text: Optional[str] = None,
-        ascending: Optional[bool] = True,
-    ) -> List[Passage]:
-        # TODO: Thread actor directly through this function, since the top level caller most likely already retrieved the user
-        actor = self.user_manager.get_user_or_default(user_id=user_id)
-
-        # iterate over records
-        records = self.agent_manager.list_passages(
-            actor=actor,
-            agent_id=agent_id,
-            after=after,
-            query_text=query_text,
-            before=before,
-            ascending=ascending,
-            limit=limit,
-        )
-        return records
-
     async def get_agent_archival_async(
         self,
         agent_id: str,
@@ -1153,7 +1130,7 @@ class SyncServer(Server):
         agent_state = await self.agent_manager.get_agent_by_id_async(agent_id=agent_id, actor=actor)
 
         # Insert passages into the archive
-        passages = await self.passage_manager.insert_passage_async(agent_state=agent_state, text=memory_contents, actor=actor)
+        passages = await self.passage_manager.insert_passage(agent_state=agent_state, text=memory_contents, actor=actor)
 
         return passages
 
@@ -1471,10 +1448,6 @@ class SyncServer(Server):
         passage_count, document_count = await load_data(connector, source, self.passage_manager, self.file_manager, actor=actor)
         return passage_count, document_count
 
-    def list_data_source_passages(self, user_id: str, source_id: str) -> List[Passage]:
-        # TODO: move this query into PassageManager
-        return self.agent_manager.list_passages(actor=self.user_manager.get_user_or_default(user_id=user_id), source_id=source_id)
-
     def list_all_sources(self, actor: User) -> List[Source]:
         # TODO: legacy: remove
         """List all sources (w/ extra metadata) belonging to a user"""
@@ -1934,12 +1907,19 @@ class SyncServer(Server):
         pip_requirements: Optional[List[PipRequirement]] = None,
     ) -> ToolReturnMessage:
         """Run a tool from source code"""
-        if tool_source_type is not None and tool_source_type != "python":
-            raise ValueError("Only Python source code is supported at this time")
+
+        if tool_source_type not in (None, ToolSourceType.python, ToolSourceType.typescript):
+            raise ValueError("Tool source type is not supported at this time. Found {tool_source_type}")
 
         # If tools_json_schema is explicitly passed in, override it on the created Tool object
         if tool_json_schema:
-            tool = Tool(name=tool_name, source_code=tool_source, json_schema=tool_json_schema, pip_requirements=pip_requirements)
+            tool = Tool(
+                name=tool_name,
+                source_code=tool_source,
+                json_schema=tool_json_schema,
+                pip_requirements=pip_requirements,
+                source_type=tool_source_type,
+            )
         else:
             # NOTE: we're creating a floating Tool object and NOT persisting to DB
             tool = Tool(
@@ -1947,6 +1927,7 @@ class SyncServer(Server):
                 source_code=tool_source,
                 args_json_schema=tool_args_json_schema,
                 pip_requirements=pip_requirements,
+                source_type=tool_source_type,
             )
 
         assert tool.name is not None, "Failed to create tool object"
@@ -2086,7 +2067,15 @@ class SyncServer(Server):
         if mcp_server_name not in self.mcp_clients:
             raise ValueError(f"No client was created for MCP server: {mcp_server_name}")
 
-        return await self.mcp_clients[mcp_server_name].list_tools()
+        tools = await self.mcp_clients[mcp_server_name].list_tools()
+
+        # Add health information to each tool
+        for tool in tools:
+            if tool.inputSchema:
+                health_status, reasons = validate_complete_json_schema(tool.inputSchema)
+                tool.health = MCPToolHealth(status=health_status.value, reasons=reasons)
+
+        return tools
 
     async def add_mcp_server_to_config(
         self, server_config: Union[SSEServerConfig, StdioServerConfig], allow_upsert: bool = True

letta/services/agent_manager.py

@@ -19,8 +19,9 @@ from letta.constants import (
     DEFAULT_MAX_FILES_OPEN,
     DEFAULT_TIMEZONE,
     DEPRECATED_LETTA_TOOLS,
-    EXCLUDED_PROVIDERS_FROM_BASE_TOOL_RULES,
+    EXCLUDE_MODEL_KEYWORDS_FROM_BASE_TOOL_RULES,
     FILES_TOOLS,
+    INCLUDE_MODEL_KEYWORDS_BASE_TOOL_RULES,
 )
 from letta.helpers import ToolRulesSolver
 from letta.helpers.datetime_helpers import get_utc_time
@@ -117,6 +118,21 @@ class AgentManager:
         self.identity_manager = IdentityManager()
         self.file_agent_manager = FileAgentManager()
 
+    @staticmethod
+    def _should_exclude_model_from_base_tool_rules(model: str) -> bool:
+        """Check if a model should be excluded from base tool rules based on model keywords."""
+        # First check if model contains any include keywords (overrides exclusion)
+        for include_keyword in INCLUDE_MODEL_KEYWORDS_BASE_TOOL_RULES:
+            if include_keyword in model:
+                return False
+
+        # Then check if model contains any exclude keywords
+        for exclude_keyword in EXCLUDE_MODEL_KEYWORDS_FROM_BASE_TOOL_RULES:
+            if exclude_keyword in model:
+                return True
+
+        return False
+
     @staticmethod
     def _resolve_tools(session, names: Set[str], ids: Set[str], org_id: str) -> Tuple[Dict[str, str], Dict[str, str]]:
         """
@@ -334,16 +350,16 @@ class AgentManager:
 
                 tool_rules = list(agent_create.tool_rules or [])
 
-                # Override include_base_tool_rules to False if provider is not in excluded set and include_base_tool_rules is not explicitly set to True
+                # Override include_base_tool_rules to False if model matches exclusion keywords and include_base_tool_rules is not explicitly set to True
                 if (
                     (
-                        agent_create.llm_config.model_endpoint_type in EXCLUDED_PROVIDERS_FROM_BASE_TOOL_RULES
+                        self._should_exclude_model_from_base_tool_rules(agent_create.llm_config.model)
                         and agent_create.include_base_tool_rules is None
                     )
                     and agent_create.agent_type != AgentType.sleeptime_agent
                 ) or agent_create.include_base_tool_rules is False:
                     agent_create.include_base_tool_rules = False
-                    logger.info(f"Overriding include_base_tool_rules to False for provider: {agent_create.llm_config.model_endpoint_type}")
+                    logger.info(f"Overriding include_base_tool_rules to False for model: {agent_create.llm_config.model}")
                 else:
                     agent_create.include_base_tool_rules = True
 
@@ -543,16 +559,16 @@ class AgentManager:
                 tool_names = set(name_to_id.keys())  # now canonical
                 tool_rules = list(agent_create.tool_rules or [])
 
-                # Override include_base_tool_rules to False if provider is not in excluded set and include_base_tool_rules is not explicitly set to True
+                # Override include_base_tool_rules to False if model matches exclusion keywords and include_base_tool_rules is not explicitly set to True
                 if (
                     (
-                        agent_create.llm_config.model_endpoint_type in EXCLUDED_PROVIDERS_FROM_BASE_TOOL_RULES
+                        self._should_exclude_model_from_base_tool_rules(agent_create.llm_config.model)
                         and agent_create.include_base_tool_rules is None
                     )
                     and agent_create.agent_type != AgentType.sleeptime_agent
                 ) or agent_create.include_base_tool_rules is False:
                     agent_create.include_base_tool_rules = False
-                    logger.info(f"Overriding include_base_tool_rules to False for provider: {agent_create.llm_config.model_endpoint_type}")
+                    logger.info(f"Overriding include_base_tool_rules to False for model: {agent_create.llm_config.model}")
                 else:
                     agent_create.include_base_tool_rules = True
 
@@ -630,6 +646,7 @@ class AgentManager:
                     [{"agent_id": aid, "identity_id": iid} for iid in identity_ids],
                 )
 
+                env_rows = []
                 if agent_create.tool_exec_environment_variables:
                     env_rows = [
                         {
@@ -640,7 +657,8 @@ class AgentManager:
                         }
                         for key, val in agent_create.tool_exec_environment_variables.items()
                     ]
-                    await session.execute(insert(AgentEnvironmentVariable).values(env_rows))
+                    result = await session.execute(insert(AgentEnvironmentVariable).values(env_rows).returning(AgentEnvironmentVariable.id))
+                    env_rows = [{**row, "id": env_var_id} for row, env_var_id in zip(env_rows, result.scalars().all())]
 
                 include_relationships = []
                 if tool_ids:
@@ -656,6 +674,9 @@ class AgentManager:
 
                 result = await new_agent.to_pydantic_async(include_relationships=include_relationships)
 
+                if agent_create.tool_exec_environment_variables and env_rows:
+                    result.tool_exec_environment_variables = [AgentEnvironmentVariable(**row) for row in env_rows]
+
                 # initial message sequence (skip if _init_with_no_messages is True)
                 if not _init_with_no_messages:
                     init_messages = await self._generate_initial_message_sequence_async(
@@ -1986,6 +2007,26 @@ class AgentManager:
     @enforce_types
     @trace_method
     async def refresh_file_blocks(self, agent_state: PydanticAgentState, actor: PydanticUser) -> PydanticAgentState:
+        """
+        Refresh the file blocks in an agent's memory with current file content.
+
+        This method synchronizes the agent's in-memory file blocks with the actual
+        file content from attached sources. It respects the per-file view window
+        limit to prevent excessive memory usage.
+
+        Args:
+            agent_state: The current agent state containing memory configuration
+            actor: The user performing this action (for permission checking)
+
+        Returns:
+            Updated agent state with refreshed file blocks
+
+        Important:
+            - File blocks are truncated based on per_file_view_window_char_limit
+            - None values are filtered out (files that couldn't be loaded)
+            - This does NOT persist changes to the database, only updates the state object
+            - Call this before agent interactions if files may have changed externally
+        """
         file_blocks = await self.file_agent_manager.list_files_for_agent(
             agent_id=agent_state.id,
             per_file_view_window_char_limit=agent_state.per_file_view_window_char_limit,
@@ -2035,6 +2076,28 @@ class AgentManager:
     @enforce_types
     @trace_method
     def append_system_message(self, agent_id: str, content: str, actor: PydanticUser):
+        """
+        Append a system message to an agent's in-context message history.
+
+        This method is typically used during agent initialization to add system prompts,
+        instructions, or context that should be treated as system-level guidance.
+        Unlike user messages, system messages directly influence the agent's behavior
+        and understanding of its role.
+
+        Args:
+            agent_id: The ID of the agent to append the message to
+            content: The system message content (e.g., instructions, context, role definition)
+            actor: The user performing this action (for permission checking)
+
+        Side Effects:
+            - Creates a new Message object in the database
+            - Updates the agent's in_context_message_ids list
+            - The message becomes part of the agent's permanent context window
+
+        Note:
+            System messages consume tokens in the context window and cannot be
+            removed without rebuilding the agent's message history.
+        """
 
         # get the agent
         agent = self.get_agent_by_id(agent_id=agent_id, actor=actor)
@@ -2048,6 +2111,15 @@ class AgentManager:
     @enforce_types
     @trace_method
     async def append_system_message_async(self, agent_id: str, content: str, actor: PydanticUser):
+        """
+        Async version of append_system_message.
+
+        Append a system message to an agent's in-context message history.
+        See append_system_message for detailed documentation.
+
+        This async version is preferred for high-throughput scenarios or when
+        called within other async operations to avoid blocking the event loop.
+        """
 
         # get the agent
         agent = await self.get_agent_by_id_async(agent_id=agent_id, actor=actor)
@@ -2354,7 +2426,7 @@ class AgentManager:
 
     @enforce_types
     @trace_method
-    def list_passages(
+    async def list_passages(
         self,
         actor: PydanticUser,
         agent_id: Optional[str] = None,
@@ -2372,8 +2444,8 @@ class AgentManager:
         agent_only: bool = False,
     ) -> List[PydanticPassage]:
         """Lists all passages attached to an agent."""
-        with db_registry.session() as session:
-            main_query = build_passage_query(
+        async with db_registry.async_session() as session:
+            main_query = await build_passage_query(
                 actor=actor,
                 agent_id=agent_id,
                 file_id=file_id,
@@ -2394,7 +2466,7 @@ class AgentManager:
                 main_query = main_query.limit(limit)
 
             # Execute query
-            result = session.execute(main_query)
+            result = await session.execute(main_query)
 
             passages = []
             for row in result:
@@ -2437,7 +2509,7 @@ class AgentManager:
     ) -> List[PydanticPassage]:
         """Lists all passages attached to an agent."""
         async with db_registry.async_session() as session:
-            main_query = build_passage_query(
+            main_query = await build_passage_query(
                 actor=actor,
                 agent_id=agent_id,
                 file_id=file_id,
@@ -2500,7 +2572,7 @@ class AgentManager:
     ) -> List[PydanticPassage]:
         """Lists all passages attached to an agent."""
         async with db_registry.async_session() as session:
-            main_query = build_source_passage_query(
+            main_query = await build_source_passage_query(
                 actor=actor,
                 agent_id=agent_id,
                 file_id=file_id,
@@ -2546,7 +2618,7 @@ class AgentManager:
     ) -> List[PydanticPassage]:
         """Lists all passages attached to an agent."""
         async with db_registry.async_session() as session:
-            main_query = build_agent_passage_query(
+            main_query = await build_agent_passage_query(
                 actor=actor,
                 agent_id=agent_id,
                 query_text=query_text,
@@ -2574,7 +2646,7 @@ class AgentManager:
 
     @enforce_types
     @trace_method
-    def passage_size(
+    async def passage_size(
         self,
         actor: PydanticUser,
         agent_id: Optional[str] = None,
@@ -2591,8 +2663,8 @@ class AgentManager:
         agent_only: bool = False,
     ) -> int:
         """Returns the count of passages matching the given criteria."""
-        with db_registry.session() as session:
-            main_query = build_passage_query(
+        async with db_registry.async_session() as session:
+            main_query = await build_passage_query(
                 actor=actor,
                 agent_id=agent_id,
                 file_id=file_id,
@@ -2610,7 +2682,7 @@ class AgentManager:
 
             # Convert to count query
             count_query = select(func.count()).select_from(main_query.subquery())
-            return session.scalar(count_query) or 0
+            return (await session.scalar(count_query)) or 0
 
     @enforce_types
     async def passage_size_async(
@@ -2630,7 +2702,7 @@ class AgentManager:
         agent_only: bool = False,
     ) -> int:
         async with db_registry.async_session() as session:
-            main_query = build_passage_query(
+            main_query = await build_passage_query(
                 actor=actor,
                 agent_id=agent_id,
                 file_id=file_id,

letta/services/agent_serialization_manager.py

@@ -2,7 +2,7 @@ from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional
 
 from letta.constants import MCP_TOOL_TAG_NAME_PREFIX
-from letta.errors import AgentFileExportError, AgentFileImportError
+from letta.errors import AgentExportIdMappingError, AgentExportProcessingError, AgentFileImportError, AgentNotFoundForExportError
 from letta.helpers.pinecone_utils import should_use_pinecone
 from letta.log import get_logger
 from letta.schemas.agent import AgentState, CreateAgent
@@ -118,10 +118,7 @@ class AgentSerializationManager:
             return self._db_to_file_ids[db_id]
 
         if not allow_new:
-            raise AgentFileExportError(
-                f"Unexpected new {entity_type} ID '{db_id}' encountered during conversion. "
-                f"All IDs should have been mapped during agent processing."
-            )
+            raise AgentExportIdMappingError(db_id, entity_type)
 
         file_id = self._generate_file_id(entity_type)
         self._db_to_file_ids[db_id] = file_id
@@ -352,7 +349,7 @@ class AgentSerializationManager:
             if len(agent_states) != len(agent_ids):
                 found_ids = {agent.id for agent in agent_states}
                 missing_ids = [agent_id for agent_id in agent_ids if agent_id not in found_ids]
-                raise AgentFileExportError(f"The following agent IDs were not found: {missing_ids}")
+                raise AgentNotFoundForExportError(missing_ids)
 
             groups = []
             group_agent_ids = []
@@ -417,7 +414,7 @@ class AgentSerializationManager:
 
         except Exception as e:
             logger.error(f"Failed to export agent file: {e}")
-            raise AgentFileExportError(f"Export failed: {e}") from e
+            raise AgentExportProcessingError(str(e), e) from e
 
     async def import_file(
         self,
@@ -657,6 +654,12 @@ class AgentSerializationManager:
                     )
                     imported_count += len(files_for_agent)
 
+            # Extract the imported agent IDs (database IDs)
+            imported_agent_ids = []
+            for agent_schema in schema.agents:
+                if agent_schema.id in file_to_db_ids:
+                    imported_agent_ids.append(file_to_db_ids[agent_schema.id])
+
             for group in schema.groups:
                 group_data = group.model_dump(exclude={"id"})
                 group_data["agent_ids"] = [file_to_db_ids[agent_id] for agent_id in group_data["agent_ids"]]
@@ -670,6 +673,7 @@ class AgentSerializationManager:
                 success=True,
                 message=f"Import completed successfully. Imported {imported_count} entities.",
                 imported_count=imported_count,
+                imported_agent_ids=imported_agent_ids,
                 id_mappings=file_to_db_ids,
             )