From 4a2ecd430e7a12c988405ba17d0c2d1840d7b6e4 Mon Sep 17 00:00:00 2001 From: Huixin615 Date: Thu, 23 Jul 2026 22:56:12 +0800 Subject: [PATCH] fix(streaming): expose custom events to astream_events (#4403) * fix(streaming): expose custom events to astream_events * test(streaming): validate real custom event emitters --------- Co-authored-by: Willem Jiang --- README.md | 2 + backend/AGENTS.md | 3 +- backend/docs/STREAMING.md | 12 +- .../llm_error_handling_middleware.py | 61 +++- .../safety_finish_reason_middleware.py | 103 ++++-- .../deerflow/tools/builtins/task_tool.py | 41 ++- .../harness/deerflow/utils/custom_events.py | 57 ++++ backend/tests/test_custom_events.py | 304 ++++++++++++++++++ .../test_llm_error_handling_middleware.py | 56 ++++ ..._safety_finish_reason_graph_integration.py | 37 +++ .../test_safety_finish_reason_middleware.py | 87 +++++ backend/tests/test_task_tool_core_logic.py | 14 + 12 files changed, 726 insertions(+), 51 deletions(-) create mode 100644 backend/packages/harness/deerflow/utils/custom_events.py create mode 100644 backend/tests/test_custom_events.py diff --git a/README.md b/README.md index c73873787..a7c4d6084 100644 --- a/README.md +++ b/README.md @@ -332,6 +332,8 @@ DeerFlow runs the agent runtime inside the Gateway API. Development mode enables Gateway owns `/api/langgraph/*` and translates those public LangGraph-compatible paths to its native `/api/*` routers behind nginx. +DeerFlow's built-in custom events are available through both LangGraph streaming interfaces: native clients can continue subscribing to `stream_mode="custom"`, while callback-based integrations can consume the same payloads as `on_custom_event` records from `astream_events(version="v2")`. The callback event name matches the payload's `type` field. + #### Docker Production Deployment `deploy.sh` supports building and starting separately: diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 9e94d11fd..d01060492 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -963,8 +963,9 @@ Gateway API endpoints and `DeerFlowClient` methods can modify MCP servers and sk - `stream(message, thread_id)` — subscribes to LangGraph `stream_mode=["values", "messages", "custom"]` and yields `StreamEvent`: - `"values"` — full state snapshot (title, messages, artifacts); AI text already delivered via `messages` mode is **not** re-synthesized here to avoid duplicate deliveries - `"messages-tuple"` — per-chunk update: for AI text this is a **delta** (concat per `id` to rebuild the full message); tool calls and tool results are emitted once each - - `"custom"` — forwarded from `StreamWriter` + - `"custom"` — forwarded from `StreamWriter`; DeerFlow-built-in custom events are dual-emitted through `deerflow.utils.custom_events`, so `astream_events(version="v2")` consumers also receive one `on_custom_event` with `name=payload["type"]` and the unchanged payload as `data` - `"end"` — stream finished (carries cumulative `usage` counted once per message id) +- **Custom-event invariant** — production DeerFlow emitters must use `emit_custom_event` / `aemit_custom_event`, not call `StreamWriter` alone. Every built-in payload must carry a non-empty string `type`; typeless payloads remain writer-only and are intentionally absent from `astream_events`. The writer runs first and remains authoritative for Gateway, Web UI, and embedded-client compatibility; callback dispatch is best-effort and must not break that path. Async graph hooks must await the async helper rather than invoking synchronous dispatch on a running event loop. - Agent created lazily via `create_agent()` + `build_middlewares()`, same as `make_lead_agent` - Supports `checkpointer` parameter for state persistence across turns - `reset_agent()` forces agent recreation (e.g. after memory or skill changes) diff --git a/backend/docs/STREAMING.md b/backend/docs/STREAMING.md index 73756a276..7b0039b96 100644 --- a/backend/docs/STREAMING.md +++ b/backend/docs/STREAMING.md @@ -7,7 +7,7 @@ ## TL;DR - DeerFlow 有**两条并行**的流式路径：**Gateway 路径**（async / HTTP SSE / JSON 序列化）服务浏览器和 IM 渠道；**DeerFlowClient 路径**（sync / in-process / 原生 LangChain 对象）服务 Jupyter、脚本、测试。它们**无法合并**——消费者模型不同。 -- 两条路径都从 `create_agent()` 工厂出发，核心都是订阅 LangGraph 的 `stream_mode=["values", "messages", "custom"]`。`values` 是节点级 state 快照，`messages` 是 LLM token 级 delta，`custom` 是显式 `StreamWriter` 事件。**这三种模式不是详细程度的梯度，是三个独立的事件源**，要 token 流就必须显式订阅 `messages`。 +- 两条路径都从 `create_agent()` 工厂出发，核心都是订阅 LangGraph 的 `stream_mode=["values", "messages", "custom"]`。`values` 是节点级 state 快照，`messages` 是 LLM token 级 delta，`custom` 是显式 `StreamWriter` 事件。DeerFlow 内置 custom 事件同时通过 callback dispatch 暴露为 `astream_events(version="v2")` 的 `on_custom_event`，供 AG-UI 等 callback 型消费者使用。**这些接口不是详细程度的梯度，而是独立事件源**，消费者必须订阅自己需要的接口。 - 嵌入式 client 为每个 `stream()` 调用维护三个 `set[str]`：`seen_ids` / `streamed_ids` / `counted_usage_ids`。三者看起来相似但管理**三个独立的不变式**，不能合并。 --- @@ -65,11 +65,14 @@ flowchart LR LG -->|"每个节点完成后"| V["values: 完整 state 快照"] Node1 -->|"LLM 每产生一个 token"| M["messages: (AIMessageChunk, meta)"] - Node1 -->|"StreamWriter.write()"| C["custom: 任意 dict"] + Node1 -->|"emit_custom_event()"| E["DeerFlow custom event helper"] + E -->|"StreamWriter.write()"| C["custom: 任意 dict"] + E -->|"dispatch_custom_event()"| A["astream_events(v2): on_custom_event"] class V values class M messages class C custom + class A custom ``` | Mode | 发射时机 | Payload | 粒度 | @@ -77,6 +80,9 @@ flowchart LR | `values` | 每个 graph 节点完成后 | 完整 state dict（title、messages、artifacts）| 节点级 | | `messages` | LLM 每次 yield 一个 chunk；tool 节点完成时 | `(AIMessageChunk \| ToolMessage, metadata_dict)` | token 级 | | `custom` | 用户代码显式调用 `StreamWriter.write()` | 任意 dict | 应用定义 | +| `on_custom_event` | 用户代码调用 `dispatch_custom_event()`；通过 `astream_events(version="v2")` 消费 | `name` + 任意 `data` | 应用定义 | + +DeerFlow 自身产生的事件必须通过 `deerflow.utils.custom_events` 的同步或异步 helper 发送，并且每个内置 payload 必须携带非空字符串 `type`；缺少合法 `type` 的 payload 只进入 `custom` stream，不会出现在 `astream_events`。helper 先写入 `custom` stream，再 best-effort dispatch callback；callback 名称取 payload 的 `type`，`data` 保留完整 payload。这样原生 Gateway / Web UI / `DeerFlowClient` 的 custom 事件不变，`astream_events` 消费者也能观察同一事件。callback dispatch 的普通异常只记 debug 日志，不允许打断原有 writer 链路；writer 的异常语义保持不变。 ### 两套命名的由来 @@ -293,6 +299,8 @@ sequenceDiagram 本文档的直接起因是 bytedance/deer-flow#1969：`DeerFlowClient.stream()` 原本只订阅 `["values", "custom"]`，**漏了 `"messages"`**。结果 `client.stream("hello")` 等价于一次性返回，视觉上和 `chat()` 没区别。 +Custom 事件还有一条独立回归边界：`get_stream_writer()` 产生的 chunk 不会自动成为 `astream_events(version="v2")` 的 `on_custom_event`，而 callback dispatch 也不会自动进入 `stream_mode="custom"`。测试必须使用真实最小 LangGraph 同时锁定两种 API，断言每个消费者各收到一次且 payload 相同；仅 mock 任一函数无法证明协议互操作。 + 这类 bug 有三个结构性原因： 1. **多协议层命名**：`messages` / `messages-tuple` / HTTP SSE `messages` 是同一概念的三个名字。在其中一层出错不会在另外两层报错。 diff --git a/backend/packages/harness/deerflow/agents/middlewares/llm_error_handling_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/llm_error_handling_middleware.py index e04ee4d5f..1b840fe09 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/llm_error_handling_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/llm_error_handling_middleware.py @@ -23,6 +23,7 @@ from langchain_core.messages import AIMessage from langgraph.errors import GraphBubbleUp from deerflow.config.app_config import AppConfig +from deerflow.utils.custom_events import aemit_custom_event, emit_custom_event logger = logging.getLogger(__name__) @@ -709,6 +710,26 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]): detail=_extract_error_detail(exc), ) + def _build_retry_event( + self, + attempt: int, + wait_ms: int, + reason: str, + *, + max_attempts: int, + ) -> dict[str, Any]: + return { + "type": "llm_retry", + "attempt": attempt, + # Effective budget for this call (burst-rate == 2), not the + # configured ceiling - the frontend renders this and the + # ``message`` below, so both must describe the loop that runs. + "max_attempts": max_attempts, + "wait_ms": wait_ms, + "reason": reason, + "message": self._build_retry_message(attempt, wait_ms, reason, max_attempts=max_attempts), + } + def _emit_retry_event( self, attempt: int, @@ -721,22 +742,36 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]): from langgraph.config import get_stream_writer writer = get_stream_writer() - writer( - { - "type": "llm_retry", - "attempt": attempt, - # Effective budget for this call (burst-rate == 2), not the - # configured ceiling - the frontend renders this and the - # ``message`` below, so both must describe the loop that runs. - "max_attempts": max_attempts, - "wait_ms": wait_ms, - "reason": reason, - "message": self._build_retry_message(attempt, wait_ms, reason, max_attempts=max_attempts), - } + emit_custom_event( + self._build_retry_event(attempt, wait_ms, reason, max_attempts=max_attempts), + writer=writer, ) + except GraphBubbleUp: + raise except Exception: logger.debug("Failed to emit llm_retry event", exc_info=True) + async def _aemit_retry_event( + self, + attempt: int, + wait_ms: int, + reason: str, + *, + max_attempts: int, + ) -> None: + try: + from langgraph.config import get_stream_writer + + writer = get_stream_writer() + await aemit_custom_event( + self._build_retry_event(attempt, wait_ms, reason, max_attempts=max_attempts), + writer=writer, + ) + except GraphBubbleUp: + raise + except Exception: + logger.debug("Failed to emit async llm_retry event", exc_info=True) + @override def wrap_model_call( self, @@ -834,7 +869,7 @@ class LLMErrorHandlingMiddleware(AgentMiddleware[AgentState]): wait_ms, _extract_error_detail(exc), ) - self._emit_retry_event(attempt, wait_ms, reason, max_attempts=max_attempts) + await self._aemit_retry_event(attempt, wait_ms, reason, max_attempts=max_attempts) await asyncio.sleep(wait_ms / 1000) attempt += 1 continue diff --git a/backend/packages/harness/deerflow/agents/middlewares/safety_finish_reason_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/safety_finish_reason_middleware.py index 4a57588ea..1a83d0764 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/safety_finish_reason_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/safety_finish_reason_middleware.py @@ -48,11 +48,13 @@ and Loop then accounts against the cleaned message. from __future__ import annotations import logging +from dataclasses import dataclass from typing import TYPE_CHECKING, override from langchain.agents import AgentState from langchain.agents.middleware import AgentMiddleware from langchain_core.messages import AIMessage +from langgraph.errors import GraphBubbleUp from langgraph.runtime import Runtime from deerflow.agents.middlewares.safety_termination_detectors import ( @@ -62,6 +64,7 @@ from deerflow.agents.middlewares.safety_termination_detectors import ( ) from deerflow.agents.middlewares.tool_call_metadata import clone_ai_message_with_tool_calls from deerflow.runtime.events.catalog import MIDDLEWARE_SAFETY_TERMINATION_TAG +from deerflow.utils.custom_events import aemit_custom_event, emit_custom_event from deerflow.utils.messages import message_content_to_text if TYPE_CHECKING: @@ -84,6 +87,15 @@ _USER_FACING_MESSAGE = ( _USER_FACING_EMPTY_MESSAGE = "The model provider stopped this response with a safety-related signal ({reason_field}={reason_value!r}, detector={detector!r}) and returned no content. Please rephrase your request or start a new conversation." +@dataclass(frozen=True) +class _SafetyIntervention: + update: dict + termination: SafetyTermination + suppressed_names: list[str] + message: AIMessage + tool_calls: list[dict] + + class SafetyFinishReasonMiddleware(AgentMiddleware[AgentState]): """Repair AIMessages flagged by a SafetyTerminationDetector: strip tool calls, or backfill an explanation when the message is otherwise empty.""" @@ -190,6 +202,25 @@ class SafetyFinishReasonMiddleware(AgentMiddleware[AgentState]): # ----- observability --------------------------------------------------- + @staticmethod + def _build_event_payload( + termination: SafetyTermination, + suppressed_names: list[str], + runtime: Runtime, + ) -> dict: + thread_id = None + if runtime is not None and getattr(runtime, "context", None): + thread_id = runtime.context.get("thread_id") if isinstance(runtime.context, dict) else None + return { + "type": "safety_termination", + "detector": termination.detector, + "reason_field": termination.reason_field, + "reason_value": termination.reason_value, + "suppressed_tool_call_count": len(suppressed_names), + "suppressed_tool_call_names": suppressed_names, + "thread_id": thread_id, + } + def _emit_event( self, termination: SafetyTermination, @@ -204,29 +235,42 @@ class SafetyFinishReasonMiddleware(AgentMiddleware[AgentState]): from langgraph.config import get_stream_writer writer = get_stream_writer() + except GraphBubbleUp: + raise except Exception: # noqa: BLE001 logger.debug("get_stream_writer unavailable; skipping safety_termination event", exc_info=True) return - thread_id = None - if runtime is not None and getattr(runtime, "context", None): - thread_id = runtime.context.get("thread_id") if isinstance(runtime.context, dict) else None - try: - writer( - { - "type": "safety_termination", - "detector": termination.detector, - "reason_field": termination.reason_field, - "reason_value": termination.reason_value, - "suppressed_tool_call_count": len(suppressed_names), - "suppressed_tool_call_names": suppressed_names, - "thread_id": thread_id, - } - ) + emit_custom_event(self._build_event_payload(termination, suppressed_names, runtime), writer=writer) + except GraphBubbleUp: + raise except Exception: # noqa: BLE001 logger.debug("Failed to emit safety_termination stream event", exc_info=True) + async def _aemit_event( + self, + termination: SafetyTermination, + suppressed_names: list[str], + runtime: Runtime, + ) -> None: + try: + from langgraph.config import get_stream_writer + + writer = get_stream_writer() + except GraphBubbleUp: + raise + except Exception: # noqa: BLE001 + logger.debug("get_stream_writer unavailable; skipping async safety_termination event", exc_info=True) + return + + try: + await aemit_custom_event(self._build_event_payload(termination, suppressed_names, runtime), writer=writer) + except GraphBubbleUp: + raise + except Exception: # noqa: BLE001 + logger.debug("Failed to emit async safety_termination stream event", exc_info=True) + def _record_audit_event( self, termination: SafetyTermination, @@ -286,7 +330,7 @@ class SafetyFinishReasonMiddleware(AgentMiddleware[AgentState]): # ----- main apply ------------------------------------------------------ - def _apply(self, state: AgentState, runtime: Runtime) -> dict | None: + def _prepare_intervention(self, state: AgentState, runtime: Runtime) -> _SafetyIntervention | None: messages = state.get("messages", []) if not messages: return None @@ -348,10 +392,23 @@ class SafetyFinishReasonMiddleware(AgentMiddleware[AgentState]): }, ) - self._emit_event(termination, [tc.get("name") or "unknown" for tc in tool_calls], runtime) - self._record_audit_event(termination, last, list(tool_calls), runtime) + tool_calls = list(tool_calls) + return _SafetyIntervention( + update={"messages": [patched]}, + termination=termination, + suppressed_names=[tc.get("name") or "unknown" for tc in tool_calls], + message=last, + tool_calls=tool_calls, + ) - return {"messages": [patched]} + def _apply(self, state: AgentState, runtime: Runtime) -> dict | None: + intervention = self._prepare_intervention(state, runtime) + if intervention is None: + return None + + self._emit_event(intervention.termination, intervention.suppressed_names, runtime) + self._record_audit_event(intervention.termination, intervention.message, intervention.tool_calls, runtime) + return intervention.update # ----- hooks ----------------------------------------------------------- @@ -361,4 +418,10 @@ class SafetyFinishReasonMiddleware(AgentMiddleware[AgentState]): @override async def aafter_model(self, state: AgentState, runtime: Runtime) -> dict | None: - return self._apply(state, runtime) + intervention = self._prepare_intervention(state, runtime) + if intervention is None: + return None + + await self._aemit_event(intervention.termination, intervention.suppressed_names, runtime) + self._record_audit_event(intervention.termination, intervention.message, intervention.tool_calls, runtime) + return intervention.update diff --git a/backend/packages/harness/deerflow/tools/builtins/task_tool.py b/backend/packages/harness/deerflow/tools/builtins/task_tool.py index 5a5014bb1..3f11592ae 100644 --- a/backend/packages/harness/deerflow/tools/builtins/task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/task_tool.py @@ -32,6 +32,7 @@ from deerflow.subagents.status_contract import ( ) from deerflow.tools.types import Runtime from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, get_current_trace_id, normalize_trace_id +from deerflow.utils.custom_events import aemit_custom_event if TYPE_CHECKING: from deerflow.config.app_config import AppConfig @@ -419,13 +420,14 @@ async def task_tool( writer = get_stream_writer() # Send Task Started message' - writer( + await aemit_custom_event( { "type": "task_started", "task_id": task_id, "description": description, "model_name": effective_model, - } + }, + writer=writer, ) try: @@ -434,7 +436,10 @@ async def task_tool( if result is None: logger.error(f"[trace={trace_id}] Task {task_id} not found in background tasks") - writer({"type": "task_failed", "task_id": task_id, "error": "Task disappeared from background tasks"}) + await aemit_custom_event( + {"type": "task_failed", "task_id": task_id, "error": "Task disappeared from background tasks"}, + writer=writer, + ) cleanup_background_task(task_id) error = f"Task {task_id} disappeared from background tasks" return _task_result_command( @@ -460,7 +465,7 @@ async def task_tool( # Send task_running event for each new message for i in range(last_message_count, current_message_count): message = ai_messages[i] - writer( + await aemit_custom_event( { "type": "task_running", "task_id": task_id, @@ -469,7 +474,8 @@ async def task_tool( "total_messages": current_message_count, "usage": usage, "model_name": effective_model, - } + }, + writer=writer, ) logger.info(f"[trace={trace_id}] Task {task_id} sent message #{i + 1}/{current_message_count}") last_message_count = current_message_count @@ -478,14 +484,15 @@ async def task_tool( if result.status == SubagentStatus.COMPLETED: _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) _report_subagent_usage(runtime, result) - writer( + await aemit_custom_event( { "type": "task_completed", "task_id": task_id, "result": result.result, "usage": usage, "model_name": effective_model, - } + }, + writer=writer, ) logger.info(f"[trace={trace_id}] Task {task_id} completed after {poll_count} polls") cleanup_background_task(task_id) @@ -503,14 +510,15 @@ async def task_tool( elif result.status == SubagentStatus.FAILED: _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) _report_subagent_usage(runtime, result) - writer( + await aemit_custom_event( { "type": "task_failed", "task_id": task_id, "error": result.error, "usage": usage, "model_name": effective_model, - } + }, + writer=writer, ) logger.error(f"[trace={trace_id}] Task {task_id} failed: {result.error}") cleanup_background_task(task_id) @@ -528,14 +536,15 @@ async def task_tool( elif result.status == SubagentStatus.CANCELLED: _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) _report_subagent_usage(runtime, result) - writer( + await aemit_custom_event( { "type": "task_cancelled", "task_id": task_id, "error": result.error, "usage": usage, "model_name": effective_model, - } + }, + writer=writer, ) logger.info(f"[trace={trace_id}] Task {task_id} cancelled: {result.error}") cleanup_background_task(task_id) @@ -549,14 +558,15 @@ async def task_tool( elif result.status == SubagentStatus.TIMED_OUT: _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) _report_subagent_usage(runtime, result) - writer( + await aemit_custom_event( { "type": "task_timed_out", "task_id": task_id, "error": result.error, "usage": usage, "model_name": effective_model, - } + }, + writer=writer, ) logger.warning(f"[trace={trace_id}] Task {task_id} timed out: {result.error}") cleanup_background_task(task_id) @@ -581,13 +591,14 @@ async def task_tool( _report_subagent_usage(runtime, result) usage = _summarize_usage(getattr(result, "token_usage_records", None)) _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) - writer( + await aemit_custom_event( { "type": "task_timed_out", "task_id": task_id, "usage": usage, "model_name": effective_model, - } + }, + writer=writer, ) # The task may still be running in the background. Signal cooperative # cancellation and schedule deferred cleanup to remove the entry from diff --git a/backend/packages/harness/deerflow/utils/custom_events.py b/backend/packages/harness/deerflow/utils/custom_events.py new file mode 100644 index 000000000..45b239ac8 --- /dev/null +++ b/backend/packages/harness/deerflow/utils/custom_events.py @@ -0,0 +1,57 @@ +"""Compatibility helpers for DeerFlow custom stream events.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import Any + +from langchain_core.callbacks import adispatch_custom_event, dispatch_custom_event +from langgraph.errors import GraphBubbleUp + +logger = logging.getLogger(__name__) + +StreamWriter = Callable[[Any], None] + + +def _event_name(payload: dict[str, Any]) -> str | None: + event_type = payload.get("type") + if isinstance(event_type, str) and event_type: + return event_type + logger.debug("Custom stream payload has no non-empty string 'type'; skipping callback dispatch") + return None + + +def emit_custom_event(payload: dict[str, Any], *, writer: StreamWriter) -> None: + """Emit one event to LangGraph's custom stream and callback APIs. + + The writer remains the primary compatibility path. Callback dispatch is + best-effort so an optional ``astream_events`` consumer cannot break an + existing DeerFlow run. + """ + + writer(payload) + event_name = _event_name(payload) + if event_name is None: + return + try: + dispatch_custom_event(event_name, payload) + except GraphBubbleUp: + raise + except Exception: + logger.debug("Failed to dispatch custom callback event %s", event_name, exc_info=True) + + +async def aemit_custom_event(payload: dict[str, Any], *, writer: StreamWriter) -> None: + """Async counterpart to :func:`emit_custom_event`.""" + + writer(payload) + event_name = _event_name(payload) + if event_name is None: + return + try: + await adispatch_custom_event(event_name, payload) + except GraphBubbleUp: + raise + except Exception: + logger.debug("Failed to dispatch async custom callback event %s", event_name, exc_info=True) diff --git a/backend/tests/test_custom_events.py b/backend/tests/test_custom_events.py new file mode 100644 index 000000000..84c977e8b --- /dev/null +++ b/backend/tests/test_custom_events.py @@ -0,0 +1,304 @@ +from __future__ import annotations + +import asyncio +import importlib +from enum import Enum +from types import SimpleNamespace +from typing import TypedDict + +import pytest +from langchain.agents import create_agent +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.outputs import ChatGeneration, ChatResult +from langgraph.config import get_stream_writer +from langgraph.errors import GraphInterrupt +from langgraph.graph import END, START, StateGraph +from langgraph.types import Interrupt + +from deerflow.subagents.config import SubagentConfig +from deerflow.utils import custom_events as custom_events_module +from deerflow.utils.custom_events import aemit_custom_event, emit_custom_event + +task_tool_module = importlib.import_module("deerflow.tools.builtins.task_tool") + + +class _State(TypedDict): + value: int + + +def _compile_graph(node): + builder = StateGraph(_State) + builder.add_node("emit", node) + builder.add_edge(START, "emit") + builder.add_edge("emit", END) + return builder.compile() + + +def _sync_node(state: _State) -> _State: + payload = {"type": "sync_probe", "value": state["value"]} + emit_custom_event(payload, writer=get_stream_writer()) + return state + + +async def _async_node(state: _State) -> _State: + payload = {"type": "async_probe", "value": state["value"]} + await aemit_custom_event(payload, writer=get_stream_writer()) + return state + + +async def _custom_events(graph) -> list[dict]: + return [chunk async for chunk in graph.astream({"value": 7}, stream_mode="custom")] + + +async def _astream_events(graph) -> list[dict]: + return [event async for event in graph.astream_events({"value": 7}, version="v2") if event["event"] == "on_custom_event"] + + +@pytest.mark.anyio +@pytest.mark.parametrize( + ("node", "event_name"), + [ + (_sync_node, "sync_probe"), + (_async_node, "async_probe"), + ], +) +async def test_custom_event_is_emitted_once_to_each_streaming_api(node, event_name): + graph = _compile_graph(node) + + custom_chunks = await _custom_events(graph) + callback_events = await _astream_events(graph) + + expected = {"type": event_name, "value": 7} + assert custom_chunks == [expected] + assert len(callback_events) == 1 + assert callback_events[0]["name"] == event_name + assert callback_events[0]["data"] == expected + + +class _TaskCallingModel(BaseChatModel): + call_count: int = 0 + + @property + def _llm_type(self) -> str: + return "fake-task-caller" + + def bind_tools(self, tools, **kwargs): + return self + + def _generate(self, messages, stop=None, run_manager=None, **kwargs): + self.call_count += 1 + if self.call_count == 1: + message = AIMessage( + content="", + tool_calls=[ + { + "id": "task-call-1", + "name": "task", + "args": { + "description": "validate streaming", + "prompt": "run the delegated task", + "subagent_type": "general-purpose", + }, + } + ], + response_metadata={"finish_reason": "tool_calls"}, + ) + else: + message = AIMessage(content="done", response_metadata={"finish_reason": "stop"}) + return ChatResult(generations=[ChatGeneration(message=message)]) + + async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs): + return self._generate(messages, stop=stop, run_manager=run_manager, **kwargs) + + +@pytest.mark.anyio +async def test_real_task_tool_events_reach_astream_events(monkeypatch): + """Exercise the real ToolNode/runtime callback context used by task_tool.""" + + class _SubagentStatus(Enum): + COMPLETED = "completed" + + config = SubagentConfig( + name="general-purpose", + description="General helper", + system_prompt="Test prompt", + model="test-model", + timeout_seconds=10, + ) + completed = SimpleNamespace( + status=_SubagentStatus.COMPLETED, + ai_messages=[], + result="delegated result", + error=None, + stop_reason=None, + token_usage_records=[], + usage_reported=False, + ) + + class _Executor: + def __init__(self, **_kwargs): + pass + + def execute_async(self, _prompt, task_id=None): + return task_id + + monkeypatch.setattr(task_tool_module, "SubagentStatus", _SubagentStatus) + monkeypatch.setattr(task_tool_module, "SubagentExecutor", _Executor) + monkeypatch.setattr(task_tool_module, "get_available_subagent_names", lambda: ["general-purpose"]) + monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _name: config) + monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _task_id: completed) + monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda _task_id: None) + monkeypatch.setattr(task_tool_module, "_token_usage_cache_enabled", lambda _config: False) + monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **_kwargs: []) + + agent = create_agent( + model=_TaskCallingModel(), + tools=[task_tool_module.task_tool], + context_schema=dict, + ) + events = [ + event + async for event in agent.astream_events( + {"messages": [HumanMessage(content="delegate this")]}, + version="v2", + context={"thread_id": "task-stream-thread"}, + ) + if event["event"] == "on_custom_event" + ] + + assert [event["name"] for event in events] == ["task_started", "task_completed"] + assert [event["data"]["type"] for event in events] == ["task_started", "task_completed"] + assert all(event["data"]["task_id"] == "task-call-1" for event in events) + assert events[0]["data"]["description"] == "validate streaming" + assert events[1]["data"]["result"] == "delegated result" + + +def test_sync_dispatch_failure_does_not_break_writer(monkeypatch): + payload = {"type": "sync_probe", "value": 1} + written: list[dict] = [] + + def fail_dispatch(*_args, **_kwargs): + raise RuntimeError("callback failed") + + monkeypatch.setattr(custom_events_module, "dispatch_custom_event", fail_dispatch) + + emit_custom_event(payload, writer=written.append) + + assert written == [payload] + + +def test_sync_dispatch_without_parent_run_does_not_break_writer(): + payload = {"type": "sync_probe", "value": 1} + written: list[dict] = [] + + emit_custom_event(payload, writer=written.append) + + assert written == [payload] + + +@pytest.mark.anyio +async def test_async_dispatch_failure_does_not_break_writer(monkeypatch): + payload = {"type": "async_probe", "value": 1} + written: list[dict] = [] + + async def fail_dispatch(*_args, **_kwargs): + raise RuntimeError("callback failed") + + monkeypatch.setattr(custom_events_module, "adispatch_custom_event", fail_dispatch) + + await aemit_custom_event(payload, writer=written.append) + + assert written == [payload] + + +@pytest.mark.anyio +async def test_async_dispatch_without_parent_run_does_not_break_writer(): + payload = {"type": "async_probe", "value": 1} + written: list[dict] = [] + + await aemit_custom_event(payload, writer=written.append) + + assert written == [payload] + + +def test_missing_event_type_preserves_writer_and_skips_dispatch(monkeypatch): + payload = {"value": 1} + written: list[dict] = [] + dispatched: list[tuple] = [] + + monkeypatch.setattr(custom_events_module, "dispatch_custom_event", lambda *args, **kwargs: dispatched.append((args, kwargs))) + + emit_custom_event(payload, writer=written.append) + + assert written == [payload] + assert dispatched == [] + + +def test_writer_failure_propagates_before_dispatch(monkeypatch): + dispatched: list[tuple] = [] + + def fail_writer(_payload): + raise RuntimeError("writer failed") + + monkeypatch.setattr(custom_events_module, "dispatch_custom_event", lambda *args, **kwargs: dispatched.append((args, kwargs))) + + with pytest.raises(RuntimeError, match="writer failed"): + emit_custom_event({"type": "sync_probe"}, writer=fail_writer) + + assert dispatched == [] + + +@pytest.mark.anyio +async def test_async_writer_failure_propagates_before_dispatch(monkeypatch): + dispatched: list[tuple] = [] + + def fail_writer(_payload): + raise RuntimeError("writer failed") + + async def record_dispatch(*args, **kwargs): + dispatched.append((args, kwargs)) + + monkeypatch.setattr(custom_events_module, "adispatch_custom_event", record_dispatch) + + with pytest.raises(RuntimeError, match="writer failed"): + await aemit_custom_event({"type": "async_probe"}, writer=fail_writer) + + assert dispatched == [] + + +@pytest.mark.anyio +async def test_async_cancellation_is_not_swallowed(monkeypatch): + async def cancel_dispatch(*_args, **_kwargs): + raise asyncio.CancelledError + + monkeypatch.setattr(custom_events_module, "adispatch_custom_event", cancel_dispatch) + + with pytest.raises(asyncio.CancelledError): + await aemit_custom_event({"type": "async_probe"}, writer=lambda _payload: None) + + +@pytest.mark.parametrize("async_dispatch", [False, True]) +def test_langgraph_control_flow_is_not_swallowed(monkeypatch, async_dispatch): + control_flow = GraphInterrupt((Interrupt(value="pause"),)) + + if async_dispatch: + + async def interrupt_dispatch(*_args, **_kwargs): + raise control_flow + + monkeypatch.setattr(custom_events_module, "adispatch_custom_event", interrupt_dispatch) + + with pytest.raises(GraphInterrupt) as raised: + asyncio.run(aemit_custom_event({"type": "async_probe"}, writer=lambda _payload: None)) + else: + + def interrupt_dispatch(*_args, **_kwargs): + raise control_flow + + monkeypatch.setattr(custom_events_module, "dispatch_custom_event", interrupt_dispatch) + + with pytest.raises(GraphInterrupt) as raised: + emit_custom_event({"type": "sync_probe"}, writer=lambda _payload: None) + + assert raised.value is control_flow diff --git a/backend/tests/test_llm_error_handling_middleware.py b/backend/tests/test_llm_error_handling_middleware.py index 08c1acf51..acf480b11 100644 --- a/backend/tests/test_llm_error_handling_middleware.py +++ b/backend/tests/test_llm_error_handling_middleware.py @@ -95,6 +95,7 @@ def test_async_model_call_retries_busy_provider_then_succeeds( attempts = 0 waits: list[float] = [] events: list[dict] = [] + dispatched_events: list[dict] = [] async def fake_sleep(delay: float) -> None: waits.append(delay) @@ -102,6 +103,10 @@ def test_async_model_call_retries_busy_provider_then_succeeds( def fake_writer(): return events.append + async def fake_emit_custom_event(payload, *, writer): + writer(payload) + dispatched_events.append(payload) + async def handler(_request) -> AIMessage: nonlocal attempts attempts += 1 @@ -114,6 +119,10 @@ def test_async_model_call_retries_busy_provider_then_succeeds( "langgraph.config.get_stream_writer", fake_writer, ) + monkeypatch.setattr( + "deerflow.agents.middlewares.llm_error_handling_middleware.aemit_custom_event", + fake_emit_custom_event, + ) result = asyncio.run(middleware.awrap_model_call(SimpleNamespace(), handler)) @@ -122,6 +131,7 @@ def test_async_model_call_retries_busy_provider_then_succeeds( assert attempts == 3 assert waits == [0.025, 0.025] assert [event["type"] for event in events] == ["llm_retry", "llm_retry"] + assert dispatched_events == events def test_async_model_call_returns_user_message_for_quota_errors() -> None: @@ -168,11 +178,17 @@ def test_async_model_call_marks_transient_retry_exhaustion_as_error_fallback( def test_sync_model_call_uses_retry_after_header(monkeypatch: pytest.MonkeyPatch) -> None: middleware = _build_middleware(retry_max_attempts=2, retry_base_delay_ms=10, retry_cap_delay_ms=10) waits: list[float] = [] + events: list[dict] = [] + dispatched_events: list[dict] = [] attempts = 0 def fake_sleep(delay: float) -> None: waits.append(delay) + def fake_emit_custom_event(payload, *, writer): + writer(payload) + dispatched_events.append(payload) + def handler(_request) -> AIMessage: nonlocal attempts attempts += 1 @@ -185,12 +201,52 @@ def test_sync_model_call_uses_retry_after_header(monkeypatch: pytest.MonkeyPatch return AIMessage(content="ok") monkeypatch.setattr("time.sleep", fake_sleep) + monkeypatch.setattr("langgraph.config.get_stream_writer", lambda: events.append) + monkeypatch.setattr( + "deerflow.agents.middlewares.llm_error_handling_middleware.emit_custom_event", + fake_emit_custom_event, + ) result = middleware.wrap_model_call(SimpleNamespace(), handler) assert isinstance(result, AIMessage) assert result.content == "ok" assert waits == [2.0] + assert dispatched_events == events + assert [event["type"] for event in events] == ["llm_retry"] + + +def test_sync_retry_event_preserves_langgraph_control_flow(monkeypatch: pytest.MonkeyPatch) -> None: + middleware = _build_middleware() + + def interrupt_dispatch(*_args, **_kwargs): + raise GraphBubbleUp + + monkeypatch.setattr("langgraph.config.get_stream_writer", lambda: lambda _payload: None) + monkeypatch.setattr( + "deerflow.agents.middlewares.llm_error_handling_middleware.emit_custom_event", + interrupt_dispatch, + ) + + with pytest.raises(GraphBubbleUp): + middleware._emit_retry_event(1, 10, "busy", max_attempts=2) + + +@pytest.mark.anyio +async def test_async_retry_event_preserves_langgraph_control_flow(monkeypatch: pytest.MonkeyPatch) -> None: + middleware = _build_middleware() + + async def interrupt_dispatch(*_args, **_kwargs): + raise GraphBubbleUp + + monkeypatch.setattr("langgraph.config.get_stream_writer", lambda: lambda _payload: None) + monkeypatch.setattr( + "deerflow.agents.middlewares.llm_error_handling_middleware.aemit_custom_event", + interrupt_dispatch, + ) + + with pytest.raises(GraphBubbleUp): + await middleware._aemit_retry_event(1, 10, "busy", max_attempts=2) def test_sync_model_call_propagates_graph_bubble_up() -> None: diff --git a/backend/tests/test_safety_finish_reason_graph_integration.py b/backend/tests/test_safety_finish_reason_graph_integration.py index 5695084c8..d13b51230 100644 --- a/backend/tests/test_safety_finish_reason_graph_integration.py +++ b/backend/tests/test_safety_finish_reason_graph_integration.py @@ -20,6 +20,7 @@ from __future__ import annotations from typing import Any +import pytest from langchain.agents import create_agent from langchain.agents.middleware import AgentMiddleware from langchain.agents.middleware.types import ModelRequest, ModelResponse @@ -140,6 +141,42 @@ def test_content_filter_with_tool_calls_does_not_invoke_tool_node(): assert final_ai.response_metadata.get("finish_reason") == "content_filter" +@pytest.mark.anyio +async def test_safety_termination_event_reaches_astream_events(): + """Exercise the middleware's real async graph hook and callback context.""" + + _TOOL_INVOCATIONS.clear() + agent = create_agent( + model=_ContentFilteredModel(), + tools=[write_file], + middleware=[SafetyFinishReasonMiddleware()], + context_schema=dict, + ) + + events = [ + event + async for event in agent.astream_events( + {"messages": [HumanMessage(content="write me a report")]}, + version="v2", + context={"thread_id": "safety-stream-thread"}, + ) + if event["event"] == "on_custom_event" + ] + + assert len(events) == 1 + assert events[0]["name"] == "safety_termination" + assert events[0]["data"] == { + "type": "safety_termination", + "detector": "openai_compatible_content_filter", + "reason_field": "finish_reason", + "reason_value": "content_filter", + "suppressed_tool_call_count": 1, + "suppressed_tool_call_names": ["write_file"], + "thread_id": "safety-stream-thread", + } + assert _TOOL_INVOCATIONS == [] + + def test_content_filter_without_tool_calls_passes_through_unchanged(): """No tool calls => issue scope says don't intervene; the partial response should be delivered as-is so the user sees what they got.""" diff --git a/backend/tests/test_safety_finish_reason_middleware.py b/backend/tests/test_safety_finish_reason_middleware.py index 5cf4af7b1..2b9618aad 100644 --- a/backend/tests/test_safety_finish_reason_middleware.py +++ b/backend/tests/test_safety_finish_reason_middleware.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock import pytest from langchain_core.messages import AIMessage, HumanMessage, SystemMessage +from langgraph.errors import GraphBubbleUp from deerflow.agents.middlewares.safety_finish_reason_middleware import SafetyFinishReasonMiddleware from deerflow.agents.middlewares.safety_termination_detectors import ( @@ -689,14 +690,23 @@ class TestAuditEvent: class TestStreamEvent: def test_emits_event_when_writer_available(self, monkeypatch): captured: list = [] + dispatched: list = [] def fake_writer(payload): captured.append(payload) + def fake_emit_custom_event(payload, *, writer): + writer(payload) + dispatched.append(payload) + # Patch get_stream_writer at the symbol-resolution site. import langgraph.config monkeypatch.setattr(langgraph.config, "get_stream_writer", lambda: fake_writer) + monkeypatch.setattr( + "deerflow.agents.middlewares.safety_finish_reason_middleware.emit_custom_event", + fake_emit_custom_event, + ) mw = SafetyFinishReasonMiddleware() state = { @@ -718,6 +728,83 @@ class TestStreamEvent: assert payload["suppressed_tool_call_count"] == 1 assert payload["suppressed_tool_call_names"] == ["write_file"] assert payload["thread_id"] == "t-stream" + assert dispatched == captured + + @pytest.mark.anyio + async def test_async_hook_uses_async_event_dispatch(self, monkeypatch): + captured: list = [] + dispatched: list = [] + + async def fake_emit_custom_event(payload, *, writer): + writer(payload) + dispatched.append(payload) + + import langgraph.config + + monkeypatch.setattr(langgraph.config, "get_stream_writer", lambda: captured.append) + monkeypatch.setattr( + "deerflow.agents.middlewares.safety_finish_reason_middleware.aemit_custom_event", + fake_emit_custom_event, + ) + + mw = SafetyFinishReasonMiddleware() + state = { + "messages": [ + _ai( + tool_calls=[_write_call()], + response_metadata={"finish_reason": "content_filter"}, + ) + ] + } + + result = await mw.aafter_model(state, _runtime("t-async-stream")) + + assert result is not None + assert result["messages"][0].tool_calls == [] + assert dispatched == captured + assert [payload["type"] for payload in captured] == ["safety_termination"] + assert captured[0]["thread_id"] == "t-async-stream" + + def test_sync_event_preserves_langgraph_control_flow(self, monkeypatch): + import langgraph.config + + def interrupt_dispatch(*_args, **_kwargs): + raise GraphBubbleUp + + monkeypatch.setattr(langgraph.config, "get_stream_writer", lambda: lambda _payload: None) + monkeypatch.setattr( + "deerflow.agents.middlewares.safety_finish_reason_middleware.emit_custom_event", + interrupt_dispatch, + ) + + termination = SafetyTermination( + detector="test", + reason_field="finish_reason", + reason_value="content_filter", + ) + with pytest.raises(GraphBubbleUp): + SafetyFinishReasonMiddleware()._emit_event(termination, ["write_file"], _runtime()) + + @pytest.mark.anyio + async def test_async_event_preserves_langgraph_control_flow(self, monkeypatch): + import langgraph.config + + async def interrupt_dispatch(*_args, **_kwargs): + raise GraphBubbleUp + + monkeypatch.setattr(langgraph.config, "get_stream_writer", lambda: lambda _payload: None) + monkeypatch.setattr( + "deerflow.agents.middlewares.safety_finish_reason_middleware.aemit_custom_event", + interrupt_dispatch, + ) + + termination = SafetyTermination( + detector="test", + reason_field="finish_reason", + reason_value="content_filter", + ) + with pytest.raises(GraphBubbleUp): + await SafetyFinishReasonMiddleware()._aemit_event(termination, ["write_file"], _runtime()) def test_writer_unavailable_does_not_break(self, monkeypatch): import langgraph.config diff --git a/backend/tests/test_task_tool_core_logic.py b/backend/tests/test_task_tool_core_logic.py index 30c46f24c..6de04fed2 100644 --- a/backend/tests/test_task_tool_core_logic.py +++ b/backend/tests/test_task_tool_core_logic.py @@ -499,9 +499,14 @@ def test_task_tool_emits_running_and_completed_events(monkeypatch): runtime = _make_runtime() runtime.context["deerflow_trace_id"] = "task-trace-1" events = [] + dispatched_events = [] captured = {} get_available_tools = MagicMock(return_value=["tool-a", "tool-b"]) + async def fake_emit_custom_event(payload, *, writer): + writer(payload) + dispatched_events.append(payload) + class DummyExecutor: def __init__(self, **kwargs): captured["executor_kwargs"] = kwargs @@ -529,6 +534,7 @@ def test_task_tool_emits_running_and_completed_events(monkeypatch): monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _: next(responses)) monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: events.append) + monkeypatch.setattr(task_tool_module, "aemit_custom_event", fake_emit_custom_event) monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) # task_tool lazily imports from deerflow.tools at call time, so patch that module-level function. monkeypatch.setattr("deerflow.tools.get_available_tools", get_available_tools) @@ -557,6 +563,7 @@ def test_task_tool_emits_running_and_completed_events(monkeypatch): event_types = [e["type"] for e in events] assert event_types == ["task_started", "task_running", "task_running", "task_completed"] + assert dispatched_events == events assert events[0]["model_name"] == "ark-model" assert events[-1]["result"] == "all done" @@ -1701,6 +1708,11 @@ def test_terminal_events_include_usage(monkeypatch, status, expected_type): config = _make_subagent_config() runtime = _make_runtime() events = [] + dispatched_events = [] + + async def fake_emit_custom_event(payload, *, writer): + writer(payload) + dispatched_events.append(payload) records = [ {"source_run_id": "r1", "caller": "subagent:general-purpose", "input_tokens": 100, "output_tokens": 50, "total_tokens": 150}, @@ -1712,6 +1724,7 @@ def test_terminal_events_include_usage(monkeypatch, status, expected_type): monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _: result) monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: events.append) + monkeypatch.setattr(task_tool_module, "aemit_custom_event", fake_emit_custom_event) monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) monkeypatch.setattr(task_tool_module, "_report_subagent_usage", lambda *_: None) monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda _: None) @@ -1727,6 +1740,7 @@ def test_terminal_events_include_usage(monkeypatch, status, expected_type): terminal_events = [e for e in events if e["type"] == expected_type] assert len(terminal_events) == 1 + assert dispatched_events == events assert terminal_events[0]["usage"] == { "input_tokens": 300, "output_tokens": 130,