mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 05:21:57 +00:00
问数据一轮:停止正好落在发起运行的那一下,也不会留下没人管的运行
发起运行(launch_run)是这一轮后台任务里的一步,停止就是取消这个任务:取消落在发起的半路上,运行可能已经建好、开始执行,这一轮 却还没记下它的 id——停止取消不到它,它照样跑、照样计费,那一行也不知道有过这次运行。现在发起那一下不让停止打断:等它发起完、 记下运行,再照常取消它,那一行收成已取消、meta.runId 是这次运行。 测试:停止落在发起的半路上(发起被拦住、放行之前先停),停下之后这次运行是已取消、记在这一轮上;去掉修复回跑会失败。顺手把 「运行中停止」那一条改成等任务上记下运行 id 再停(stage=run 在发起之前就记下了,满载跑全量时读到的 run_id 偶尔还是空的)。
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user