engine.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796
  1. """Agent 执行内核:产出结构化事件,统一处理落库、计费、取消与收尾。
  2. updates 通道处理工具、待办和子 Agent;messages 通道处理正文和用量。
  3. 本模块不依赖 HTTP、SSE 或 API 请求模型。
  4. """
  5. from __future__ import annotations
  6. import asyncio
  7. from contextlib import aclosing
  8. import time
  9. from dataclasses import dataclass, field
  10. from datetime import datetime
  11. import anyio
  12. from typing import Any, AsyncIterator
  13. import config
  14. from runtime import concurrency
  15. from runtime.interrupts import pack, interrupt_event
  16. from runtime.context_meter import ContextMeter
  17. from runtime.events import EventBuilder, ExecutionEvent
  18. from runtime.source_tracker import SourceTracker
  19. from runtime.subagent_tracker import SubagentTracker
  20. from billing.meter import UsageMeter
  21. from db import chat_store, db_conn, run_store
  22. from obs import get_logger, get_trace_id
  23. logger = get_logger(__name__)
  24. MAIN_AGENT_NAME = "deep-agent"
  25. MAIN_AGENT_CN = "深度智能体"
  26. # 摘要中间件的模型调用也在 messages 通道里推 token,靠这个标记过滤掉,
  27. # 否则摘要正文会混进给用户看的回答里
  28. _SUMMARIZATION_SOURCE = "summarization"
  29. @dataclass
  30. class ExecutionRequest:
  31. """一次流式请求的全部入参。"""
  32. agent: Any
  33. session_id: str
  34. thread_id: str
  35. config: dict[str, Any]
  36. message: str
  37. user_name: str = "default"
  38. kind: str = "chat"
  39. model: str = ""
  40. original_user_message: str = ""
  41. trace_id: str = ""
  42. slot: concurrency.TaskSlot | None = None
  43. started_at: float = field(default_factory=time.time)
  44. agent_name: str = MAIN_AGENT_NAME
  45. agent_cn_name: str = MAIN_AGENT_CN
  46. # 非 None 表示这是一次中断恢复:不组装消息,直接把 Command(resume=...) 喂回图
  47. resume_payload: Any = None
  48. session_lease: Any = None
  49. run_id: str = field(default_factory=db_conn.new_id)
  50. @dataclass(frozen=True)
  51. class ExecutionResult:
  52. run_id: str
  53. session_id: str
  54. status: str
  55. output: str
  56. error: str
  57. usage: dict[str, Any]
  58. citations: list[dict[str, Any]]
  59. interrupts: list[Any]
  60. def normalize_chunk(chunk: Any) -> tuple[str | None, tuple[Any, ...], Any]:
  61. """把不同版本的 astream chunk 归一成 ``(mode, ns, data)``。
  62. 兼容三种形态:v2 的 ``{"type","ns","data"}`` 字典、旧版 ``(mode, data)``
  63. 元组、以及无法识别的原样返回(mode=None,上层据此忽略)。
  64. """
  65. if isinstance(chunk, dict) and "data" in chunk and "type" in chunk:
  66. return str(chunk["type"]), tuple(chunk.get("ns") or ()), chunk["data"]
  67. if isinstance(chunk, tuple) and len(chunk) == 2 and isinstance(chunk[0], str):
  68. return chunk[0], (), chunk[1]
  69. return None, (), chunk
  70. def _message_text(message: Any) -> str:
  71. """把消息 content 拍平成纯文本(content 可能是字符串或内容块列表)。"""
  72. content = getattr(message, "content", None)
  73. if isinstance(content, str):
  74. return content
  75. if isinstance(content, list):
  76. parts: list[str] = []
  77. for block in content:
  78. if isinstance(block, str):
  79. parts.append(block)
  80. elif isinstance(block, dict):
  81. text = block.get("text")
  82. if isinstance(text, str):
  83. parts.append(text)
  84. return "".join(parts)
  85. return ""
  86. def _reasoning_text(message: Any) -> str:
  87. """抽取思考链文本(deepseek 的 reasoning_content)。"""
  88. extra = getattr(message, "additional_kwargs", None)
  89. if isinstance(extra, dict):
  90. value = extra.get("reasoning_content")
  91. if isinstance(value, str):
  92. return value
  93. return ""
  94. class Execution:
  95. """一次任务执行,可逐事件消费,也可等待最终结果。"""
  96. def __init__(self, request: ExecutionRequest) -> None:
  97. self.req = request
  98. self.writer = EventBuilder(request.session_id, request.trace_id or get_trace_id(), encode=lambda event: event)
  99. self.meter = _meter_for_request()
  100. self.tracker = SubagentTracker()
  101. self.sources = SourceTracker()
  102. self.billing = UsageMeter(
  103. user_name=request.user_name,
  104. session_id=request.session_id,
  105. kind=request.kind,
  106. model=request.model or config.MODEL,
  107. trace_id=request.trace_id,
  108. usage_id=request.run_id,
  109. usage_day=db_conn.day_key(datetime.fromtimestamp(request.started_at)),
  110. )
  111. # 累计状态
  112. self.answer_parts: list[str] = []
  113. self.tool_calls: list[dict[str, Any]] = []
  114. self.todos: list[dict[str, Any]] = []
  115. self.interrupt_payloads: list[Any] = []
  116. self.finished = False
  117. self.status = "running"
  118. self.error = ""
  119. self._billing_added = False
  120. self._started = False
  121. self._snapshot = None
  122. self._answer_saved = False
  123. self._settlement_task: asyncio.Task | None = None
  124. self._consumed = False
  125. async def run(self) -> ExecutionResult:
  126. """消费事件并完成资源清理后返回结果,供定时任务、批处理等入口复用。"""
  127. async with aclosing(self.events()) as events:
  128. async for _ in events:
  129. pass
  130. return self.result()
  131. def result(self) -> ExecutionResult:
  132. if not self.finished:
  133. raise RuntimeError("任务尚未完成,不能读取最终结果")
  134. return ExecutionResult(self.req.run_id, self.req.session_id, self.status,
  135. "".join(self.answer_parts), self.error, self._usage_payload(),
  136. self.sources.build_citations(), list(self.interrupt_payloads))
  137. # ========================================================
  138. # 主入口
  139. # ========================================================
  140. async def events(self) -> AsyncIterator[ExecutionEvent]:
  141. """产出结构化执行事件。保证:计费幂等入账、并发位释放、图游标关闭。
  142. 这里显式绑定 trace_id 而不是依赖 HTTP 中间件设的 ContextVar:
  143. 响应体的迭代在 Starlette 的中间件调用栈之外执行,中间件里绑的值
  144. 在生成器里不保证可见。绑定发生在生成器内部,生命周期与本次流一致。
  145. """
  146. from obs import bind_trace_id, unbind_trace_id
  147. if self._consumed:
  148. raise RuntimeError("同一任务只能执行一次,请创建新的 Execution")
  149. self._consumed = True
  150. token = bind_trace_id(self.req.trace_id or None)
  151. stream = None
  152. try:
  153. await asyncio.to_thread(run_store.start, self.req.run_id, self.req.session_id,
  154. self.req.user_name, self.req.kind)
  155. self._started = True
  156. yield self.writer.agent_start(self.req.agent_name, self.req.agent_cn_name)
  157. yield self.writer.thinking("正在理解您的问题...", self.req.agent_cn_name)
  158. graph_input = await self._prepare_input()
  159. stream = self.req.agent.astream(
  160. graph_input,
  161. stream_mode=["updates", "messages"],
  162. config=self.req.config,
  163. subgraphs=True,
  164. version="v2",
  165. )
  166. async for frame in self._drive(stream):
  167. yield frame
  168. # 中断检测必须在流结束后做:interrupt 会让流自然结束
  169. async for frame in self._emit_interrupts():
  170. yield frame
  171. async for frame in self._completion():
  172. yield frame
  173. except (asyncio.CancelledError, GeneratorExit):
  174. self.status = "cancelled"
  175. raise
  176. except Exception as exc:
  177. self.status = "failed"
  178. self.error = type(exc).__name__
  179. async for frame in self._error_events(exc):
  180. yield frame
  181. finally:
  182. # 清理与发帧独立。ASGI 断连的取消作用域不能打断保存和释放。
  183. with anyio.CancelScope(shield=True):
  184. cleanup = asyncio.create_task(self._cleanup(stream))
  185. try:
  186. await asyncio.shield(cleanup)
  187. except asyncio.CancelledError:
  188. await cleanup
  189. unbind_trace_id(token)
  190. async def _cleanup(self, stream: Any) -> None:
  191. try:
  192. if stream is not None:
  193. close = getattr(stream, "aclose", None)
  194. if close is not None:
  195. try:
  196. await close()
  197. except Exception:
  198. logger.exception("关闭图事件流失败 session=%s", self.req.session_id)
  199. await self._settle(aborted=self.status in {"cancelled", "failed"})
  200. finally:
  201. try:
  202. await asyncio.to_thread(self._finalize)
  203. finally:
  204. if self.req.session_lease is not None:
  205. self.req.session_lease.release()
  206. # ========================================================
  207. # 输入准备
  208. # ========================================================
  209. async def _prepare_input(self) -> Any:
  210. """组装喂给图的输入。
  211. 两种情形:
  212. - **恢复**(``resume_payload`` 非 None):喂 ``Command(resume=...)``,
  213. 不组装消息——图会从挂起的那个节点继续。
  214. - **新请求**:热线程(checkpoint 里已有状态)只喂当轮新消息;冷线程
  215. 需要从会话历史补种。
  216. 补种这一步不能省,也不能对所有线程都做:热线程若把历史重复追加,
  217. 上下文会随轮次按 O(N²) 膨胀——这是长会话最容易被忽略的膨胀来源。
  218. """
  219. if self.req.resume_payload is not None:
  220. from langgraph.types import Command
  221. logger.info("恢复中断 session=%s", self.req.session_id)
  222. return Command(resume=self.req.resume_payload)
  223. from langchain_core.messages import AIMessage, HumanMessage
  224. from agents import agent_memory
  225. messages: list[Any] = []
  226. warm = await agent_memory.has_checkpoint(self.req.thread_id)
  227. if not warm:
  228. history = await asyncio.to_thread(chat_store.get_messages,
  229. self.req.session_id, limit=config.CTX_RECENT_LIMIT
  230. )
  231. for row in history:
  232. content = row.get("content") or ""
  233. if not content:
  234. continue
  235. if row.get("msg_role") == "user":
  236. messages.append(HumanMessage(content=content))
  237. else:
  238. messages.append(AIMessage(content=content))
  239. if messages:
  240. logger.info(
  241. "冷线程补种历史 session=%s 条数=%d", self.req.session_id, len(messages)
  242. )
  243. else:
  244. logger.info("热线程,仅喂当轮消息 session=%s", self.req.session_id)
  245. # 用户消息先落库:即使这一轮流到一半挂了,历史里也得有用户问了什么
  246. await asyncio.to_thread(chat_store.save_message, self.req.session_id, "user",
  247. self.req.original_user_message or self.req.message)
  248. self.meter.note_text(self.req.message)
  249. messages.append(HumanMessage(content=self.req.message))
  250. # 必须是 {"messages": [...]} 字典:直接传 list 会被当成 __start__ 节点的
  251. # 返回值,图会报 InvalidUpdateError: Expected dict
  252. return {"messages": messages}
  253. # ========================================================
  254. # 流驱动
  255. # ========================================================
  256. async def _drive(self, stream: Any) -> AsyncIterator[ExecutionEvent]:
  257. """消费 astream,派发事件帧。带墙钟上限。"""
  258. remaining = config.TASK_WALL_TIMEOUT - (time.time() - self.req.started_at)
  259. deadline = time.monotonic() + max(remaining, 0)
  260. timed_out = False
  261. iterator = stream.__aiter__()
  262. while True:
  263. if time.monotonic() >= deadline:
  264. timed_out = True
  265. break
  266. try:
  267. chunk = await asyncio.wait_for(
  268. iterator.__anext__(), timeout=max(deadline - time.monotonic(), 0.001)
  269. )
  270. except StopAsyncIteration:
  271. break
  272. except asyncio.TimeoutError:
  273. timed_out = True
  274. break
  275. mode, ns, data = normalize_chunk(chunk)
  276. if mode == "updates":
  277. for frame in self._handle_updates(data, ns):
  278. yield frame
  279. elif mode == "messages":
  280. for frame in self._handle_messages(data, ns):
  281. yield frame
  282. # mode 为 None:无法识别的 chunk,忽略而不是猜
  283. if timed_out:
  284. raise concurrency.TaskWallTimeoutError(
  285. f"任务超过墙钟上限 {config.TASK_WALL_TIMEOUT}s 已被中止"
  286. )
  287. # 子智能体可能因为异常/中断没收到结果,收尾时统一关掉,避免前端卡片长转
  288. for run in self.tracker.finish_all():
  289. yield self.writer.progress(
  290. f"子智能体 {run.subagent_type} 未正常返回,已中止",
  291. subagent=run.subagent_type, subagent_id=run.call_id,
  292. )
  293. # -------- updates 通道 --------
  294. def _handle_updates(self, data: Any, ns: tuple[Any, ...]) -> list[ExecutionEvent]:
  295. """处理 updates 通道:工具事件、待办、子智能体台账。"""
  296. from langchain_core.messages import AIMessage, ToolMessage
  297. frames: list[ExecutionEvent] = []
  298. if not isinstance(data, dict):
  299. return frames
  300. # 中断会以 __interrupt__ 键出现在 updates 里,比 aget_state 更及时
  301. if "__interrupt__" in data:
  302. for item in data.get("__interrupt__") or ():
  303. self.interrupt_payloads.append(pack(item))
  304. return frames
  305. # 子智能体身份:updates 的 ns 指向它所属的 task 调用
  306. run = self.tracker.note_namespace(ns) if ns else None
  307. subagent_type = run.subagent_type if run else ""
  308. subagent_id = run.call_id if run else ""
  309. for node, update in data.items():
  310. if not isinstance(update, dict):
  311. continue
  312. todos = update.get("todos")
  313. if isinstance(todos, list):
  314. self.todos = [t for t in todos if isinstance(t, dict)]
  315. frames.append(self.writer.todo_list(self.todos))
  316. for message in update.get("messages") or ():
  317. if isinstance(message, ToolMessage):
  318. frames.extend(self._on_tool_message(message, subagent_type, subagent_id))
  319. elif isinstance(message, AIMessage):
  320. frames.extend(self._on_ai_message(message, subagent_type, subagent_id))
  321. return frames
  322. def _on_ai_message(
  323. self, message: Any, subagent_type: str, subagent_id: str
  324. ) -> list[ExecutionEvent]:
  325. """处理一条完整的 AIMessage:登记 task 调用、发工具调用卡。"""
  326. from tools.registry import get_label
  327. frames: list[ExecutionEvent] = []
  328. tool_calls = list(getattr(message, "tool_calls", None) or [])
  329. # 先让跟踪器登记 task 调用,这样后续子智能体事件能立刻带上身份
  330. for run in self.tracker.note_task_calls(tool_calls):
  331. frames.append(
  332. self.writer.progress(
  333. f"委派子智能体:{run.subagent_type}",
  334. subagent=run.subagent_type,
  335. subagent_id=run.call_id,
  336. description=run.description,
  337. )
  338. )
  339. names = [c.get("name", "") for c in tool_calls if isinstance(c, dict)]
  340. if names:
  341. frames.append(
  342. self.writer.executing(
  343. f"调用工具:{'、'.join(get_label(n) for n in names)}",
  344. cn_agent=self.req.agent_cn_name,
  345. tools=names,
  346. cn_tools=[get_label(n) for n in names],
  347. )
  348. )
  349. for call in tool_calls:
  350. if not isinstance(call, dict):
  351. continue
  352. name = call.get("name") or ""
  353. if not name:
  354. continue
  355. args = call.get("args")
  356. self.tool_calls.append({"name": name, "args": args})
  357. frames.append(
  358. self.writer.tool_call(
  359. name,
  360. args=args,
  361. cn_tool=get_label(name),
  362. agent_id=subagent_id,
  363. tool_call_id=str(call.get("id") or ""),
  364. )
  365. )
  366. return frames
  367. def _on_tool_message(
  368. self, message: Any, subagent_type: str, subagent_id: str
  369. ) -> list[ExecutionEvent]:
  370. """处理一条 ToolMessage:发工具结果卡、抽取来源、结子智能体。"""
  371. from tools.registry import get_label
  372. frames: list[ExecutionEvent] = []
  373. name = getattr(message, "name", None) or "tool"
  374. content = _message_text(message)
  375. # task 的返回即子智能体结束
  376. call_id = getattr(message, "tool_call_id", None) or ""
  377. if name == "task" and call_id:
  378. run = self.tracker.note_result(call_id, ok=True)
  379. if run is not None:
  380. frames.append(
  381. self.writer.agent_done(
  382. run.subagent_type,
  383. cn_agent=run.subagent_type,
  384. duration_ms=run.duration_ms,
  385. progress="已完成",
  386. agent_id=run.call_id,
  387. )
  388. )
  389. self.meter.note_text(content)
  390. self.sources.note_tool_result(name, content)
  391. frames.append(
  392. self.writer.tool_result(
  393. name,
  394. content=content[:4000],
  395. cn_tool=get_label(name),
  396. agent_id=subagent_id,
  397. error=str(getattr(message, "status", "") or "") == "error",
  398. tool_call_id=call_id,
  399. )
  400. )
  401. return frames
  402. # -------- messages 通道 --------
  403. def _handle_messages(self, data: Any, ns: tuple[Any, ...]) -> list[ExecutionEvent]:
  404. """处理 messages 通道:正文 token、思考 token、用量计量。"""
  405. from langchain_core.messages import AIMessage, AIMessageChunk
  406. frames: list[ExecutionEvent] = []
  407. if not isinstance(data, tuple) or len(data) != 2:
  408. return frames
  409. chunk, metadata = data
  410. metadata = metadata if isinstance(metadata, dict) else {}
  411. # 摘要中间件的模型调用也在这条通道上:它的 token 不能给用户看
  412. source = metadata.get("lc_source")
  413. # 汇总全部模型费用;只有主线程更新当前上下文窗口占用。
  414. self.meter.note_message(chunk, is_main_context=not ns and source != _SUMMARIZATION_SOURCE)
  415. if source == _SUMMARIZATION_SOURCE or ns:
  416. return frames
  417. # 只有 AIMessageChunk 的内容才是「模型正在说的话」。ToolMessage 也走
  418. # 这条通道,它的 content 是工具返回的原始数据——不加这个判断,工具
  419. # 结果的 JSON 会混进给用户看的正文里。
  420. if not isinstance(chunk, (AIMessage, AIMessageChunk)):
  421. return frames
  422. reasoning = _reasoning_text(chunk)
  423. if reasoning:
  424. frames.append(self.writer.thinking_token(reasoning))
  425. text = _message_text(chunk)
  426. if text:
  427. self.meter.note_text(text)
  428. self.answer_parts.append(text)
  429. frames.append(self.writer.token(text))
  430. return frames
  431. # ========================================================
  432. # 中断
  433. # ========================================================
  434. async def _emit_interrupts(self) -> AsyncIterator[ExecutionEvent]:
  435. """把中断转成 1.5 形状的 interrupt 事件。
  436. 两个来源:流里捕获到的 ``__interrupt__``,以及流结束后从图状态里
  437. 读出的 pending interrupt。优先用前者(更及时),两者都空则无中断。
  438. """
  439. payloads = list(self.interrupt_payloads)
  440. if not payloads:
  441. payloads = await self._interrupts_from_state()
  442. if not payloads:
  443. return
  444. self.interrupt_payloads = payloads
  445. for payload in payloads:
  446. frame = self._adapt_interrupt(payload)
  447. if frame:
  448. yield frame
  449. async def _interrupts_from_state(self) -> list[Any]:
  450. """从图状态里读出挂起的中断载荷。"""
  451. get_state = getattr(self.req.agent, "aget_state", None)
  452. if get_state is None:
  453. return []
  454. try:
  455. state = await get_state(self.req.config)
  456. except Exception as exc: # pragma: no cover - 取决于底层库
  457. logger.warning("读取图状态失败:%s: %s", type(exc).__name__, exc)
  458. return []
  459. found: list[Any] = []
  460. for task in getattr(state, "tasks", None) or ():
  461. for item in getattr(task, "interrupts", None) or ():
  462. found.append(pack(item))
  463. return found
  464. def _adapt_interrupt(self, payload: Any) -> ExecutionEvent:
  465. return interrupt_event(self.writer, payload)
  466. # ========================================================
  467. # 收尾
  468. # ========================================================
  469. async def _settle(self, aborted: bool = False) -> None:
  470. # 断连恰逢数据库写入时继续等待同一次收尾,避免在线程仍写入时重做。
  471. if self._settlement_task is None:
  472. self._settlement_task = asyncio.create_task(self._settle_once(aborted))
  473. await asyncio.shield(self._settlement_task)
  474. async def _settle_once(self, aborted: bool = False) -> None:
  475. """落库 → 计费 → 来源 → 上下文用量 → agent_done/done。
  476. 无论正常结束、异常还是断连都走这里(幂等,重复调用只生效一次)。
  477. """
  478. if self.finished:
  479. return
  480. duration_ms = int((time.time() - self.req.started_at) * 1000)
  481. answer = "".join(self.answer_parts)
  482. if aborted and answer:
  483. answer += "\n\n(本次回复被中断,内容可能不完整)"
  484. # 1) 上下文用量快照。**只取一次**:snapshot() 会消费掉压缩存档的待办标记,
  485. # 取两次会导致第二次拿不到存档内容。
  486. self._apply_estimate_fallback()
  487. if self._snapshot is None:
  488. self._snapshot = self.meter.snapshot()
  489. snapshot = self._snapshot
  490. # 2) 计费:把上下文计量器的数字转给计费计量器。
  491. # 这是两个计量器唯一的连接点——漏了这一步,交互式对话全程不会计费,
  492. # 而 /usage/today 会显示 0,表面上「一切正常」。
  493. if not self._billing_added:
  494. self.billing.add(snapshot.input_tokens, snapshot.output_tokens, snapshot.usage_source)
  495. self._billing_added = True
  496. # 3) 落库:assistant 消息是会话历史的事实源
  497. citations = self.sources.build_citations()
  498. if (answer or citations) and not self._answer_saved:
  499. try:
  500. await asyncio.to_thread(
  501. chat_store.save_message, self.req.session_id,
  502. "assistant",
  503. answer,
  504. tool_calls=self.tool_calls or None,
  505. citations=citations or None,
  506. )
  507. self._answer_saved = True
  508. except Exception as exc:
  509. self.status = "failed"
  510. self.error = "message_persistence_failed"
  511. logger.error("落库失败 session=%s:%s: %s", self.req.session_id, type(exc).__name__, exc)
  512. # 4) 上下文用量 + 压缩存档
  513. await self._persist_context(snapshot)
  514. # 5) 计费入账(幂等)
  515. await asyncio.to_thread(self.billing.settle)
  516. if self.status == "running":
  517. self.status = "waiting_input" if self.interrupt_payloads else "success"
  518. if self._started:
  519. await asyncio.to_thread(run_store.finish, self.req.run_id, self.status, self.error)
  520. # 用户在普通恢复接口完成定时任务审批后,同步其执行记录。
  521. if self.req.kind == "resume":
  522. await asyncio.to_thread(db_conn.execute,
  523. "UPDATE schedule_runs SET status = ?, finished_at = ? WHERE thread_id = ? AND status = 'waiting_input'",
  524. (self.status, None if self.status == "waiting_input" else db_conn.fmt_dt(), self.req.session_id))
  525. self.finished = True
  526. async def _completion(self, aborted: bool = False) -> AsyncIterator[ExecutionEvent]:
  527. await self._settle(aborted)
  528. duration_ms = int((time.time() - self.req.started_at) * 1000)
  529. citations = self.sources.build_citations()
  530. # 6) 来源标注
  531. yield self.writer.citations(citations, model_only=not citations)
  532. if aborted or self.status in {"failed", "cancelled"}:
  533. if self.error == "message_persistence_failed":
  534. yield self.writer.error("回复未能保存,请稍后检查历史记录", code=self.error)
  535. yield self.writer.done(duration_ms=duration_ms, aborted=True)
  536. return
  537. if self.status == "waiting_input":
  538. yield self.writer.done(duration_ms=duration_ms, waiting_input=True, run_id=self.req.run_id, usage=self._usage_payload())
  539. return
  540. yield self.writer.agent_done(
  541. self.req.agent_name,
  542. cn_agent=self.req.agent_cn_name,
  543. duration_ms=duration_ms,
  544. progress="已完成",
  545. **self.meter.summary_fields(),
  546. )
  547. yield self.writer.done(duration_ms=duration_ms, run_id=self.req.run_id, usage=self._usage_payload())
  548. def _usage_payload(self) -> dict[str, Any]:
  549. """给前端的用量摘要(便于测试页显示,1.5 前端会忽略未知字段)。"""
  550. return {
  551. "input_tokens": self.meter.input_tokens,
  552. "output_tokens": self.meter.output_tokens,
  553. "total_tokens": self.meter.total_tokens,
  554. "usage_source": self.meter.usage_source,
  555. "context_tokens": self.meter.context_tokens,
  556. }
  557. def _apply_estimate_fallback(self) -> None:
  558. """provider 没返回任何 usage 时,用缓冲区文本估算兜底。
  559. 不兜底的话,这类请求会以 0 token 入账——用户白用、额度不受控,
  560. 而且从账面上完全看不出来。兜底后至少 usage_source=estimate 能
  561. 标示这批数字不可尽信。
  562. """
  563. if self.meter.total_tokens > 0:
  564. return
  565. estimated = self.meter.estimate_from_buffer()
  566. if estimated > 0:
  567. logger.warning(
  568. "本次未取到 provider usage,按缓冲区估算 %d tokens(usage_source=estimate)"
  569. " session=%s",
  570. estimated, self.req.session_id,
  571. )
  572. self.meter.note_estimated(estimated)
  573. async def _persist_context(self, snapshot: Any) -> None:
  574. """落 session_context_usage,并在发生压缩时登记原文存档位置。
  575. 存档优先用**框架自己的产物**:摘要中间件会把被淘汰的历史追加写到
  576. ``/conversation_history/{thread_id}.md``,并把路径嵌进摘要正文——
  577. 那份文件是权威的,agent 也能自己读回来。只有框架产物不存在时才
  578. 回落写我们自己的缓冲区,避免「以为有存档其实没有」。
  579. """
  580. archive_path = ""
  581. try:
  582. if snapshot.compacted:
  583. archive_path = await asyncio.to_thread(
  584. self._archive_compaction, snapshot.archive_text, snapshot.archive_tokens
  585. )
  586. await asyncio.to_thread(chat_store.save_context_usage,
  587. session_id=self.req.session_id,
  588. turn=chat_store.count_messages(self.req.session_id) // 2,
  589. input_tokens=snapshot.input_tokens,
  590. output_tokens=snapshot.output_tokens,
  591. cache_read_tokens=snapshot.cache_read_tokens,
  592. usage_source=snapshot.usage_source,
  593. compacted=snapshot.compacted,
  594. compact_count=snapshot.compact_count,
  595. snapshot={
  596. "context_tokens": snapshot.context_tokens,
  597. "context_window": self.meter.window,
  598. "model_calls": snapshot.model_calls,
  599. "archive_path": archive_path,
  600. },
  601. )
  602. except Exception as exc:
  603. logger.error(
  604. "上下文用量落库失败 session=%s:%s: %s",
  605. self.req.session_id, type(exc).__name__, exc,
  606. )
  607. def _archive_compaction(self, buffered_text: str | None, buffered_tokens: int) -> str:
  608. """登记(必要时写出)压缩前原文,返回存档路径或空串。"""
  609. from agents.workspace import session_root
  610. history_dir = session_root(self.req.user_name, self.req.session_id) / "conversation_history"
  611. framework_file = history_dir / f"{self.req.thread_id}.md"
  612. if framework_file.exists():
  613. path = str(framework_file)
  614. chat_store.save_context_archive(
  615. self.req.session_id, 0, path, buffered_tokens
  616. )
  617. logger.info("压缩原文由框架落盘:%s", path)
  618. return path
  619. if not buffered_text:
  620. logger.warning(
  621. "检测到压缩但既无框架落盘也无本地缓冲,无法存档 session=%s",
  622. self.req.session_id,
  623. )
  624. return ""
  625. archive_dir = history_dir
  626. archive_dir.mkdir(parents=True, exist_ok=True)
  627. stamp = time.strftime("%Y%m%d-%H%M%S")
  628. target = archive_dir / f"{self.req.session_id}-{stamp}.md"
  629. target.write_text(buffered_text, encoding="utf-8")
  630. chat_store.save_context_archive(
  631. self.req.session_id, 0, str(target), buffered_tokens
  632. )
  633. logger.info("压缩原文由本进程落盘:%s(%d tokens)", target, buffered_tokens)
  634. return str(target)
  635. # ========================================================
  636. # 异常
  637. # ========================================================
  638. async def _error_events(self, exc: Exception) -> AsyncIterator[ExecutionEvent]:
  639. """把异常翻译成 error 帧 + 收尾帧。"""
  640. from langgraph.errors import GraphRecursionError
  641. if isinstance(exc, GraphRecursionError):
  642. logger.warning("图递归上限触发 session=%s:%s", self.req.session_id, exc)
  643. yield self.writer.error(
  644. "任务步骤过多已被中止,请把任务拆成更小的步骤重试", code="recursion_limit"
  645. )
  646. async for frame in self._completion(aborted=True):
  647. yield frame
  648. return
  649. if isinstance(exc, concurrency.TaskWallTimeoutError):
  650. logger.warning("任务墙钟超时 session=%s:%s", self.req.session_id, exc)
  651. yield self.writer.error(str(exc), code="task_timeout")
  652. async for frame in self._completion(aborted=True):
  653. yield frame
  654. return
  655. logger.exception("流式执行异常 session=%s", self.req.session_id)
  656. yield self.writer.error(
  657. f"服务内部错误:{type(exc).__name__}", code="internal_error"
  658. )
  659. async for frame in self._completion(aborted=True):
  660. yield frame
  661. # ========================================================
  662. # 资源
  663. # ========================================================
  664. def _finalize(self) -> None:
  665. """释放并发位并兜底结算。幂等。"""
  666. try:
  667. if not self.billing.settled:
  668. self.billing.settle()
  669. finally:
  670. if self.req.slot is not None:
  671. self.req.slot.release()
  672. def _meter_for_request() -> ContextMeter:
  673. """给本轮请求建一个上下文计量器,窗口取自模型配置。"""
  674. from agents import model_factory
  675. try:
  676. window = model_factory.context_window()
  677. except Exception: # pragma: no cover - 配置异常不该阻断对话
  678. window = 0
  679. return ContextMeter(window=window)