ag-ui-protocol / ag-ui-protocol/ag-ui

[Bug]: LangGraphAGUIAgent is a stateful singleton — concurrent requests corrupt shared state causing KeyError / TypeError crashes

未關閉
#1,277 0 則留言 0 個 reaction 已指派 0 人 在 GitHub 檢視
bug
主要語言
Python
星號
15.9k
分支
1.4k
平均合併
1 天 17 小時
30 天內合併 PR
163

描述

### 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 {}
```

貢獻指南

開啟貢獻指南

評估

這個 Issue 還沒有評估資料。

把新 issue 寄到你的電子郵件信箱

精選適合新手參與的 GitHub issue 摘要。