agentengine-sdk-python 0.8.4__py3-none-any.whl → 0.8.5__py3-none-any.whl

ksadk/runners/langgraph_runner.py

@@ -6,7 +6,6 @@ LangGraphRunner - LangGraph 框架运行时
 
 import asyncio
 import base64
-import inspect
 import os
 import re
 import sqlite3
@@ -17,13 +16,21 @@ from langgraph.types import Command
 
 from ksadk.conversations.attachments import classify_attachment_kind, read_attachment_uri_bytes
 from ksadk.runners._langgraph_runner_streams import _LangGraphStreamMixin
+from ksadk.runners._langgraph_state_reader import (
+    interrupt_info_from_state,
+    known_subgraph_namespaces,
+    read_graph_state,
+)
+from ksadk.runners._session_identity import LangGraphSessionIdentityMixin
 from ksadk.runners.base_runner import BaseRunner
 from ksadk.runners.utils import load_agent_module
+from ksadk.runtime.skill_eval_result import skill_eval_response_fields
+from ksadk.runtime.timing import extract_timing
 from ksadk.sessions import resolve_persistence_topology
 from ksadk.sessions.continuity import LangGraphSessionAdapter
 
 
-class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
+class LangGraphRunner(LangGraphSessionIdentityMixin, _LangGraphStreamMixin, BaseRunner):
     """LangGraph 框架运行时
 
     透传原生 LangGraph 功能，支持任意 State 格式
@@ -35,12 +42,14 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
 
     def __init__(self, detection_result: Any, project_dir: str):
         super().__init__(detection_result, project_dir)
+        self._module: Any = None
         self._managed_checkpoint_lock = asyncio.Lock()
         self._managed_checkpoint_prepared = False
         self._managed_checkpoint_state = "uninitialized"
         self._managed_checkpoint_error: tuple[str, str] | None = None
         self._managed_checkpoint_pool: Any = None
-        self._managed_checkpoint_namespace = ""
+        self._identity_thread_bindings: dict[str, str] = {}
+        self._identity_thread_lock = asyncio.Lock()
 
     def load_agent(self) -> None:
         self._load_agent(force_reload=False)
@@ -133,8 +142,7 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
             return "unknown"
         try:
             closed_values = {
-                name: cell.cell_contents
-                for name, cell in zip(code.co_freevars, closure)
+                name: cell.cell_contents for name, cell in zip(code.co_freevars, closure)
             }
         except (AttributeError, ValueError):
             return "unknown"
@@ -199,8 +207,7 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
                 "ResumeMode": "none",
                 "ReasonCode": "CHECKPOINTER_NOT_DURABLE",
                 "Reason": (
-                    "In-memory checkpoint cannot be recovered after process restart "
-                    "or across pods"
+                    "In-memory checkpoint cannot be recovered after process restart or across pods"
                 ),
             }
         return {
@@ -326,18 +333,6 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
             "on",
         }
 
-    @staticmethod
-    def _resolve_checkpoint_namespace() -> str:
-        session_namespace = str(os.getenv("KSADK_SESSION_NAMESPACE") or "").strip()
-        if session_namespace:
-            return session_namespace
-        agent_id = str(
-            os.getenv("AGENTENGINE_AGENT_ID")
-            or os.getenv("KSADK_AGENT_ID")
-            or "default"
-        ).strip()
-        return f"agent:{agent_id}"
-
     async def _create_managed_postgres_saver(self, dsn: str) -> tuple[Any, Any]:
         from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
         from psycopg.rows import dict_row
@@ -390,20 +385,12 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
         )
         checkpoint_id = str(
             native_ref.get("checkpoint_id")
-            or (
-                langgraph_ref.get("checkpoint_id")
-                if isinstance(langgraph_ref, Mapping)
-                else ""
-            )
+            or (langgraph_ref.get("checkpoint_id") if isinstance(langgraph_ref, Mapping) else "")
             or ""
         ).strip()
         thread_id = str(
             native_ref.get("thread_id")
-            or (
-                langgraph_ref.get("thread_id")
-                if isinstance(langgraph_ref, Mapping)
-                else ""
-            )
+            or (langgraph_ref.get("thread_id") if isinstance(langgraph_ref, Mapping) else "")
             or ""
         ).strip()
         if not checkpoint_id or not thread_id:
@@ -424,11 +411,7 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
         }
         checkpoint_ns = str(
             native_ref.get("checkpoint_ns")
-            or (
-                langgraph_ref.get("checkpoint_ns")
-                if isinstance(langgraph_ref, Mapping)
-                else ""
-            )
+            or (langgraph_ref.get("checkpoint_ns") if isinstance(langgraph_ref, Mapping) else "")
             or ""
         ).strip()
         if checkpoint_ns:
@@ -436,8 +419,7 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
         config = {"configurable": configurable}
 
         if callable(
-            getattr(self._agent, "aget_state", None)
-            or getattr(self._agent, "get_state", None)
+            getattr(self._agent, "aget_state", None) or getattr(self._agent, "get_state", None)
         ):
             return await self._checkpoint_state_resolves(
                 config,
@@ -470,14 +452,9 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
     ) -> bool:
         """Confirm one exact persisted checkpoint through the active graph."""
 
-        get_state = getattr(self._agent, "aget_state", None) or getattr(
-            self._agent, "get_state", None
-        )
-        if not callable(get_state):
+        state, _error = await self._read_graph_state(config)
+        if state is None:
             return False
-        state = get_state(config)
-        if inspect.isawaitable(state):
-            state = await state
         state_created_at = (
             state.get("created_at")
             if isinstance(state, Mapping)
@@ -509,19 +486,13 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
         async with self._managed_checkpoint_lock:
             if self._managed_checkpoint_state in {"ready", "terminal_failure"}:
                 return
-            if (
-                self._managed_checkpoint_state == "transient_failure"
-                and not allow_transient_retry
-            ):
+            if self._managed_checkpoint_state == "transient_failure" and not allow_transient_retry:
                 return
 
             checkpointer = getattr(self._agent, "checkpointer", None)
             if checkpointer is None:
                 checkpointer = getattr(self._agent, "_checkpointer", None)
             if self._checkpoint_backend_from_saver(checkpointer) == "postgres":
-                self._managed_checkpoint_namespace = (
-                    self._resolve_checkpoint_namespace()
-                )
                 self._managed_checkpoint_prepared = True
                 self._managed_checkpoint_state = "ready"
                 return
@@ -537,6 +508,13 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
                 self._managed_checkpoint_state = "terminal_failure"
                 return
 
+            # 探测路径可能先于首个 chat 请求 probe（本地调试与手动部署形态的
+            # capability adapter 均未 preflight）。此时驻留图还没加载，先完成
+            # 模块加载：既避免把"未加载"误判为"无 durable checkpointer"，
+            # 也避免对未初始化的 _module 抛 AttributeError 被静默吞掉。
+            if self._agent is None:
+                self.load_agent()
+
             factory = getattr(self._module, "ksadk_graph_factory", None)
             if not callable(factory):
                 self._managed_checkpoint_error = (
@@ -548,6 +526,23 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
                 self._managed_checkpoint_state = "terminal_failure"
                 return
 
+            # 托管 checkpointer 从 checkpoint 恢复 messages（append-only reducer）。
+            # 若接入方未导出 ksadk_prepare_state，runner 走 _to_state 把 history 再注入
+            # 一遍，与 checkpointer 恢复的 messages 双重叠加。托管 checkpoint 必须同时
+            # 提供 hook，否则 fail-closed，避免生产 auto checkpoint 下静默双重注入。
+            prepare_state_hook = getattr(self._module, "ksadk_prepare_state", None)
+            if not callable(prepare_state_hook):
+                self._managed_checkpoint_error = (
+                    "PREPARE_STATE_HOOK_REQUIRED",
+                    "Managed PostgreSQL checkpoint requires ksadk_prepare_state hook to "
+                    "avoid double history injection (checkpointer messages restore + "
+                    "_to_state history). Export ksadk_prepare_state(payload, "
+                    "session_context) in the agent module.",
+                )
+                self._managed_checkpoint_prepared = True
+                self._managed_checkpoint_state = "terminal_failure"
+                return
+
             pool = None
             try:
                 saver, pool = await self._create_managed_postgres_saver(checkpoint_target.dsn)
@@ -556,9 +551,6 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
                     raise TypeError("ksadk_graph_factory must return a compiled LangGraph graph")
                 self._agent = managed_graph
                 self._managed_checkpoint_pool = pool
-                self._managed_checkpoint_namespace = (
-                    self._resolve_checkpoint_namespace()
-                )
                 self._managed_checkpoint_error = None
                 self._managed_checkpoint_state = "ready"
             except (ModuleNotFoundError, ImportError):
@@ -577,8 +569,10 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
                 permission_failure = any(
                     marker in error_name for marker in ("privilege", "permission")
                 )
-                terminal = authentication_failure or permission_failure or isinstance(
-                    exc, (TypeError, ValueError)
+                terminal = (
+                    authentication_failure
+                    or permission_failure
+                    or isinstance(exc, (TypeError, ValueError))
                 )
                 reason_code = (
                     "AUTH_FAILED"
@@ -612,11 +606,14 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
         await super().close()
 
     def _get_config(self, session_id: str) -> dict:
-        """获取运行配置"""
-        config: dict[str, Any] = {"configurable": {"thread_id": session_id}}
-        if self._managed_checkpoint_namespace:
-            config["configurable"]["checkpoint_ns"] = self._managed_checkpoint_namespace
-        return config
+        """获取运行配置
+
+        只写 thread_id：checkpoint_ns 是 LangGraph 的子图寻址字段，不是租户
+        隔离字段。把租户/agent scope 塞进 configurable.checkpoint_ns 会让
+        Pregel 的 aget_state 走子图重定向并抛 "Subgraph <scope> not found"，
+        静默打掉流式路径的待审批探测。租户隔离由 thread_id 承担。
+        """
+        return {"configurable": {"thread_id": session_id}}
 
     @staticmethod
     def _extract_langgraph_checkpoint_ref(payload: Dict[str, Any]) -> dict[str, Any]:
@@ -628,19 +625,28 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
             return {}
         return dict(langgraph_ref)
 
-    @classmethod
+    def _known_subgraph_namespaces(self) -> set[str]:
+        """当前图里真实存在的子图 namespace 集合（无法探测时返回空集）。"""
+        return known_subgraph_namespaces(self._agent)
+
     def _apply_checkpoint_resume_config(
-        cls,
+        self,
         config: dict[str, Any],
         *,
         session_id: str,
         checkpoint_ref: dict[str, Any],
+        enforce_bound_thread: bool = False,
     ) -> dict[str, Any]:
         checkpoint_id = str(checkpoint_ref.get("checkpoint_id") or "").strip()
         if not checkpoint_id:
             raise ValueError("checkpoint_resume requires framework_ref.langgraph.checkpoint_id")
 
-        thread_id = str(checkpoint_ref.get("thread_id") or session_id or "").strip()
+        configured = dict(config.get("configurable") or {})
+        bound_thread_id = str(configured.get("thread_id") or session_id or "").strip()
+        requested_thread_id = str(checkpoint_ref.get("thread_id") or "").strip()
+        if enforce_bound_thread and requested_thread_id and requested_thread_id != bound_thread_id:
+            raise ValueError("checkpoint_resume thread does not belong to this session")
+        thread_id = requested_thread_id or bound_thread_id
         if not thread_id:
             raise ValueError(
                 "checkpoint_resume requires session_id or framework_ref.langgraph.thread_id"
@@ -649,7 +655,23 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
         next_config = dict(config)
         configurable = dict(next_config.get("configurable") or {})
         configurable["thread_id"] = thread_id
-        configurable["checkpoint_ns"] = str(checkpoint_ref.get("checkpoint_ns") or "")
+        # checkpoint_ns 只承载 LangGraph 子图寻址语义。真实运行时 namespace 的
+        # 格式是 "<子图名>:<task_path>"（如 inner:<task-id>），get_subgraphs()
+        # 返回的是裸子图名——所以校验必须只比对首段，合法时保留原始完整
+        # namespace（带任务路径；否则 aget_state 定位不到对应 checkpoint）。
+        # 不匹配的地址（含历史租户/agent scope 残留）显性报错，绝不静默改写
+        # 成根地址——静默改写曾让恢复悄悄指向不存在的 checkpoint。
+        requested_namespace = str(checkpoint_ref.get("checkpoint_ns") or "").strip()
+        if requested_namespace:
+            known = self._known_subgraph_namespaces()
+            leading = requested_namespace.split(":", 1)[0]
+            if leading not in known:
+                raise ValueError(
+                    "checkpoint_resume checkpoint_ns does not address any subgraph "
+                    f"of this graph: {requested_namespace!r} "
+                    f"(known subgraphs: {sorted(known) or 'none'})"
+                )
+        configurable["checkpoint_ns"] = requested_namespace
         configurable["checkpoint_id"] = checkpoint_id
         next_config["configurable"] = configurable
         return next_config
@@ -754,13 +776,8 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
         }
 
     async def _latest_checkpoint_metadata(self, config: dict[str, Any]) -> dict[str, Any]:
-        state = None
-        try:
-            if callable(getattr(self._agent, "aget_state", None)):
-                state = await self._agent.aget_state(config)
-            elif callable(getattr(self._agent, "get_state", None)):
-                state = self._agent.get_state(config)
-        except Exception:
+        state, _error = await self._read_graph_state(config)
+        if state is None:
             return {}
         framework_ref = self._checkpoint_ref_from_state(state)
         if not framework_ref:
@@ -796,13 +813,8 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
         }
 
     async def _latest_state_usage(self, config: dict[str, Any]) -> dict[str, Any]:
-        state = None
-        try:
-            if callable(getattr(self._agent, "aget_state", None)):
-                state = await self._agent.aget_state(config)
-            elif callable(getattr(self._agent, "get_state", None)):
-                state = self._agent.get_state(config)
-        except Exception:
+        state, _error = await self._read_graph_state(config)
+        if state is None:
             return {}
         values = getattr(state, "values", None)
         if values is not None:
@@ -1091,10 +1103,15 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
             final_metadata = final_chunk.get("metadata")
             if isinstance(final_metadata, Mapping):
                 metadata = {**dict(final_metadata), **metadata}
-            result: dict[str, Any] = {"output": output_text, "raw": {"chunks": chunks}}
+            result: dict[str, Any] = {
+                "output": output_text, "raw": {"chunks": chunks},
+                **skill_eval_response_fields(final_chunk),
+            }
             usage = final_chunk.get("usage")
             if isinstance(usage, Mapping) and usage:
                 result["usage"] = dict(usage)
+            if timing := extract_timing(final_chunk):
+                result["timing"] = timing
             if metadata:
                 result["metadata"] = metadata
             return result
@@ -1126,12 +1143,13 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
             payload.get("platform_context"),
             context_schema=getattr(self._agent, "context_schema", None),
         )
-        config = self._get_config(session_id)
+        config = await self._get_session_config(session_id)
         if is_checkpoint_resume:
             config = self._apply_checkpoint_resume_config(
                 config,
                 session_id=session_id,
                 checkpoint_ref=checkpoint_ref,
+                enforce_bound_thread=bool(self._invocation_identity_scope_ref()),
             )
 
         # 判断输入格式 / resume
@@ -1172,7 +1190,12 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
                     context=native_context,
                 )
 
-            output = {"output": self._extract_output(result), "raw": result}
+            output = {
+                "output": self._extract_output(result), "raw": result,
+                **skill_eval_response_fields(result),
+            }
+            if timing := extract_timing(result):
+                output["timing"] = timing
             usage = self._extract_usage(result)
             if usage:
                 output["usage"] = usage
@@ -1186,7 +1209,14 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
 
         except Exception as e:
             if "Interrupt" in type(e).__name__:
-                interrupt_info = self._get_interrupt_info(self._agent.get_state(config))
+                state, error = await self._read_graph_state(config, include_subgraphs=True)
+                if state is None:
+                    # 图已经抛出 interrupt，但状态读不回来：这里必须显性失败，
+                    # 不能返回空 interrupt_info，更不能伪装成正常完成。
+                    raise RuntimeError(
+                        "graph interrupted but its checkpoint state could not be read"
+                    ) from error
+                interrupt_info = self._get_interrupt_info(state)
                 return {
                     "type": "interrupt",
                     "interrupt_info": interrupt_info,
@@ -1221,35 +1251,20 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
                     return str(content)
         return str(result) if result else ""
 
+    async def _read_graph_state(
+        self,
+        config: Mapping[str, Any],
+        *,
+        include_subgraphs: bool = False,
+    ) -> tuple[Any | None, Exception | None]:
+        """统一的图状态读取入口（实现在 `_langgraph_state_reader`，全仓唯一）。"""
+        return await read_graph_state(
+            self._agent, config, include_subgraphs=include_subgraphs
+        )
+
     def _get_interrupt_info(self, state) -> dict:
-        """从 state 中获取 interrupt 信息"""
-        if hasattr(state, "tasks") and state.tasks:
-            for task in state.tasks:
-                if hasattr(task, "interrupts") and task.interrupts:
-                    for intr in task.interrupts:
-                        if hasattr(intr, "value"):
-                            value = intr.value
-                            info = dict(value) if isinstance(value, Mapping) else {"value": value}
-                            interrupt_id = str(getattr(intr, "id", "") or "")
-                            if interrupt_id:
-                                info.setdefault("approval_request_id", interrupt_id)
-                            # HumanInTheLoopMiddleware 的 interrupt value 把
-                            # tool_name/arguments 嵌在 action_requests[0] 里；
-                            # 下游期望顶层字段，这里提取并保留原 action_requests。
-                            action_requests = info.get("action_requests")
-                            if isinstance(action_requests, list) and action_requests:
-                                first = action_requests[0]
-                                if isinstance(first, Mapping):
-                                    info.setdefault("tool_name", str(first.get("name") or ""))
-                                    raw_args = first.get("args") or first.get("arguments")
-                                    if raw_args is not None:
-                                        info.setdefault("arguments", raw_args)
-                                    if first.get("description") is not None:
-                                        info.setdefault(
-                                            "description", str(first.get("description"))
-                                        )
-                            return info
-        return {}
+        """从 state 中获取 interrupt 信息（含子图嵌套任务）。"""
+        return interrupt_info_from_state(state)
 
     async def _stream_checkpoint_resume_updates(
         self,
@@ -1293,6 +1308,7 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
         latest_config = self._latest_checkpoint_config(config)
         state_usage = await self._latest_state_usage(latest_config)
         state_output = ""
+        result_fields: dict[str, Any] = {}
         try:
             state = None
             if callable(getattr(self._agent, "aget_state", None)):
@@ -1302,12 +1318,14 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
             values = getattr(state, "values", None)
             if values is not None:
                 state_output = self._extract_output(values)
+                result_fields = skill_eval_response_fields(values)
         except Exception:
             state_output = ""
         final_output = state_output or self._extract_output(latest_output)
         yield {
             "type": "final",
             "output": final_output,
+            **result_fields,
             **({"usage": state_usage} if state_usage else {}),
             **({"resume_noop": True} if not emitted_update else {}),
         }
@@ -1386,7 +1404,6 @@ class LangGraphRunner(_LangGraphStreamMixin, BaseRunner):
                 )
         return events
 
-
     def _filter_tool_tags(self, content: str) -> str:
         """过滤 <tool_call> 标签"""
         if not isinstance(content, str):

ksadk/runners/patch_langchain.py

@@ -26,6 +26,12 @@ logger = logging.getLogger(__name__)
 
 # Original function reference (to avoid recursion if patched multiple times)
 _original_convert_delta = None
+_original_convert_chunk_to_generation_chunk = None
+
+# Dedup cache for upstream streams that return identical usage in both the
+# finish chunk and the trailing choices=[] chunk. Keyed by response ID so
+# concurrent streams don't interfere.
+_usage_dedup_cache: dict[str, dict] = {}
 
 
 def _patched_convert_delta_to_message_chunk(
@@ -154,10 +160,43 @@ def _patched_chat_get_request_payload(self, input_, *, stop=None, **kwargs):
     return _merge_payload_options(payload, model_options_for_chat_completions(model_options))
 
 
+def _patched_convert_chunk_to_generation_chunk(
+    self,
+    chunk: dict,
+    default_chunk_class: type,
+    base_generation_info: dict | None = None,
+):
+    """Suppress duplicate usage from upstream streams that return identical
+    usage in both the finish chunk and the trailing choices=[] chunk.
+
+    LangChain-core sums usage_metadata across chunks, so two identical usage
+    values produce a 2x result in Langfuse. This patch strips the second
+    occurrence before LangChain sees it, keyed by response ID.
+    """
+    token_usage = chunk.get("usage")
+    response_id = str(chunk.get("id") or "")
+    if token_usage is not None and response_id:
+        cached = _usage_dedup_cache.get(response_id)
+        if cached is not None and cached == token_usage:
+            chunk = dict(chunk)
+            chunk.pop("usage", None)
+        else:
+            _usage_dedup_cache[response_id] = dict(token_usage)
+            if len(_usage_dedup_cache) > 1024:
+                _usage_dedup_cache.clear()
+    choices = chunk.get("choices") or []
+    if len(choices) == 0 and response_id:
+        _usage_dedup_cache.pop(response_id, None)
+    return _original_convert_chunk_to_generation_chunk(
+        self, chunk, default_chunk_class, base_generation_info
+    )
+
+
 def apply_patch():
     """Apply the monkey patch to langchain_openai."""
     global _original_convert_delta, _original_convert_message_to_dict
     global _original_base_get_request_payload, _original_chat_get_request_payload
+    global _original_convert_chunk_to_generation_chunk
     try:
         import langchain_openai.chat_models.base as base_module
 
@@ -178,6 +217,14 @@ def apply_patch():
         _original_chat_get_request_payload = base_module.ChatOpenAI._get_request_payload
         base_module.ChatOpenAI._get_request_payload = _patched_chat_get_request_payload
 
+        # Patch 3: Deduplicate streaming usage to prevent Langfuse token doubling
+        _original_convert_chunk_to_generation_chunk = (
+            base_module.BaseChatOpenAI._convert_chunk_to_generation_chunk
+        )
+        base_module.BaseChatOpenAI._convert_chunk_to_generation_chunk = (
+            _patched_convert_chunk_to_generation_chunk
+        )
+
         base_module._ksadk_patched = True
 
         # logger.info("Applied langchain-openai patch for reasoning_content support")

ksadk/runners/usage_accumulator.py

@@ -1,4 +1,4 @@
-"""usage 累加工具:逐字段求和(input/output/total + details 子键)。
+"""usage 累加工具:逐字段求和(主计数/canonical 明细 + details 子键)。
 
 用于 runner 本轮内多次 LLM 调用的 usage 累积。input_tokens 各 provider 均含 cache
 (Gemini prompt_token_count 含 cached_content_token_count;OpenAI prompt_tokens 含
@@ -11,7 +11,13 @@ from __future__ import annotations
 
 from typing import Any
 
-_MAIN_FIELDS = ("input_tokens", "output_tokens", "total_tokens")
+_SCALAR_FIELDS = (
+    "input_tokens",
+    "output_tokens",
+    "total_tokens",
+    "cached_tokens",
+    "reasoning_tokens",
+)
 _DETAIL_FIELDS = ("input_token_details", "output_token_details")
 
 
@@ -20,7 +26,9 @@ def accumulate_usage(acc: dict[str, Any], delta: dict[str, Any]) -> dict[str, An
     if not delta:
         return dict(acc)
     result = dict(acc)
-    for key in _MAIN_FIELDS:
+    for key in _SCALAR_FIELDS:
+        if key not in result and key not in delta:
+            continue
         result[key] = int(result.get(key) or 0) + int(delta.get(key) or 0)
     for detail_key in _DETAIL_FIELDS:
         delta_details = delta.get(detail_key)