fix(samples): serialize chat state writes in the FastAPI sample

python/samples/agentchat_fastapi/app_agent.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import os
 from typing import Any
@@ -36,6 +37,19 @@ async def root():
 state_path = "agent_state.json"
 history_path = "agent_history.json"
 
+# /chat requests share these files, so their read-modify-write cycles must not
+# interleave. The lock serializes them and the atomic replace below keeps a
+# killed writer from leaving a truncated JSON document behind.
+_state_lock = asyncio.Lock()
+
+
+async def write_json_atomic(path: str, data: Any) -> None:
+    """Write JSON next to the target, then replace it in one step."""
+    tmp_path = f"{path}.tmp"
+    async with aiofiles.open(tmp_path, "w") as file:
+        await file.write(json.dumps(data))
+    os.replace(tmp_path, path)
+
 
 async def get_agent() -> AssistantAgent:
     """Get the assistant agent, load state from file."""
@@ -77,24 +91,25 @@ async def history() -> list[dict[str, Any]]:
 @app.post("/chat", response_model=TextMessage)
 async def chat(request: TextMessage) -> TextMessage:
     try:
-        # Get the agent and respond to the message.
-        agent = await get_agent()
-        response = await agent.on_messages(messages=[request], cancellation_token=CancellationToken())
-
-        # Save agent state to file.
-        state = await agent.save_state()
-        async with aiofiles.open(state_path, "w") as file:
-            await file.write(json.dumps(state))
-
-        # Save chat history to file.
-        history = await get_history()
-        history.append(request.model_dump())
-        history.append(response.chat_message.model_dump())
-        async with aiofiles.open(history_path, "w") as file:
-            await file.write(json.dumps(history))
-
-        assert isinstance(response.chat_message, TextMessage)
-        return response.chat_message
+        # Serialize the whole read-modify-write cycle: the agent state and the
+        # history file are shared by every request.
+        async with _state_lock:
+            # Get the agent and respond to the message.
+            agent = await get_agent()
+            response = await agent.on_messages(messages=[request], cancellation_token=CancellationToken())
+
+            # Save agent state to file.
+            state = await agent.save_state()
+            await write_json_atomic(state_path, state)
+
+            # Save chat history to file.
+            history = await get_history()
+            history.append(request.model_dump(mode="json"))
+            history.append(response.chat_message.model_dump(mode="json"))
+            await write_json_atomic(history_path, history)
+
+            assert isinstance(response.chat_message, TextMessage)
+            return response.chat_message
     except Exception as e:
         error_message = {
             "type": "error",

python/samples/agentchat_fastapi/test_app_agent.py

@@ -0,0 +1,59 @@
+import asyncio
+import json
+
+import app_agent
+import httpx
+import pytest
+from autogen_agentchat.messages import TextMessage
+
+
+class _FakeAgent:
+    async def on_messages(self, messages, cancellation_token):
+        await asyncio.sleep(0)
+        return _FakeResponse(TextMessage(content="ok", source="assistant"))
+
+    async def save_state(self):
+        return {"turns": 1}
+
+
+class _FakeResponse:
+    def __init__(self, chat_message):
+        self.chat_message = chat_message
+
+
+@pytest.fixture
+def chat_client(tmp_path, monkeypatch):
+    history_path = tmp_path / "agent_history.json"
+    state_path = tmp_path / "agent_state.json"
+    monkeypatch.setattr(app_agent, "history_path", str(history_path))
+    monkeypatch.setattr(app_agent, "state_path", str(state_path))
+
+    async def fake_get_agent():
+        return _FakeAgent()
+
+    monkeypatch.setattr(app_agent, "get_agent", fake_get_agent)
+
+    transport = httpx.ASGITransport(app=app_agent.app)
+    return httpx.AsyncClient(transport=transport, base_url="http://test"), history_path
+
+
+@pytest.mark.asyncio
+async def test_concurrent_chats_keep_state_and_history_valid(chat_client):
+    client, history_path = chat_client
+    async with client:
+        responses = await asyncio.gather(
+            *[
+                client.post(
+                    "/chat",
+                    json={"source": "user", "content": f"message {i}", "type": "TextMessage"},
+                )
+                for i in range(10)
+            ]
+        )
+
+    for response in responses:
+        assert response.status_code == 200
+
+    history = json.loads(history_path.read_text())
+    assert len(history) == 20
+    assert all(isinstance(item, dict) for item in history)