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 <willem.jiang@gmail.com>

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:

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)

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` 是同一概念的三个名字。在其中一层出错不会在另外两层报错。

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

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

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

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)

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

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:

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."""

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

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,