diff --git a/python/samples/agentchat_fastapi/app_agent.py b/python/samples/agentchat_fastapi/app_agent.py index 9f114a651655..aec4e2a9b9f6 100644 --- a/python/samples/agentchat_fastapi/app_agent.py +++ b/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", diff --git a/python/samples/agentchat_fastapi/test_app_agent.py b/python/samples/agentchat_fastapi/test_app_agent.py new file mode 100644 index 000000000000..947b11f8b28e --- /dev/null +++ b/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)