Вопрос по 2FA-авторизации в aiogram

У меня есть код на aiogram, который авторизовывает пользователя, и, если у него включена 2-факторная аутентификация, запрашивает код доступа.
Проблема в том, что это реализовано в with-менеджере, и, если я буду переключать состояние в случае запроса кода, with автоматически завершится, и выпадет ошибка о том, что страница закрыта.
Как можно реализовать логику ожидания ввода кода подтверждения?
Код:

class AuthState(StatesGroup):
    enter_login = State()
    enter_password = State()
    enter_2fa = State()

...

@dp.message(AuthState.enter_login)
async def proceed_login(message: Message, state: FSMContext):
    await state.update_data(login=message.text)
    await state.set_state(AuthState.enter_password)
    await message.answer("Введите пароль: ")

@dp.message(AuthState.enter_password)
async def proceed_pass(message: Message, state: FSMContext):
    await state.update_data(password=message.text)
    await start_auth(message, state)

async def start_auth(message: Message, state: FSMContext):
    data = await state.get_data()
    async with async_playwright() as p:
        auth = await my_lib.AUTH(p)
        value = await auth.obtain_token(data["login"], data["password"])
        if value == '2FA_NEEDED':
            await message.answer("Введите код подтверждения: ")
            value = ...  # ?
    user = await my_lib.Client(value)