letta-nightly 0.9.0.dev20250726104256__py3-none-any.whl → 0.9.1.dev20250727104258__py3-none-any.whl

letta/server/rest_api/routers/v1/tools.py

@@ -1,4 +1,6 @@
+import asyncio
 import json
+from collections.abc import AsyncGenerator
 from typing import Any, Dict, List, Optional, Union
 
 from composio.client import ComposioClientError, HTTPError, NoItemsFound
@@ -11,27 +13,38 @@ from composio.exceptions import (
     EnumStringNotFound,
 )
 from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query
+from fastapi.responses import HTMLResponse
 from pydantic import BaseModel, Field
+from starlette.responses import StreamingResponse
 
 from letta.errors import LettaToolCreateError
 from letta.functions.functions import derive_openai_json_schema
 from letta.functions.mcp_client.exceptions import MCPTimeoutError
-from letta.functions.mcp_client.types import MCPServerType, MCPTool, SSEServerConfig, StdioServerConfig, StreamableHTTPServerConfig
+from letta.functions.mcp_client.types import MCPTool, SSEServerConfig, StdioServerConfig, StreamableHTTPServerConfig
 from letta.helpers.composio_helpers import get_composio_api_key
+from letta.helpers.decorators import deprecated
 from letta.llm_api.llm_client import LLMClient
 from letta.log import get_logger
 from letta.orm.errors import UniqueConstraintViolationError
+from letta.orm.mcp_oauth import OAuthSessionStatus
 from letta.schemas.enums import MessageRole
 from letta.schemas.letta_message import ToolReturnMessage
 from letta.schemas.letta_message_content import TextContent
-from letta.schemas.mcp import UpdateSSEMCPServer, UpdateStdioMCPServer, UpdateStreamableHTTPMCPServer
+from letta.schemas.mcp import MCPOAuthSessionCreate, UpdateSSEMCPServer, UpdateStdioMCPServer, UpdateStreamableHTTPMCPServer
 from letta.schemas.message import Message
 from letta.schemas.tool import Tool, ToolCreate, ToolRunFromSource, ToolUpdate
+from letta.server.rest_api.streaming_response import StreamingResponseWithStatusCode
 from letta.server.rest_api.utils import get_letta_server
 from letta.server.server import SyncServer
-from letta.services.mcp.sse_client import AsyncSSEMCPClient
+from letta.services.mcp.oauth_utils import (
+    MCPOAuthSession,
+    create_oauth_provider,
+    drill_down_exception,
+    get_oauth_success_html,
+    oauth_stream_event,
+)
 from letta.services.mcp.stdio_client import AsyncStdioMCPClient
-from letta.services.mcp.streamable_http_client import AsyncStreamableHTTPMCPClient
+from letta.services.mcp.types import OauthStreamEvent
 from letta.settings import tool_settings
 
 router = APIRouter(prefix="/tools", tags=["tools"])
@@ -475,7 +488,11 @@ async def add_mcp_tool(
             )
 
         tool_create = ToolCreate.from_mcp(mcp_server_name=mcp_server_name, mcp_tool=mcp_tool)
-        return await server.tool_manager.create_mcp_tool_async(tool_create=tool_create, mcp_server_name=mcp_server_name, actor=actor)
+        # For config-based servers, use the server name as ID since they don't have database IDs
+        mcp_server_id = mcp_server_name
+        return await server.tool_manager.create_mcp_tool_async(
+            tool_create=tool_create, mcp_server_name=mcp_server_name, mcp_server_id=mcp_server_id, actor=actor
+        )
 
     else:
         return await server.mcp_manager.add_tool_from_mcp_server(mcp_server_name=mcp_server_name, mcp_tool_name=mcp_tool_name, actor=actor)
