Asyncio：如何向add_done_callback传递ContextVar或上下文

你遇到的核心问题是：回调函数运行的上下文与任务协程设置ContextVar的上下文不统一。当任务协程执行完毕，其对应的上下文会被销毁，而回调是在事件循环的默认上下文中触发的，无法获取之前设置的msg_code值，因此抛出LookupError。

利用add_done_callback的context参数可以完美解决这个问题——通过将任务协程的上下文传递给回调，确保回调在包含目标ContextVar值的上下文中运行。

修改后的代码示例 #

import asyncio
from contextvars import ContextVar, copy_context

msg_dict = {
    'code1': 'msg1 by code1',
    'code2': 'msg2 by code2',
    'code3': 'msg3 by code3'
}

msg_code = ContextVar('msg_code')


async def print_msg(code):
    await asyncio.sleep(0.5)
    print(f'Message: {msg_dict[code]}')


def callback_code(*args):
    code = msg_code.get()
    print(f'Code: {code}')


async def main():
    tasks = []
    for code in msg_dict.keys():
        # 临时设置ContextVar，让任务继承该上下文
        token = msg_code.set(code)
        try:
            task = asyncio.create_task(print_msg(code))
            # 复制当前包含msg_code的上下文快照
            task_context = copy_context()
            # 将上下文传递给回调，确保回调在该上下文中运行
            task.add_done_callback(callback_code, context=task_context)
            tasks.append(task)
        finally:
            # 重置当前上下文的ContextVar，避免污染后续任务逻辑
            msg_code.reset(token)
    await asyncio.gather(*tasks)


asyncio.run(main())

代码说明 #

1. 在创建任务前，临时设置msg_code为当前消息码，让任务继承这个包含目标值的上下文；
2. 调用copy_context()复制当前上下文快照，确保后续回调能访问到msg_code的值；
3. 将上下文快照作为context参数传入add_done_callback，指定回调运行的上下文；
4. 用finally块重置当前上下文的msg_code，避免该值影响其他任务或后续逻辑。

运行这段代码后，会依次输出消息文本和对应的消息码，不会再触发LookupError。

另一种思路（协程内部设置ContextVar） #

如果必须在print_msg内部设置msg_code，可以在协程中保存上下文快照到任务的自定义属性，再在回调中使用该上下文执行逻辑：

import asyncio
from contextvars import ContextVar, copy_context

msg_dict = {
    'code1': 'msg1 by code1',
    'code2': 'msg2 by code2',
    'code3': 'msg3 by code3'
}

msg_code = ContextVar('msg_code')


async def print_msg(code, task):
    await asyncio.sleep(0.5)
    msg_code.set(code)
    # 保存上下文快照到任务属性
    task.msg_context = copy_context()
    print(f'Message: {msg_dict[code]}')


def callback_code(task):
    # 使用任务保存的上下文执行回调逻辑
    task.msg_context.run(lambda: print(f'Code: {msg_code.get()}'))


async def main():
    tasks = []
    for code in msg_dict.keys():
        task = asyncio.create_task(print_msg(code, task=None))
        # 更新协程参数，将任务对象自身传入
        task._coro = print_msg(code, task)
        task.add_done_callback(callback_code)
        tasks.append(task)
    await asyncio.gather(*tasks)


asyncio.run(main())

这种方式稍显繁琐，更推荐第一种方案。

内容的提问来源于stack exchange，提问作者IzaeDA