letta-nightly 0.9.1.dev20250731104458__py3-none-any.whl → 0.10.0.dev20250801010504__py3-none-any.whl

letta/schemas/sandbox_config.py

@@ -1,21 +1,17 @@
 import hashlib
 import json
-from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Union
 
 from pydantic import BaseModel, Field, model_validator
 
 from letta.constants import LETTA_TOOL_EXECUTION_DIR
 from letta.schemas.agent import AgentState
+from letta.schemas.enums import SandboxType
 from letta.schemas.letta_base import LettaBase, OrmMetadataBase
 from letta.schemas.pip_requirement import PipRequirement
 from letta.settings import tool_settings
 
-
 # Sandbox Config
-class SandboxType(str, Enum):
-    E2B = "e2b"
-    LOCAL = "local"
 
 
 class SandboxRunResult(BaseModel):
@@ -83,6 +79,15 @@ class E2BSandboxConfig(BaseModel):
         return data
 
 
+class ModalSandboxConfig(BaseModel):
+    timeout: int = Field(5 * 60, description="Time limit for the sandbox (in seconds).")
+    pip_requirements: Optional[List[str]] = Field(None, description="A list of pip packages to install in the Modal sandbox")
+
+    @property
+    def type(self) -> "SandboxType":
+        return SandboxType.MODAL
+
+
 class SandboxConfigBase(OrmMetadataBase):
     __id_prefix__ = "sandbox"
 
@@ -99,6 +104,9 @@ class SandboxConfig(SandboxConfigBase):
     def get_local_config(self) -> LocalSandboxConfig:
         return LocalSandboxConfig(**self.config)
 
+    def get_modal_config(self) -> ModalSandboxConfig:
+        return ModalSandboxConfig(**self.config)
+
     def fingerprint(self) -> str:
         # Only take into account type, org_id, and the config items
         # Canonicalize input data into JSON with sorted keys
@@ -120,10 +128,12 @@ class SandboxConfig(SandboxConfigBase):
 
 
 class SandboxConfigCreate(LettaBase):
-    config: Union[LocalSandboxConfig, E2BSandboxConfig] = Field(..., description="The configuration for the sandbox.")
+    config: Union[LocalSandboxConfig, E2BSandboxConfig, ModalSandboxConfig] = Field(..., description="The configuration for the sandbox.")
 
 
 class SandboxConfigUpdate(LettaBase):
     """Pydantic model for updating SandboxConfig fields."""
 
-    config: Union[LocalSandboxConfig, E2BSandboxConfig] = Field(None, description="The JSON configuration data for the sandbox.")
+    config: Union[LocalSandboxConfig, E2BSandboxConfig, ModalSandboxConfig] = Field(
+        None, description="The JSON configuration data for the sandbox."
+    )

letta/server/rest_api/app.py

@@ -28,6 +28,7 @@ from letta.schemas.letta_message_content import (
     create_letta_message_content_union_schema,
     create_letta_user_message_content_union_schema,
 )
+from letta.schemas.letta_ping import create_letta_ping_schema
 from letta.server.constants import REST_DEFAULT_PORT
 from letta.server.db import db_registry
 
@@ -67,6 +68,7 @@ def generate_openapi_schema(app: FastAPI):
     letta_docs["components"]["schemas"]["LettaMessageContentUnion"] = create_letta_message_content_union_schema()
     letta_docs["components"]["schemas"]["LettaAssistantMessageContentUnion"] = create_letta_assistant_message_content_union_schema()
     letta_docs["components"]["schemas"]["LettaUserMessageContentUnion"] = create_letta_user_message_content_union_schema()
+    letta_docs["components"]["schemas"]["LettaPing"] = create_letta_ping_schema()
 
     # Update the app's schema with our modified version
     app.openapi_schema = letta_docs

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

@@ -41,7 +41,7 @@ from letta.server.server import SyncServer
 from letta.services.summarizer.enums import SummarizationMode
 from letta.services.telemetry_manager import NoopTelemetryManager
 from letta.settings import settings
-from letta.utils import safe_create_task
+from letta.utils import safe_create_task, truncate_file_visible_content
 
 # These can be forward refs, but because Fastapi needs them at runtime the must be imported normally
 
