How to handle session management in FastAPI with Okta OIDC & PKCE

Question

I have a python FastAPI service that also has a front-end component. I am integrating the service with Okta and I am having difficulty with the session management.

I can do the initial login to the app through okta, but on subsequent attempts, one of the two below happen:

1. The state is not preserved in the callback and I get a CSRF error for mismatching state.
2. The session is storing up to 5 states (new state each time I use the /login/okta URL and then stops storing the states and I get the same error.

I'm still very novice to FastAPI/Okta and need some guidance in ironing this out.

import os

from authlib.integrations.base_client import MismatchingStateError
from authlib.integrations.starlette_client import OAuth
from fastapi import Request, HTTPException
from starlette.datastructures import URL
from starlette.responses import RedirectResponse

# Set up OAuth client
oauth = OAuth()
oauth.register(
    name="okta",
    client_id=os.environ.get("OKTA_CLIENT_ID"),
    client_kwargs={"scope": "openid profile email"},
    code_challenge_method="S256",
    server_metadata_url=os.environ.get("OKTA_SERVER_METADATA_URL"),
)

app = FastAPI(
    title="My API",
    version="1.0.0",
)

app.mount("/static", StaticFiles(directory="etdge/static", html=True), name="static")
app.add_middleware(SessionMiddleware, secret_key="!secret")

@app.get("/okta")
async def login_okta(request: Request):
    okta = oauth.create_client("okta")
    state = "foo"
    redirect_uri = request.url_for("login_callback")
    if not os.environ.get("LOCAL"):
        redirect_uri = URL(scheme="https", netloc=redirect_uri.netloc, path=redirect_uri.path, query=redirect_uri.query)
    return await okta.authorize_redirect(request, redirect_uri=redirect_uri, state=state)


@app.get("/callback")
async def login_callback(request: Request):
    okta = oauth.create_client("okta")
    try:
        token = await okta.authorize_access_token(request)
    except MismatchingStateError as exc:
        raise HTTPException(401, f"Authentication was not successful: {exc}")
    request.session["token"] = token
    return RedirectResponse(url="/")

I've tried managing the state myself, but even though the state gets added to the session, it is sometimes unavailable in the callback.