Error while using API and add_deploy

[matof]

After updating Pyinfra to version 3.5.1 I get error "Cannot call operations outside of Prepare/Execute stages"(pyinfra\api\operation.py", line 270) while using Pyinfra API. It is working without problem with version 3.3.X.

Our setup a is a little bit complicated but I will paste here all I think is important.

import pyinfra.api
import pyinfra.api.connect
from pyinfra.api.deploy import add_deploy
import pyinfra.local

class PyinfraStateCallback(pyinfra.api.state.BaseStateCallback):
  ...

class PyinfraRunner:

  def run_pyinfra_deploy(self):
    working_dir = ...
    hosts = [...]
    inventory = pyinfra.api.Inventory((hosts, {"ssh_user": "root"}))
    config = pyinfra.api.config.Config(
      CWD=working_dir
    )

    state = pyinfra.api.state.State(inventory, config)
    state.add_callback_handler(PyinfraStateCallback(self))

    pyinfra.api.connect.connect_all(state)

    module_name = "my.custom.deploy"
    deploy_module = importlib.import_module(module_name)
    add_deploy(state, deploy_module.task, task_context={...})

    pyinfra.api.operations.run_ops(state)

The custom deploy script (my.custom.deploy) looks like:

from pyinfra.api import deploy
from pyinfra.operations import files, python
from pyinfra import host

@deploy("Download file")
def task(task_context):
    src_path = ...
    dest_path = ...

    files.get(
        name="Download file from remote",
        src=src_path,
        dest=dest_path
    )

The line with files.get is triggering the exception.
Can you advice me if I am doing something wrong? Or did something change in version 3.5 and I shouldn't be using it in this way anymore?
Thank you.

[matof]

Adding minimal test case to reproduce the issue. Just put the valid host IP address in the main.php and run it.

reproduce-3.5.1.zip

Run result is:

e:\pyinfra-temp>python main.py
Traceback (most recent call last):
  File "e:\pyinfra-temp\main.py", line 29, in <module>
    main()
  File "e:\pyinfra-temp\main.py", line 23, in main
    add_deploy(state, download_file.task)
  File "C:\Program Files\Python311\Lib\site-packages\pyinfra\api\deploy.py", line 51, in add_deploy
    deploy_func(*args, **kwargs)
  File "C:\Program Files\Python311\Lib\site-packages\pyinfra\api\deploy.py", line 94, in decorated_func
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "e:\pyinfra-temp\download_file.py", line 6, in task
    files.get(
  File "C:\Program Files\Python311\Lib\site-packages\pyinfra\api\operation.py", line 270, in decorated_func
    raise Exception("Cannot call operations outside of Prepare/Execute stages")
Exception: Cannot call operations outside of Prepare/Execute stages