Asyncio: pass context or contextvar to add_done_callback

I am learning asyncio callbacks. My task is- I have a message dict, message codes are keys, message texts are values. In coro main I have to create a number of asynchronous tasks (in my case 3 tasks), each task wraps a coro which prints one message. Also I have to add a callback to each task. Callback must print a code associated with a message printed by the coroutine wrapped by the task. The question is- how to pass code to callback? The staright solution is to add name to each task with the value of a code, but I dont want to go this way. I decided to use ContextVar for this purpose. So I create a global context variable, set() the value to the variable equal to code. Then I try to get() the context variable value from a callback but receive an Exception LookupError: <ContextVar name='msg_code' at 0x000001C596D94F40>. That's my code:

import asyncio
from contextvars import ContextVar


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)
    msg_code.set(code)
    print(f'Message: {msg_dict[code]}')


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


async def main():
    tasks = [asyncio.create_task(print_msg(code)) for code in msg_dict.keys()]
    [task.add_done_callback(callback_code) for task in tasks]
    await asyncio.gather(*tasks)


asyncio.run(main())

I found that add_done_callback() also has keyword argument context= but I can't find any examples of how to pass task's context to a callback.

No tricks are needed, just specify the context.

async def main():
    tasks = []
    for code in msg_dict:
        ctx = copy_context()    # note: import copy_context from contexvars
        task = asyncio.create_task(print_msg(code), context=ctx)
        task.add_done_callback(callback_code, context=ctx)
        tasks.append(task)
    await asyncio.gather(*tasks)