googleapis / googleapis/python-aiplatform

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

Open
#7,136 1 comment 0 reactions 0 assignees View on GitHub
api: vertex-ai
Dominant language
Python
Stars
905
Forks
465
Avg merge
1d 13h
Merged PRs (30d)
44

Description

`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:

```python
# 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:
```python
import asyncio
from vertexai import agent_engines

agent = agent_engines.get("projects/

/locations//reasoningEngines/")

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 took X.XX seconds" because the thread is blocked on socket reads without yielding.
```

Contributor guide

Open the contributing guide

Research direction

Start in vertexai/agent_engines/_agent_engines.py at _wrap_async_stream_query_operation and compare the synchronous execution_api_client call with self.execution_async_client. Reproduce the stream with asyncio.run(..., debug=True); done means async_stream_query no longer blocks the event-loop thread during socket reads while still yielding parsed chunks.

Written by the indexing model from the issue text.

Assessment

Tech stack
grpc, python
Domain
api, backend
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.