问数据一轮:停止正好落在发起运行的那一下,也不会留下没人管的运行

发起运行(launch_run)是这一轮后台任务里的一步,停止就是取消这个任务:取消落在发起的半路上,运行可能已经建好、开始执行,这一轮
却还没记下它的 id——停止取消不到它,它照样跑、照样计费,那一行也不知道有过这次运行。现在发起那一下不让停止打断:等它发起完、
记下运行,再照常取消它,那一行收成已取消、meta.runId 是这次运行。

测试:停止落在发起的半路上(发起被拦住、放行之前先停),停下之后这次运行是已取消、记在这一轮上;去掉修复回跑会失败。顺手把
「运行中停止」那一条改成等任务上记下运行 id 再停(stage=run 在发起之前就记下了,满载跑全量时读到的 run_id 偶尔还是空的)。
This commit is contained in:
ThearchyHelios
2026-10-18 17:10:01 +08:00
parent 8b74c0d532
commit 124131b393
2 changed files with 69 additions and 3 deletions
+12
View File
@@ -36,6 +36,7 @@
from __future__ import annotations
import asyncio
import contextlib
import logging
import time
from collections.abc import Mapping
@@ -173,6 +174,17 @@ class ChatDriver(TurnDriver):
raise LaunchRefused(403, str(e)) from e
return args
async def _launch(self, graph: dict[str, Any], payload: dict[str, Any] | None = None) -> str | None:
"""发起那一下不让停止打断到一半:运行建好了、这一轮却不知道它,就成了没人管、照样计费的运行。停止正好落在这时,
等它发起完、记下运行(on_launched),再照常去取消它(_stopped)。"""
task = asyncio.ensure_future(super()._launch(graph, payload))
try:
return await asyncio.shield(task)
except asyncio.CancelledError:
with contextlib.suppress(Exception):
await task
raise
async def on_launched(self, run: Any) -> None:
await self.execution.launched(run)
+57 -3
View File
@@ -104,6 +104,17 @@ async def _stage(turn_id: str, stage: str) -> None:
await _until(hit)
async def _launched(turn_id: str) -> str:
"""等这一轮发起了运行,交回运行 id(stage=run 在发起之前就记下了,要等的是任务上的 run_id)。"""
async def hit() -> bool:
async with SessionLocal() as session:
job = await session.get(ChatTurnJob, turn_id)
return job is not None and bool(job.run_id)
await _until(hit)
async with SessionLocal() as session:
return (await session.get(ChatTurnJob, turn_id)).run_id
def _sse(text: str) -> list[dict[str, Any]]:
return [json.loads(line[6:]) for line in text.splitlines() if line.startswith("data: ")]
@@ -303,9 +314,7 @@ async def test_stop_while_running_cancels_the_run(client, monkeypatch, source):
monkeypatch.setattr(mock_model.MockChatModel, "_astream", slow_stream)
conversation_id = await _conversation(client)
turn_id = await _ask(client, conversation_id, datasource_ids=[source["id"]])
await _stage(turn_id, "run")
async with SessionLocal() as session:
run_id = (await session.get(ChatTurnJob, turn_id)).run_id
run_id = await _launched(turn_id)
assert (await client.post(f"/api/chat/turns/{turn_id}/cancel")).status_code == 200
job = await _done(turn_id)
assert job.status == "cancelled"
@@ -881,3 +890,48 @@ async def test_several_stranded_turns_of_one_conversation_can_be_resumed_togethe
await _done(turn_id)
row = await _row(turn_id)
assert row.status == "done" and row.answer == answer
async def test_a_stop_that_lands_while_launching_still_cancels_the_run(client, monkeypatch, source):
"""停止正好落在发起运行的那一下:发起不能被打断到一半(运行建好了、这一轮却不知道它,就成了没人管的运行)——等它发起完、
记下运行,再照常取消它,那一行收成已取消。"""
import app.engine.launch as launch
from app.providers import mock_model
es.use_copilot(monkeypatch, es.Copilot(source["name"]))
hold = asyncio.Event()
async def slow(self, messages, stop=None, run_manager=None, **kwargs):
await hold.wait()
raise AssertionError("不该放行")
async def slow_stream(self, messages, stop=None, run_manager=None, **kwargs):
await hold.wait()
raise AssertionError("不该放行")
yield # noqa: unreachable - 让它成为异步生成器
monkeypatch.setattr(mock_model.MockChatModel, "_agenerate", slow)
monkeypatch.setattr(mock_model.MockChatModel, "_astream", slow_stream)
entered, release = asyncio.Event(), asyncio.Event()
real = launch.launch_run
async def gated(session, **kw):
entered.set()
await release.wait()
return await real(session, **kw)
monkeypatch.setattr(launch, "launch_run", gated)
conversation_id = await _conversation(client)
turn_id = await _ask(client, conversation_id, datasource_ids=[source["id"]])
await asyncio.wait_for(entered.wait(), timeout=20)
stopping = asyncio.ensure_future(client.post(f"/api/chat/turns/{turn_id}/cancel"))
await asyncio.sleep(0.2)
release.set()
assert (await stopping).status_code == 200
job = await _done(turn_id)
assert job.status == "cancelled" and job.run_id, "发起完了才停:运行记在这一轮上"
await run_manager.wait_idle(job.run_id)
async with SessionLocal() as session:
assert (await session.get(Run, job.run_id)).status == "cancelled", "这次运行跟着取消了,没有没人管的运行"
row = await _row(turn_id)
assert row.status == "error" and row.error == "已取消" and row.meta["runId"] == job.run_id