@@ -608,35 +625,26 @@ async def delete_mcp_server_from_config(
         return [server.to_config() for server in all_servers]
 
 
-@router.post("/mcp/servers/test", response_model=List[MCPTool], operation_id="test_mcp_server")
+@deprecated("Deprecated in favor of /mcp/servers/connect which handles OAuth flow via SSE stream")
+@router.post("/mcp/servers/test", operation_id="test_mcp_server")
 async def test_mcp_server(
     request: Union[StdioServerConfig, SSEServerConfig, StreamableHTTPServerConfig] = Body(...),
+    server: SyncServer = Depends(get_letta_server),
+    actor_id: Optional[str] = Header(None, alias="user_id"),
 ):
     """
     Test connection to an MCP server without adding it.
-    Returns the list of available tools if successful.
+    Returns the list of available tools if successful, or OAuth information if OAuth is required.
     """
     client = None
     try:
-        # create a temporary MCP client based on the server type
-        if request.type == MCPServerType.SSE:
-            if not isinstance(request, SSEServerConfig):
-                request = SSEServerConfig(**request.model_dump())
-            client = AsyncSSEMCPClient(request)
-        elif request.type == MCPServerType.STREAMABLE_HTTP:
-            if not isinstance(request, StreamableHTTPServerConfig):
-                request = StreamableHTTPServerConfig(**request.model_dump())
-            client = AsyncStreamableHTTPMCPClient(request)
-        elif request.type == MCPServerType.STDIO:
-            if not isinstance(request, StdioServerConfig):
-                request = StdioServerConfig(**request.model_dump())
-            client = AsyncStdioMCPClient(request)
-        else:
-            raise ValueError(f"Invalid MCP server type: {request.type}")
+        actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
+        client = await server.mcp_manager.get_mcp_client(request, actor)
 
         await client.connect_to_server()
         tools = await client.list_tools()
-        return tools
+
+        return {"status": "success", "tools": tools}
     except ConnectionError as e:
         raise HTTPException(
             status_code=400,
@@ -672,6 +680,160 @@ async def test_mcp_server(
                 logger.warning(f"Error during MCP client cleanup: {cleanup_error}")
 
 
+@router.post(
+    "/mcp/servers/connect",
+    response_model=None,
+    responses={
+        200: {
+            "description": "Successful response",
+            "content": {
+                "text/event-stream": {"description": "Server-Sent Events stream"},
+            },
+        }
+    },
+    operation_id="connect_mcp_server",
+)
+async def connect_mcp_server(
+    request: Union[StdioServerConfig, SSEServerConfig, StreamableHTTPServerConfig] = Body(...),
+    server: SyncServer = Depends(get_letta_server),
+    actor_id: Optional[str] = Header(None, alias="user_id"),
+) -> StreamingResponse:
+    """
+    Connect to an MCP server with support for OAuth via SSE.
+    Returns a stream of events handling authorization state and exchange if OAuth is required.
+    """
+
+    async def oauth_stream_generator(
+        request: Union[StdioServerConfig, SSEServerConfig, StreamableHTTPServerConfig],
+    ) -> AsyncGenerator[str, None]:
+        client = None
+        oauth_provider = None
+        temp_client = None
+        connect_task = None
+
+        try:
+            # Acknolwedge connection attempt
+            yield oauth_stream_event(OauthStreamEvent.CONNECTION_ATTEMPT, server_name=request.server_name)
+
+            actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
+
+            # Create MCP client with respective transport type
+            try:
+                client = await server.mcp_manager.get_mcp_client(request, actor)
+            except ValueError as e:
+                yield oauth_stream_event(OauthStreamEvent.ERROR, message=str(e))
+                return
+
+            # Try normal connection first for flows that don't require OAuth
+            try:
+                await client.connect_to_server()
+                tools = await client.list_tools(serialize=True)
+                yield oauth_stream_event(OauthStreamEvent.SUCCESS, tools=tools)
+                return
+            except ConnectionError:
+                # TODO: jnjpng make this connection error check more specific to the 401 unauthorized error
+                if isinstance(client, AsyncStdioMCPClient):
+                    logger.warning(f"OAuth not supported for stdio")
+                    yield oauth_stream_event(OauthStreamEvent.ERROR, message=f"OAuth not supported for stdio")
+                    return
+                # Continue to OAuth flow
+                logger.info(f"Attempting OAuth flow for {request}...")
+            except Exception as e:
+                yield oauth_stream_event(OauthStreamEvent.ERROR, message=f"Connection failed: {str(e)}")
+                return
+
+            # OAuth required, yield state to client to prepare to handle authorization URL
+            yield oauth_stream_event(OauthStreamEvent.OAUTH_REQUIRED, message="OAuth authentication required")
+
+            # Create OAuth session to persist the state of the OAuth flow
+            session_create = MCPOAuthSessionCreate(
+                server_url=request.server_url,
+                server_name=request.server_name,
+                user_id=actor.id,
+                organization_id=actor.organization_id,
+            )
+            oauth_session = await server.mcp_manager.create_oauth_session(session_create, actor)
+            session_id = oauth_session.id
+
+            # Create OAuth provider for the instance of the stream connection
+            # Note: Using the correct API path for the callback
+            # do not edit this this is the correct url
+            redirect_uri = f"http://localhost:8283/v1/tools/mcp/oauth/callback/{session_id}"
+            oauth_provider = await create_oauth_provider(session_id, request.server_url, redirect_uri, server.mcp_manager, actor)
+
+            # Get authorization URL by triggering OAuth flow
+            temp_client = None
+            try:
+                temp_client = await server.mcp_manager.get_mcp_client(request, actor, oauth_provider)
+
+                # Run connect_to_server in background to avoid blocking
+                # This will trigger the OAuth flow and the redirect_handler will save the authorization URL to database
+                connect_task = asyncio.create_task(temp_client.connect_to_server())
+
+                # Give the OAuth flow time to trigger and save the URL
+                await asyncio.sleep(1.0)
+
+                # Fetch the authorization URL from database and yield state to client to proceed with handling authorization URL
+                auth_session = await server.mcp_manager.get_oauth_session_by_id(session_id, actor)
+                if auth_session and auth_session.authorization_url:
+                    yield oauth_stream_event(OauthStreamEvent.AUTHORIZATION_URL, url=auth_session.authorization_url, session_id=session_id)
+
+            except Exception as e:
+                logger.error(f"Error triggering OAuth flow: {e}")
+                yield oauth_stream_event(OauthStreamEvent.ERROR, message=f"Failed to trigger OAuth: {str(e)}")
+
+                # Clean up active resources
+                if connect_task and not connect_task.done():
+                    connect_task.cancel()
+                    try:
+                        await connect_task
+                    except asyncio.CancelledError:
+                        pass
+                if temp_client:
+                    try:
+                        await temp_client.cleanup()
+                    except Exception as cleanup_error:
+                        logger.warning(f"Error during temp MCP client cleanup: {cleanup_error}")
+                return
+
+            # Wait for user authorization (with timeout), client should render loading state until user completes the flow and /mcp/oauth/callback/{session_id} is hit
+            yield oauth_stream_event(OauthStreamEvent.WAITING_FOR_AUTH, message="Waiting for user authorization...")
+
+            # Callback handler will poll for authorization code and state and update the OAuth session
+            await connect_task
+
+            tools = await temp_client.list_tools(serialize=True)
+
+            yield oauth_stream_event(OauthStreamEvent.SUCCESS, tools=tools)
+            return
+        except Exception as e:
+            detailed_error = drill_down_exception(e)
+            logger.error(f"Error in OAuth stream:\n{detailed_error}")
+            yield oauth_stream_event(OauthStreamEvent.ERROR, message=f"Internal error: {detailed_error}")
+        finally:
+            if connect_task and not connect_task.done():
+                connect_task.cancel()
+                try:
+                    await connect_task
+                except asyncio.CancelledError:
+                    pass
+            if client:
+                try:
+                    await client.cleanup()
+                except Exception as cleanup_error:
+                    detailed_error = drill_down_exception(cleanup_error)
+                    logger.warning(f"Error during MCP client cleanup: {detailed_error}")
+            if temp_client:
+                try:
+                    await temp_client.cleanup()
+                except Exception as cleanup_error:
+                    # TODO: @jnjpng fix async cancel scope issue
+                    # detailed_error = drill_down_exception(cleanup_error)
+                    logger.warning(f"Aysnc cleanup confict during temp MCP client cleanup: {cleanup_error}")
+
+    return StreamingResponseWithStatusCode(oauth_stream_generator(request), media_type="text/event-stream")
+
+
 class CodeInput(BaseModel):
     code: str = Field(..., description="Python source code to parse for JSON schema")
 
@@ -693,6 +855,45 @@ async def generate_json_schema(
         raise HTTPException(status_code=400, detail=f"Failed to generate schema: {str(e)}")
 
 
+# TODO: @jnjpng need to route this through cloud API for production
+@router.get("/mcp/oauth/callback/{session_id}", operation_id="mcp_oauth_callback", response_class=HTMLResponse)
+async def mcp_oauth_callback(
+    session_id: str,
+    code: Optional[str] = Query(None, description="OAuth authorization code"),
+    state: Optional[str] = Query(None, description="OAuth state parameter"),
+    error: Optional[str] = Query(None, description="OAuth error"),
+    error_description: Optional[str] = Query(None, description="OAuth error description"),
+):
+    """
+    Handle OAuth callback for MCP server authentication.
+    """
+    try:
+        oauth_session = MCPOAuthSession(session_id)
+
+        if error:
+            error_msg = f"OAuth error: {error}"
+            if error_description:
+                error_msg += f" - {error_description}"
+            await oauth_session.update_session_status(OAuthSessionStatus.ERROR)
+            return {"status": "error", "message": error_msg}
+
+        if not code or not state:
+            await oauth_session.update_session_status(OAuthSessionStatus.ERROR)
+            return {"status": "error", "message": "Missing authorization code or state"}
+
+        # Store authorization code
+        success = await oauth_session.store_authorization_code(code, state)
+        if not success:
+            await oauth_session.update_session_status(OAuthSessionStatus.ERROR)
+            return {"status": "error", "message": "Invalid state parameter"}
+
+        return HTMLResponse(content=get_oauth_success_html(), status_code=200)
+
+    except Exception as e:
+        logger.error(f"OAuth callback error: {e}")
+        return {"status": "error", "message": f"OAuth callback failed: {str(e)}"}
+
+
 class GenerateToolInput(BaseModel):
     tool_name: str = Field(..., description="Name of the tool to generate code for")
     prompt: str = Field(..., description="User prompt to generate code")

letta/services/agent_manager.py

@@ -15,6 +15,8 @@ from letta.constants import (
     BASE_TOOLS,
     BASE_VOICE_SLEEPTIME_CHAT_TOOLS,
     BASE_VOICE_SLEEPTIME_TOOLS,
+    DEFAULT_CORE_MEMORY_SOURCE_CHAR_LIMIT,
+    DEFAULT_MAX_FILES_OPEN,
     DEFAULT_TIMEZONE,
     DEPRECATED_LETTA_TOOLS,
     FILES_TOOLS,
@@ -1644,7 +1646,7 @@ class AgentManager:
 
         # note: we only update the system prompt if the core memory is changed
         # this means that the archival/recall memory statistics may be someout out of date
-        curr_memory_str = agent_state.memory.compile(
+        curr_memory_str = await agent_state.memory.compile_async(
             sources=agent_state.sources,
             tool_usage_rules=tool_rules_solver.compile_tool_rule_prompts(),
             max_files_open=agent_state.max_files_open,
@@ -1834,14 +1836,12 @@ class AgentManager:
         agent_state = await self.get_agent_by_id_async(agent_id=agent_id, actor=actor, include_relationships=["memory", "sources"])
         system_message = await self.message_manager.get_message_by_id_async(message_id=agent_state.message_ids[0], actor=actor)
         temp_tool_rules_solver = ToolRulesSolver(agent_state.tool_rules)
-        if (
-            new_memory.compile(
-                sources=agent_state.sources,
-                tool_usage_rules=temp_tool_rules_solver.compile_tool_rule_prompts(),
-                max_files_open=agent_state.max_files_open,
-            )
-            not in system_message.content[0].text
-        ):
+        new_memory_str = await new_memory.compile_async(
+            sources=agent_state.sources,
+            tool_usage_rules=temp_tool_rules_solver.compile_tool_rule_prompts(),
+            max_files_open=agent_state.max_files_open,
+        )
+        if new_memory_str not in system_message.content[0].text:
             # update the blocks (LRW) in the DB
             for label in agent_state.memory.list_block_labels():
                 updated_value = new_memory.get_block(label).value
@@ -3169,6 +3169,12 @@ class AgentManager:
                 if max_files is None:
                     max_files = default_max_files
 
+            # FINAL fallback: ensure neither is None (should never happen, but just in case)
+            if per_file_limit is None:
+                per_file_limit = DEFAULT_CORE_MEMORY_SOURCE_CHAR_LIMIT
+            if max_files is None:
+                max_files = DEFAULT_MAX_FILES_OPEN
+
             return per_file_limit, max_files
 
     @enforce_types

letta/services/agent_serialization_manager.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timezone
 from typing import Dict, List
 
+from letta.constants import MCP_TOOL_TAG_NAME_PREFIX
 from letta.errors import AgentFileExportError, AgentFileImportError
 from letta.helpers.pinecone_utils import should_use_pinecone
 from letta.log import get_logger
@@ -13,6 +14,7 @@ from letta.schemas.agent_file import (
     FileSchema,
     GroupSchema,
     ImportResult,
+    MCPServerSchema,
     MessageSchema,
     SourceSchema,
     ToolSchema,
@@ -20,6 +22,7 @@ from letta.schemas.agent_file import (
 from letta.schemas.block import Block
 from letta.schemas.enums import FileProcessingStatus
 from letta.schemas.file import FileMetadata
+from letta.schemas.mcp import MCPServer
 from letta.schemas.message import Message
 from letta.schemas.source import Source
 from letta.schemas.tool import Tool
@@ -92,7 +95,7 @@ class AgentSerializationManager:
             ToolSchema.__id_prefix__: 0,
             MessageSchema.__id_prefix__: 0,
             FileAgentSchema.__id_prefix__: 0,
-            # MCPServerSchema.__id_prefix__: 0,
+            MCPServerSchema.__id_prefix__: 0,
         }
 
     def _reset_state(self):
@@ -258,6 +261,53 @@ class AgentSerializationManager:
         file_schema.source_id = self._map_db_to_file_id(file_metadata.source_id, SourceSchema.__id_prefix__, allow_new=False)
         return file_schema
 
+    async def _extract_unique_mcp_servers(self, tools: List, actor: User) -> List:
+        """Extract unique MCP servers from tools based on metadata, using server_id if available, otherwise falling back to server_name."""
+        mcp_server_ids = set()
+        mcp_server_names = set()
+        for tool in tools:
+            # Check if tool has MCP metadata
+            if tool.metadata_ and MCP_TOOL_TAG_NAME_PREFIX in tool.metadata_:
+                mcp_metadata = tool.metadata_[MCP_TOOL_TAG_NAME_PREFIX]
+                # TODO: @jnjpng clean this up once we fully migrate to server_id being the main identifier
+                if "server_id" in mcp_metadata:
+                    mcp_server_ids.add(mcp_metadata["server_id"])
+                elif "server_name" in mcp_metadata:
+                    mcp_server_names.add(mcp_metadata["server_name"])
+
+        # Fetch MCP servers by ID
+        mcp_servers = []
+        fetched_server_ids = set()
+        if mcp_server_ids:
+            try:
+                mcp_servers = await self.mcp_manager.get_mcp_servers_by_ids(list(mcp_server_ids), actor)
+                fetched_server_ids.update([mcp_server.id for mcp_server in mcp_servers])
+            except Exception as e:
+                logger.warning(f"Failed to fetch MCP servers by IDs {mcp_server_ids}: {e}")
+
+        # Fetch MCP servers by name if not already fetched by ID
+        if mcp_server_names:
+            for server_name in mcp_server_names:
+                try:
+                    mcp_server = await self.mcp_manager.get_mcp_server(server_name, actor)
+                    if mcp_server and mcp_server.id not in fetched_server_ids:
+                        mcp_servers.append(mcp_server)
+                except Exception as e:
+                    logger.warning(f"Failed to fetch MCP server by name {server_name}: {e}")
+
+        return mcp_servers
+
+    def _convert_mcp_server_to_schema(self, mcp_server: MCPServer) -> MCPServerSchema:
+        """Convert MCPServer to MCPServerSchema with ID remapping and auth scrubbing"""
+        try:
+            mcp_file_id = self._map_db_to_file_id(mcp_server.id, MCPServerSchema.__id_prefix__, allow_new=False)
+            mcp_schema = MCPServerSchema.from_mcp_server(mcp_server)
+            mcp_schema.id = mcp_file_id
+            return mcp_schema
+        except Exception as e:
+            logger.error(f"Failed to convert MCP server {mcp_server.id}: {e}")
+            raise
+
     async def export(self, agent_ids: List[str], actor: User) -> AgentFileSchema:
         """
         Export agents and their related entities to AgentFileSchema format.
@@ -289,6 +339,13 @@ class AgentSerializationManager:
             tool_set = self._extract_unique_tools(agent_states)
             block_set = self._extract_unique_blocks(agent_states)
 
+            # Extract MCP servers from tools BEFORE conversion (must be done before ID mapping)
+            mcp_server_set = await self._extract_unique_mcp_servers(tool_set, actor)
+
+            # Map MCP server IDs before converting schemas
+            for mcp_server in mcp_server_set:
+                self._map_db_to_file_id(mcp_server.id, MCPServerSchema.__id_prefix__)
+
             # Extract sources and files from agent states BEFORE conversion (with caching)
             source_set, file_set = await self._extract_unique_sources_and_files_from_agents(agent_states, actor, files_agents_cache)
 
@@ -301,6 +358,7 @@ class AgentSerializationManager:
             block_schemas = [self._convert_block_to_schema(block) for block in block_set]
             source_schemas = [self._convert_source_to_schema(source) for source in source_set]
             file_schemas = [self._convert_file_to_schema(file_metadata) for file_metadata in file_set]
+            mcp_server_schemas = [self._convert_mcp_server_to_schema(mcp_server) for mcp_server in mcp_server_set]
 
             logger.info(f"Exporting {len(agent_ids)} agents to agent file format")
 
@@ -312,7 +370,7 @@ class AgentSerializationManager:
                 files=file_schemas,
                 sources=source_schemas,
                 tools=tool_schemas,
-                # mcp_servers=[],  # TODO: Extract and convert MCP servers
+                mcp_servers=mcp_server_schemas,
                 metadata={"revision_id": await get_latest_alembic_revision()},
                 created_at=datetime.now(timezone.utc),
             )
@@ -359,7 +417,20 @@ class AgentSerializationManager:
             # in-memory cache for file metadata to avoid repeated db calls
             file_metadata_cache = {}  # Maps database file ID to FileMetadata
 
-            # 1. Create tools first (no dependencies) - using bulk upsert for efficiency
+            # 1. Create MCP servers first (tools depend on them)
+            if schema.mcp_servers:
+                for mcp_server_schema in schema.mcp_servers:
+                    server_data = mcp_server_schema.model_dump(exclude={"id"})
+                    filtered_server_data = self._filter_dict_for_model(server_data, MCPServer)
+                    create_schema = MCPServer(**filtered_server_data)
+
+                    # Note: We don't have auth info from export, so the user will need to re-configure auth.
+                    # TODO: @jnjpng store metadata about obfuscated metadata to surface to the user
+                    created_mcp_server = await self.mcp_manager.create_or_update_mcp_server(create_schema, actor)
+                    file_to_db_ids[mcp_server_schema.id] = created_mcp_server.id
+                    imported_count += 1
+
+            # 2. Create tools (may depend on MCP servers) - using bulk upsert for efficiency
             if schema.tools:
                 # convert tool schemas to pydantic tools
                 pydantic_tools = []
@@ -559,6 +630,7 @@ class AgentSerializationManager:
             (schema.files, FileSchema.__id_prefix__),
             (schema.sources, SourceSchema.__id_prefix__),
             (schema.tools, ToolSchema.__id_prefix__),
+            (schema.mcp_servers, MCPServerSchema.__id_prefix__),
         ]
 
         for entities, expected_prefix in entity_checks:
@@ -601,6 +673,7 @@ class AgentSerializationManager:
             ("files", schema.files),
             ("sources", schema.sources),
             ("tools", schema.tools),
+            ("mcp_servers", schema.mcp_servers),
         ]
 
         for entity_type, entities in entity_collections:
@@ -705,3 +778,11 @@ class AgentSerializationManager:
             raise AgentFileImportError(f"Schema validation failed: {'; '.join(errors)}")
 
         logger.info("Schema validation passed")
+
+    def _filter_dict_for_model(self, data: dict, model_cls):
+        """Filter a dictionary to only include keys that are in the model fields"""
+        try:
+            allowed = model_cls.model_fields.keys()  # Pydantic v2
+        except AttributeError:
+            allowed = model_cls.__fields__.keys()  # Pydantic v1
+        return {k: v for k, v in data.items() if k in allowed}

letta/services/block_manager.py

@@ -178,6 +178,7 @@ class BlockManager:
         template_name: Optional[str] = None,
         identity_id: Optional[str] = None,
         identifier_keys: Optional[List[str]] = None,
+        project_id: Optional[str] = None,
         before: Optional[str] = None,
         after: Optional[str] = None,
         limit: Optional[int] = 50,
@@ -210,6 +211,9 @@ class BlockManager:
             if template_name:
                 query = query.where(BlockModel.template_name == template_name)
 
+            if project_id:
+                query = query.where(BlockModel.project_id == project_id)
+
             needs_distinct = False
 
             if identifier_keys:

letta/services/file_manager.py

@@ -151,7 +151,8 @@ class FileManager:
         Enforces state transition rules (when enforce_state_transitions=True):
         - PENDING -> PARSING -> EMBEDDING -> COMPLETED (normal flow)
         - Any non-terminal state -> ERROR
-        - ERROR and COMPLETED are terminal (no transitions allowed)
+        - Same-state transitions are allowed (e.g., EMBEDDING -> EMBEDDING)
+        - ERROR and COMPLETED are terminal (no status transitions allowed, metadata updates blocked)
 
         Args:
             file_id: ID of the file to update
@@ -196,22 +197,31 @@ class FileManager:
             ]
 
             # only add state transition validation if enforce_state_transitions is True
-            if enforce_state_transitions:
-                # prevent updates to terminal states (ERROR, COMPLETED)
+            if enforce_state_transitions and processing_status is not None:
+                # enforce specific transitions based on target status
+                if processing_status == FileProcessingStatus.PARSING:
+                    where_conditions.append(
+                        FileMetadataModel.processing_status.in_([FileProcessingStatus.PENDING, FileProcessingStatus.PARSING])
+                    )
+                elif processing_status == FileProcessingStatus.EMBEDDING:
+                    where_conditions.append(
+                        FileMetadataModel.processing_status.in_([FileProcessingStatus.PARSING, FileProcessingStatus.EMBEDDING])
+                    )
+                elif processing_status == FileProcessingStatus.COMPLETED:
+                    where_conditions.append(
+                        FileMetadataModel.processing_status.in_([FileProcessingStatus.EMBEDDING, FileProcessingStatus.COMPLETED])
+                    )
+                elif processing_status == FileProcessingStatus.ERROR:
+                    # ERROR can be set from any non-terminal state
+                    where_conditions.append(
+                        FileMetadataModel.processing_status.notin_([FileProcessingStatus.ERROR, FileProcessingStatus.COMPLETED])
+                    )
+            elif enforce_state_transitions and processing_status is None:
+                # If only updating metadata fields (not status), prevent updates to terminal states
                 where_conditions.append(
                     FileMetadataModel.processing_status.notin_([FileProcessingStatus.ERROR, FileProcessingStatus.COMPLETED])
                 )
 
-                if processing_status is not None:
-                    # enforce specific transitions based on target status
-                    if processing_status == FileProcessingStatus.PARSING:
-                        where_conditions.append(FileMetadataModel.processing_status == FileProcessingStatus.PENDING)
-                    elif processing_status == FileProcessingStatus.EMBEDDING:
-                        where_conditions.append(FileMetadataModel.processing_status == FileProcessingStatus.PARSING)
-                    elif processing_status == FileProcessingStatus.COMPLETED:
-                        where_conditions.append(FileMetadataModel.processing_status == FileProcessingStatus.EMBEDDING)
-                    # ERROR can be set from any non-terminal state (already handled by terminal check above)
-
             # fast in-place update with state validation
             stmt = (
                 update(FileMetadataModel)

letta/services/file_processor/file_processor.py

@@ -69,6 +69,15 @@ class FileProcessor:
                     raise ValueError("No chunks created from text")
                 all_chunks.extend(chunks)
 
+            # Update with chunks length
+            file_metadata = await self.file_manager.update_file_status(
+                file_id=file_metadata.id,
+                actor=self.actor,
+                processing_status=FileProcessingStatus.EMBEDDING,
+                total_chunks=len(all_chunks),
+                chunks_embedded=0,
+            )
+
             all_passages = await self.embedder.generate_embedded_passages(
                 file_id=file_metadata.id,
                 source_id=source_id,
@@ -177,9 +186,7 @@ class FileProcessor:
                 "file_processor.ocr_completed",
                 {"filename": filename, "pages_extracted": len(ocr_response.pages), "text_length": len(raw_markdown_text)},
             )
-            file_metadata = await self.file_manager.update_file_status(
-                file_id=file_metadata.id, actor=self.actor, processing_status=FileProcessingStatus.EMBEDDING
-            )
+
             file_metadata = await self.file_manager.upsert_file_content(file_id=file_metadata.id, text=raw_markdown_text, actor=self.actor)
 
             await self.agent_manager.insert_file_into_context_windows(
@@ -241,13 +248,7 @@ class FileProcessor:
                     file_id=file_metadata.id,
                     actor=self.actor,
                     processing_status=FileProcessingStatus.COMPLETED,
-                )
-            else:
-                await self.file_manager.update_file_status(
-                    file_id=file_metadata.id,
-                    actor=self.actor,
-                    total_chunks=len(all_passages),
-                    chunks_embedded=0,
+                    chunks_embedded=len(all_passages),
                 )
 
             return all_passages
@@ -286,6 +287,7 @@ class FileProcessor:
             document_annotation=None,
         )
 
+    # TODO: The file state machine here is kind of out of date, we need to match with the correct one above
     @trace_method
     async def process_imported_file(self, file_metadata: FileMetadata, source_id: str) -> List[Passage]:
         """Process an imported file that already has content - skip OCR, do chunking/embedding"""