feat: headers support

This commit is contained in:
Timothy Jaeryang Baek 2025-06-04 19:57:13 +04:00
parent a33ab53331
commit dc4c4abe08
3 changed files with 39 additions and 4 deletions

View File

@ -46,6 +46,12 @@ To use an SSE-compatible MCP server, simply specify the server type and endpoint
mcpo --port 8000 --api-key "top-secret" --server-type "sse" -- http://127.0.0.1:8001/sse
```
You can also provide headers for the SSE connection:
```bash
mcpo --port 8000 --api-key "top-secret" --server-type "sse" --headers '{"Authorization": "Bearer token", "X-Custom-Header": "value"}' -- http://127.0.0.1:8001/sse
```
To use a Streamable HTTP-compatible MCP server, specify the server type and endpoint:
```bash
@ -93,7 +99,11 @@ Example config.json:
},
"mcp_sse": {
"type": "sse", // Explicitly define type
"url": "http://127.0.0.1:8001/sse"
"url": "http://127.0.0.1:8001/sse",
"headers": {
"Authorization": "Bearer token",
"X-Custom-Header": "value"
}
},
"mcp_streamable_http": {
"type": "streamable_http",

View File

@ -28,7 +28,9 @@ def main(
] = None,
strict_auth: Annotated[
Optional[bool],
typer.Option("--strict-auth", help="API key protects all endpoints and documentation"),
typer.Option(
"--strict-auth", help="API key protects all endpoints and documentation"
),
] = False,
env: Annotated[
Optional[List[str]], typer.Option("--env", "-e", help="Environment variables")
@ -61,6 +63,9 @@ def main(
path_prefix: Annotated[
Optional[str], typer.Option("--path-prefix", help="URL prefix")
] = None,
headers: Annotated[
Optional[str], typer.Option("--header", "-H", help="Headers in JSON format")
] = None,
):
server_command = None
if not config_path:
@ -131,6 +136,7 @@ def main(
ssl_certfile=ssl_certfile,
ssl_keyfile=ssl_keyfile,
path_prefix=path_prefix,
headers=headers,
)
)

View File

@ -116,7 +116,10 @@ async def lifespan(app: FastAPI):
await create_dynamic_endpoints(app, api_dependency=api_dependency)
yield
if server_type == "sse":
async with sse_client(url=args[0], sse_read_timeout=None) as (
headers = getattr(app.state, "headers", None)
async with sse_client(
url=args[0], sse_read_timeout=None, headers=headers
) as (
reader,
writer,
):
@ -125,13 +128,15 @@ async def lifespan(app: FastAPI):
await create_dynamic_endpoints(app, api_dependency=api_dependency)
yield
if server_type == "streamablehttp" or server_type == "streamable_http":
headers = getattr(app.state, "headers", None)
# Ensure URL has trailing slash to avoid redirects
url = args[0]
if not url.endswith("/"):
url = f"{url}/"
# Connect using streamablehttp_client from the SDK, similar to sse_client
async with streamablehttp_client(url=url) as (
async with streamablehttp_client(url=url, headers=headers) as (
reader,
writer,
_, # get_session_id callback not needed for ClientSession
@ -212,6 +217,14 @@ async def run(
if api_key and strict_auth:
main_app.add_middleware(APIKeyMiddleware, api_key=api_key)
headers = kwargs.get("headers")
if headers and isinstance(headers, str):
try:
headers = json.loads(headers)
except json.JSONDecodeError:
print("Warning: Invalid JSON format for headers. Headers will be ignored.")
headers = None
if server_type == "sse":
logger.info(
f"Configuring for a single SSE MCP Server with URL {server_command[0]}"
@ -219,6 +232,7 @@ async def run(
main_app.state.server_type = "sse"
main_app.state.args = server_command[0] # Expects URL as the first element
main_app.state.api_dependency = api_dependency
main_app.state.headers = headers
elif server_type == "streamablehttp" or server_type == "streamable_http":
logger.info(
f"Configuring for a single StreamableHTTP MCP Server with URL {server_command[0]}"
@ -226,6 +240,7 @@ async def run(
main_app.state.server_type = "streamablehttp"
main_app.state.args = server_command[0] # Expects URL as the first element
main_app.state.api_dependency = api_dependency
main_app.state.headers = headers
elif server_command: # This handles stdio
logger.info(
f"Configuring for a single Stdio MCP Server with command: {' '.join(server_command)}"
@ -306,6 +321,7 @@ async def run(
if server_config_type == "sse" and server_cfg.get("url"):
sub_app.state.server_type = "sse"
sub_app.state.args = server_cfg["url"]
sub_app.state.headers = server_cfg.get("headers")
elif (
server_config_type == "streamablehttp"
or server_config_type == "streamable_http"
@ -316,11 +332,14 @@ async def run(
url = f"{url}/"
sub_app.state.server_type = "streamablehttp"
sub_app.state.args = url
sub_app.state.headers = server_cfg.get("headers")
elif not server_config_type and server_cfg.get(
"url"
): # Fallback for old SSE config
sub_app.state.server_type = "sse"
sub_app.state.args = server_cfg["url"]
sub_app.state.headers = server_cfg.get("headers")
# Add middleware to protect also documentation and spec
if api_key and strict_auth: