Skip to content

Commit 7a0dd69

Browse files
R-RudolfR-Rudolf
committed
working-sse-with-auth
1 parent ec193be commit 7a0dd69

1 file changed

Lines changed: 82 additions & 32 deletions

File tree

‎spendee/spendee_mcp.py‎

Lines changed: 82 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -41,51 +41,101 @@ def get_wallets():
4141
]
4242

4343

44-
# def server_with_authentication():
45-
# session_manager = StreamableHTTPSessionManager(
46-
# app=mcp._mcp_server, # Use the underlying MCPServer instance
47-
# )
44+
def server_with_authentication():
45+
session_manager = StreamableHTTPSessionManager(
46+
app=mcp._mcp_server, # Use the underlying MCPServer instance
47+
)
4848

49-
# async def auth_middleware(scope, receive, send):
50-
# request = Request(scope, receive)
51-
# # Log request method, url, and headers for troubleshooting
52-
# if DEBUG_MODE:
53-
# logger.debug(f"Incoming request: method={request.method}, url={request.url}")
54-
# logger.debug(f"Request headers: {dict(request.headers)}")
49+
async def auth_middleware(scope, receive, send):
50+
request = Request(scope, receive)
51+
# Log request method, url, and headers for troubleshooting
52+
if DEBUG_MODE:
53+
logger.debug(f"Incoming request: method={request.method}, url={request.url}")
54+
logger.debug(f"Request headers: {dict(request.headers)}")
5555

56-
# auth_header = request.headers.get("authorization")
56+
auth_header = request.headers.get("authorization")
5757

58-
# if not auth_header or not auth_header.lower().startswith("bearer "):
59-
# logger.warning("Missing or invalid Authorization header.")
60-
# raise HTTPException(401, "Missing or invalid Authorization header.")
58+
if not auth_header or not auth_header.lower().startswith("bearer "):
59+
logger.warning("Missing or invalid Authorization header.")
60+
raise HTTPException(401, "Missing or invalid Authorization header.")
6161

62-
# token = auth_header.split(" ", 1)[1]
63-
# if token != ACCEPTED_TOKEN:
64-
# logger.warning("Invalid or expired token.")
65-
# raise HTTPException(401, "Invalid or expired token.")
62+
token = auth_header.split(" ", 1)[1]
63+
if token != ACCEPTED_TOKEN:
64+
logger.warning("Invalid or expired token.")
65+
raise HTTPException(401, "Invalid or expired token.")
6666

67-
# await session_manager.handle_request(scope, receive, send)
67+
await session_manager.handle_request(scope, receive, send)
6868

69-
# @contextlib.asynccontextmanager
70-
# async def lifespan(app: FastAPI):
71-
# async with session_manager.run():
72-
# logger.info("StreamableHTTPSessionManager started.")
73-
# yield
74-
# logger.info("StreamableHTTPSessionManager stopped.")
69+
@contextlib.asynccontextmanager
70+
async def lifespan(app: FastAPI):
71+
async with session_manager.run():
72+
logger.info("StreamableHTTPSessionManager started.")
73+
yield
74+
logger.info("StreamableHTTPSessionManager stopped.")
7575

76-
# app = FastAPI(lifespan=lifespan)
77-
# app.mount("/mcp", auth_middleware)
76+
app = FastAPI(lifespan=lifespan)
77+
app.mount("/mcp", auth_middleware)
7878

79-
# logger.info(f"Starting Spendee MCP Server with Bearer Token Authentication on port {PORT}")
80-
# logger.info(f"Access the MCP endpoint at http://0.0.0.0:{PORT}/mcp")
81-
# uvicorn.run(app, host="0.0.0.0", port=PORT)
79+
logger.info(f"Starting Spendee MCP Server with Bearer Token Authentication on port {PORT}")
80+
logger.info(f"Access the MCP endpoint at http://0.0.0.0:{PORT}/mcp")
81+
uvicorn.run(app, host="0.0.0.0", port=PORT)
8282

8383
def server_with_sse():
8484

85-
# Create the Starlette app and mount the MCP servers
85+
# Persistent SSE transport instance
86+
sse_transport = mcp._sse_transport if hasattr(mcp, "_sse_transport") else None
87+
if not sse_transport:
88+
# Create and cache the transport instance
89+
from mcp.server.sse import SseServerTransport
90+
sse_transport = SseServerTransport("/messages/")
91+
mcp._sse_transport = sse_transport
92+
93+
from starlette.requests import Request
94+
from starlette.responses import Response
95+
from starlette.types import Scope, Receive, Send
96+
97+
async def sse_auth_middleware(scope: Scope, receive: Receive, send: Send):
98+
request = Request(scope, receive)
99+
if DEBUG_MODE:
100+
logger.debug(f"Incoming request: method={request.method}, url={request.url}")
101+
logger.debug(f"Request headers: {dict(request.headers)}")
102+
103+
auth_header = request.headers.get("authorization")
104+
if not auth_header or not auth_header.lower().startswith("bearer "):
105+
logger.warning("Missing or invalid Authorization header.")
106+
response = Response("Missing or invalid Authorization header.", status_code=401)
107+
return await response(scope, receive, send)
108+
109+
token = auth_header.split(" ", 1)[1]
110+
if token != ACCEPTED_TOKEN:
111+
logger.warning("Invalid or expired token.")
112+
response = Response("Invalid or expired token.", status_code=401)
113+
return await response(scope, receive, send)
114+
115+
# If auth passes, route to the persistent SSE transport
116+
# Mount /sse and /messages/ endpoints
117+
from starlette.routing import Route, Mount
118+
from starlette.applications import Starlette
119+
120+
# Only create the app once
121+
if not hasattr(mcp, "_starlette_sse_app"):
122+
async def handle_sse(request):
123+
async with sse_transport.connect_sse(request.scope, request.receive, request._send) as streams:
124+
await mcp._mcp_server.run(streams[0], streams[1], mcp._mcp_server.create_initialization_options())
125+
return Response()
126+
127+
routes = [
128+
Route("/sse", endpoint=handle_sse, methods=["GET"]),
129+
Mount("/messages/", app=sse_transport.handle_post_message),
130+
]
131+
mcp._starlette_sse_app = Starlette(routes=routes)
132+
133+
await mcp._starlette_sse_app(scope, receive, send)
134+
135+
# Mount the auth middleware at root
86136
app = Starlette(
87137
routes=[
88-
Mount("/", mcp.sse_app()),
138+
Mount("/", sse_auth_middleware),
89139
],
90140
)
91141
uvicorn.run(app, host="0.0.0.0", port=PORT)

0 commit comments

Comments
 (0)