langchain-agentx-python 0.9.2__py3-none-any.whl → 0.9.4__py3-none-any.whl

langchain_agentx/__init__.py

@@ -1,64 +1,64 @@
-"""
-LangChain-based agent utilities for CodeBaseX.
-
-This package provides:
-- LangChain-AgentX LangGraph agent factory (`loop`)
-
-Typical usage:
-
-```python
-from langchain_agentx import create_loop_agent
-```
-"""
-
-__version__ = "0.9.2"
-
-from .loop import (  # noqa: F401
-    create_loop_agent,
-)
-
-from .loop.hook import (  # noqa: F401
-    ATTACHMENT_HOOK_BLOCKING_ERROR,
-    ATTACHMENT_HOOK_CANCELLED,
-    ATTACHMENT_HOOK_NON_BLOCKING_ERROR,
-    ATTACHMENT_HOOK_STOPPED_CONTINUATION,
-    ATTACHMENT_HOOK_SUCCESS,
-)
-from .observability import (  # noqa: F401
-    EventAdapterState,
-    LangchainAgentEvent,
-    LangchainAgentEventType,
-    LangGraphToLangchainAgentEventAdapter,
-    TraceCollector,
-    TraceSession,
-    TraceSpan,
-)
-from .provider import CompatibleChatOpenAI  # noqa: F401
-from .session import AgentSession, create_agent_session  # noqa: F401
-from .health import HealthCheck, HealthReport, PreflightResult, preflight, sdk_health  # noqa: F401
-from .workflow import BaseWorkflow  # noqa: F401
-
-__all__ = [
-    "ATTACHMENT_HOOK_BLOCKING_ERROR",
-    "ATTACHMENT_HOOK_CANCELLED",
-    "ATTACHMENT_HOOK_NON_BLOCKING_ERROR",
-    "ATTACHMENT_HOOK_STOPPED_CONTINUATION",
-    "ATTACHMENT_HOOK_SUCCESS",
-    "create_loop_agent",
-    "create_agent_session",
-    "AgentSession",
-    "BaseWorkflow",
-    "EventAdapterState",
-    "LangchainAgentEvent",
-    "LangchainAgentEventType",
-    "LangGraphToLangchainAgentEventAdapter",
-    "TraceCollector",
-    "TraceSession",
-    "TraceSpan",
-    "CompatibleChatOpenAI",
-    "HealthCheck",
-    "HealthReport",
-    "PreflightResult",
-    "preflight",
-    "sdk_health",
-]
+"""
+LangChain-based agent utilities for CodeBaseX.
+
+This package provides:
+- LangChain-AgentX LangGraph agent factory (`loop`)
+
+Typical usage:
+
+```python
+from langchain_agentx import create_loop_agent
+```
+"""
+
+__version__ = "0.9.4"
+
+from .loop import (  # noqa: F401
+    create_loop_agent,
+)
+
+from .loop.hook import (  # noqa: F401
+    ATTACHMENT_HOOK_BLOCKING_ERROR,
+    ATTACHMENT_HOOK_CANCELLED,
+    ATTACHMENT_HOOK_NON_BLOCKING_ERROR,
+    ATTACHMENT_HOOK_STOPPED_CONTINUATION,
+    ATTACHMENT_HOOK_SUCCESS,
+)
+from .observability import (  # noqa: F401
+    EventAdapterState,
+    LangchainAgentEvent,
+    LangchainAgentEventType,
+    LangGraphToLangchainAgentEventAdapter,
+    TraceCollector,
+    TraceSession,
+    TraceSpan,
+)
+from .provider import CompatibleChatOpenAI  # noqa: F401
+from .session import AgentSession, create_agent_session  # noqa: F401
+from .health import HealthCheck, HealthReport, PreflightResult, preflight, sdk_health  # noqa: F401
+from .workflow import BaseWorkflow  # noqa: F401
+
+__all__ = [
+    "ATTACHMENT_HOOK_BLOCKING_ERROR",
+    "ATTACHMENT_HOOK_CANCELLED",
+    "ATTACHMENT_HOOK_NON_BLOCKING_ERROR",
+    "ATTACHMENT_HOOK_STOPPED_CONTINUATION",
+    "ATTACHMENT_HOOK_SUCCESS",
+    "create_loop_agent",
+    "create_agent_session",
+    "AgentSession",
+    "BaseWorkflow",
+    "EventAdapterState",
+    "LangchainAgentEvent",
+    "LangchainAgentEventType",
+    "LangGraphToLangchainAgentEventAdapter",
+    "TraceCollector",
+    "TraceSession",
+    "TraceSpan",
+    "CompatibleChatOpenAI",
+    "HealthCheck",
+    "HealthReport",
+    "PreflightResult",
+    "preflight",
+    "sdk_health",
+]

