Unexpected ContextVar handling for lifespan context managers

[rmorshea]

Description

A ContextVar set within a lifespan context manager is not available inside request handlers. By contrast, doing the same thing via an application level dependency does appear set in the require handler.

MCVE

import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from contextvars import ContextVar

from litestar import Litestar
from litestar import get
from litestar.testing import AsyncTestClient

VAR = ContextVar[int]("VAR", default=0)


@asynccontextmanager
async def set_var() -> AsyncIterator[None]:
    token = VAR.set(1)
    try:
        yield
    finally:
        VAR.reset(token)


@get("/example")
async def example() -> int:
    return VAR.get()


app = Litestar([example], lifespan=[set_var()])


async def test() -> None:
    async with AsyncTestClient(app) as client:
        response = await client.get("/example")
        assert (value := response.json()) == 1, f"Expected 1, got {value}"


if __name__ == "__main__":
    asyncio.run(test())

Steps to reproduce

Run the example above. Either via python example.py or uvicorn example:app and check the /example route.

We should expect the response to be 1 since that's what's set during the lifespan context manager, but it's always 0.

If you instead swap the lifespan context manager for a dependency it works as expected:

import asyncio
from collections.abc import AsyncIterator
from contextvars import ContextVar

from litestar import Litestar
from litestar import get
from litestar.testing import AsyncTestClient

VAR = ContextVar[int]("VAR", default=0)


async def set_var() -> AsyncIterator[None]:
    token = VAR.set(1)
    try:
        yield
    finally:
        VAR.reset(token)


@get("/example")
async def example(set_var: None) -> int:
    return VAR.get()


app = Litestar([example], dependencies={"set_var": set_var})


async def test() -> None:
    async with AsyncTestClient(app) as client:
        response = await client.get("/example")
        assert (value := response.json()) == 1, f"Expected 1, got {value}"


if __name__ == "__main__":
    asyncio.run(test())

Litestar Version

2.12.1

Platform

• Linux
• Mac
• Windows
• Other (Please specify in the description above)

---

Note

While we are open for sponsoring on GitHub Sponsors and
OpenCollective, we also utilize Polar.sh to engage in pledge-based sponsorship.

Check out all issues funded or available for funding on our Polar.sh dashboard

• If you would like to see an issue prioritized, make a pledge towards it!
• We receive the pledge once the issue is completed & verified
• This, along with engagement in the community, helps us know which features are a priority to our users.

[euri10]

vey interesting, I played a litttle with it, changing lifespan to the old way with on_startup / on_sutdown and the result is even weirder.

I didn't go through all this thread but that seems related pytest-dev/pytest-asyncio#127

note on my "tests" I didnt use pytest-asyncio but anyio

[rmorshea]

It looks like this is expected behavior in Starlette/FastAPI. The response there is pretty turse so it's hard to see exactly why.

[rmorshea]

If this turns out to be a fundamental limitation it might be nice to support dependencies that have side effects but don't actually return a value. I could imagine something similar to Pytest's auto use fixtures:

VAR = ContextVar("VAR", default=0)


async def set_var() -> AsyncIterator[None]:
    token = VAR.set(1)
    try:
        yield
    finally:
        VAR.reset(token)


@get("/", dependencies={"set_var": Provide(set_var, always=True)})
async def get_thing() -> None:
    assert VAR.get() == 1

[euri10]

yep, pytest-asyncio might add another layer of complexity but this fails without regardless so the issue lies elsewhere,

now adding some goold old-style prints to this:

from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from contextvars import ContextVar

from litestar import Litestar
from litestar import get



VAR = ContextVar("VAR", default=0)
print(f"1: {VAR}")


@asynccontextmanager
async def set_var() -> AsyncIterator[None]:
    token = VAR.set(1)
    try:
        yield
    # except Exception as exc:
    #     print(exc)
    finally:
        print("Finally")
        VAR.reset(token)


@get("/example")
async def example() -> int:
    print(VAR)
    return VAR.get()


app = Litestar([example], lifespan=[set_var()], debug=True)

if __name__ == "__main__":
    # asyncio.run(test())
    print(f"1: {VAR}")
    import uvicorn
    uvicorn.run("cv:app", reload=True)

I get this log, note that there is something going on if I understand it clearly with the way the var is created / initialized as the handler seems to not use the same var that was created in lifespan

/home/lotso/PycharmProjects/abdul/.venv/bin/python /home/lotso/PycharmProjects/abdul/cv.py 
1: <ContextVar name='VAR' default=0 at 0x7fb24500be20>
1: <ContextVar name='VAR' default=0 at 0x7fb24500be20>
INFO:     Will watch for changes in these directories: ['/home/lotso/PycharmProjects/abdul']
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
INFO:     Started reloader process [436898] using WatchFiles
1: <ContextVar name='VAR' default=0 at 0x7fabc7fedfd0>
1: <ContextVar name='VAR' default=0 at 0x7fabc57972e0>
INFO:     Started server process [436904]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO - 2024-10-09 08:08:12,131 - watchfiles.main - main - 3 changes detected
<ContextVar name='VAR' default=0 at 0x7fabc57972e0>
INFO:     127.0.0.1:57262 - "GET /example HTTP/1.1" 200 OK

[rmorshea]

Here's a longer running issue: encode/uvicorn#1151

[rmorshea]

For additional context (😉) here's related Pytest issue.

[provinzkraut]

This is working as intended and expected actually. ContextVars in asyncio are scoped to a task. While executing an ASGI app, many such tasks may be created, and there are no guarantees either by Litestar or the ASGI server that the lifespan will be executed in the same task context as the request.

Example

This will fail because VAR is unset inside the newly created task:

import asyncio
import contextvars

VAR = contextvars.ContextVar("VAR")

async def set_var():
    VAR.set("hello")


async def get_var():
    print(VAR.get())


async def main():
    await asyncio.create_task(set_var())
    await asyncio.create_task(get_var())


asyncio.run(main())

---

If you want to share application wide state, that is what app.state is for.

So you would do something like:

async def set_var(app: Litestar) -> AsyncIterator[None]:
    app.state.token = 1
    try:
        yield
    finally:
        del app.state.token

def get_token(state: State) -> int:
  return state.token

@get("/example")
async def example(state: State) -> int:
    return state.token

app = Litestar([example], lifespan=[set_var], dependencies={"token": get_token})

[rmorshea]

I don't think state is a full solution since (IIRC) it isn't context aware. Thus if you spawn tasks as part of a request handler and mutate state those changes won't be appropriately isolated - users would need to use a combination of state and contextvars to deal with this. Additionally, state is a feature of Litestar, not of ASGI apps so someone developing a generic ASGI middleware wouldn't be able to rely on it.

All that said, I understand that this is due to the fact that the lifespan context and request context are fundamentally distinct and that this is true in almost any ASGI server you could imagine implementing. Thus it probably doesn't make sense for Litestar to supply a workaround.

For posterity I wanted to link the solution I came to in my ASGI middleware and underlying utility that copy contextual values set by lifespan context managers into the context of request handlers.

One could imagine generalizing this ASGI middleware to take a list of context variables that should be copied from the lifespan context to the request context. That, or the default behavior could be to copy all contextual values set by lifespan context managers using Context.items().