ag-ui-protocol / ag-ui-protocol/ag-ui
[Bug]: LangGraphAGUIAgent is a stateful singleton — concurrent requests corrupt shared state causing KeyError / TypeError crashes
- Langage dominant
- Python
- Étoiles
- 15.9k
- Forks
- 1.4k
- Merge moyen
- 1 j 17 h
- PR mergées (30 j)
- 163
Description
### Pre-flight Checklist
- [x] I have searched [existing issues](https://github.com/ag-ui-protocol/ag-ui/issues) and this hasn't been reported yet.
- [x] I am using the **latest** version AG-UI.
### Describe the Bug
`LangGraphAGUIAgent` stores per-request working state in two instance-level
attributes (`self.active_run` and `self.messages_in_process`). Because one
agent instance is registered once at startup and shared across all requests, and
because asyncio yields at every `await` point, concurrent requests overwrite
each other's state mid-execution.
This produces **three distinct crash types** under any concurrent load:
**1. `KeyError: 'schema_keys'`** — `agent.py` line 487
`self.active_run` is assigned a fresh dict (no `schema_keys` key) at line 126.
`schema_keys` is only added after the first `await` at line 286. A second
concurrent request replaces `self.active_run` with its own fresh dict at line
126 before the first request reaches line 286. When the first request resumes
and its event handler calls `get_state_snapshot()` (line 487), it reads the
second request's dict, which has no `schema_keys` yet.
**2. `KeyError: 'mode'`** — `agent.py` lines 330, 342
Same race. `mode` is written to `self.active_run` at lines 140–142 (also after
an `await`), and read back in `prepare_stream()` at lines 330 and 342. A
concurrent request replaces `self.active_run` with a fresh dict that has no
`mode` key before those reads occur.
**3. `TypeError: 'NoneType' object is not a mapping`** — `agent.py` line 407
When a message finishes streaming, the library cleans up by setting the tracking
slot to `None` (lines 555, 563, 629, 636):
```python
self.messages_in_process[self.active_run["id"]] = None
```
When the next message in the same run starts, `set_message_in_progress()` (line
405) reads that slot back with `.get(run_id, {})`. However, `.get(key, default)`
only returns the default when the key is **absent** — if the key is present with
a `None` value, it returns `None`. The subsequent `{**None}` raises `TypeError`.
This third bug can technically occur in a single-request scenario (any
multi-step response), but the concurrency race makes it far more frequent because
`self.active_run["id"]` in the cleanup line can point to a different request's ID.
### Steps to Reproduce
1. Install `ag-ui-langgraph` (tested on v0.0.25, Python 3.12)
2. Create a FastAPI app with `add_langgraph_fastapi_endpoint`, registering the
agent once at startup via `LangGraphAGUIAgent(name=..., graph=...)`
3. Run uvicorn with default settings (1 worker, asyncio event loop)
4. Send two `POST /agent/` requests to the same endpoint simultaneously
(overlapping, not sequential) — e.g. using `asyncio.gather` or `httpx` async
client with two requests in flight at the same time
```python
import asyncio, httpx
async def main():
async with httpx.AsyncClient() as client:
await asyncio.gather(
client.post("http://localhost:8000/agent/math",
headers={"x-thread-id": "t1", "Accept": "text/event-stream"},
content=payload_1),
client.post("http://localhost:8000/agent/math",
headers={"x-thread-id": "t2", "Accept": "text/event-stream"},
content=payload_2),
)
asyncio.run(main())
```
1. Observe `KeyError: 'schema_keys'` or `KeyError: 'mode'` in the server log for
one of the two requests. Also observe `TypeError: 'NoneType' object is not a
mapping` for any agent response that involves a tool call followed by a text reply.
### Expected Behavior
Each request should be fully isolated. Concurrent requests to the same agent
endpoint should not corrupt each other's state. A multi-step response (tool
call + text reply) should not crash.
### Environment
```text
- **AG-UI package(s) & version(s):** `ag-ui-langgraph==0.0.25`
- **Runtime:** Python 3.12, FastAPI, uvicorn (asyncio event loop)
- **LangGraph version:** see pyproject.toml
```
### Screenshots
_No response_
### Logs & Errors
```shell
ERROR uvicorn.error Exception in ASGI application
File "ag_ui_langgraph/agent.py", line 487, in get_state_snapshot
schema_keys = self.active_run["schema_keys"]
KeyError: 'schema_keys'
ERROR uvicorn.error Exception in ASGI application
File "ag_ui_langgraph/agent.py", line 342, in prepare_stream
mode=self.active_run["mode"],
KeyError: 'mode'
ERROR uvicorn.error Exception in ASGI application
File "ag_ui_langgraph/agent.py", line 407, in set_message_in_progress
self.messages_in_process[run_id] = {
**current_message_in_progress,
TypeError: 'NoneType' object is not a mapping
```
### Additional Context
`self.active_run` and `self.messages_in_process` are instance-level attributes
mutated mid-flight, but the agent instance is a singleton shared across all
concurrent requests in the same event loop.
**Some fix directions:**
- For the concurrency bugs (KeyErrors): make `active_run` a local variable
passed as a parameter through the call chain, or store it in a
`contextvars.ContextVar` so each asyncio Task has its own isolated copy.
- For the `TypeError` bug: change the cleanup from setting `None` to actually
deleting the key (`del self.messages_in_process[run_id]`), or guard the reader:
```python
current = self.messages_in_process.get(run_id) or {}
```
Guide de contribution
Ouvrir le guide de contribution
Évaluation
Cette issue n'a pas encore été évaluée.