fix(journal): dedup llm.ai.response persistence on re-fired on_llm_end (#5187)

* fix(journal): dedup llm.ai.response persistence on re-fired on_llm_end

LangChain may deliver on_llm_end more than once for the same run_id.
RunJournal already dedups token accounting and the run summary
(_record_message_summary) on that premise via _counted_message_llm_run_ids,
but the durable llm.ai.response self._put() call was left unguarded.

The event store is append-only and count_messages/list_messages read raw
rows without read-time dedup, so a replayed callback persists a second
llm.ai.response row for one logical response while the run's own
message_count counts it once. This inflates count_messages, duplicates a
message in list_messages pagination, and leaves the durable feed
inconsistent with the run summary.

Gate the persistence + summary block by the existing per-run_id guard so a
replayed callback is a no-op, keeping the durable message feed and the run
summary in agreement. Distinct run_ids are unaffected.

Adds regression tests: a re-fired callback for one run_id persists exactly
one row (red on main), and distinct run_ids each still persist a message.

* fix(journal): preserve canonical response on late usage

* fix(journal): preserve late usage while deduplicating responses

* fix(journal): keep first callback response canonical

* fix(journal): snapshot canonical response summaries

---------

Co-authored-by: CorgiBoyG <CorgiBoyG@users.noreply.github.com>
This commit is contained in:
RongJie G 2026-09-06 10:16:17 +08:00 committed by GitHub
parent ab9c1719ee
commit a4ff4b0b3b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 622 additions and 49 deletions

View File

@ -61,6 +61,25 @@ fetch-and-decode of every message row's tool outputs on long threads.
client input, because a welded-in seq goes stale when a fork re-seeds the feed client input, because a welded-in seq goes stale when a fork re-seeds the feed
(#4380). (#4380).
**LLM response callback coalescing** (`runtime/journal.py`): a provider may fire
`on_llm_end` twice for one LangChain run id, first without usage (or with all token
counts zero) and immediately again with usage populated. The first callback's generation
set is always canonical: `RunJournal` stages only its response events and immutable
message-summary fields while retaining the first caller, and applies that callback's
fallback state and tool-call bookkeeping immediately; those effects remain canonical.
It must not retain provider-owned message objects because a provider may mutate and
reuse the same response for the usage replay. Usage metadata is deep-snapshotted,
including nested token-detail mappings, before it enters a staged or buffered event.
An adjacent same-id positive-usage replay may enrich only each corresponding staged
event's metadata/content usage fields. Replay
generation-count differences never add, remove, or replace canonical messages. The next
unrelated event, an effective buffer size (committed plus pending events) reaching the
flush threshold, or an explicit flush commits the staged unit and updates the message
summary. Once that ordering boundary is crossed, a late usage replay can still update the
authoritative run token summary, but it cannot mutate the append-only message event,
caller attribution, fallback state, or tool-call bookkeeping. Closed journals return
from `on_llm_end` before inspecting the response or touching any run state.
**Run delivery receipts** (`runtime/journal.py` + `runs/worker.py`): **Run delivery receipts** (`runtime/journal.py` + `runs/worker.py`):
`RunJournal` records each non-empty artifact update once per tool `Command` for `RunJournal` records each non-empty artifact update once per tool `Command` for
the terminal `run.delivery` event. When a command contains multiple messages, a the terminal `run.delivery` event. When a command contains multiple messages, a

View File

@ -22,6 +22,8 @@ import logging
import threading import threading
import time import time
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from datetime import UTC, datetime from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, cast from typing import TYPE_CHECKING, Any, cast
from uuid import UUID from uuid import UUID
@ -53,6 +55,14 @@ _LEGACY_SUMMARY_MESSAGE_NAME = "summary"
_PERSISTED_HIDDEN_HUMAN_INPUT_RESPONSE_SOURCES = frozenset({"ask_clarification", "sandbox_network"}) _PERSISTED_HIDDEN_HUMAN_INPUT_RESPONSE_SOURCES = frozenset({"ask_clarification", "sandbox_network"})
@dataclass
class _PendingLlmResponse:
llm_run_id: str
events: list[dict]
message_count: int
last_ai_message: str | None
def _should_persist_human_input_message(message: BaseMessage) -> bool: def _should_persist_human_input_message(message: BaseMessage) -> bool:
if not isinstance(message, HumanMessage): if not isinstance(message, HumanMessage):
return False return False
@ -243,6 +253,7 @@ class RunJournal(BaseCallbackHandler):
# Write buffer # Write buffer
self._buffer: list[dict] = [] self._buffer: list[dict] = []
self._pending_llm_response: _PendingLlmResponse | None = None
self._pending_flush_tasks: set[asyncio.Task[None]] = set() self._pending_flush_tasks: set[asyncio.Task[None]] = set()
self._pending_progress_task: asyncio.Task[None] | None = None self._pending_progress_task: asyncio.Task[None] | None = None
self._pending_progress_delayed = False self._pending_progress_delayed = False
@ -267,6 +278,7 @@ class RunJournal(BaseCallbackHandler):
self._counted_llm_run_ids: set[str] = set() self._counted_llm_run_ids: set[str] = set()
self._counted_external_source_ids: set[str] = set() self._counted_external_source_ids: set[str] = set()
self._counted_message_llm_run_ids: set[str] = set() self._counted_message_llm_run_ids: set[str] = set()
self._llm_response_callers: dict[str, str] = {}
self._memory_context_recorded = False self._memory_context_recorded = False
self._tool_promotion_claim_lock = threading.Lock() self._tool_promotion_claim_lock = threading.Lock()
self._claimed_tool_promotions: set[str] = set() self._claimed_tool_promotions: set[str] = set()
@ -306,6 +318,14 @@ class RunJournal(BaseCallbackHandler):
"""Extract displayable text from a message's mixed content shape.""" """Extract displayable text from a message's mixed content shape."""
return message_to_text(message, text_attribute_fallback=True) return message_to_text(message, text_attribute_fallback=True)
def _message_summary_text(self, message: BaseMessage, *, caller: str | None = None) -> str | None:
"""Return the bounded user-facing AI summary text for one message."""
is_ai_message = isinstance(message, AIMessage) or getattr(message, "type", None) == "ai"
if not is_ai_message or (caller is not None and caller != "lead_agent"):
return None
text = self._message_text(message).strip()
return text[:2000] if text else None
def _record_message_summary(self, message: BaseMessage, *, caller: str | None = None) -> None: def _record_message_summary(self, message: BaseMessage, *, caller: str | None = None) -> None:
"""Update run-level convenience fields for persisted run rows.""" """Update run-level convenience fields for persisted run rows."""
self._msg_count += 1 self._msg_count += 1
@ -313,11 +333,9 @@ class RunJournal(BaseCallbackHandler):
# ``last_ai_message`` should represent the lead agent's user-facing # ``last_ai_message`` should represent the lead agent's user-facing
# answer. Middleware/subagent model calls and empty tool-call-only # answer. Middleware/subagent model calls and empty tool-call-only
# AI messages must not overwrite the last useful assistant text. # AI messages must not overwrite the last useful assistant text.
is_ai_message = isinstance(message, AIMessage) or getattr(message, "type", None) == "ai" summary_text = self._message_summary_text(message, caller=caller)
if is_ai_message and (caller is None or caller == "lead_agent"): if summary_text is not None:
text = self._message_text(message).strip() self._last_ai_msg = summary_text
if text:
self._last_ai_msg = text[:2000]
def on_chain_start( def on_chain_start(
self, self,
@ -433,7 +451,16 @@ class RunJournal(BaseCallbackHandler):
tags: list[str] | None = None, tags: list[str] | None = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
if self._closed:
return
messages: list[AnyMessage] = [] messages: list[AnyMessage] = []
response_events: list[dict] = []
should_schedule_progress = False
rid = str(run_id)
callback_caller = self._identify_caller(tags)
is_canonical_callback = rid not in self._counted_message_llm_run_ids
caller = self._llm_response_callers.get(rid, callback_caller)
logger.debug("on_llm_end %s: tags=%s", run_id, tags) logger.debug("on_llm_end %s: tags=%s", run_id, tags)
for generation in response.generations: for generation in response.generations:
for gen in generation: for gen in generation:
@ -443,19 +470,20 @@ class RunJournal(BaseCallbackHandler):
logger.warning(f"on_llm_end {run_id}: generation has no message attribute: {gen}") logger.warning(f"on_llm_end {run_id}: generation has no message attribute: {gen}")
for message in messages: for message in messages:
caller = self._identify_caller(tags) if is_canonical_callback:
self._remember_current_run_tool_calls(message, caller=caller) self._remember_current_run_tool_calls(message, caller=caller)
# Latency # Latency
rid = str(run_id)
start = self._llm_start_times.pop(rid, None) start = self._llm_start_times.pop(rid, None)
latency_ms = int((time.monotonic() - start) * 1000) if start else None latency_ms = int((time.monotonic() - start) * 1000) if start else None
# Token usage from message # Token usage from message
usage = getattr(message, "usage_metadata", None) usage = getattr(message, "usage_metadata", None)
usage_dict = dict(usage) if usage else {} # Providers may mutate and reuse the same response object after the
# callback returns, including nested token-detail mappings.
usage_dict = deepcopy(dict(usage)) if usage else {}
additional_kwargs = getattr(message, "additional_kwargs", None) or {} additional_kwargs = getattr(message, "additional_kwargs", None) or {}
if isinstance(additional_kwargs, dict) and additional_kwargs.get("deerflow_error_fallback"): if is_canonical_callback and isinstance(additional_kwargs, dict) and additional_kwargs.get("deerflow_error_fallback"):
self._had_llm_error_fallback = True self._had_llm_error_fallback = True
detail = additional_kwargs.get("error_detail") detail = additional_kwargs.get("error_detail")
reason = additional_kwargs.get("error_reason") reason = additional_kwargs.get("error_reason")
@ -475,20 +503,19 @@ class RunJournal(BaseCallbackHandler):
call_index = self._llm_call_index call_index = self._llm_call_index
self._seen_llm_starts.add(rid) self._seen_llm_starts.add(rid)
# Message event: checkpoint-aligned llm.ai.response payload. response_events.append(
self._put( self._make_event(
event_type=LLM_AI_RESPONSE_EVENT.event_type, event_type=LLM_AI_RESPONSE_EVENT.event_type,
category=LLM_AI_RESPONSE_EVENT.category, category=LLM_AI_RESPONSE_EVENT.category,
content=message.model_dump(), content=message.model_dump(),
metadata={ metadata={
"caller": caller, "caller": caller,
"usage": usage_dict, "usage": usage_dict,
"latency_ms": latency_ms, "latency_ms": latency_ms,
"llm_call_index": call_index, "llm_call_index": call_index,
}, },
)
) )
if rid not in self._counted_message_llm_run_ids:
self._record_message_summary(message, caller=caller)
# Token accumulation (dedup by langchain run_id to avoid double-counting # Token accumulation (dedup by langchain run_id to avoid double-counting
# when the callback fires more than once for the same response) # when the callback fires more than once for the same response)
@ -519,10 +546,18 @@ class RunJournal(BaseCallbackHandler):
per_call_model = response_metadata.get("model_name") or response_metadata.get("model") per_call_model = response_metadata.get("model_name") or response_metadata.get("model")
self._record_model_usage(per_call_model, input_tk, output_tk, total_tk, self._extract_cache_read(usage_dict)) self._record_model_usage(per_call_model, input_tk, output_tk, total_tk, self._extract_cache_read(usage_dict))
self._schedule_progress_flush() should_schedule_progress = True
if messages: if messages:
self._counted_message_llm_run_ids.add(str(run_id)) self._queue_llm_response_events(
str(run_id),
response_events,
messages,
caller=caller,
)
if should_schedule_progress:
self._schedule_progress_flush()
def on_llm_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None: def on_llm_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None:
self._llm_start_times.pop(str(run_id), None) self._llm_start_times.pop(str(run_id), None)
@ -661,20 +696,123 @@ class RunJournal(BaseCallbackHandler):
if self._should_reconcile_tool_message(message): if self._should_reconcile_tool_message(message):
self._persist_tool_result_message(message) self._persist_tool_result_message(message)
def _make_event(self, *, event_type: str, category: str, content: str | dict = "", metadata: dict | None = None) -> dict:
return {
"thread_id": self.thread_id,
"run_id": self.run_id,
"event_type": event_type,
"category": category,
"content": content,
"metadata": metadata or {},
"created_at": datetime.now(UTC).isoformat(),
}
def _commit_pending_llm_response(self) -> None:
pending = self._pending_llm_response
if pending is None:
return
self._pending_llm_response = None
self._buffer.extend(pending.events)
self._msg_count += pending.message_count
if pending.last_ai_message is not None:
self._last_ai_msg = pending.last_ai_message
def _snapshot_message_summary(self, messages: Sequence[AnyMessage], *, caller: str) -> tuple[int, str | None]:
"""Freeze summary fields before a provider can mutate replayed messages."""
last_ai_message: str | None = None
for message in messages:
summary_text = self._message_summary_text(message, caller=caller)
if summary_text is not None:
last_ai_message = summary_text
return len(messages), last_ai_message
@staticmethod
def _has_positive_usage(events: list[dict]) -> bool:
for event in events:
usage = event["metadata"].get("usage")
if not isinstance(usage, Mapping):
continue
for key in ("input_tokens", "output_tokens", "total_tokens"):
try:
if int(usage.get(key) or 0) > 0:
return True
except (TypeError, ValueError):
continue
return False
@staticmethod
def _merge_response_event_usage(canonical_events: list[dict], replay_events: list[dict]) -> None:
"""Enrich canonical generation events with only replayed usage fields."""
for canonical, replay in zip(canonical_events, replay_events, strict=False):
replay_metadata = replay.get("metadata")
if isinstance(replay_metadata, Mapping):
replay_usage = replay_metadata.get("usage")
if isinstance(replay_usage, Mapping):
canonical["metadata"]["usage"] = deepcopy(dict(replay_usage))
canonical_content = canonical.get("content")
replay_content = replay.get("content")
if isinstance(canonical_content, dict) and isinstance(replay_content, Mapping) and "usage_metadata" in replay_content:
replay_content_usage = replay_content.get("usage_metadata")
canonical_content["usage_metadata"] = deepcopy(dict(replay_content_usage)) if isinstance(replay_content_usage, Mapping) else replay_content_usage
def _flush_if_threshold_reached(self) -> None:
pending_count = len(self._pending_llm_response.events) if self._pending_llm_response is not None else 0
if len(self._buffer) + pending_count >= self._flush_threshold:
self._flush_sync()
def _queue_llm_response_events(
self,
llm_run_id: str,
events: list[dict],
messages: list[AnyMessage],
*,
caller: str,
) -> None:
"""Queue one logical response and merge usage into its canonical callback."""
if self._closed:
return
has_usage = self._has_positive_usage(events)
pending = self._pending_llm_response
if pending is not None and pending.llm_run_id == llm_run_id:
if has_usage:
# The first callback's generation set, immutable summary,
# caller, and non-usage payload are canonical. A provider's
# immediate replay may enrich only corresponding usage fields.
self._merge_response_event_usage(pending.events, events)
self._commit_pending_llm_response()
self._flush_if_threshold_reached()
return
if llm_run_id in self._counted_message_llm_run_ids:
return
# A different event is the ordering boundary for an earlier no-usage
# callback. Commit it before accepting this response.
self._commit_pending_llm_response()
self._flush_if_threshold_reached()
message_count, last_ai_message = self._snapshot_message_summary(messages, caller=caller)
pending_response = _PendingLlmResponse(
llm_run_id=llm_run_id,
events=events,
message_count=message_count,
last_ai_message=last_ai_message,
)
self._counted_message_llm_run_ids.add(llm_run_id)
self._llm_response_callers[llm_run_id] = caller
self._pending_llm_response = pending_response
if has_usage:
self._commit_pending_llm_response()
self._flush_if_threshold_reached()
# Some providers immediately re-fire on_llm_end with usage filled in.
# Defer an incomplete copy until the next event or flush.
def _put(self, *, event_type: str, category: str, content: str | dict = "", metadata: dict | None = None) -> None: def _put(self, *, event_type: str, category: str, content: str | dict = "", metadata: dict | None = None) -> None:
if self._closed: if self._closed:
return return
self._buffer.append( self._commit_pending_llm_response()
{ self._buffer.append(self._make_event(event_type=event_type, category=category, content=content, metadata=metadata))
"thread_id": self.thread_id,
"run_id": self.run_id,
"event_type": event_type,
"category": category,
"content": content,
"metadata": metadata or {},
"created_at": datetime.now(UTC).isoformat(),
}
)
if len(self._buffer) >= self._flush_threshold: if len(self._buffer) >= self._flush_threshold:
self._flush_sync() self._flush_sync()
@ -686,6 +824,7 @@ class RunJournal(BaseCallbackHandler):
stay in the buffer and are flushed later by the async ``flush()`` stay in the buffer and are flushed later by the async ``flush()``
call in the worker's ``finally`` block. call in the worker's ``finally`` block.
""" """
self._commit_pending_llm_response()
if not self._buffer: if not self._buffer:
return return
# Skip if a flush is already in flight — avoids concurrent writes # Skip if a flush is already in flight — avoids concurrent writes
@ -929,6 +1068,7 @@ class RunJournal(BaseCallbackHandler):
"""Force flush remaining buffer. Called in worker's finally block.""" """Force flush remaining buffer. Called in worker's finally block."""
if self._closed: if self._closed:
return return
self._commit_pending_llm_response()
if self._pending_flush_tasks: if self._pending_flush_tasks:
await asyncio.gather(*tuple(self._pending_flush_tasks), return_exceptions=True) await asyncio.gather(*tuple(self._pending_flush_tasks), return_exceptions=True)
while self._pending_progress_task is not None: while self._pending_progress_task is not None:
@ -968,6 +1108,7 @@ class RunJournal(BaseCallbackHandler):
self._store = None self._store = None
self._progress_reporter = None self._progress_reporter = None
self._buffer.clear() self._buffer.clear()
self._pending_llm_response = None
self._pending_flush_tasks.clear() self._pending_flush_tasks.clear()
self._pending_progress_task = None self._pending_progress_task = None
self._pending_progress_delayed = False self._pending_progress_delayed = False
@ -976,6 +1117,7 @@ class RunJournal(BaseCallbackHandler):
self._counted_llm_run_ids.clear() self._counted_llm_run_ids.clear()
self._counted_external_source_ids.clear() self._counted_external_source_ids.clear()
self._counted_message_llm_run_ids.clear() self._counted_message_llm_run_ids.clear()
self._llm_response_callers.clear()
self._llm_start_times.clear() self._llm_start_times.clear()
self._seen_llm_starts.clear() self._seen_llm_starts.clear()
self._current_run_tool_call_names.clear() self._current_run_tool_call_names.clear()

View File

@ -11,7 +11,8 @@ from unittest.mock import MagicMock
from uuid import uuid4 from uuid import uuid4
import pytest import pytest
from langchain_core.messages import HumanMessage from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, LLMResult
from deerflow.runtime.events.store.memory import MemoryRunEventStore from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.journal import RunJournal from deerflow.runtime.journal import RunJournal
@ -68,6 +69,24 @@ async def test_close_flushes_and_detaches_runtime_dependencies():
assert reporter_ref() is None assert reporter_ref() is None
@pytest.mark.anyio
async def test_closed_on_llm_end_returns_before_touching_response_or_state():
store = MemoryRunEventStore()
journal = RunJournal("r-closed-callback", "t-closed-callback", store)
await journal.close()
completion_before = journal.get_completion_data()
# A plain object has no generations attribute, so this also pins the
# early return ahead of response inspection.
journal.on_llm_end(object(), run_id=uuid4(), tags=["lead_agent"])
assert journal.get_completion_data() == completion_before
assert journal._pending_llm_response is None
assert journal._buffer == []
assert journal._counted_message_llm_run_ids == set()
assert journal._counted_llm_run_ids == set()
@pytest.mark.anyio @pytest.mark.anyio
async def test_close_preserves_buffer_and_dependencies_when_flush_fails(): async def test_close_preserves_buffer_and_dependencies_when_flush_fails():
class FailOnceRunEventStore(MemoryRunEventStore): class FailOnceRunEventStore(MemoryRunEventStore):
@ -101,6 +120,73 @@ async def test_close_preserves_buffer_and_dependencies_when_flush_fails():
assert [event["event_type"] for event in events] == ["middleware:test"] assert [event["event_type"] for event in events] == ["middleware:test"]
@pytest.mark.anyio
async def test_close_retries_pending_no_usage_response_without_duplication():
class FailOnceRunEventStore(MemoryRunEventStore):
def __init__(self) -> None:
super().__init__()
self.put_batch_calls = 0
async def put_batch(self, events):
self.put_batch_calls += 1
if self.put_batch_calls == 1:
raise RuntimeError("transient store failure")
return await super().put_batch(events)
async def progress_reporter(snapshot):
del snapshot
store = FailOnceRunEventStore()
journal = RunJournal(
"r-close-pending-retry",
"t-close-pending-retry",
store,
flush_threshold=100,
progress_reporter=progress_reporter,
)
journal.record_middleware("before", name="test", hook="after", action="record", changes={})
journal.on_llm_end(
_make_llm_response("Canonical without usage"),
run_id=uuid4(),
parent_run_id=None,
tags=["lead_agent"],
)
assert journal._pending_llm_response is not None
assert journal.get_completion_data()["message_count"] == 0
with pytest.raises(RuntimeError, match="transient store failure"):
await journal.close()
assert journal._closed is False
assert journal._store is store
assert journal._progress_reporter is progress_reporter
assert journal._pending_llm_response is None
assert [event["event_type"] for event in journal._buffer] == [
"middleware:before",
"llm.ai.response",
]
assert journal.get_completion_data()["message_count"] == 1
assert journal.get_completion_data()["last_ai_message"] == "Canonical without usage"
await journal.close()
events = await store.list_events("t-close-pending-retry", "r-close-pending-retry")
assert [event["event_type"] for event in events] == [
"middleware:before",
"llm.ai.response",
]
responses = [event for event in events if event["event_type"] == "llm.ai.response"]
assert len(responses) == 1
assert responses[0]["content"]["content"] == "Canonical without usage"
assert responses[0]["content"]["usage_metadata"] is None
assert responses[0]["metadata"]["usage"] == {}
assert journal.get_completion_data()["message_count"] == 1
assert journal._closed is True
assert journal._store is None
assert journal._progress_reporter is None
@pytest.mark.anyio @pytest.mark.anyio
async def test_close_without_flush_discards_buffer_and_detaches_runtime_dependencies(): async def test_close_without_flush_discards_buffer_and_detaches_runtime_dependencies():
class TrackingRunEventStore(MemoryRunEventStore): class TrackingRunEventStore(MemoryRunEventStore):
@ -197,6 +283,12 @@ def _make_llm_response(content="Hello", usage=None, tool_calls=None, additional_
return response return response
def _combine_llm_responses(*responses):
response = MagicMock()
response.generations = [generation for item in responses for generation in item.generations]
return response
class TestLlmCallbacks: class TestLlmCallbacks:
@pytest.mark.anyio @pytest.mark.anyio
async def test_on_chat_model_start_persists_original_user_input_without_mutating_model_message(self, journal_setup): async def test_on_chat_model_start_persists_original_user_input_without_mutating_model_message(self, journal_setup):
@ -568,14 +660,28 @@ class TestBufferFlush:
j, store = journal_setup j, store = journal_setup
j._flush_threshold = 2 j._flush_threshold = 2
# Each on_llm_end emits 1 event # Each on_llm_end emits 1 event
j.on_llm_end(_make_llm_response("A"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"]) usage = {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
j.on_llm_end(_make_llm_response("A", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
assert len(j._buffer) == 1 assert len(j._buffer) == 1
j.on_llm_end(_make_llm_response("B"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"]) j.on_llm_end(_make_llm_response("B", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
# At threshold the buffer should have been flushed asynchronously # At threshold the buffer should have been flushed asynchronously
await asyncio.sleep(0.1) await asyncio.sleep(0.1)
events = await store.list_events("t1", "r1") events = await store.list_events("t1", "r1")
assert len(events) >= 2 assert len(events) >= 2
@pytest.mark.anyio
async def test_pending_response_counts_toward_flush_threshold(self, journal_setup):
j, store = journal_setup
j._flush_threshold = 2
j.record_middleware("before", name="BeforeMiddleware", hook="after_model", action="record", changes={})
j.on_llm_end(_make_llm_response("Pending"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
await asyncio.sleep(0.1)
assert j._pending_llm_response is None
events = await store.list_events("t1", "r1")
assert [event["event_type"] for event in events] == ["middleware:before", "llm.ai.response"]
@pytest.mark.anyio @pytest.mark.anyio
async def test_events_retained_when_no_loop(self, journal_setup): async def test_events_retained_when_no_loop(self, journal_setup):
"""Events buffered in a sync (no-loop) context should survive """Events buffered in a sync (no-loop) context should survive
@ -611,11 +717,12 @@ class TestFeedGeneration:
""" """
@pytest.mark.anyio @pytest.mark.anyio
async def test_buffering_alone_does_not_advance_it(self, journal_setup): async def test_pending_response_alone_does_not_advance_it(self, journal_setup):
j, _store = journal_setup j, _store = journal_setup
j.on_llm_end(_make_llm_response("A"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"]) j.on_llm_end(_make_llm_response("A"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
assert len(j._buffer) == 1 assert j._buffer == []
assert j._pending_llm_response is not None
assert j.feed_generation == 0 assert j.feed_generation == 0
@pytest.mark.anyio @pytest.mark.anyio
@ -623,7 +730,8 @@ class TestFeedGeneration:
j, _store = journal_setup j, _store = journal_setup
j._flush_threshold = 1 j._flush_threshold = 1
j.on_llm_end(_make_llm_response("A"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"]) usage = {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
j.on_llm_end(_make_llm_response("A", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
await asyncio.sleep(0.1) await asyncio.sleep(0.1)
assert j.feed_generation == 1 assert j.feed_generation == 1
@ -734,6 +842,7 @@ class TestConvenienceFields:
parent_run_id=None, parent_run_id=None,
tags=["lead_agent"], tags=["lead_agent"],
) )
await j.flush()
data = j.get_completion_data() data = j.get_completion_data()
@ -756,6 +865,7 @@ class TestConvenienceFields:
parent_run_id=None, parent_run_id=None,
tags=["lead_agent"], tags=["lead_agent"],
) )
await j.flush()
data = j.get_completion_data() data = j.get_completion_data()
@ -766,6 +876,7 @@ class TestConvenienceFields:
async def test_last_ai_message_extracts_mapping_content(self, journal_setup): async def test_last_ai_message_extracts_mapping_content(self, journal_setup):
j, _ = journal_setup j, _ = journal_setup
j.on_llm_end(_make_llm_response({"content": "Nested answer"}), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"]) j.on_llm_end(_make_llm_response({"content": "Nested answer"}), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
await j.flush()
data = j.get_completion_data() data = j.get_completion_data()
@ -796,6 +907,7 @@ class TestConvenienceFields:
j, _ = journal_setup j, _ = journal_setup
j.on_llm_end(_make_llm_response("Lead answer"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"]) j.on_llm_end(_make_llm_response("Lead answer"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
j.on_llm_end(_make_llm_response("Subagent detail"), run_id=uuid4(), parent_run_id=None, tags=["subagent:research"]) j.on_llm_end(_make_llm_response("Subagent detail"), run_id=uuid4(), parent_run_id=None, tags=["subagent:research"])
await j.flush()
data = j.get_completion_data() data = j.get_completion_data()
@ -950,17 +1062,317 @@ class TestCallerBucketing:
assert j._lead_agent_tokens == 15 assert j._lead_agent_tokens == 15
assert j._llm_call_count == 1 assert j._llm_call_count == 1
def test_first_no_usage_second_with_usage(self, journal_setup): @pytest.mark.anyio
"""First callback with no usage must not block second callback with usage for same run_id.""" async def test_dedup_same_run_id_persists_single_message(self, journal_setup):
j, _ = journal_setup """A re-fired on_llm_end for one run_id must persist the message once.
LangChain can deliver on_llm_end more than once for the same run_id.
Token accounting already dedups on that; the durable llm.ai.response
row must be deduped on the same premise, or count_messages and message
pagination (which read append-only rows without dedup) inflate.
"""
j, store = journal_setup
run_id = uuid4()
response = _make_llm_response("Answer")
j.on_llm_end(response, run_id=run_id, parent_run_id=None, tags=["lead_agent"])
j.on_llm_end(response, run_id=run_id, parent_run_id=None, tags=["lead_agent"])
await j.flush()
messages = await store.list_messages("t1")
assert [m["event_type"] for m in messages] == ["llm.ai.response"]
assert await store.count_messages("t1") == 1
# The run summary counts the message exactly once as well.
assert j._msg_count == 1
@pytest.mark.anyio
async def test_adjacent_late_usage_enriches_canonical_response_only(self, journal_setup):
j, store = journal_setup
run_id = uuid4()
usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
original_tool_calls = [{"id": "call-original", "name": "search", "args": {}}]
replay_tool_calls = [{"id": "call-replay", "name": "write_file", "args": {}}]
j.on_llm_end(
_make_llm_response(
"Canonical",
tool_calls=original_tool_calls,
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "canonical fallback",
},
),
run_id=run_id,
parent_run_id=None,
tags=["lead_agent"],
)
j.on_llm_end(
_make_llm_response(
"Replay",
usage=usage,
tool_calls=replay_tool_calls,
additional_kwargs={
"deerflow_error_fallback": True,
"error_detail": "replay fallback",
},
),
run_id=run_id,
parent_run_id=None,
tags=["subagent:research"],
)
await j.flush()
messages = await store.list_messages("t1")
assert len(messages) == 1
assert messages[0]["content"]["content"] == "Canonical"
assert messages[0]["content"]["tool_calls"] == original_tool_calls
assert messages[0]["content"]["additional_kwargs"]["error_detail"] == "canonical fallback"
assert messages[0]["content"]["usage_metadata"] == usage
assert messages[0]["metadata"]["caller"] == "lead_agent"
assert messages[0]["metadata"]["usage"] == usage
assert j._current_run_tool_call_names == {"call-original": "search"}
assert j.had_llm_error_fallback is True
assert j.llm_error_fallback_message == "canonical fallback"
assert j.get_completion_data()["last_ai_message"] == "Canonical"
assert j.get_completion_data()["lead_agent_tokens"] == 15
assert j.get_completion_data()["subagent_tokens"] == 0
@pytest.mark.anyio
async def test_same_message_object_replay_cannot_mutate_canonical_summary(self, journal_setup):
j, store = journal_setup
run_id = uuid4()
message = AIMessage(content="Canonical answer")
response = LLMResult(generations=[[ChatGeneration(message=message)]])
j.on_llm_end(response, run_id=run_id, parent_run_id=None, tags=["lead_agent"])
message.content = "Replay answer"
message.usage_metadata = {
"input_tokens": 10,
"output_tokens": 5,
"total_tokens": 15,
"input_token_details": {"cache_read": 3},
}
j.on_llm_end(response, run_id=run_id, parent_run_id=None, tags=["lead_agent"])
message.usage_metadata["input_token_details"]["cache_read"] = 999
await j.flush()
messages = await store.list_messages("t1")
assert len(messages) == 1
assert messages[0]["content"]["content"] == "Canonical answer"
expected_usage = {
"input_tokens": 10,
"output_tokens": 5,
"total_tokens": 15,
"input_token_details": {"cache_read": 3},
}
assert messages[0]["metadata"]["usage"] == expected_usage
assert messages[0]["content"]["usage_metadata"] == expected_usage
assert j.get_completion_data()["message_count"] == 1
assert j.get_completion_data()["last_ai_message"] == "Canonical answer"
@pytest.mark.anyio
async def test_positive_usage_event_does_not_retain_nested_provider_metadata(self, journal_setup):
j, store = journal_setup
usage = {
"input_tokens": 8,
"output_tokens": 3,
"total_tokens": 11,
"output_token_details": {"reasoning": 2},
}
message = AIMessage(content="Canonical with usage", usage_metadata=usage)
response = LLMResult(generations=[[ChatGeneration(message=message)]])
j.on_llm_end(response, run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
message.usage_metadata["output_token_details"]["reasoning"] = 999
await j.flush()
messages = await store.list_messages("t1")
assert len(messages) == 1
assert messages[0]["metadata"]["usage"]["output_token_details"] == {"reasoning": 2}
assert messages[0]["content"]["usage_metadata"]["output_token_details"] == {"reasoning": 2}
@pytest.mark.anyio
async def test_mutating_staged_message_before_flush_cannot_mutate_canonical_summary(self, journal_setup):
j, store = journal_setup
message = AIMessage(content="Canonical before flush")
response = LLMResult(generations=[[ChatGeneration(message=message)]])
j.on_llm_end(response, run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
message.content = "Mutation before flush"
await j.flush()
messages = await store.list_messages("t1")
assert len(messages) == 1
assert messages[0]["content"]["content"] == "Canonical before flush"
assert j.get_completion_data()["message_count"] == 1
assert j.get_completion_data()["last_ai_message"] == "Canonical before flush"
@pytest.mark.anyio
async def test_nested_same_message_object_replay_cannot_mutate_canonical_summary(self, journal_setup):
j, store = journal_setup
run_id = uuid4()
message = AIMessage(content=[{"type": "text", "text": "Canonical nested answer"}])
response = LLMResult(generations=[[ChatGeneration(message=message)]])
j.on_llm_end(response, run_id=run_id, parent_run_id=None, tags=["lead_agent"])
message.content[0]["text"] = "Replay nested answer"
message.usage_metadata = {"input_tokens": 8, "output_tokens": 3, "total_tokens": 11}
j.on_llm_end(response, run_id=run_id, parent_run_id=None, tags=["lead_agent"])
await j.flush()
messages = await store.list_messages("t1")
assert len(messages) == 1
assert messages[0]["content"]["content"] == [{"type": "text", "text": "Canonical nested answer"}]
assert messages[0]["content"]["usage_metadata"] == message.usage_metadata
assert j.get_completion_data()["message_count"] == 1
assert j.get_completion_data()["last_ai_message"] == "Canonical nested answer"
@pytest.mark.anyio
async def test_all_zero_usage_remains_pending_and_positive_usage_enriches_it(self, journal_setup):
j, store = journal_setup
run_id = uuid4()
zero_usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
positive_usage = {"input_tokens": 4, "output_tokens": 2, "total_tokens": 6}
j.on_llm_end(_make_llm_response("Zero usage", usage=zero_usage), run_id=run_id, parent_run_id=None, tags=["lead_agent"])
assert j._buffer == []
assert j._pending_llm_response is not None
assert j.get_completion_data()["message_count"] == 0
j.on_llm_end(_make_llm_response("Replay payload", usage=positive_usage), run_id=run_id, parent_run_id=None, tags=["lead_agent"])
await j.flush()
messages = await store.list_messages("t1")
assert len(messages) == 1
assert messages[0]["content"]["content"] == "Zero usage"
assert messages[0]["content"]["usage_metadata"] == positive_usage
assert messages[0]["metadata"]["usage"] == positive_usage
assert j.get_completion_data()["message_count"] == 1
assert j.get_completion_data()["last_ai_message"] == "Zero usage"
@pytest.mark.anyio
async def test_replay_generation_length_cannot_change_canonical_set(self, journal_setup):
j, store = journal_setup
short_usage = {"input_tokens": 8, "output_tokens": 3, "total_tokens": 11}
extra_usage = {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
first_run_id = uuid4()
second_run_id = uuid4()
j.on_llm_end(
_combine_llm_responses(_make_llm_response("Canonical one"), _make_llm_response("Canonical two")),
run_id=first_run_id,
parent_run_id=None,
tags=["lead_agent"],
)
j.on_llm_end(
_make_llm_response("Short replay", usage=short_usage),
run_id=first_run_id,
parent_run_id=None,
tags=["lead_agent"],
)
j.on_llm_end(
_make_llm_response("Single canonical"),
run_id=second_run_id,
parent_run_id=None,
tags=["lead_agent"],
)
j.on_llm_end(
_combine_llm_responses(
_make_llm_response("Long replay one", usage=extra_usage),
_make_llm_response("Long replay two"),
),
run_id=second_run_id,
parent_run_id=None,
tags=["lead_agent"],
)
await j.flush()
messages = await store.list_messages("t1")
assert [message["content"]["content"] for message in messages] == [
"Canonical one",
"Canonical two",
"Single canonical",
]
assert messages[0]["metadata"]["usage"] == short_usage
assert messages[0]["content"]["usage_metadata"] == short_usage
assert messages[1]["metadata"]["usage"] == {}
assert messages[1]["content"]["usage_metadata"] is None
assert messages[2]["metadata"]["usage"] == extra_usage
assert messages[2]["content"]["usage_metadata"] == extra_usage
assert j.get_completion_data()["message_count"] == 3
assert j.get_completion_data()["last_ai_message"] == "Single canonical"
@pytest.mark.anyio
async def test_interleaved_late_usage_updates_summary_only(self, journal_setup):
j, store = journal_setup
first_run_id = uuid4()
second_run_id = uuid4()
usage = {"input_tokens": 9, "output_tokens": 4, "total_tokens": 13}
j.on_llm_end(_make_llm_response("First canonical"), run_id=first_run_id, parent_run_id=None, tags=["lead_agent"])
j.on_llm_end(_make_llm_response("Second canonical"), run_id=second_run_id, parent_run_id=None, tags=["lead_agent"])
j.on_llm_end(
_make_llm_response(
"Late replay",
usage=usage,
tool_calls=[{"id": "late-call", "name": "write_file", "args": {}}],
additional_kwargs={"deerflow_error_fallback": True, "error_detail": "late fallback"},
),
run_id=first_run_id,
parent_run_id=None,
tags=["subagent:research"],
)
await j.flush()
messages = await store.list_messages("t1")
assert [message["content"]["content"] for message in messages] == ["First canonical", "Second canonical"]
assert messages[0]["metadata"]["usage"] == {}
assert messages[0]["content"]["usage_metadata"] is None
assert j.get_completion_data()["total_tokens"] == 13
assert j.get_completion_data()["lead_agent_tokens"] == 13
assert j.get_completion_data()["subagent_tokens"] == 0
assert j.get_completion_data()["message_count"] == 2
assert j.get_completion_data()["last_ai_message"] == "Second canonical"
assert "late-call" not in j._current_run_tool_call_names
assert j.had_llm_error_fallback is False
@pytest.mark.anyio
async def test_single_no_usage_response_persists_once_at_flush(self, journal_setup):
j, store = journal_setup
j.on_llm_end(_make_llm_response("No usage"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
assert j._buffer == []
assert j._pending_llm_response is not None
await j.flush()
messages = await store.list_messages("t1")
assert len(messages) == 1
assert messages[0]["content"]["content"] == "No usage"
assert messages[0]["metadata"]["usage"] == {}
@pytest.mark.anyio
async def test_distinct_run_ids_each_persist_a_message(self, journal_setup):
"""The dedup guard is per run_id and must not drop distinct responses."""
j, store = journal_setup
j.on_llm_end(_make_llm_response("First"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
j.on_llm_end(_make_llm_response("Second"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
await j.flush()
assert await store.count_messages("t1") == 2
@pytest.mark.anyio
async def test_first_no_usage_second_with_usage(self, journal_setup):
"""Late usage enriches the single canonical event and the run summary."""
j, store = journal_setup
run_id = uuid4() run_id = uuid4()
j.on_llm_end(_make_llm_response("A", usage=None), run_id=run_id, parent_run_id=None, tags=["lead_agent"]) j.on_llm_end(_make_llm_response("A", usage=None), run_id=run_id, parent_run_id=None, tags=["lead_agent"])
assert str(run_id) not in j._counted_llm_run_ids
# Second callback for the same run_id with actual usage must still count
usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15} usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
j.on_llm_end(_make_llm_response("A", usage=usage), run_id=run_id, parent_run_id=None, tags=["lead_agent"]) j.on_llm_end(_make_llm_response("A", usage=usage), run_id=run_id, parent_run_id=None, tags=["lead_agent"])
assert j._total_tokens == 15 await j.flush()
assert j._lead_agent_tokens == 15
messages = await store.list_messages("t1")
assert len(messages) == 1
assert messages[0]["metadata"]["usage"] == usage
assert messages[0]["content"]["usage_metadata"] == usage
assert j.get_completion_data()["total_tokens"] == 15
def test_track_token_usage_false_skips_buckets(self): def test_track_token_usage_false_skips_buckets(self):
"""When token tracking is disabled, caller buckets stay at 0.""" """When token tracking is disabled, caller buckets stay at 0."""