@@ -65,7 +65,7 @@ async def list_agents(
     after: str | None = Query(None, description="Cursor for pagination"),
     limit: int | None = Query(50, description="Limit for pagination"),
     query_text: str | None = Query(None, description="Search agents by name"),
-    project_id: str | None = Query(None, description="Search agents by project ID"),
+    project_id: str | None = Query(None, description="Search agents by project ID - this will default to your default project on cloud"),
     template_id: str | None = Query(None, description="Search agents by template ID"),
     base_template_id: str | None = Query(None, description="Search agents by base template ID"),
     identity_id: str | None = Query(None, description="Search agents by identity ID"),
@@ -86,6 +86,11 @@ async def list_agents(
         "created_at",
         description="Field to sort by. Options: 'created_at' (default), 'last_run_completion'",
     ),
+    show_hidden_agents: bool | None = Query(
+        False,
+        include_in_schema=False,
+        description="If set to True, include agents marked as hidden in the results.",
+    ),
 ):
     """
     List all agents associated with a given user.
@@ -115,6 +120,7 @@ async def list_agents(
         include_relationships=include_relationships,
         ascending=ascending,
         sort_by=sort_by,
+        show_hidden_agents=show_hidden_agents,
     )
 
 
@@ -478,14 +484,23 @@ async def open_file(
     if not file_metadata:
         raise HTTPException(status_code=404, detail=f"File with id={file_id} not found")
 
+    # Process file content with line numbers using LineChunker
+    from letta.services.file_processor.chunker.line_chunker import LineChunker
+
+    content_lines = LineChunker().chunk_text(file_metadata=file_metadata, validate_range=False)
+    visible_content = "\n".join(content_lines)
+
+    # Truncate if needed
+    visible_content = truncate_file_visible_content(visible_content, True, per_file_view_window_char_limit)
+
     # Use enforce_max_open_files_and_open for efficient LRU handling
-    closed_files, was_already_open = await server.file_agent_manager.enforce_max_open_files_and_open(
+    closed_files, was_already_open, _ = await server.file_agent_manager.enforce_max_open_files_and_open(
         agent_id=agent_id,
         file_id=file_id,
         file_name=file_metadata.file_name,
         source_id=file_metadata.source_id,
         actor=actor,
-        visible_content=file_metadata.content[:per_file_view_window_char_limit] if file_metadata.content else "",
+        visible_content=visible_content,
         max_files_open=max_files_open,
     )
 
@@ -850,7 +865,15 @@ async def send_message(
     # TODO: This is redundant, remove soon
     agent = await server.agent_manager.get_agent_by_id_async(agent_id, actor, include_relationships=["multi_agent_group"])
     agent_eligible = agent.multi_agent_group is None or agent.multi_agent_group.manager_type in ["sleeptime", "voice_sleeptime"]
-    model_compatible = agent.llm_config.model_endpoint_type in ["anthropic", "openai", "together", "google_ai", "google_vertex", "bedrock"]
+    model_compatible = agent.llm_config.model_endpoint_type in [
+        "anthropic",
+        "openai",
+        "together",
+        "google_ai",
+        "google_vertex",
+        "bedrock",
+        "ollama",
+    ]
 
     # Create a new run for execution tracking
     if settings.track_agent_run:
@@ -984,7 +1007,15 @@ async def send_message_streaming(
     # TODO: This is redundant, remove soon
     agent = await server.agent_manager.get_agent_by_id_async(agent_id, actor, include_relationships=["multi_agent_group"])
     agent_eligible = agent.multi_agent_group is None or agent.multi_agent_group.manager_type in ["sleeptime", "voice_sleeptime"]
-    model_compatible = agent.llm_config.model_endpoint_type in ["anthropic", "openai", "together", "google_ai", "google_vertex", "bedrock"]
+    model_compatible = agent.llm_config.model_endpoint_type in [
+        "anthropic",
+        "openai",
+        "together",
+        "google_ai",
+        "google_vertex",
+        "bedrock",
+        "ollama",
+    ]
     model_compatible_token_streaming = agent.llm_config.model_endpoint_type in ["anthropic", "openai", "bedrock"]
     not_letta_endpoint = agent.llm_config.model_endpoint != LETTA_MODEL_ENDPOINT
 
@@ -1052,28 +1083,42 @@ async def send_message_streaming(
                         else SummarizationMode.PARTIAL_EVICT_MESSAGE_BUFFER
                     ),
                 )
-            from letta.server.rest_api.streaming_response import StreamingResponseWithStatusCode
+            from letta.server.rest_api.streaming_response import StreamingResponseWithStatusCode, add_keepalive_to_stream
 
             if request.stream_tokens and model_compatible_token_streaming and not_letta_endpoint:
+                raw_stream = agent_loop.step_stream(
+                    input_messages=request.messages,
+                    max_steps=request.max_steps,
+                    use_assistant_message=request.use_assistant_message,
+                    request_start_timestamp_ns=request_start_timestamp_ns,
+                    include_return_message_types=request.include_return_message_types,
+                )
+                # Conditionally wrap with keepalive based on request parameter
+                if request.include_pings and settings.enable_keepalive:
+                    stream = add_keepalive_to_stream(raw_stream, keepalive_interval=settings.keepalive_interval)
+                else:
+                    stream = raw_stream
+
                 result = StreamingResponseWithStatusCode(
-                    agent_loop.step_stream(
-                        input_messages=request.messages,
-                        max_steps=request.max_steps,
-                        use_assistant_message=request.use_assistant_message,
-                        request_start_timestamp_ns=request_start_timestamp_ns,
-                        include_return_message_types=request.include_return_message_types,
-                    ),
+                    stream,
                     media_type="text/event-stream",
                 )
             else:
+                raw_stream = agent_loop.step_stream_no_tokens(
+                    request.messages,
+                    max_steps=request.max_steps,
+                    use_assistant_message=request.use_assistant_message,
+                    request_start_timestamp_ns=request_start_timestamp_ns,
+                    include_return_message_types=request.include_return_message_types,
+                )
+                # Conditionally wrap with keepalive based on request parameter
+                if request.include_pings and settings.enable_keepalive:
+                    stream = add_keepalive_to_stream(raw_stream, keepalive_interval=settings.keepalive_interval)
+                else:
+                    stream = raw_stream
+
                 result = StreamingResponseWithStatusCode(
-                    agent_loop.step_stream_no_tokens(
-                        request.messages,
-                        max_steps=request.max_steps,
-                        use_assistant_message=request.use_assistant_message,
-                        request_start_timestamp_ns=request_start_timestamp_ns,
-                        include_return_message_types=request.include_return_message_types,
-                    ),
+                    stream,
                     media_type="text/event-stream",
                 )
         else:
@@ -1165,6 +1210,7 @@ async def _process_message_background(
             "google_ai",
             "google_vertex",
             "bedrock",
+            "ollama",
         ]
         if agent_eligible and model_compatible:
             if agent.enable_sleeptime and agent.agent_type != AgentType.voice_convo_agent:
@@ -1344,7 +1390,15 @@ async def preview_raw_payload(
     actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
     agent = await server.agent_manager.get_agent_by_id_async(agent_id, actor, include_relationships=["multi_agent_group"])
     agent_eligible = agent.multi_agent_group is None or agent.multi_agent_group.manager_type in ["sleeptime", "voice_sleeptime"]
-    model_compatible = agent.llm_config.model_endpoint_type in ["anthropic", "openai", "together", "google_ai", "google_vertex", "bedrock"]
+    model_compatible = agent.llm_config.model_endpoint_type in [
+        "anthropic",
+        "openai",
+        "together",
+        "google_ai",
+        "google_vertex",
+        "bedrock",
+        "ollama",
+    ]
 
     if agent_eligible and model_compatible:
         if agent.enable_sleeptime:
@@ -1386,7 +1440,7 @@ async def preview_raw_payload(
         )
 
 
-@router.post("/{agent_id}/summarize", response_model=AgentState, operation_id="summarize_agent_conversation")
+@router.post("/{agent_id}/summarize", status_code=204, operation_id="summarize_agent_conversation")
 async def summarize_agent_conversation(
     agent_id: str,
     request_obj: Request,  # FastAPI Request
@@ -1404,7 +1458,15 @@ async def summarize_agent_conversation(
     actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
     agent = await server.agent_manager.get_agent_by_id_async(agent_id, actor, include_relationships=["multi_agent_group"])
     agent_eligible = agent.multi_agent_group is None or agent.multi_agent_group.manager_type in ["sleeptime", "voice_sleeptime"]
-    model_compatible = agent.llm_config.model_endpoint_type in ["anthropic", "openai", "together", "google_ai", "google_vertex", "bedrock"]
+    model_compatible = agent.llm_config.model_endpoint_type in [
+        "anthropic",
+        "openai",
+        "together",
+        "google_ai",
+        "google_vertex",
+        "bedrock",
+        "ollama",
+    ]
 
     if agent_eligible and model_compatible:
         agent = LettaAgent(
@@ -1419,9 +1481,10 @@ async def summarize_agent_conversation(
             telemetry_manager=server.telemetry_manager if settings.llm_api_logging else NoopTelemetryManager(),
             message_buffer_min=max_message_length,
         )
-        return await agent.summarize_conversation_history()
-
-    raise HTTPException(
-        status_code=status.HTTP_403_FORBIDDEN,
-        detail="Summarization is not currently supported for this agent configuration. Please contact Letta support.",
-    )
+        await agent.summarize_conversation_history()
+        # Summarization completed, return 204 No Content
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Summarization is not currently supported for this agent configuration. Please contact Letta support.",
+        )

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

@@ -24,6 +24,50 @@ async def list_blocks(
     identifier_keys: Optional[List[str]] = Query(None, description="Search agents by identifier keys"),
     project_id: Optional[str] = Query(None, description="Search blocks by project id"),
     limit: Optional[int] = Query(50, description="Number of blocks to return"),
+    before: Optional[str] = Query(
+        None,
+        description="Cursor for pagination. If provided, returns blocks before this cursor.",
+    ),
+    after: Optional[str] = Query(
+        None,
+        description="Cursor for pagination. If provided, returns blocks after this cursor.",
+    ),
+    label_search: Optional[str] = Query(
+        None,
+        description=("Search blocks by label. If provided, returns blocks that match this label. " "This is a full-text search on labels."),
+    ),
+    description_search: Optional[str] = Query(
+        None,
+        description=(
+            "Search blocks by description. If provided, returns blocks that match this description. "
+            "This is a full-text search on block descriptions."
+        ),
+    ),
+    value_search: Optional[str] = Query(
+        None,
+        description=("Search blocks by value. If provided, returns blocks that match this value."),
+    ),
+    connected_to_agents_count_gt: Optional[int] = Query(
+        None,
+        description=(
+            "Filter blocks by the number of connected agents. "
+            "If provided, returns blocks that have more than this number of connected agents."
+        ),
+    ),
+    connected_to_agents_count_lt: Optional[int] = Query(
+        None,
+        description=(
+            "Filter blocks by the number of connected agents. "
+            "If provided, returns blocks that have less than this number of connected agents."
+        ),
+    ),
+    connected_to_agents_count_eq: Optional[List[int]] = Query(
+        None,
+        description=(
+            "Filter blocks by the exact number of connected agents. "
+            "If provided, returns blocks that have exactly this number of connected agents."
+        ),
+    ),
     server: SyncServer = Depends(get_letta_server),
     actor_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
 ):
@@ -32,11 +76,19 @@ async def list_blocks(
         actor=actor,
         label=label,
         is_template=templates_only,
+        value_search=value_search,
+        label_search=label_search,
+        description_search=description_search,
         template_name=name,
         identity_id=identity_id,
         identifier_keys=identifier_keys,
         project_id=project_id,
+        before=before,
+        connected_to_agents_count_gt=connected_to_agents_count_gt,
+        connected_to_agents_count_lt=connected_to_agents_count_lt,
+        connected_to_agents_count_eq=connected_to_agents_count_eq,
         limit=limit,
+        after=after,
     )
 
 

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

@@ -5,11 +5,12 @@ from typing import List, Optional
 from fastapi import APIRouter, Depends, HTTPException, Query
 
 from letta.log import get_logger
+from letta.schemas.enums import SandboxType
 from letta.schemas.environment_variables import SandboxEnvironmentVariable as PydanticEnvVar
 from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate, SandboxEnvironmentVariableUpdate
 from letta.schemas.sandbox_config import LocalSandboxConfig
 from letta.schemas.sandbox_config import SandboxConfig as PydanticSandboxConfig
-from letta.schemas.sandbox_config import SandboxConfigCreate, SandboxConfigUpdate, SandboxType
+from letta.schemas.sandbox_config import SandboxConfigCreate, SandboxConfigUpdate
 from letta.server.rest_api.utils import get_letta_server, get_user_id
 from letta.server.server import SyncServer
 from letta.services.helpers.tool_execution_helper import create_venv_for_local_sandbox, install_pip_requirements_for_sandbox

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

@@ -1,4 +1,3 @@
-import asyncio
 import json
 from collections.abc import AsyncGenerator
 from typing import Any, Dict, List, Optional, Union
@@ -12,12 +11,12 @@ from composio.exceptions import (
     EnumMetadataNotFound,
     EnumStringNotFound,
 )
-from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query
-from fastapi.responses import HTMLResponse
+from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query, Request
+from httpx import HTTPStatusError
 from pydantic import BaseModel, Field
 from starlette.responses import StreamingResponse
 
-from letta.errors import LettaToolCreateError
+from letta.errors import LettaToolCreateError, LettaToolNameConflictError
 from letta.functions.functions import derive_openai_json_schema
 from letta.functions.mcp_client.exceptions import MCPTimeoutError
 from letta.functions.mcp_client.types import MCPTool, SSEServerConfig, StdioServerConfig, StreamableHTTPServerConfig
@@ -27,22 +26,18 @@ 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.prompts.gpt_system import get_system_text
 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 MCPOAuthSessionCreate, UpdateSSEMCPServer, UpdateStdioMCPServer, UpdateStreamableHTTPMCPServer
+from letta.schemas.mcp import UpdateSSEMCPServer, UpdateStdioMCPServer, UpdateStreamableHTTPMCPServer
 from letta.schemas.message import Message
+from letta.schemas.pip_requirement import PipRequirement
 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.oauth_utils import (
-    MCPOAuthSession,
-    create_oauth_provider,
-    drill_down_exception,
-    get_oauth_success_html,
-    oauth_stream_event,
-)
+from letta.services.mcp.oauth_utils import MCPOAuthSession, drill_down_exception, oauth_stream_event
 from letta.services.mcp.stdio_client import AsyncStdioMCPClient
 from letta.services.mcp.types import OauthStreamEvent
 from letta.settings import tool_settings
@@ -196,6 +191,10 @@ async def modify_tool(
     try:
         actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
         return await server.tool_manager.update_tool_by_id_async(tool_id=tool_id, tool_update=request, actor=actor)
+    except LettaToolNameConflictError as e:
+        # HTTP 409 == Conflict
+        print(f"Tool name conflict during update: {e}")
+        raise HTTPException(status_code=409, detail=str(e))
     except LettaToolCreateError as e:
         # HTTP 400 == Bad Request
         print(f"Error occurred during tool update: {e}")
@@ -394,7 +393,7 @@ async def list_mcp_servers(server: SyncServer = Depends(get_letta_server), user_
     else:
         actor = await server.user_manager.get_actor_or_default_async(actor_id=user_id)
         mcp_servers = await server.mcp_manager.list_mcp_servers(actor=actor)
-        return {server.server_name: server.to_config() for server in mcp_servers}
+        return {server.server_name: server.to_config(resolve_variables=False) for server in mcp_servers}
 
 
 # NOTE: async because the MCP client/session calls are async
@@ -634,11 +633,12 @@ async def test_mcp_server(
 ):
     """
     Test connection to an MCP server without adding it.
