[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.

[ArushKhasru]

Hi maintainers, I'd like to work on this issue. Could you assign it to me if it's available?

I noticed that #5773 and the open PR #6295 cover related async streaming behavior. I'd be happy to help with any remaining work, including regression tests to verify that streaming doesn't block the asyncio event loop. Please let me know whether you'd prefer contributions to the existing fix or a separate PR for any remaining gaps.

Thanks!