langchain_agentx/loop/context/compaction_service.py

@@ -7,7 +7,8 @@
 
 设计要点：
 - `LoopContextCompaction` 构造注入 `ContextPipeline`，便于单测替换阶段序列；
-- `default_loop_context_compaction()` 返回进程内单例，避免每轮重复装配。
+- `default_loop_context_compaction()` 返回进程内单例，避免每轮重复装配；
+- 支持 `mode` 参数区分 PROACTIVE（常规）与 REACTIVE_CONTEXT_WINDOW（溢出恢复）。
 
 注意：
 - 不修改 `should_end` / `finish_reason` 等 exit 字段；只产出 `messages` 增量补丁。
@@ -19,6 +20,7 @@ from __future__ import annotations
 from typing import Any, Dict, Mapping
 
 from .pipeline import ContextPipeline
+from .recovery_mode import CompactionMode
 from .settings import CompactionSettings
 from .types import CompactProgressCallback, ContextCompactionContext
 
@@ -38,9 +40,22 @@ class LoopContextCompaction:
         """使用 CC 默认五段顺序装配管道。"""
         return cls(ContextPipeline.with_default_stages(settings=settings))
 
-    def run(self, state: Mapping[str, Any], model: Any | None = None) -> Dict[str, Any]:
-        """返回 `Command(update=...)` 可用的补丁（常见为 ``{}`` 或 ``{"messages": ...}``）。"""
-        ctx = ContextCompactionContext.from_state(state, model)
+    def run(
+        self,
+        state: Mapping[str, Any],
+        model: Any | None = None,
+        *,
+        mode: str | CompactionMode = CompactionMode.PROACTIVE,
+    ) -> Dict[str, Any]:
+        """返回 `Command(update=...)` 可用的补丁（常见为 ``{}`` 或 ``{"messages": ...}``）。
+
+        参数:
+            state: LangGraph state
+            model: 可选的 LLM model（供 autocompact 使用）
+            mode: 压缩模式，默认 PROACTIVE
+        """
+        mode_str = mode if isinstance(mode, str) else mode.value
+        ctx = ContextCompactionContext.from_state(state, model, mode=mode_str)
         return self._pipeline.run(state, ctx)
 
     async def arun(
@@ -48,6 +63,8 @@ class LoopContextCompaction:
         state: Mapping[str, Any],
         model: Any | None = None,
         progress_callback: CompactProgressCallback = None,
+        *,
+        mode: str | CompactionMode = CompactionMode.PROACTIVE,
     ) -> Dict[str, Any]:
         """Async 版本（供 /compact 命令使用），支持进度回调。
 
@@ -55,11 +72,13 @@ class LoopContextCompaction:
             state: LangGraph state
             model: 可选的 LLM model（供 autocompact 使用）
             progress_callback: 进度回调函数，接收 CompactProgress 对象
+            mode: 压缩模式，默认 PROACTIVE
 
         返回:
             与 run() 相同的补丁格式
         """
