Connect to Litmus over MCP

from fastapi import FastAPI, Request
from starlette.middleware.sessions import SessionMiddleware
from authlib.integrations.starlette_client import OAuth

app = FastAPI()
# Required — gives requests a signed-cookie session, which is where
# Authlib stashes state + the PKCE verifier between this route and the
# callback. SESSION_SECRET is your own app secret (e.g. secrets.token_urlsafe(32)),
# not anything Litmus issues — same role as Django's SECRET_KEY.
app.add_middleware(SessionMiddleware, secret_key=SESSION_SECRET)

oauth = OAuth()
oauth.register(
    name="litmus",
    client_id=CLIENT_ID,  # no client_secret — Litmus issues public, PKCE-only clients
    authorize_url=f"{LITMUS_OAUTH_SERVER_DOMAIN}/oauth2/authorize",
    access_token_url=f"{LITMUS_OAUTH_SERVER_DOMAIN}/oauth2/token",
    client_kwargs={
        "scope": "openid offline_access",  # both required — see below
        "code_challenge_method": "S256",  # turns on PKCE end to end, automatically
    },
)

@app.get("/connect/litmus")
async def connect_litmus(request: Request):
    # REDIRECT_URI is the exact URI you gave Litmus — see "What we need
    # from you first" — not something derived at request time. Building it
    # from request.url_for() instead can silently drift from what's
    # registered (e.g. http vs. https behind a proxy), which Litmus will
    # then reject as a mismatch.
    return await oauth.litmus.authorize_redirect(request, REDIRECT_URI, resource=LITMUS_MCP_SERVER_URL)

from fastapi.responses import RedirectResponse

# The path here is REDIRECT_URI's path — this route IS your redirect URI.
@app.get("/auth/litmus/callback")
async def litmus_callback(request: Request):
    # Reads code + state from the query string and checks state against
    # what step A stashed in the session — automatically, nothing to
    # compare by hand.
    tokens = await oauth.litmus.authorize_access_token(request)
    await save_tokens_for_user(current_user_id(request), tokens)
    # Send them back into your product, not Litmus — wherever you want
    # a "Litmus connected" state to show. This URL is illustrative.
    return RedirectResponse("/settings/integrations?connected=litmus")

# The line that does it — inside the callback shown in step B:
tokens = await oauth.litmus.authorize_access_token(request)
# tokens: { access_token, token_type, expires_in, refresh_token? }

async def refresh_litmus_tokens(refresh_token: str) -> dict:
    # Not automatic — call this yourself on a 401 from the MCP server.
    return await oauth.litmus.fetch_access_token(
        grant_type="refresh_token",
        refresh_token=refresh_token,
    )