-    Returns the list of available tools if successful, or OAuth information if OAuth is required.
+    Returns the list of available tools if successful.
     """
     client = None
     try:
         actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
+        request.resolve_environment_variables()
         client = await server.mcp_manager.get_mcp_client(request, actor)
 
         await client.connect_to_server()
@@ -697,6 +697,7 @@ 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"),
+    http_request: Request = None,
 ) -> StreamingResponse:
     """
     Connect to an MCP server with support for OAuth via SSE.
@@ -705,12 +706,11 @@ async def connect_mcp_server(
 
     async def oauth_stream_generator(
         request: Union[StdioServerConfig, SSEServerConfig, StreamableHTTPServerConfig],
+        http_request: Request,
     ) -> AsyncGenerator[str, None]:
         client = None
-        oauth_provider = None
-        temp_client = None
-        connect_task = None
 
+        oauth_flow_attempted = False
         try:
             # Acknolwedge connection attempt
             yield oauth_stream_event(OauthStreamEvent.CONNECTION_ATTEMPT, server_name=request.server_name)
@@ -719,6 +719,7 @@ async def connect_mcp_server(
 
             # Create MCP client with respective transport type
             try:
+                request.resolve_environment_variables()
                 client = await server.mcp_manager.get_mcp_client(request, actor)
             except ValueError as e:
                 yield oauth_stream_event(OauthStreamEvent.ERROR, message=str(e))
@@ -741,97 +742,35 @@ async def connect_mcp_server(
             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:
+            finally:
+                if 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)
+                        await client.cleanup()
+                    # This is a workaround to catch the expected 401 Unauthorized from the official MCP SDK, see their streamable_http.py
+                    # For SSE transport types, we catch the ConnectionError above, but Streamable HTTP doesn't bubble up the exception
+                    except* HTTPStatusError:
+                        oauth_flow_attempted = True
+                        async for event in server.mcp_manager.handle_oauth_flow(request=request, actor=actor, http_request=http_request):
+                            yield event
+
+            # Failsafe to make sure we don't try to handle OAuth flow twice
+            if not oauth_flow_attempted:
+                async for event in server.mcp_manager.handle_oauth_flow(request=request, actor=actor, http_request=http_request):
+                    yield event
             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}")
+                    logger.warning(f"Error during temp MCP client cleanup: {cleanup_error}")
 
-    return StreamingResponseWithStatusCode(oauth_stream_generator(request), media_type="text/event-stream")
+    return StreamingResponseWithStatusCode(oauth_stream_generator(request, http_request), media_type="text/event-stream")
 
 
 class CodeInput(BaseModel):
@@ -856,7 +795,7 @@ async def generate_json_schema(
 
 
 # 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)
+@router.get("/mcp/oauth/callback/{session_id}", operation_id="mcp_oauth_callback")
 async def mcp_oauth_callback(
     session_id: str,
     code: Optional[str] = Query(None, description="OAuth authorization code"),
@@ -869,7 +808,6 @@ async def mcp_oauth_callback(
     """
     try:
         oauth_session = MCPOAuthSession(session_id)
-
         if error:
             error_msg = f"OAuth error: {error}"
             if error_description:
@@ -887,7 +825,7 @@ async def mcp_oauth_callback(
             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)
+        return {"status": "success", "message": "Authorization successful", "server_url": success.server_url}
 
     except Exception as e:
         logger.error(f"OAuth callback error: {e}")
@@ -932,9 +870,11 @@ async def generate_tool_from_prompt(
         )
         assert llm_client is not None
 
+        assistant_message_ack = "Understood, I will respond with generated python source code and sample arguments that can be used to test the functionality once I receive the user prompt. I'm ready."
+
         input_messages = [
-            Message(role=MessageRole.system, content=[TextContent(text="Placeholder system message")]),
-            Message(role=MessageRole.assistant, content=[TextContent(text="Placeholder assistant message")]),
+            Message(role=MessageRole.system, content=[TextContent(text=get_system_text("memgpt_generate_tool"))]),
+            Message(role=MessageRole.assistant, content=[TextContent(text=assistant_message_ack)]),
             Message(role=MessageRole.user, content=[TextContent(text=formatted_prompt)]),
         ]
 
@@ -965,11 +905,13 @@ async def generate_tool_from_prompt(
         response_data = await llm_client.request_async(request_data, llm_config)
         response = llm_client.convert_response_to_chat_completion(response_data, input_messages, llm_config)
         output = json.loads(response.choices[0].message.tool_calls[0].function.arguments)
+        pip_requirements = [PipRequirement(name=k, version=v or None) for k, v in json.loads(output["pip_requirements_json"]).items()]
         return GenerateToolOutput(
             tool=Tool(
                 name=request.tool_name,
                 source_type="python",
                 source_code=output["raw_source_code"],
+                pip_requirements=pip_requirements,
             ),
             sample_args=json.loads(output["sample_args_json"]),
             response=response.choices[0].message.content,