-        ctx = ContextCompactionContext.from_state(state, model)
+        mode_str = mode if isinstance(mode, str) else mode.value
+        ctx = ContextCompactionContext.from_state(state, model, mode=mode_str)
         ctx.progress_callback = progress_callback
         return await self._pipeline.arun(state, ctx)
 

langchain_agentx/loop/context/pipeline.py

@@ -114,6 +114,7 @@ class ContextPipeline:
                 "snip_tokens_freed_carry_final": ctx.snip_tokens_freed_carry,
                 "token_warning_level": tw.level,
                 "stages": stages_meta,
+                "mode": ctx.mode,
             }
 
         if not patch:
@@ -194,6 +195,7 @@ class ContextPipeline:
                 "snip_tokens_freed_carry_final": ctx.snip_tokens_freed_carry,
                 "token_warning_level": tw.level,
                 "stages": stages_meta,
+                "mode": ctx.mode,
             }
 
         if not patch:

langchain_agentx/loop/context/recovery_mode.py

@@ -0,0 +1,37 @@
+"""
+recovery_mode.py — 上下文压缩模式枚举。
+
+职责：
+    定义 PROACTIVE（常规 pre-model）与 REACTIVE_CONTEXT_WINDOW（溢出恢复）两种模式，
+    供 ContextPipeline 与 ReactiveCompactNode 使用。
+
+链路位置：
+    - CompactionMode 传入 ContextCompactionContext.mode
+    - 写入 compaction_pipeline_meta.mode 字段
+    - ReactiveCompactNode 使用 REACTIVE_CONTEXT_WINDOW 模式
+
+当前裁剪范围：
+    当前仅定义枚举类型；未来可扩展 mode 切换检测逻辑。
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class CompactionMode(str, Enum):
+    """上下文压缩模式。"""
+
+    PROACTIVE = "proactive"
+    """常规 pre-model 治理模式（prune_and_compress 节点默认）。"""
+
+    REACTIVE_CONTEXT_WINDOW = "reactive_context_window"
+    """上下文窗口溢出恢复模式（reactive_compact 节点专用）。"""
+
+    @classmethod
+    def is_valid(cls, value: str) -> bool:
+        """检查字符串是否为有效 CompactionMode。"""
+        return value in {m.value for m in cls}
+
+
+__all__ = ["CompactionMode"]

langchain_agentx/loop/context/types.py

@@ -40,16 +40,19 @@ class ContextCompactionContext:
     query_source: str | None = None
     config: CompactionConfig = field(default_factory=CompactionConfig)
     progress_callback: CompactProgressCallback = None
+    mode: str = "proactive"
+    """压缩模式：proactive（常规）或 reactive_context_window（溢出恢复）。"""
 
     @classmethod
     def from_state(
-        cls, state: Mapping[str, Any], model: Any | None
+        cls, state: Mapping[str, Any], model: Any | None, mode: str = "proactive"
     ) -> ContextCompactionContext:
         return cls(
             model=model,
             query_source=state.get("query_source")
             if isinstance(state.get("query_source"), str)
             else None,
+            mode=mode,
         )
 
 

langchain_agentx/loop/graph/factory.py

@@ -52,6 +52,7 @@ from ..exit.exit_logic import (
     LoopAgentStateFields,
 )
 from .graph_edges import _make_model_to_model_edge, _make_tools_to_model_edge, _make_loop_controller_edge
+from .reactive_compact_node import ReactiveCompactNode
 from ..model.schema_and_format import (
     STRUCTURED_OUTPUT_ERROR_TEMPLATE,
     FALLBACK_MODELS_WITH_STRUCTURED_OUTPUT,
@@ -1372,7 +1373,27 @@ class LoopGraphBuilder:
         graph_inner.add_edge("hook_after_model", "hook_stop")
         graph_inner.add_edge("hook_stop", "loop_controller")
 
-        loop_controller_destinations = ["tools", loop_entry_node, exit_node]
+        # 添加 reactive_compact 节点（上下文窗口溢出恢复专用节点）
+        def _reactive_compact_node(state, runtime):
+            node = ReactiveCompactNode(
+                compaction_service=default_loop_context_compaction(),
+                loop_entry_destination=loop_entry_node,
+                end_destination=exit_node,
+            )
+            return node(state)
+
+        graph_inner.add_node(
+            "reactive_compact",
+            RunnableCallable(_reactive_compact_node, None, trace=False),
+        )
+        # reactive_compact 成功路径 → loop_entry_node
+        graph_inner.add_conditional_edges(
+            "reactive_compact",
+            lambda state: state.get("destination", loop_entry_node),
+            [loop_entry_node, exit_node],
+        )
+
+        loop_controller_destinations = ["tools", loop_entry_node, exit_node, "reactive_compact"]
         if provider_chain and "task_event_commit_ack" not in loop_controller_destinations:
             loop_controller_destinations.append("task_event_commit_ack")
         graph_inner.add_conditional_edges(

langchain_agentx/loop/graph/graph_edges.py

@@ -328,10 +328,12 @@ class LoopControllerEdge:
         withheld_error = state.get("withheld_error")
         if not withheld_error:
             return None
+        # 先检查是否为 context_window 溢出（CC 对齐：prompt_too_long 或 max_tokens 都可能）
+        overflow_kind = resolve_overflow_kind(withheld_error)
+        if overflow_kind == "context_window":
+            return self._handle_context_window_recovery(state, withheld_error)
+        # 再检查是否为 max_tokens（输出限制）
         if withheld_error.get("reason") == TERMINAL_REASON_MAX_TOKENS:
-            overflow_kind = resolve_overflow_kind(withheld_error)
-            if overflow_kind == "context_window":
-                return self._handle_context_window_recovery(state, withheld_error)
             return self._handle_max_tokens_recovery(state, withheld_error)
         return self._handle_generic_withheld(state, withheld_error)
 
@@ -340,25 +342,27 @@ class LoopControllerEdge:
         state: dict[str, Any],
         withheld_error: dict[str, Any],
     ) -> str:
+        """处理上下文窗口溢出恢复。
+
+        CC 对齐设计：
+        - 路由到 reactive_compact 节点，由节点决定成功或 exhausted
+        - transition 在 reactive_compact 成功后才写入
+        - 已尝试过时直接走 exhausted 路径
+        """
         if not state.get("has_attempted_reactive_compact", False):
-            state["has_attempted_reactive_compact"] = True
-            state["transition"] = {
-                "reason": TRANSITION_REACTIVE_COMPACT_RETRY,
-                "error_reason": TERMINAL_REASON_MAX_TOKENS,
-                "overflow_kind": "context_window",
-            }
             emit_loop_decision_event(
                 state=state,
                 decision="continue",
-                selected_edge=self._model_destination,
-                reason_code="reactive_compact_retry",
+                selected_edge="reactive_compact",
+                reason_code="context_window_recovery_to_compact",
                 evidence={
                     "withheld_error": withheld_error,
                     "overflow_kind": "context_window",
                     "step": state.get("step", 0),
                 },
             )
-            return self._model_destination
+            return "reactive_compact"
+        # 已尝试过，走 exhausted 路径
         state["withheld_error"] = None
         state["should_end"] = True
         state["terminal_reason"] = TERMINAL_REASON_MAX_TOKENS

langchain_agentx/loop/graph/reactive_compact_node.py

@@ -0,0 +1,230 @@
+"""
+reactive_compact_node.py — 上下文窗口溢出恢复节点。
+
+职责：
+    作为专用 overflow recovery 节点，仅在 context_window 溢出时触发，
+    执行 recovery compact 成功后才写入 reactive_compact_retry transition。
+
+链路位置：
+    loop_controller (检测 context_window 溢出) → reactive_compact → loop_entry
+    失败路径: reactive_compact → exhausted → END
+
+当前裁剪范围：
+    - 实现 _should_attempt_recovery 判断
+    - 实现 _execute_recovery_compact 核心逻辑
+    - 成功路径：写 transition、回 loop_entry
+    - 失败路径：exhausted、走 END
+    - 暂不实现 LLM 摘要（mode=REACTIVE 时只走 prune 路径）
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ..context.compaction_service import LoopContextCompaction
+from ..context.recovery_mode import CompactionMode
+from ..exit.reason_codes import (
+    TERMINAL_REASON_MAX_TOKENS,
+    TRANSITION_CONTEXT_WINDOW_RECOVERY_EXHAUSTED,
+    TRANSITION_REACTIVE_COMPACT_RETRY,
+)
+from ...observability.trace.event_emitter import emit_loop_decision_event
+
+logger = logging.getLogger(__name__)
+
+# 状态键常量
+STATE_HAS_ATTEMPTED_REACTIVE_COMPACT = "has_attempted_reactive_compact"
+STATE_MESSAGES = "messages"
+STATE_WITHHELD_ERROR = "withheld_error"
+STATE_SHOULD_END = "should_end"
+STATE_TERMINAL_REASON = "terminal_reason"
+STATE_TRANSITION = "transition"
+
+
+class ReactiveCompactNode:
+    """上下文窗口溢出恢复专用节点。
+
+    与 prune_and_compress 的区别：
+    - prune_and_compress：通用 pre-model 治理节点，PROACTIVE 模式
+    - reactive_compact：仅在 context_window 溢出时触发，REACTIVE_CONTEXT_WINDOW 模式
+
+    关键契约：
+    - 只在 compact 成功（messages 实际被压缩）后才写 transition
+    - 失败时设置 exhausted=True，走 terminal 分支
+    """
+
+    def __init__(
+        self,
+        compaction_service: LoopContextCompaction | None = None,
+        *,
+        loop_entry_destination: str = "loop_entry_node",
+        end_destination: str = "__end__",
+    ) -> None:
+        """初始化 ReactiveCompactNode。
+
+        参数:
+            compaction_service: 压缩服务，默认使用 default_loop_context_compaction()
+            loop_entry_destination: 成功后跳转目标（默认 loop_entry_node）
+            end_destination: 失败后跳转目标（默认 __end__）
+        """
+        if compaction_service is None:
+            from ..context.compaction_service import default_loop_context_compaction
+
+            compaction_service = default_loop_context_compaction()
+        self._compaction_service = compaction_service
+        self._loop_entry_destination = loop_entry_destination
+        self._end_destination = end_destination
+
+    def _should_attempt_recovery(self, state: dict[str, Any]) -> bool:
+        """判断是否应该尝试恢复。
+
+        检查 state.has_attempted_reactive_compact 标记：
+        - False：首次溢出，应尝试 recovery
+        - True：已尝试过，应走 exhausted 分支
+
+        参数:
+            state: LangGraph state
+
+        返回:
+            True 表示应尝试 recovery，False 表示应直接 exhausted
+        """
+        return not state.get(STATE_HAS_ATTEMPTED_REACTIVE_COMPACT, False)
+
+    def _execute_recovery_compact(
+        self, state: dict[str, Any]
+    ) -> dict[str, Any]:
+        """执行 recovery compact 并返回状态补丁。
+
+        参数:
+            state: LangGraph state
+
+        返回:
+            状态补丁，可能包含：
+            - 成功：{messages: [...], transition: {...}, has_attempted_reactive_compact: True}
+            - 失败：{exhausted: True, error_kind: "context_window", should_end: True, ...}
+        """
+        messages = state.get(STATE_MESSAGES, [])
+        if not isinstance(messages, list):
+            logger.warning("Recovery compact: messages 不是列表，跳过 compact")
+            return self._build_exhausted_patch(state)
+
+        # 记录原始 messages 用于比较
+        original_id = id(messages)
+        original_len = len(messages)
+
+        try:
+            # 使用 REACTIVE_CONTEXT_WINDOW 模式执行 compact
+            patch = self._compaction_service.run(
+                state,
+                model=None,  # 暂不支持 LLM 摘要
+                mode=CompactionMode.REACTIVE_CONTEXT_WINDOW,
+            )
+        except Exception as e:
+            logger.exception(f"Recovery compact 异常: {e}")
+            return self._build_exhausted_patch(state)
+
+        # 检查 compact 是否实际压缩了 messages
+        # 成功标准：messages 被替换（id 变化）或数量减少
+        if not patch:
+            logger.warning("Recovery compact: 无变化，视为失败")
+            return self._build_exhausted_patch(state)
+
+        new_messages = patch.get(STATE_MESSAGES, messages)
+        if not isinstance(new_messages, list):
+            logger.warning("Recovery compact: 返回 messages 不是列表")
+            return self._build_exhausted_patch(state)
+
+        # 检查是否实际压缩
+        is_compacted = (
+            id(new_messages) != original_id or len(new_messages) < original_len
+        )
+
+        if not is_compacted:
+            logger.warning("Recovery compact: messages 未实际压缩，视为失败")
+            return self._build_exhausted_patch(state)
+
+        # 成功：构建返回补丁
+        result = dict(patch)
+        result[STATE_HAS_ATTEMPTED_REACTIVE_COMPACT] = True
+        result[STATE_TRANSITION] = {
+            "reason": TRANSITION_REACTIVE_COMPACT_RETRY,
+            "error_reason": TERMINAL_REASON_MAX_TOKENS,
+            "overflow_kind": "context_window",
+        }
+        result["destination"] = self._loop_entry_destination
+
+        logger.info(f"Recovery compact 成功: {original_len} → {len(new_messages)} 条消息")
+        return result
+
+    def _build_exhausted_patch(
+        self, state: dict[str, Any]
+    ) -> dict[str, Any]:
+        """构建 exhausted 状态补丁。
+
+        参数:
+            state: LangGraph state
+
+        返回:
+            exhausted 状态补丁
+        """
+        return {
+            STATE_HAS_ATTEMPTED_REACTIVE_COMPACT: True,
+            STATE_WITHHELD_ERROR: None,
+            STATE_SHOULD_END: True,
+            STATE_TERMINAL_REASON: TERMINAL_REASON_MAX_TOKENS,
+            STATE_TRANSITION: {
+                "reason": TRANSITION_CONTEXT_WINDOW_RECOVERY_EXHAUSTED,
+                "error_reason": TERMINAL_REASON_MAX_TOKENS,
+                "overflow_kind": "context_window",
+            },
+            "error_kind": "context_window",
+            "exhausted": True,
+            "destination": self._end_destination,
+        }
+
+    def __call__(self, state: dict[str, Any]) -> dict[str, Any]:
+        """节点入口（LangGraph 调用约定）。
+
+        参数:
+            state: LangGraph state
+
+        返回:
+            状态补丁
+        """
+        if not self._should_attempt_recovery(state):
+            logger.info("Recovery compact: 已尝试过，直接 exhausted")
+            result = self._build_exhausted_patch(state)
+            emit_loop_decision_event(
+                state=state,
+                decision="terminal",
+                selected_edge=self._end_destination,
+                reason_code="context_window_recovery_exhausted",
+                evidence={"step": state.get("step", 0), "already_attempted": True},
+            )
+            return result
+
+        logger.info("Recovery compact: 首次溢出，尝试 recovery")
+        result = self._execute_recovery_compact(state)
+
+        if result.get("exhausted"):
+            emit_loop_decision_event(
+                state=state,
+                decision="terminal",
+                selected_edge=self._end_destination,
+                reason_code="context_window_recovery_exhausted",
+                evidence={"step": state.get("step", 0), "compact_failed": True},
+            )
+        else:
+            emit_loop_decision_event(
+                state=state,
+                decision="continue",
+                selected_edge=self._loop_entry_destination,
+                reason_code="reactive_compact_retry",
+                evidence={"step": state.get("step", 0)},
+            )
+
+        return result
+
+
+__all__ = ["ReactiveCompactNode"]

