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

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)

ksadk/runtime/_runner_adapter/stream_mapping.py

@@ -15,6 +15,7 @@ import logging
 import time
 from collections.abc import Mapping
 from contextlib import nullcontext
+from dataclasses import replace
 from typing import TYPE_CHECKING, Any, AsyncIterator, Optional, cast
 
 from pydantic import JsonValue
@@ -25,6 +26,7 @@ from ksadk.conversations.runtime_observability import (
     _set_conversation_output_attributes,
     _set_conversation_span_attributes,
     _set_conversation_usage_attributes,
+    _set_skill_eval_result_attributes,
 )
 from ksadk.events.canonical import (
     ApprovalRequest,
@@ -36,6 +38,7 @@ from ksadk.events.canonical import (
     ItemCompleted,
     ItemStarted,
     ItemUpdated,
+    RunCompleted,
     RunFailed,
     RunProgress,
     RuntimeEvent,
@@ -46,6 +49,9 @@ from ksadk.events.content import DataContent, TextContent, ToolCallContent, Tool
 from ksadk.events.identity import stable_event_id, stable_item_id, stable_scope_id
 from ksadk.runtime.adapter import RunHandle
 from ksadk.runtime.preprocessing import PreparedRuntimeStart
+from ksadk.runtime.skill_eval_result import extract_skill_eval_result
+from ksadk.runtime.timing import extract_timing
+from ksadk.runtime.usage import canonical_usage_payload
 from ksadk.runtime_context import platform_invocation_scope
 from ksadk.tools.gateway import approval_interrupt_info_from_result
 
@@ -65,6 +71,38 @@ async def _anext_or_stop(gen: AsyncIterator[Any]) -> Any:
         return _STREAM_STOP
 
 
+async def _pump_stream_in_one_context(
+    gen: AsyncIterator[Any],
+    requests: asyncio.Queue[None],
+    results: asyncio.Queue[tuple[str, Any]],
+) -> None:
+    """Own one async generator in one Task without reading ahead."""
+    outcome: tuple[str, Any] = ("stop", _STREAM_STOP)
+    try:
+        while True:
+            await requests.get()
+            try:
+                chunk = await gen.__anext__()
+            except StopAsyncIteration:
+                break
+            await results.put(("chunk", chunk))
+    except asyncio.CancelledError:
+        raise
+    except BaseException as exc:
+        outcome = ("error", exc)
+    finally:
+        aclose = getattr(gen, "aclose", None)
+        if callable(aclose):
+            try:
+                await aclose()
+            except asyncio.CancelledError:
+                raise
+            except BaseException as exc:
+                if outcome[0] != "error":
+                    outcome = ("error", exc)
+    await results.put(outcome)
+
+
 def _a2ui_surface_event(
     self: Any,
     handle: RunHandle,
@@ -192,14 +230,27 @@ class _RunnerStreamMappingMixin:
     """``_map_runner_stream`` / ``_chunk_to_event`` 的实现载体(纯移动自 runner_adapter)。"""
 
     async def _map_runner_stream(
-        self, handle: RunHandle, runner_input: dict
+        self,
+        handle: RunHandle,
+        runner_input: dict,
+        *,
+        active_run: Optional[_ActiveRun] = None,
     ) -> AsyncIterator[RuntimeEvent]:
-        run: Optional[_ActiveRun] = self._active_runs.get(handle.run_id)  # type: ignore[attr-defined]
+        run: Optional[_ActiveRun] = active_run
+        if run is None:
+            run = self._active_runs.get(handle.run_id)  # type: ignore[attr-defined]
         interrupt = run.interrupt_event if run is not None else None
         prepared_start = run.__dict__.get("_prepared_start") if run is not None else None
         invocation_context = (
             prepared_start.context if isinstance(prepared_start, PreparedRuntimeStart) else None
         )
+        if run is not None and run.__dict__.get("_tag_context") is not None:
+            tag_context = run.__dict__["_tag_context"]
+            invocation_context = (
+                replace(invocation_context, session=tag_context.session)
+                if invocation_context is not None
+                else tag_context
+            )
         scope = (
             platform_invocation_scope(invocation_context)
             if invocation_context is not None
@@ -209,6 +260,9 @@ class _RunnerStreamMappingMixin:
         accumulated_output = ""
         usage: dict[str, Any] = {}
         runner_gen: Optional[AsyncIterator[Any]] = None
+        pump_requests: asyncio.Queue[None] | None = None
+        pump_queue: asyncio.Queue[tuple[str, Any]] | None = None
+        pump_task: asyncio.Task[None] | None = None
         # span scope 经 runner_adapter 模块属性间接解析,保持既有 monkeypatch
         # patch 点(tests/agui/test_runtime_preprocessing.py)继续生效。
         from ksadk.runtime import runner_adapter as _runner_adapter_module
@@ -234,7 +288,9 @@ class _RunnerStreamMappingMixin:
                     # _chunk_to_events 的 tool_result 分支,canonical 快速路径
                     # (stream_canonical_events)不覆盖该语义。
                     if getattr(
-                        self._runner, "supports_gateway_approval_semantic_resume", False  # type: ignore[attr-defined]
+                        self._runner,
+                        "supports_gateway_approval_semantic_resume",
+                        False,  # type: ignore[attr-defined]
                     ):
                         canonical_stream = None
                     stream_result = (
@@ -247,10 +303,26 @@ class _RunnerStreamMappingMixin:
                         # 调用返回 coroutine,需 await 得到迭代器。
                         stream_result = await stream_result
                     runner_gen = cast(AsyncIterator[Any], stream_result)
+                    if self._requires_stable_stream_context():
+                        pump_requests = asyncio.Queue(maxsize=1)
+                        pump_queue = asyncio.Queue()
+                        pump_task = asyncio.create_task(
+                            _pump_stream_in_one_context(
+                                runner_gen, pump_requests, pump_queue
+                            )
+                        )
+                        if run is not None:
+                            run.chunk_task = pump_task
                     while True:
                         # 竞速:下一个 runner chunk vs cancel 中断事件。
-                        chunk_task = asyncio.ensure_future(_anext_or_stop(runner_gen))
-                        if run is not None:
+                        if pump_requests is not None:
+                            pump_requests.put_nowait(None)
+                        chunk_task = asyncio.ensure_future(
+                            pump_queue.get()
+                            if pump_queue is not None
+                            else _anext_or_stop(runner_gen)
+                        )
+                        if run is not None and pump_task is None:
                             run.chunk_task = chunk_task
                         wait_set = {chunk_task}
                         interrupt_task = (
@@ -269,23 +341,33 @@ class _RunnerStreamMappingMixin:
                             await asyncio.gather(*pending, return_exceptions=True)
                         if interrupt_task is not None and interrupt_task in done:
                             # cancel 中断:安全关闭 runner 流(同一 task)并停止。
-                            chunk_task.cancel()
-                            await asyncio.gather(chunk_task, return_exceptions=True)
+                            owner_task = pump_task or chunk_task
+                            owner_task.cancel()
+                            await asyncio.gather(owner_task, return_exceptions=True)
                             return
                         try:
                             chunk = chunk_task.result()
+                            if pump_task is not None:
+                                outcome, payload = cast(tuple[str, Any], chunk)
+                                if outcome == "error":
+                                    raise payload
+                                chunk = payload
                         except asyncio.CancelledError:
                             if interrupt is not None and interrupt.is_set():
                                 return
                             raise
                         finally:
-                            if run is not None:
+                            if run is not None and pump_task is None:
                                 run.chunk_task = None
                         if chunk is _STREAM_STOP:
                             return
                         if isinstance(chunk, EventEnvelope):
                             # canonical 事件(来自 stream_canonical_events):直接转发,
                             # 追踪 output/usage 供 span 属性。
+                            if isinstance(chunk, RunCompleted):
+                                evidence = extract_skill_eval_result(chunk.source.metadata.get("metrics"))
+                                if evidence:
+                                    _set_skill_eval_result_attributes(span, evidence)
                             if isinstance(chunk, ItemCompleted) and chunk.item_kind == "message":
                                 accumulated_output = "".join(
                                     part.text
@@ -305,6 +387,11 @@ class _RunnerStreamMappingMixin:
                         if isinstance(chunk, dict):
                             chunk_type = str(chunk.get("type") or "")
                             if chunk_type == "final" and run is not None:
+                                if evidence := extract_skill_eval_result(chunk):
+                                    run.completion_metrics["skill_eval_result"] = evidence
+                                    _set_skill_eval_result_attributes(span, evidence)
+                                if timing := extract_timing(chunk):
+                                    run.completion_metrics["timing"] = timing
                                 for source_key, target_key in (
                                     ("duration_ms", "duration_ms"),
                                     ("started_at", "started_at"),
@@ -333,7 +420,11 @@ class _RunnerStreamMappingMixin:
                 if accumulated_output:
                     _set_conversation_output_attributes(span, accumulated_output)
                 _set_conversation_usage_attributes(span, usage)
-                if runner_gen is not None:
+                if pump_task is not None:
+                    if not pump_task.done():
+                        pump_task.cancel()
+                    await asyncio.gather(pump_task, return_exceptions=True)
+                elif runner_gen is not None:
                     aclose = getattr(runner_gen, "aclose", None)
                     if callable(aclose):
                         try:
@@ -341,6 +432,8 @@ class _RunnerStreamMappingMixin:
                         except Exception:  # noqa: BLE001
                             pass
                 if run is not None:
+                    if pump_task is not None:
+                        run.chunk_task = None
                     run.cancellation_ack.set()
 
     def _chunk_to_event(
@@ -387,6 +480,24 @@ class _RunnerStreamMappingMixin:
                 )
             ]
 
+        def tool_item_id(effective_call_id: str, item_kind: str) -> str:
+            components = (effective_call_id, item_kind)
+            if run is None:
+                return stable_item_id(framework, run_id, *components)
+            occurrences = run.__dict__.setdefault("_tool_item_occurrences", {})
+            occurrence_key = (scope_id, item_kind, effective_call_id)
+            occurrence = int(occurrences.get(occurrence_key, 0))
+            occurrences[occurrence_key] = occurrence + 1
+            if occurrence == 0:
+                return stable_item_id(framework, run_id, *components)
+            return stable_item_id(
+                framework,
+                run_id,
+                *components,
+                "occurrence",
+                occurrence,
+            )
+
         chunk_type = chunk.get("type")
 
         # ---- reasoning ----
@@ -449,7 +560,7 @@ class _RunnerStreamMappingMixin:
             )
             name = str(chunk.get("tool_name") or chunk.get("name") or "tool")
             effective_call_id = call_id or name
-            item_id = stable_item_id(framework, run_id, effective_call_id, "tool_call")
+            item_id = tool_item_id(effective_call_id, "tool_call")
             part_id = "tool_call"
             tc_content = ToolCallContent(
                 part_id=part_id,
@@ -496,7 +607,7 @@ class _RunnerStreamMappingMixin:
             )
             name = str(chunk.get("tool_name") or chunk.get("name") or "tool")
             effective_call_id = call_id or name
-            item_id = stable_item_id(framework, run_id, effective_call_id, "tool_result")
+            item_id = tool_item_id(effective_call_id, "tool_result")
             part_id = "tool_result"
             result_data = chunk.get("tool_output", chunk.get("output"))
             # ToolGateway 审批可能出现在"本已终态"的 tool result 里;识别后转为
@@ -618,26 +729,55 @@ class _RunnerStreamMappingMixin:
             if isinstance(runtime_ref, dict):
                 handle.native_ref.update(runtime_ref)
             item_id = stable_item_id(framework, run_id, "$run")
+            # Extract checkpoint capability from agentengine metadata so the
+            # canonical event carries backend/scope/next_node facts that the
+            # projection layer (projection.py) needs for the REST payload.
+            capability: dict[str, Any] = {}
+            for cap_key in (
+                # Capability / resume fields
+                "backend", "scope", "durable", "next_node",
+                "is_terminal", "is_resumable", "resume_status",
+                "resume_disabled_reason",
+                # Checkpoint display fields
+                "phase", "stage", "stage_name", "stage_key",
+                "stage_index", "total_stages", "summary",
+                "next_action", "status", "artifact_preview",
+                "tool_name", "receipt_key", "artifact_path",
+            ):
+                cap_val = agentengine.get(cap_key)
+                if cap_val is not None:
+                    capability[cap_key] = cap_val
+            canonical_kwargs = self._canonical_kwargs(  # type: ignore[attr-defined]
+                handle,
+                scope_id=scope_id,
+                item_id=item_id,
+                event_type="continuation.created",
+                part_id="continuation",
+            )
+            if capability:
+                base_source = canonical_kwargs["source"]
+                canonical_kwargs["source"] = SourceRef(
+                    framework=base_source.framework,
+                    protocol=base_source.protocol,
+                    native_run_id=base_source.native_run_id,
+                    metadata={**base_source.metadata, "capability": capability},
+                )
             ref_value = cast(
                 JsonValue,
                 framework_ref if isinstance(framework_ref, dict) else {},
             )
+            next_node_val = agentengine.get("next_node")
             return [
                 ContinuationCreated(
-                    **self._canonical_kwargs(  # type: ignore[attr-defined]
-                        handle,
-                        scope_id=scope_id,
-                        item_id=item_id,
-                        event_type="continuation.created",
-                        part_id="continuation",
-                    ),
+                    **canonical_kwargs,
                     continuation_id=checkpoint_id,
                     continuation_kind="graph_checkpoint",
-                    resumable=True,
+                    resumable=bool(agentengine.get("is_resumable", True)),
                     ref={
                         "framework": ckpt_framework,
                         "framework_ref": ref_value,
                         "resume_target": ref_value,
+                        **({"next_node": str(next_node_val)} if next_node_val else {}),
                     },
                 )
             ]
@@ -662,7 +802,11 @@ class _RunnerStreamMappingMixin:
         # ---- usage ----
         if chunk_type == "usage":
             raw_usage = chunk.get("usage")
-            usage_dict: dict[str, Any] = raw_usage if isinstance(raw_usage, dict) else {}
+            usage_dict = (
+                canonical_usage_payload(raw_usage, runtime_type=framework)
+                if isinstance(raw_usage, Mapping)
+                else {}
+            )
             item_id = stable_item_id(framework, run_id, "$run")
             return [
                 UsageReported(
@@ -749,6 +893,36 @@ class _RunnerStreamMappingMixin:
             events.extend(
                 ensure_started(item_id=item_id, item_kind="message", phase="final_answer")
             )
+            # The dict-chunk path is still required for gateway approval flows.
+            # It must project runner-final usage into the canonical stream too;
+            # otherwise non-streaming Runs aggregate zero even when the runner
+            # returned provider token accounting.
+            raw_usage = chunk.get("usage")
+            if isinstance(raw_usage, Mapping):
+                usage_dict = canonical_usage_payload(
+                    raw_usage,
+                    runtime_type=framework,
+                )
+                input_tokens = usage_dict["input_tokens"]
+                output_tokens = usage_dict["output_tokens"]
+                total_tokens = usage_dict["total_tokens"]
+                if input_tokens or output_tokens or total_tokens:
+                    events.append(
+                        UsageReported(
+                            **self._canonical_kwargs(  # type: ignore[attr-defined]
+                                handle,
+                                scope_id=scope_id,
+                                item_id=stable_item_id(framework, run_id, "$run"),
+                                event_type="usage.reported",
+                                part_id="usage",
+                            ),
+                            input_tokens=input_tokens,
+                            output_tokens=output_tokens,
+                            total_tokens=total_tokens,
+                            cached_tokens=usage_dict["cached_tokens"],
+                            reasoning_tokens=usage_dict["reasoning_tokens"],
+                        )
+                    )
             events.append(
                 ItemCompleted(
                     **self._canonical_kwargs(  # type: ignore[attr-defined]

ksadk/runtime/adapter.py

@@ -114,6 +114,8 @@ class ResumePayload(BaseModel):
     call_id: Optional[str] = None
     """对应的 ``approval.requested`` / ``tool.call`` 的 id(如有)。"""
     data: Any = None
+    session_context: dict[str, Any] | None = None
+    """Platform-owned snapshot for a new resume invocation; never model input."""
 
 
 # ---------------------------------------------------------------------------