Serialize chat resets with message generation

This commit is contained in:
eezd committed 2026-10-20 18:50:05 +08:00
1 parent 6bcff9c572
commit ab6c875b1f
2 files changed
+59 -7

No files matched your search

+8 -7
View File
@@ -29,7 +29,7 @@ MAX_TOOL_ROUNDS = 8
TOOL_PREVIEW_CHARS = 2000
STOPPED_RESULT = json.dumps({"error": "用户已停止生成,未完成的工具结果尚未确认,请核对后再重试"}, ensure_ascii=False)
# api 只有一个进程:同一时间只允许一轮生成,Reset 也要等生成结束
# api 只有一个进程:生成与 Reset 共用互斥锁,不能同时修改当前会话。
generating = asyncio.Lock()
logger = logging.getLogger("alphapulse.api.chat")
@@ -127,7 +127,7 @@ async def get_chat_session(session: Annotated[AsyncSession, Depends(get_session)
async def send_message(body: SendIn) -> AsyncIterator[ServerSentEvent]:
"""事件:prompt、thinking、delta、tool_call、tool_result、skill_draft、context_updated、usage、error、done"""
if generating.locked():
yield event("error", {"message": "上一轮还在生成,请等待完成或先停止"})
yield event("error", {"message": "正在生成或重置会话,请等待完成"})
return
async with generating:
async with SessionLocal() as session:
@@ -233,11 +233,12 @@ async def send_message(body: SendIn) -> AsyncIterator[ServerSentEvent]:
async def reset(body: ResetIn, session: Annotated[AsyncSession, Depends(get_session)]) -> None:
"""结束当前会话(消息保留在库里,不再展示),可选同时清空 CONTEXT.md(写入空版本,可找回)"""
if generating.locked():
raise HTTPException(status.HTTP_409_CONFLICT, "正在生成,请先停止")
await session.execute(update(ChatSession).where(ChatSession.ended_at.is_(None)).values(ended_at=datetime.now(UTC)))
if body.clear_context:
session.add(ContextVersion(content="", source="reset"))
await session.commit()
raise HTTPException(status.HTTP_409_CONFLICT, "正在生成或重置会话,请稍后重试")
async with generating:
await session.execute(update(ChatSession).where(ChatSession.ended_at.is_(None)).values(ended_at=datetime.now(UTC)))
if body.clear_context:
session.add(ContextVersion(content="", source="reset"))
await session.commit()
@router.get("/context")
+51
View File
@@ -179,3 +179,54 @@ def test_skill_rename_rejects_whitespace_without_writing(trading):
assert error.value.status_code == 400
assert node.name == "原名称"
asyncio.run(check())
def test_chat_reset_blocks_new_generation_until_commit(trading, monkeypatch):
from alphapulse.api import chat
async def check():
monkeypatch.setattr(chat, "generating", asyncio.Lock())
entered, finish = asyncio.Event(), asyncio.Event()
monkeypatch.setattr(chat, "SessionLocal", lambda: pytest.fail("重置期间不能进入生成的数据库路径"))
saved = []
class Session:
async def execute(self, statement):
entered.set()
await finish.wait()
def add(self, row):
saved.append(row)
async def commit(self):
assert chat.generating.locked()
async def reset_and_send():
task = asyncio.create_task(chat.reset(chat.ResetIn(clear_context=True), Session()))
try:
await entered.wait()
events = [event async for event in chat.send_message(chat.SendIn(text="重置期间的新消息"))]
assert len(events) == 1 and events[0].event == "error"
assert not task.done()
finally:
finish.set()
await task
await asyncio.wait_for(reset_and_send(), timeout=2)
assert len(saved) == 1 and saved[0].content == "" and saved[0].source == "reset"
assert not chat.generating.locked()
asyncio.run(check())
def test_chat_reset_rejects_busy_generation_and_releases_lock_on_failure(trading, monkeypatch):
from fastapi import HTTPException
from alphapulse.api import chat
async def check():
monkeypatch.setattr(chat, "generating", asyncio.Lock())
class Session:
async def execute(self, statement):
raise RuntimeError("重置写库失败")
async with chat.generating:
with pytest.raises(HTTPException) as error:
await chat.reset(chat.ResetIn(), Session())
assert error.value.status_code == 409
with pytest.raises(RuntimeError, match="重置写库失败"):
await chat.reset(chat.ResetIn(), Session())
assert not chat.generating.locked()
asyncio.run(check())