langchain_agentx/observability/events/finish_event_payload_composer.py

@@ -22,8 +22,10 @@ from langchain_core.messages import AIMessage
 
 from langchain_agentx.loop.exit.reason_codes import (
     ERROR_TERMINAL_REASONS,
+    TERMINAL_REASON_COMPLETED,
     TERMINAL_REASON_ERROR,
     TERMINAL_REASON_MAX_TOKENS,
+    TRANSITION_API_ERROR_TERMINAL,
     TRANSITION_CONTEXT_WINDOW_RECOVERY_EXHAUSTED,
     TRANSITION_MAX_OUTPUT_TOKENS_RECOVERY_EXHAUSTED,
     TRANSITION_WITHHOLD_RECOVERY_EXHAUSTED,
@@ -117,7 +119,9 @@ class FinishEventPayloadComposer:
 
         from langchain_agentx.loop.exit.exit_logic import RECOVERABLE_TERMINAL_REASONS
 
-        is_error = self._resolve_is_error(terminal_reason, withheld_error, snap)
+        is_error = self._resolve_is_error(
+            terminal_reason, withheld_error, snap, messages=messages
+        )
         recoverable = (
             terminal_reason in RECOVERABLE_TERMINAL_REASONS if terminal_reason else False
         )
@@ -125,6 +129,14 @@ class FinishEventPayloadComposer:
         recovery_exhausted = (
             transition_reason in _RECOVERY_EXHAUSTED_TRANSITIONS if is_error else False
         )
+        if (
+            not recovery_exhausted
+            and is_error
+            and terminal_reason == TERMINAL_REASON_MAX_TOKENS
+            and resolve_overflow_kind(withheld_error) == "context_window"
+            and bool(snap.get("has_attempted_reactive_compact"))
+        ):
+            recovery_exhausted = True
         error_kind = self._infer_error_kind(
             terminal_reason=terminal_reason,
             is_error=is_error,
@@ -166,7 +178,14 @@ class FinishEventPayloadComposer:
         terminal_reason: str | None,
         withheld_error: dict[str, Any] | None,
         snap: dict[str, Any],
+        *,
+        messages: list[Any] | None = None,
     ) -> bool:
+        transition_reason = _transition_reason(snap)
+        if transition_reason == TRANSITION_API_ERROR_TERMINAL:
+            return True
+        if _is_api_error_message(_last_ai_message(messages)):
+            return True
         if terminal_reason in ERROR_TERMINAL_REASONS:
             return True
         if isinstance(withheld_error, dict) and withheld_error and snap.get("should_end"):
@@ -185,6 +204,13 @@ class FinishEventPayloadComposer:
         if not is_error:
             return None
 
+        last_ai = _last_ai_message(messages)
+        if (
+            transition_reason == TRANSITION_API_ERROR_TERMINAL
+            or _is_api_error_message(last_ai)
+        ):
+            return "api_error"
+
         overflow_kind = resolve_overflow_kind(withheld_error)
         if terminal_reason == TERMINAL_REASON_MAX_TOKENS:
             if overflow_kind == "context_window" or (
@@ -198,9 +224,9 @@ class FinishEventPayloadComposer:
             return "unknown"
 
         if terminal_reason == TERMINAL_REASON_ERROR:
-            last_ai = _last_ai_message(messages)
-            if _is_api_error_message(last_ai):
-                return "api_error"
+            return "unknown"
+
+        if terminal_reason == TERMINAL_REASON_COMPLETED:
             return "unknown"
 
         return "unknown"