[Agent Engines] `_wrap_async_stream_query_operation` blocks asyncio event loop due to synchronous gRPC iteration

[steffanianigro]

agent_engine.async_stream_query(...) blocks the calling asyncio event loop thread during stream generation.

In vertexai/agent_engines/_agent_engines.py, _wrap_async_stream_query_operation wraps a synchronous client call in an async def and uses a blocking for loop:

# CURRENT IMPLEMENTATION (vertexai/agent_engines/_agent_engines.py)
def _wrap_async_stream_query_operation(*, method_name: str):
    async def _method(self, **kwargs):
        # 1. Uses sync client instead of self.execution_async_client
        response = self.execution_api_client.stream_query_reasoning_engine(...)
        # 2. Synchronous iteration blocks the asyncio event loop on socket reads
        for chunk in response:
            for parsed_json in _utils.yield_parsed_json(chunk):
                if parsed_json is not None:
                    yield parsed_json
    return _method

Reproduction:

import asyncio
from vertexai import agent_engines

agent = agent_engines.get("projects/<P>/locations/<L>/reasoningEngines/<ID>")

async def main():
    async for chunk in agent.async_stream_query(user_id="u", message="Long prompt"):
        pass

asyncio.run(main(), debug=True)
# Result: Emits "Executing <Task ...> took X.XX seconds" because the thread is blocked on socket reads without yielding.