From 86f7e2036df85cfac97289367f5b572d830d6485 Mon Sep 17 00:00:00 2001 From: Charles Packer Date: Thu, 21 Aug 2025 13:42:19 -0700 Subject: [PATCH] fix: patch support for custom headers + auth in MCP connects (was getting a weird silent error) Co-authored-by: jnjpng Co-authored-by: Jin Peng --- letta/functions/mcp_client/types.py | 36 +++++++++++++++++++---- letta/server/rest_api/routers/v1/tools.py | 4 +-- 2 files changed, 32 insertions(+), 8 deletions(-) diff --git a/letta/functions/mcp_client/types.py b/letta/functions/mcp_client/types.py index 20a35996..b1cd27e2 100644 --- a/letta/functions/mcp_client/types.py +++ b/letta/functions/mcp_client/types.py @@ -148,9 +148,21 @@ class SSEServerConfig(BaseServerConfig): custom_headers: Optional[dict[str, str]] = Field(None, description="Custom HTTP headers to include with SSE requests") def resolve_token(self) -> Optional[str]: - if self.auth_token and self.auth_token.startswith(f"{MCP_AUTH_TOKEN_BEARER_PREFIX} "): - return self.auth_token[len(f"{MCP_AUTH_TOKEN_BEARER_PREFIX} ") :] - return self.auth_token + """ + Extract token for storage if auth_header/auth_token are provided + and not already in custom_headers. + + Returns: + The resolved token (without Bearer prefix) if it should be stored separately, None otherwise + """ + if self.auth_token and self.auth_header: + # Check if custom_headers already has the auth header + if not self.custom_headers or self.auth_header not in self.custom_headers: + # Strip Bearer prefix if present + if self.auth_token.startswith(f"{MCP_AUTH_TOKEN_BEARER_PREFIX} "): + return self.auth_token[len(f"{MCP_AUTH_TOKEN_BEARER_PREFIX} ") :] + return self.auth_token + return None def resolve_environment_variables(self, environment_variables: Optional[Dict[str, str]] = None) -> None: if self.auth_token and super().is_templated_tool_variable(self.auth_token): @@ -217,9 +229,21 @@ class StreamableHTTPServerConfig(BaseServerConfig): custom_headers: Optional[dict[str, str]] = Field(None, description="Custom HTTP headers to include with streamable HTTP requests") def resolve_token(self) -> Optional[str]: - if self.auth_token and self.auth_token.startswith(f"{MCP_AUTH_TOKEN_BEARER_PREFIX} "): - return self.auth_token[len(f"{MCP_AUTH_TOKEN_BEARER_PREFIX} ") :] - return self.auth_token + """ + Extract token for storage if auth_header/auth_token are provided + and not already in custom_headers. + + Returns: + The resolved token (without Bearer prefix) if it should be stored separately, None otherwise + """ + if self.auth_token and self.auth_header: + # Check if custom_headers already has the auth header + if not self.custom_headers or self.auth_header not in self.custom_headers: + # Strip Bearer prefix if present + if self.auth_token.startswith(f"{MCP_AUTH_TOKEN_BEARER_PREFIX} "): + return self.auth_token[len(f"{MCP_AUTH_TOKEN_BEARER_PREFIX} ") :] + return self.auth_token + return None def resolve_environment_variables(self, environment_variables: Optional[Dict[str, str]] = None) -> None: if self.auth_token and super().is_templated_tool_variable(self.auth_token): diff --git a/letta/server/rest_api/routers/v1/tools.py b/letta/server/rest_api/routers/v1/tools.py index 34d4bc72..a1a0230f 100644 --- a/letta/server/rest_api/routers/v1/tools.py +++ b/letta/server/rest_api/routers/v1/tools.py @@ -547,7 +547,7 @@ async def add_mcp_server_to_config( server_name=request.server_name, server_type=request.type, server_url=request.server_url, - token=request.resolve_token() if not request.custom_headers else None, + token=request.resolve_token(), custom_headers=request.custom_headers, ) elif isinstance(request, StreamableHTTPServerConfig): @@ -555,7 +555,7 @@ async def add_mcp_server_to_config( server_name=request.server_name, server_type=request.type, server_url=request.server_url, - token=request.resolve_token() if not request.custom_headers else None, + token=request.resolve_token(), custom_headers=request.custom_headers, )