"""Tests for RunJournal callback handler.

Uses MemoryRunEventStore as the backend for direct event inspection.
"""

import asyncio
from unittest.mock import MagicMock
from uuid import uuid4

import pytest
from langchain_core.messages import HumanMessage

from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.journal import RunJournal
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY


@pytest.fixture
def journal_setup():
    store = MemoryRunEventStore()
    j = RunJournal("r1", "t1", store, flush_threshold=100)
    return j, store


def _make_llm_response(content="Hello", usage=None, tool_calls=None, additional_kwargs=None):
    """Create a mock LLM response with a message.

    model_dump() returns checkpoint-aligned format matching real AIMessage.
    """
    msg = MagicMock()
    msg.type = "ai"
    msg.content = content
    msg.id = f"msg-{id(msg)}"
    msg.tool_calls = tool_calls or []
    msg.invalid_tool_calls = []
    msg.response_metadata = {"model_name": "test-model"}
    msg.usage_metadata = usage
    msg.additional_kwargs = additional_kwargs or {}
    msg.name = None
    # model_dump returns checkpoint-aligned format
    msg.model_dump.return_value = {
        "content": content,
        "additional_kwargs": additional_kwargs or {},
        "response_metadata": {"model_name": "test-model"},
        "type": "ai",
        "name": None,
        "id": msg.id,
        "tool_calls": tool_calls or [],
        "invalid_tool_calls": [],
        "usage_metadata": usage,
    }

    gen = MagicMock()
    gen.message = msg

    response = MagicMock()
    response.generations = [[gen]]
    return response


class TestLlmCallbacks:
    @pytest.mark.anyio
    async def test_on_chat_model_start_persists_original_user_input_without_mutating_model_message(self, journal_setup):
        j, store = journal_setup
        wrapped_content = "--- BEGIN USER INPUT ---\nShow revenue\n--- END USER INPUT ---"
        model_message = HumanMessage(
            content=wrapped_content,
            id="human-1",
            additional_kwargs={ORIGINAL_USER_CONTENT_KEY: "Show revenue", "channel": "web"},
        )

        j.on_chat_model_start({}, [[model_message]], run_id=uuid4(), tags=["lead_agent"])
        await j.flush()

        assert j._first_human_msg == "Show revenue"
        events = await store.list_events("t1", "r1")
        human_event = next(event for event in events if event["event_type"] == "llm.human.input")
        assert human_event["content"]["content"] == "Show revenue"
        assert human_event["content"]["id"] == "human-1"
        assert human_event["content"]["additional_kwargs"] == {"channel": "web"}
        assert model_message.content == wrapped_content
        assert model_message.additional_kwargs[ORIGINAL_USER_CONTENT_KEY] == "Show revenue"

    @pytest.mark.anyio
    async def test_on_llm_end_produces_trace_event(self, journal_setup):
        j, store = journal_setup
        run_id = uuid4()
        j.on_llm_start({}, [], run_id=run_id, tags=["lead_agent"])
        j.on_llm_end(_make_llm_response("Hi"), run_id=run_id, parent_run_id=None, tags=["lead_agent"])
        await j.flush()
        events = await store.list_events("t1", "r1")
        trace_events = [e for e in events if e["event_type"] == "llm.ai.response"]
        assert len(trace_events) == 1
        assert trace_events[0]["category"] == "message"

    @pytest.mark.anyio
    async def test_on_llm_end_lead_agent_produces_ai_message(self, journal_setup):
        j, store = journal_setup
        run_id = uuid4()
        j.on_llm_start({}, [], run_id=run_id, tags=["lead_agent"])
        j.on_llm_end(_make_llm_response("Answer"), 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]["event_type"] == "llm.ai.response"
        # Content is checkpoint-aligned model_dump format
        assert messages[0]["content"]["type"] == "ai"
        assert messages[0]["content"]["content"] == "Answer"

    @pytest.mark.anyio
    async def test_on_llm_end_with_tool_calls_produces_ai_tool_call(self, journal_setup):
        """LLM response with pending tool_calls emits llm.ai.response with tool_calls in content."""
        j, store = journal_setup
        run_id = uuid4()
        j.on_llm_end(
            _make_llm_response("Let me search", tool_calls=[{"id": "call_1", "name": "search", "args": {}}]),
            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]["event_type"] == "llm.ai.response"
        assert len(messages[0]["content"]["tool_calls"]) == 1

    @pytest.mark.anyio
    async def test_on_llm_end_subagent_no_ai_message(self, journal_setup):
        j, store = journal_setup
        run_id = uuid4()
        j.on_llm_start({}, [], run_id=run_id, tags=["subagent:research"])
        j.on_llm_end(_make_llm_response("Sub answer"), run_id=run_id, parent_run_id=None, tags=["subagent:research"])
        await j.flush()
        messages = await store.list_messages("t1")
        # subagent responses still emit llm.ai.response with category="message"
        assert len(messages) == 1

    @pytest.mark.anyio
    async def test_token_accumulation(self, journal_setup):
        j, store = journal_setup
        usage1 = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
        usage2 = {"input_tokens": 20, "output_tokens": 10, "total_tokens": 30}
        j.on_llm_end(_make_llm_response("A", usage=usage1), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
        j.on_llm_end(_make_llm_response("B", usage=usage2), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
        assert j._total_input_tokens == 30
        assert j._total_output_tokens == 15
        assert j._total_tokens == 45
        assert j._llm_call_count == 2

    @pytest.mark.anyio
    async def test_total_tokens_computed_from_input_output(self, journal_setup):
        """If total_tokens is 0, it should be computed from input + output."""
        j, store = journal_setup
        j.on_llm_end(
            _make_llm_response("Hi", usage={"input_tokens": 100, "output_tokens": 50, "total_tokens": 0}),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        assert j._total_tokens == 150

    @pytest.mark.anyio
    async def test_caller_token_classification(self, journal_setup):
        j, store = journal_setup
        usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
        j.on_llm_end(_make_llm_response("A", usage=usage), 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=["subagent:research"])
        j.on_llm_end(_make_llm_response("C", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["middleware:summarization"])
        # token tracking not broken by caller type
        assert j._total_tokens == 45
        assert j._llm_call_count == 3

    @pytest.mark.anyio
    async def test_usage_metadata_none_no_crash(self, journal_setup):
        j, store = journal_setup
        j.on_llm_end(_make_llm_response("No usage", usage=None), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
        await j.flush()

    @pytest.mark.anyio
    async def test_latency_tracking(self, journal_setup):
        j, store = journal_setup
        run_id = uuid4()
        j.on_llm_start({}, [], run_id=run_id, tags=["lead_agent"])
        j.on_llm_end(_make_llm_response("Fast"), run_id=run_id, parent_run_id=None, tags=["lead_agent"])
        await j.flush()
        events = await store.list_events("t1", "r1")
        llm_resp = [e for e in events if e["event_type"] == "llm.ai.response"][0]
        assert "latency_ms" in llm_resp["metadata"]
        assert llm_resp["metadata"]["latency_ms"] is not None


class TestLifecycleCallbacks:
    @pytest.mark.anyio
    async def test_chain_start_end_produce_trace_events(self, journal_setup):
        j, store = journal_setup
        j.on_chain_start({}, {}, run_id=uuid4(), parent_run_id=None)
        j.on_chain_end({}, run_id=uuid4())
        await asyncio.sleep(0.05)
        await j.flush()
        events = await store.list_events("t1", "r1")
        types = {e["event_type"] for e in events}
        assert "run.start" in types
        assert "run.end" in types

    @pytest.mark.anyio
    async def test_nested_chain_no_run_lifecycle_events(self, journal_setup):
        """Nested chains (parent_run_id set) should NOT produce root run lifecycle events."""
        j, store = journal_setup
        parent_id = uuid4()
        j.on_chain_start({}, {}, run_id=uuid4(), parent_run_id=parent_id)
        j.on_chain_end({}, run_id=uuid4(), parent_run_id=parent_id)
        await j.flush()
        events = await store.list_events("t1", "r1")
        assert not any(e["event_type"] == "run.start" for e in events)
        assert not any(e["event_type"] == "run.end" for e in events)


class TestToolCallbacks:
    @pytest.mark.anyio
    async def test_tool_end_with_tool_message(self, journal_setup):
        """on_tool_end with a ToolMessage stores it as llm.tool.result."""
        from langchain_core.messages import ToolMessage

        j, store = journal_setup
        tool_msg = ToolMessage(content="results", tool_call_id="call_1", name="web_search")
        j.on_tool_end(tool_msg, run_id=uuid4())
        await j.flush()
        messages = await store.list_messages("t1")
        assert len(messages) == 1
        assert messages[0]["event_type"] == "llm.tool.result"
        assert messages[0]["content"]["type"] == "tool"

    @pytest.mark.anyio
    async def test_tool_end_with_command_unwraps_tool_message(self, journal_setup):
        """on_tool_end with Command(update={'messages':[ToolMessage]}) unwraps inner message."""
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        j, store = journal_setup
        inner = ToolMessage(content="file list", tool_call_id="call_2", name="present_files")
        cmd = Command(update={"messages": [inner]})
        j.on_tool_end(cmd, run_id=uuid4())
        await j.flush()
        messages = await store.list_messages("t1")
        assert len(messages) == 1
        assert messages[0]["event_type"] == "llm.tool.result"
        assert messages[0]["content"]["content"] == "file list"

    @pytest.mark.anyio
    async def test_on_tool_error_no_crash(self, journal_setup):
        """on_tool_error should not crash (no event emitted by default)."""
        j, store = journal_setup
        j.on_tool_error(TimeoutError("timeout"), run_id=uuid4(), name="web_fetch")
        await j.flush()
        # Base implementation does not emit tool_error — just verify no crash
        events = await store.list_events("t1", "r1")
        assert isinstance(events, list)


class TestFinalToolMessageReconciliation:
    @pytest.mark.anyio
    async def test_root_chain_end_reconciles_missing_ask_clarification_tool_message(self, journal_setup):
        from langchain_core.messages import ToolMessage

        j, store = journal_setup
        j.on_llm_end(
            _make_llm_response("", tool_calls=[{"id": "call_clarify", "name": "ask_clarification", "args": {"question": "Which format?"}}]),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        tool_msg = ToolMessage(
            content="Which format?",
            tool_call_id="call_clarify",
            name="ask_clarification",
            artifact={"human_input": {"kind": "human_input_request", "request_id": "clarification:call_clarify"}},
        )

        j.on_chain_end({"messages": [tool_msg]}, run_id=uuid4())
        await j.flush()

        messages = await store.list_messages("t1")
        tool_results = [m for m in messages if m["event_type"] == "llm.tool.result"]
        assert len(tool_results) == 1
        assert tool_results[0]["content"]["name"] == "ask_clarification"
        assert tool_results[0]["content"]["artifact"]["human_input"]["request_id"] == "clarification:call_clarify"

    @pytest.mark.anyio
    async def test_root_chain_end_does_not_duplicate_tool_message_captured_by_on_tool_end(self, journal_setup):
        from langchain_core.messages import ToolMessage

        j, store = journal_setup
        j.on_llm_end(
            _make_llm_response("", tool_calls=[{"id": "call_clarify", "name": "ask_clarification", "args": {"question": "Which format?"}}]),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        tool_msg = ToolMessage(content="Which format?", tool_call_id="call_clarify", name="ask_clarification")

        j.on_tool_end(tool_msg, run_id=uuid4())
        j.on_chain_end({"messages": [tool_msg]}, run_id=uuid4())
        await j.flush()

        messages = await store.list_messages("t1")
        tool_results = [m for m in messages if m["event_type"] == "llm.tool.result"]
        assert len(tool_results) == 1

    @pytest.mark.anyio
    async def test_root_chain_end_ignores_retained_old_tool_message_from_previous_run(self, journal_setup):
        from langchain_core.messages import ToolMessage

        j, store = journal_setup
        j.on_llm_end(
            _make_llm_response("", tool_calls=[{"id": "call_current", "name": "ask_clarification", "args": {"question": "Current?"}}]),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        retained_old_tool_msg = ToolMessage(content="Old question", tool_call_id="call_old", name="ask_clarification")

        j.on_chain_end({"messages": [retained_old_tool_msg]}, run_id=uuid4())
        await j.flush()

        messages = await store.list_messages("t1")
        assert not any(m["event_type"] == "llm.tool.result" for m in messages)

    @pytest.mark.anyio
    async def test_root_chain_end_ignores_non_allowlisted_tool_message(self, journal_setup):
        from langchain_core.messages import ToolMessage

        j, store = journal_setup
        j.on_llm_end(
            _make_llm_response("", tool_calls=[{"id": "call_search", "name": "web_search", "args": {"query": "deerflow"}}]),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        tool_msg = ToolMessage(content="Search result", tool_call_id="call_search", name="web_search")

        j.on_chain_end({"messages": [tool_msg]}, run_id=uuid4())
        await j.flush()

        messages = await store.list_messages("t1")
        assert not any(m["event_type"] == "llm.tool.result" for m in messages)

    @pytest.mark.anyio
    async def test_root_chain_end_ignores_hidden_ask_clarification_tool_message(self, journal_setup):
        from langchain_core.messages import ToolMessage

        j, store = journal_setup
        j.on_llm_end(
            _make_llm_response("", tool_calls=[{"id": "call_clarify", "name": "ask_clarification", "args": {"question": "Hidden?"}}]),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        tool_msg = ToolMessage(
            content="Hidden?",
            tool_call_id="call_clarify",
            name="ask_clarification",
            additional_kwargs={"hide_from_ui": True},
        )

        j.on_chain_end({"messages": [tool_msg]}, run_id=uuid4())
        await j.flush()

        messages = await store.list_messages("t1")
        assert not any(m["event_type"] == "llm.tool.result" for m in messages)


class TestCustomEvents:
    @pytest.mark.anyio
    async def test_on_custom_event_not_implemented(self, journal_setup):
        """RunJournal does not implement on_custom_event — no crash expected."""
        j, store = journal_setup
        # BaseCallbackHandler.on_custom_event is a no-op by default
        j.on_custom_event("task_running", {"task_id": "t1"}, run_id=uuid4())
        await j.flush()
        events = await store.list_events("t1", "r1")
        assert isinstance(events, list)


class TestBufferFlush:
    @pytest.mark.anyio
    async def test_flush_threshold(self, journal_setup):
        j, store = journal_setup
        j._flush_threshold = 2
        # 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"])
        assert len(j._buffer) == 1
        j.on_llm_end(_make_llm_response("B"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
        # At threshold the buffer should have been flushed asynchronously
        await asyncio.sleep(0.1)
        events = await store.list_events("t1", "r1")
        assert len(events) >= 2

    @pytest.mark.anyio
    async def test_events_retained_when_no_loop(self, journal_setup):
        """Events buffered in a sync (no-loop) context should survive
        until the async flush() in the finally block."""
        j, store = journal_setup
        j._flush_threshold = 1

        original = asyncio.get_running_loop

        def no_loop():
            raise RuntimeError("no running event loop")

        asyncio.get_running_loop = no_loop
        try:
            j._put(event_type="llm.ai.response", category="message", content="test")
        finally:
            asyncio.get_running_loop = original

        assert len(j._buffer) == 1
        await j.flush()
        events = await store.list_events("t1", "r1")
        assert any(e["event_type"] == "llm.ai.response" for e in events)


class TestIdentifyCaller:
    def test_lead_agent_tag(self, journal_setup):
        j, _ = journal_setup
        assert j._identify_caller(["lead_agent"]) == "lead_agent"

    def test_subagent_tag(self, journal_setup):
        j, _ = journal_setup
        assert j._identify_caller(["subagent:research"]) == "subagent:research"

    def test_middleware_tag(self, journal_setup):
        j, _ = journal_setup
        assert j._identify_caller(["middleware:summarization"]) == "middleware:summarization"

    def test_no_tags_returns_lead_agent(self, journal_setup):
        j, _ = journal_setup
        assert j._identify_caller([]) == "lead_agent"
        assert j._identify_caller(None) == "lead_agent"


class TestChainErrorCallback:
    @pytest.mark.anyio
    async def test_on_chain_error_writes_run_error(self, journal_setup):
        j, store = journal_setup
        j.on_chain_error(ValueError("boom"), run_id=uuid4())
        await asyncio.sleep(0.05)
        await j.flush()
        events = await store.list_events("t1", "r1")
        error_events = [e for e in events if e["event_type"] == "run.error"]
        assert len(error_events) == 1
        assert "boom" in error_events[0]["content"]
        assert error_events[0]["metadata"]["error_type"] == "ValueError"


class TestTokenTrackingDisabled:
    @pytest.mark.anyio
    async def test_track_token_usage_false(self):
        store = MemoryRunEventStore()
        j = RunJournal("r1", "t1", store, track_token_usage=False, flush_threshold=100)
        j.on_llm_end(
            _make_llm_response("X", usage={"input_tokens": 50, "output_tokens": 50, "total_tokens": 100}),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        data = j.get_completion_data()
        assert data["total_tokens"] == 0
        assert data["llm_call_count"] == 0


class TestConvenienceFields:
    @pytest.mark.anyio
    async def test_first_human_message_via_set(self, journal_setup):
        j, _ = journal_setup
        j.set_first_human_message("What is AI?")
        data = j.get_completion_data()
        assert data["first_human_message"] == "What is AI?"

    @pytest.mark.anyio
    async def test_completion_data_counts_human_ai_and_tool_messages(self, journal_setup):
        from langchain_core.messages import HumanMessage, ToolMessage

        j, _ = journal_setup
        j.on_chat_model_start({}, [[HumanMessage(content="Question")]], run_id=uuid4(), tags=["lead_agent"])
        j.on_llm_end(_make_llm_response("Answer"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
        j.on_tool_end(ToolMessage(content="Tool result", tool_call_id="call_1", name="search"), run_id=uuid4())

        data = j.get_completion_data()

        assert data["message_count"] == 3
        assert data["first_human_message"] == "Question"
        assert data["last_ai_message"] == "Answer"

    @pytest.mark.anyio
    async def test_tool_call_only_ai_does_not_clear_last_ai_message(self, journal_setup):
        j, _ = journal_setup
        j.on_llm_end(_make_llm_response("Useful answer"), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
        j.on_llm_end(
            _make_llm_response("", tool_calls=[{"id": "call_1", "name": "search", "args": {}}]),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )

        data = j.get_completion_data()

        assert data["message_count"] == 2
        assert data["last_ai_message"] == "Useful answer"

    @pytest.mark.anyio
    async def test_last_ai_message_extracts_mixed_content_without_extra_newlines(self, journal_setup):
        j, _ = journal_setup
        j.on_llm_end(
            _make_llm_response(
                [
                    {"type": "text", "text": "First "},
                    {"type": "text", "content": "second"},
                    " third",
                    {"type": "image", "url": "ignored"},
                ]
            ),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )

        data = j.get_completion_data()

        assert data["message_count"] == 1
        assert data["last_ai_message"] == "First second third"

    @pytest.mark.anyio
    async def test_last_ai_message_extracts_mapping_content(self, 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"])

        data = j.get_completion_data()

        assert data["message_count"] == 1
        assert data["last_ai_message"] == "Nested answer"

    @pytest.mark.anyio
    async def test_duplicate_llm_run_id_does_not_double_count_message_summary(self, journal_setup):
        j, _ = journal_setup
        run_id = uuid4()

        j.on_llm_end(_make_llm_response("Answer", usage=None), run_id=run_id, parent_run_id=None, tags=["lead_agent"])
        j.on_llm_end(
            _make_llm_response("Answer", usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}),
            run_id=run_id,
            parent_run_id=None,
            tags=["lead_agent"],
        )

        data = j.get_completion_data()

        assert data["message_count"] == 1
        assert data["last_ai_message"] == "Answer"
        assert data["total_tokens"] == 15

    @pytest.mark.anyio
    async def test_subagent_ai_does_not_overwrite_lead_last_ai_message(self, 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("Subagent detail"), run_id=uuid4(), parent_run_id=None, tags=["subagent:research"])

        data = j.get_completion_data()

        assert data["message_count"] == 2
        assert data["last_ai_message"] == "Lead answer"

    @pytest.mark.anyio
    async def test_get_completion_data(self, journal_setup):
        j, _ = journal_setup
        j._total_tokens = 100
        j._msg_count = 5
        data = j.get_completion_data()
        assert data["total_tokens"] == 100
        assert data["message_count"] == 5


class TestMiddlewareEvents:
    @pytest.mark.anyio
    async def test_record_middleware_uses_middleware_category(self, journal_setup):
        j, store = journal_setup
        j.record_middleware(
            "title",
            name="TitleMiddleware",
            hook="after_model",
            action="generate_title",
            changes={"title": "Test Title", "thread_id": "t1"},
        )
        await j.flush()
        events = await store.list_events("t1", "r1")
        mw_events = [e for e in events if e["event_type"] == "middleware:title"]
        assert len(mw_events) == 1
        assert mw_events[0]["category"] == "middleware"
        assert mw_events[0]["content"]["name"] == "TitleMiddleware"
        assert mw_events[0]["content"]["hook"] == "after_model"
        assert mw_events[0]["content"]["action"] == "generate_title"
        assert mw_events[0]["content"]["changes"]["title"] == "Test Title"

    @pytest.mark.anyio
    async def test_middleware_tag_variants(self, journal_setup):
        """Different middleware tags produce distinct event_types."""
        j, store = journal_setup
        j.record_middleware("title", name="TitleMiddleware", hook="after_model", action="generate_title", changes={})
        j.record_middleware("guardrail", name="GuardrailMiddleware", hook="before_tool", action="deny", changes={})
        await j.flush()
        events = await store.list_events("t1", "r1")
        event_types = {e["event_type"] for e in events}
        assert "middleware:title" in event_types
        assert "middleware:guardrail" in event_types


class TestContextEvents:
    @pytest.mark.anyio
    async def test_record_memory_context_is_readable_from_public_store_contract(self, journal_setup):
        j, store = journal_setup

        j.record_memory_context(
            content_sha256="a" * 64,
        )
        # Goal continuations may enter the graph more than once under the same
        # run-scoped journal; the effective frozen memory event stays singular.
        j.record_memory_context(
            content_sha256="a" * 64,
        )
        await j.flush()

        events = await store.list_events("t1", "r1", event_types=["context:memory"])
        assert len(events) == 1
        assert events[0]["category"] == "context"
        assert events[0]["content"] == {"content_sha256": "a" * 64}

    @pytest.mark.anyio
    async def test_record_memory_context_can_retry_after_buffer_failure(self, journal_setup, monkeypatch):
        j, store = journal_setup
        original_put = j._put
        attempts = 0

        def fail_once(**kwargs):
            nonlocal attempts
            attempts += 1
            if attempts == 1:
                raise RuntimeError("buffer unavailable")
            return original_put(**kwargs)

        monkeypatch.setattr(j, "_put", fail_once)

        with pytest.raises(RuntimeError, match="buffer unavailable"):
            j.record_memory_context(content_sha256="a" * 64)
        j.record_memory_context(content_sha256="a" * 64)
        await j.flush()

        events = await store.list_events("t1", "r1", event_types=["context:memory"])
        assert len(events) == 1
        assert events[0]["content"] == {"content_sha256": "a" * 64}


class TestCallerBucketing:
    """Tests for caller-bucketed token accumulation (lead_agent / subagent / middleware)."""

    def test_lead_agent_bucketing(self, journal_setup):
        j, _ = journal_setup
        usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
        j.on_llm_end(_make_llm_response("A", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
        assert j._lead_agent_tokens == 15
        assert j._subagent_tokens == 0
        assert j._middleware_tokens == 0

    def test_subagent_bucketing(self, journal_setup):
        j, _ = journal_setup
        usage = {"input_tokens": 20, "output_tokens": 10, "total_tokens": 30}
        j.on_llm_end(_make_llm_response("B", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["subagent:research"])
        assert j._subagent_tokens == 30
        assert j._lead_agent_tokens == 0
        assert j._middleware_tokens == 0

    def test_middleware_bucketing(self, journal_setup):
        j, _ = journal_setup
        usage = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}
        j.on_llm_end(_make_llm_response("C", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["middleware:summarize"])
        assert j._middleware_tokens == 7
        assert j._lead_agent_tokens == 0
        assert j._subagent_tokens == 0

    def test_mixed_callers_sum_independently(self, journal_setup):
        j, _ = journal_setup
        usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
        j.on_llm_end(_make_llm_response("A", usage=usage), 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=["subagent:bash"])
        j.on_llm_end(_make_llm_response("C", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["middleware:title"])
        assert j._lead_agent_tokens == 15
        assert j._subagent_tokens == 15
        assert j._middleware_tokens == 15
        assert j._total_tokens == 45

    def test_get_completion_data_includes_buckets(self, journal_setup):
        j, _ = journal_setup
        j._lead_agent_tokens = 100
        j._subagent_tokens = 200
        j._middleware_tokens = 50
        data = j.get_completion_data()
        assert data["lead_agent_tokens"] == 100
        assert data["subagent_tokens"] == 200
        assert data["middleware_tokens"] == 50

    def test_dedup_same_run_id(self, journal_setup):
        """Same langchain run_id in on_llm_end must not double-count."""
        j, _ = journal_setup
        run_id = uuid4()
        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
        assert j._lead_agent_tokens == 15
        assert j._llm_call_count == 1

    def test_first_no_usage_second_with_usage(self, journal_setup):
        """First callback with no usage must not block second callback with usage for same run_id."""
        j, _ = journal_setup
        run_id = uuid4()
        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}
        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
        assert j._lead_agent_tokens == 15

    def test_track_token_usage_false_skips_buckets(self):
        """When token tracking is disabled, caller buckets stay at 0."""
        store = MemoryRunEventStore()
        j = RunJournal("r1", "t1", store, track_token_usage=False, flush_threshold=100)
        usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
        j.on_llm_end(_make_llm_response("X", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["subagent:research"])
        assert j._subagent_tokens == 0
        assert j._lead_agent_tokens == 0

    def test_default_no_tags_buckets_as_lead_agent(self, journal_setup):
        """LLM calls without explicit tags default to lead_agent bucket."""
        j, _ = journal_setup
        usage = {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10}
        j.on_llm_end(_make_llm_response("Hi", usage=usage), run_id=uuid4(), parent_run_id=None)
        assert j._lead_agent_tokens == 10
        assert j._subagent_tokens == 0
        assert j._middleware_tokens == 0

    def test_unknown_tag_buckets_as_lead_agent(self, journal_setup):
        """Calls with unrecognized tags (not lead_agent/subagent:/middleware:) go to lead_agent."""
        j, _ = journal_setup
        usage = {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10}
        j.on_llm_end(_make_llm_response("Hi", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["some_random_tag"])
        assert j._lead_agent_tokens == 10


class TestExternalUsageRecords:
    """Tests for record_external_llm_usage_records."""

    def test_records_added_to_subagent_bucket(self, journal_setup):
        j, _ = journal_setup
        records = [
            {
                "source_run_id": "ext-1",
                "caller": "subagent:general-purpose",
                "input_tokens": 100,
                "output_tokens": 50,
                "total_tokens": 150,
            }
        ]
        j.record_external_llm_usage_records(records)
        assert j._subagent_tokens == 150
        assert j._total_tokens == 150
        assert j._total_input_tokens == 100
        assert j._total_output_tokens == 50

    def test_records_added_to_middleware_bucket(self, journal_setup):
        j, _ = journal_setup
        records = [
            {
                "source_run_id": "ext-2",
                "caller": "middleware:summarize",
                "input_tokens": 30,
                "output_tokens": 10,
                "total_tokens": 40,
            }
        ]
        j.record_external_llm_usage_records(records)
        assert j._middleware_tokens == 40
        assert j._lead_agent_tokens == 0
        assert j._subagent_tokens == 0

    def test_records_added_to_lead_agent_bucket(self, journal_setup):
        j, _ = journal_setup
        records = [
            {
                "source_run_id": "ext-3",
                "caller": "lead_agent",
                "input_tokens": 10,
                "output_tokens": 5,
                "total_tokens": 15,
            }
        ]
        j.record_external_llm_usage_records(records)
        assert j._lead_agent_tokens == 15

    def test_dedup_same_source_run_id(self, journal_setup):
        """Same source_run_id must not be double-counted."""
        j, _ = journal_setup
        records = [
            {
                "source_run_id": "dup-1",
                "caller": "subagent:research",
                "input_tokens": 50,
                "output_tokens": 25,
                "total_tokens": 75,
            }
        ]
        j.record_external_llm_usage_records(records)
        j.record_external_llm_usage_records(records)
        assert j._subagent_tokens == 75
        assert j._total_tokens == 75

    def test_total_tokens_missing_computed_from_input_output(self, journal_setup):
        j, _ = journal_setup
        records = [
            {
                "source_run_id": "ext-4",
                "caller": "subagent:bash",
                "input_tokens": 200,
                "output_tokens": 100,
                "total_tokens": 0,
            }
        ]
        j.record_external_llm_usage_records(records)
        assert j._subagent_tokens == 300
        assert j._total_tokens == 300

    def test_total_tokens_zero_no_count(self, journal_setup):
        """Records with zero total and zero input+output must not be counted."""
        j, _ = journal_setup
        records = [
            {
                "source_run_id": "ext-5",
                "caller": "subagent:research",
                "input_tokens": 0,
                "output_tokens": 0,
                "total_tokens": 0,
            }
        ]
        j.record_external_llm_usage_records(records)
        assert j._total_tokens == 0
        assert j._subagent_tokens == 0

    def test_empty_source_run_id_skipped(self, journal_setup):
        j, _ = journal_setup
        records = [
            {
                "source_run_id": "",
                "caller": "subagent:research",
                "input_tokens": 50,
                "output_tokens": 25,
                "total_tokens": 75,
            }
        ]
        j.record_external_llm_usage_records(records)
        assert j._total_tokens == 0

    def test_multiple_records_in_single_call(self, journal_setup):
        j, _ = journal_setup
        records = [
            {"source_run_id": "r1", "caller": "subagent:gp", "input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
            {"source_run_id": "r2", "caller": "subagent:bash", "input_tokens": 20, "output_tokens": 10, "total_tokens": 30},
        ]
        j.record_external_llm_usage_records(records)
        assert j._subagent_tokens == 45
        assert j._total_tokens == 45

    def test_external_records_coexist_with_inline_callbacks(self, journal_setup):
        """External records and inline on_llm_end must not interfere."""
        j, _ = journal_setup
        usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
        j.on_llm_end(_make_llm_response("A", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
        j.record_external_llm_usage_records([{"source_run_id": "ext-6", "caller": "subagent:gp", "input_tokens": 100, "output_tokens": 50, "total_tokens": 150}])
        assert j._lead_agent_tokens == 15
        assert j._subagent_tokens == 150
        assert j._total_tokens == 165

    def test_track_token_usage_false_skips_external_records(self):
        """When token tracking is disabled, external records must not accumulate."""
        store = MemoryRunEventStore()
        j = RunJournal("r1", "t1", store, track_token_usage=False, flush_threshold=100)
        j.record_external_llm_usage_records([{"source_run_id": "ext-7", "caller": "subagent:gp", "input_tokens": 100, "output_tokens": 50, "total_tokens": 150}])
        assert j._total_tokens == 0
        assert j._subagent_tokens == 0


class TestProgressSnapshots:
    @pytest.mark.anyio
    async def test_on_llm_end_reports_progress_snapshot(self):
        snapshots: list[dict] = []

        async def reporter(snapshot: dict) -> None:
            snapshots.append(snapshot)

        store = MemoryRunEventStore()
        j = RunJournal(
            "r1",
            "t1",
            store,
            flush_threshold=100,
            progress_reporter=reporter,
            progress_flush_interval=0,
        )
        usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
        j.on_llm_end(_make_llm_response("Answer", usage=usage), run_id=uuid4(), parent_run_id=None, tags=["lead_agent"])
        await j.flush()

        assert snapshots
        assert snapshots[-1]["total_tokens"] == 15
        assert snapshots[-1]["llm_call_count"] == 1
        assert snapshots[-1]["message_count"] == 1
        assert snapshots[-1]["last_ai_message"] == "Answer"

    @pytest.mark.anyio
    async def test_throttled_progress_flush_emits_trailing_snapshot(self):
        snapshots: list[dict] = []
        trailing_seen = asyncio.Event()

        async def reporter(snapshot: dict) -> None:
            snapshots.append(snapshot)
            if snapshot["total_tokens"] == 45:
                trailing_seen.set()

        store = MemoryRunEventStore()
        j = RunJournal(
            "r1",
            "t1",
            store,
            flush_threshold=100,
            progress_reporter=reporter,
            progress_flush_interval=0.01,
        )
        j.on_llm_end(
            _make_llm_response("First", usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        j.on_llm_end(
            _make_llm_response("Second", usage={"input_tokens": 20, "output_tokens": 10, "total_tokens": 30}),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        await asyncio.wait_for(trailing_seen.wait(), timeout=1.0)
        await j.flush()

        assert len(snapshots) >= 2
        assert snapshots[-1]["total_tokens"] == 45
        assert snapshots[-1]["llm_call_count"] == 2
        assert snapshots[-1]["last_ai_message"] == "Second"

    @pytest.mark.anyio
    async def test_flush_cancels_delayed_progress_without_final_progress_write(self):
        snapshots: list[dict] = []

        async def reporter(snapshot: dict) -> None:
            snapshots.append(snapshot)

        store = MemoryRunEventStore()
        j = RunJournal(
            "r1",
            "t1",
            store,
            flush_threshold=100,
            progress_reporter=reporter,
            progress_flush_interval=10.0,
        )
        j.on_llm_end(
            _make_llm_response("First", usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )
        await asyncio.sleep(0)
        assert snapshots[-1]["total_tokens"] == 15
        j.on_llm_end(
            _make_llm_response("Second", usage={"input_tokens": 20, "output_tokens": 10, "total_tokens": 30}),
            run_id=uuid4(),
            parent_run_id=None,
            tags=["lead_agent"],
        )

        await asyncio.wait_for(j.flush(), timeout=0.2)

        assert snapshots[-1]["total_tokens"] == 15
        assert snapshots[-1]["llm_call_count"] == 1
        assert snapshots[-1]["last_ai_message"] == "First"


class TestChatModelStartHumanMessage:
    """Tests for on_chat_model_start extracting the first human message."""

    @staticmethod
    def _human_input_response(source: str = "ask_clarification") -> dict:
        return {
            "version": 1,
            "kind": "human_input_response",
            "source": source,
            "request_id": "clarification:call-abc",
            "response_kind": "option",
            "option_id": "option-2",
            "value": "staging",
        }

    @pytest.mark.anyio
    async def test_extracts_first_human_message(self, journal_setup):
        """on_chat_model_start captures the first HumanMessage from prompts."""
        from langchain_core.messages import AIMessage, HumanMessage

        j, store = journal_setup
        messages_batch = [
            [HumanMessage(content="What is AI?"), AIMessage(content="Hi there")],
        ]
        j.on_chat_model_start({}, messages_batch, run_id=uuid4(), tags=["lead_agent"])
        await j.flush()

        assert j._first_human_msg == "What is AI?"
        events = await store.list_events("t1", "r1")
        human_events = [e for e in events if e["event_type"] == "llm.human.input"]
        assert len(human_events) == 1
        assert human_events[0]["content"]["content"] == "What is AI?"

    @pytest.mark.anyio
    async def test_skips_hidden_human_messages(self, journal_setup):
        """HumanMessages hidden from the UI are internal context, not user input."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        messages_batch = [
            [
                HumanMessage(content="What is the weather today?"),
                HumanMessage(
                    content="Your todo list from earlier...",
                    name="todo_reminder",
                    additional_kwargs={"hide_from_ui": True},
                ),
            ],
        ]
        j.on_chat_model_start({}, messages_batch, run_id=uuid4(), tags=["lead_agent"])
        await j.flush()

        assert j._first_human_msg == "What is the weather today?"
        assert j.get_completion_data()["message_count"] == 1
        events = await store.list_events("t1", "r1")
        human_events = [e for e in events if e["event_type"] == "llm.human.input"]
        assert len(human_events) == 1
        assert human_events[0]["content"]["content"] == "What is the weather today?"

    @pytest.mark.anyio
    async def test_only_hidden_human_messages_are_not_captured(self, journal_setup):
        """A prompt containing only internal HumanMessages has no user input."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        hidden_message = HumanMessage(
            content="Internal context",
            additional_kwargs={"hide_from_ui": True},
        )
        j.on_chat_model_start({}, [[hidden_message]], run_id=uuid4(), tags=["lead_agent"])
        await j.flush()

        assert j._first_human_msg is None
        assert j.get_completion_data()["message_count"] == 0
        events = await store.list_events("t1", "r1")
        assert not any(e["event_type"] == "llm.human.input" for e in events)

    @pytest.mark.anyio
    async def test_hidden_human_input_response_is_captured(self, journal_setup):
        """Hidden HumanInputCard replies are user-authored and must survive compaction."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        hidden_response = HumanMessage(
            content='For your clarification "Which environment?", my answer is: staging',
            additional_kwargs={
                "hide_from_ui": True,
                "human_input_response": self._human_input_response(),
            },
        )
        j.on_chat_model_start({}, [[hidden_response]], run_id=uuid4(), tags=["lead_agent"])
        await j.flush()

        assert j._first_human_msg == 'For your clarification "Which environment?", my answer is: staging'
        assert j.get_completion_data()["message_count"] == 1
        events = await store.list_events("t1", "r1")
        human_events = [e for e in events if e["event_type"] == "llm.human.input"]
        assert len(human_events) == 1
        assert human_events[0]["content"]["additional_kwargs"]["hide_from_ui"] is True
        assert human_events[0]["content"]["additional_kwargs"]["human_input_response"]["request_id"] == "clarification:call-abc"

    @pytest.mark.anyio
    async def test_hidden_human_input_response_wins_over_older_visible_prompt(self, journal_setup):
        """The latest hidden card reply is the run input, not an older visible prompt."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        older_prompt = HumanMessage(content="Write a quicksort PDF")
        hidden_response = HumanMessage(
            content='For your clarification "Which format?", my answer is: tutorial',
            additional_kwargs={
                "hide_from_ui": True,
                "human_input_response": self._human_input_response(),
            },
        )
        j.on_chat_model_start({}, [[older_prompt, hidden_response]], run_id=uuid4(), tags=["lead_agent"])
        await j.flush()

        assert j._first_human_msg == 'For your clarification "Which format?", my answer is: tutorial'
        events = await store.list_events("t1", "r1")
        human_events = [e for e in events if e["event_type"] == "llm.human.input"]
        assert len(human_events) == 1
        assert human_events[0]["content"]["content"] == 'For your clarification "Which format?", my answer is: tutorial'

    @pytest.mark.anyio
    async def test_hidden_human_input_response_ignores_non_allowlisted_source(self, journal_setup):
        """Only explicit HumanInputCard sources are persisted while hidden."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        hidden_response = HumanMessage(
            content="Internal approval response",
            additional_kwargs={
                "hide_from_ui": True,
                "human_input_response": self._human_input_response(source="future_approval"),
            },
        )
        j.on_chat_model_start({}, [[hidden_response]], run_id=uuid4(), tags=["lead_agent"])
        await j.flush()

        assert j._first_human_msg is None
        assert j.get_completion_data()["message_count"] == 0
        events = await store.list_events("t1", "r1")
        assert not any(e["event_type"] == "llm.human.input" for e in events)

    @pytest.mark.anyio
    async def test_legacy_summary_message_is_not_captured_as_user_input(self, journal_setup):
        """Legacy synthetic summaries are internal context even if hide_from_ui is absent."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        legacy_summary = HumanMessage(content="Older compressed conversation state", name="summary")
        j.on_chat_model_start({}, [[legacy_summary]], run_id=uuid4(), tags=["lead_agent"])
        await j.flush()

        assert j._first_human_msg is None
        assert j.get_completion_data()["message_count"] == 0
        events = await store.list_events("t1", "r1")
        assert not any(e["event_type"] == "llm.human.input" for e in events)

    @pytest.mark.anyio
    async def test_visible_human_message_after_hidden_only_prompt_is_captured(self, journal_setup):
        """Skipping an internal-only prompt does not block later user input."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        hidden_message = HumanMessage(
            content="Internal context",
            additional_kwargs={"hide_from_ui": True},
        )
        j.on_chat_model_start({}, [[hidden_message]], run_id=uuid4(), tags=["lead_agent"])
        j.on_chat_model_start(
            {},
            [[HumanMessage(content="Real question")]],
            run_id=uuid4(),
            tags=["lead_agent"],
        )
        await j.flush()

        assert j._first_human_msg == "Real question"
        assert j.get_completion_data()["message_count"] == 1
        events = await store.list_events("t1", "r1")
        human_events = [e for e in events if e["event_type"] == "llm.human.input"]
        assert len(human_events) == 1
        assert human_events[0]["content"]["content"] == "Real question"

    @pytest.mark.anyio
    async def test_summarization_prompt_does_not_capture_first_human_message(self, journal_setup):
        """Internal summarization prompts must not replace the run's real user input."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        summarization_prompt = HumanMessage(
            content="<role>\nContext Extraction Assistant\n</role>\n\n<primary_objective>\nExtract context...",
        )
        j.on_chat_model_start(
            {},
            [[summarization_prompt]],
            run_id=uuid4(),
            tags=["middleware:summarize"],
        )
        j.on_chat_model_start(
            {},
            [[HumanMessage(content="Real user follow-up")]],
            run_id=uuid4(),
            tags=["lead_agent"],
        )
        await j.flush()

        assert j._first_human_msg == "Real user follow-up"
        assert j.get_completion_data()["message_count"] == 1
        events = await store.list_events("t1", "r1")
        human_events = [e for e in events if e["event_type"] == "llm.human.input"]
        assert len(human_events) == 1
        assert human_events[0]["content"]["content"] == "Real user follow-up"
        assert human_events[0]["metadata"]["caller"] == "lead_agent"

    @pytest.mark.anyio
    @pytest.mark.parametrize("tags", [["middleware:summarize"], ["subagent:research"]])
    async def test_non_lead_human_prompts_are_not_captured_as_user_input(self, journal_setup, tags):
        """Only lead-agent LLM starts create UI-facing human input events."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        j.on_chat_model_start(
            {},
            [[HumanMessage(content="Internal prompt")]],
            run_id=uuid4(),
            tags=tags,
        )
        await j.flush()

        assert j._first_human_msg is None
        assert j.get_completion_data()["message_count"] == 0
        events = await store.list_events("t1", "r1")
        assert not any(e["event_type"] == "llm.human.input" for e in events)

    @pytest.mark.anyio
    async def test_only_first_human_message_captured(self, journal_setup):
        """Subsequent on_chat_model_start calls do not overwrite the first message."""
        from langchain_core.messages import HumanMessage

        j, store = journal_setup
        j.on_chat_model_start({}, [[HumanMessage(content="First question")]], run_id=uuid4(), tags=["lead_agent"])
        j.on_chat_model_start({}, [[HumanMessage(content="Second question")]], run_id=uuid4(), tags=["lead_agent"])
        await j.flush()

        assert j._first_human_msg == "First question"
        events = await store.list_events("t1", "r1")
        human_events = [e for e in events if e["event_type"] == "llm.human.input"]
        assert len(human_events) == 1

    @pytest.mark.anyio
    async def test_empty_messages_no_crash(self, journal_setup):
        """on_chat_model_start with empty messages does not crash."""
        j, store = journal_setup
        j.on_chat_model_start({}, [], run_id=uuid4(), tags=["lead_agent"])
        await j.flush()
        assert j._first_human_msg is None


class TestDeliveryTracking:
    """Slice 1 (#4272): journal records artifact production for run.delivery."""

    @staticmethod
    def _register_tool_call(j: RunJournal, tool_call_id: str, name: str) -> None:
        from langchain_core.messages import AIMessage

        ai = AIMessage(content="", tool_calls=[{"id": tool_call_id, "name": name, "args": {}}])
        j._remember_current_run_tool_calls(ai, caller="lead_agent")

    def test_callbacks_run_inline_to_serialize_parallel_mutations(self, journal_setup):
        j, _ = journal_setup

        # LangChain dispatches synchronous handlers with run_inline=False via
        # run_in_executor, allowing parallel tool callbacks to mutate one
        # journal from different threads.
        assert j.run_inline is True

    @pytest.mark.anyio
    async def test_concurrent_callbacks_on_one_journal_are_serialized(self, journal_setup):
        from langchain_core.callbacks.manager import ahandle_event
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        j, _ = journal_setup
        commands = []
        for index, path in enumerate(("report.md", "report.md", "appendix.md"), start=1):
            tool_call_id = f"call_{index}"
            self._register_tool_call(j, tool_call_id, "present_files")
            commands.append(
                Command(
                    update={
                        "artifacts": [f"/mnt/user-data/outputs/{path}"],
                        "messages": [ToolMessage("Successfully presented files", tool_call_id=tool_call_id)],
                    }
                )
            )

        # This is the real LangChain async callback dispatcher. Because the
        # journal is run_inline, each synchronous mutation completes on the
        # event-loop thread instead of racing in executor threads.
        await asyncio.gather(
            *(
                ahandle_event(
                    [j],
                    "on_tool_end",
                    "ignore_agent",
                    command,
                    run_id=uuid4(),
                )
                for command in commands
            )
        )

        content = j.get_delivery_content()
        assert content["presented"] == 2
        assert set(content["paths"]) == {
            "/mnt/user-data/outputs/report.md",
            "/mnt/user-data/outputs/appendix.md",
        }
        assert set(content["by_tool"]["present_files"]) == set(content["paths"])

    @pytest.mark.anyio
    async def test_concurrent_runs_keep_delivery_accumulators_isolated(self):
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        store = MemoryRunEventStore()
        journals = [RunJournal(run_id, "t1", store, flush_threshold=100) for run_id in ("r1", "r2")]

        async def finish_run(journal: RunJournal, index: int) -> None:
            tool_call_id = f"call_run_{index}"
            self._register_tool_call(journal, tool_call_id, "present_files")
            journal.on_tool_end(
                Command(
                    update={
                        "artifacts": [f"/mnt/user-data/outputs/report-{index}.md"],
                        "messages": [ToolMessage("Successfully presented files", tool_call_id=tool_call_id)],
                    }
                ),
                run_id=uuid4(),
            )
            await asyncio.sleep(0)
            journal.record_delivery()
            await journal.flush()

        await asyncio.gather(*(finish_run(journal, index) for index, journal in enumerate(journals, start=1)))

        for index in (1, 2):
            events = await store.list_events("t1", f"r{index}")
            content = next(e for e in events if e["event_type"] == "run.delivery")["content"]
            assert content == {
                "presented": 1,
                "paths": [f"/mnt/user-data/outputs/report-{index}.md"],
                "by_tool": {"present_files": [f"/mnt/user-data/outputs/report-{index}.md"]},
            }

    @pytest.mark.anyio
    async def test_present_files_success_command_recorded_with_attribution(self, journal_setup):
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        j, store = journal_setup
        self._register_tool_call(j, "call_1", "present_files")
        cmd = Command(
            update={
                "artifacts": ["/mnt/user-data/outputs/report.md"],
                "messages": [ToolMessage("Successfully presented files", tool_call_id="call_1")],
            }
        )
        j.on_tool_end(cmd, run_id=uuid4())
        j.record_delivery()
        await j.flush()

        events = await store.list_events("t1", "r1")
        delivery = [e for e in events if e["event_type"] == "run.delivery"]
        assert len(delivery) == 1
        content = delivery[0]["content"]
        assert content["presented"] == 1
        assert content["paths"] == ["/mnt/user-data/outputs/report.md"]
        assert content["by_tool"] == {"present_files": ["/mnt/user-data/outputs/report.md"]}
        assert delivery[0]["category"] == "outputs"

    @pytest.mark.anyio
    async def test_command_with_multiple_messages_records_artifacts_once(self, journal_setup):
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        j, store = journal_setup
        self._register_tool_call(j, "call_multi", "present_files")
        cmd = Command(
            update={
                "artifacts": ["/mnt/user-data/outputs/report.md"],
                "messages": [
                    ToolMessage("Successfully presented files", tool_call_id="call_multi"),
                    HumanMessage("Additional command message"),
                ],
            }
        )
        j.on_tool_end(cmd, run_id=uuid4())
        j.record_delivery()
        await j.flush()

        events = await store.list_events("t1", "r1")
        content = next(e for e in events if e["event_type"] == "run.delivery")["content"]
        assert content == {
            "presented": 1,
            "paths": ["/mnt/user-data/outputs/report.md"],
            "by_tool": {"present_files": ["/mnt/user-data/outputs/report.md"]},
        }

    @pytest.mark.anyio
    async def test_command_with_multiple_tool_names_leaves_artifacts_unattributed(self, journal_setup):
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        j, store = journal_setup
        self._register_tool_call(j, "call_present", "present_files")
        self._register_tool_call(j, "call_browser", "browser_screenshot")
        cmd = Command(
            update={
                "artifacts": [
                    "/mnt/user-data/outputs/report.md",
                    "/mnt/user-data/outputs/shot.png",
                ],
                "messages": [
                    ToolMessage("Successfully presented files", tool_call_id="call_present"),
                    ToolMessage("Saved browser screenshot", tool_call_id="call_browser"),
                ],
            }
        )
        j.on_tool_end(cmd, run_id=uuid4())
        j.record_delivery()
        await j.flush()

        events = await store.list_events("t1", "r1")
        content = next(e for e in events if e["event_type"] == "run.delivery")["content"]
        assert content == {
            "presented": 2,
            "paths": [
                "/mnt/user-data/outputs/report.md",
                "/mnt/user-data/outputs/shot.png",
            ],
            "by_tool": {},
        }

    @pytest.mark.anyio
    async def test_error_command_without_artifacts_not_recorded(self, journal_setup):
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        j, store = journal_setup
        self._register_tool_call(j, "call_2", "present_files")
        cmd = Command(update={"messages": [ToolMessage("Error: Only files in /mnt/user-data/outputs can be presented", tool_call_id="call_2")]})
        j.on_tool_end(cmd, run_id=uuid4())
        j.record_delivery()
        await j.flush()

        events = await store.list_events("t1", "r1")
        delivery = [e for e in events if e["event_type"] == "run.delivery"]
        assert len(delivery) == 1
        assert delivery[0]["content"] == {"presented": 0, "paths": [], "by_tool": {}}

    @pytest.mark.anyio
    async def test_browser_tool_artifacts_recorded_under_producing_tool(self, journal_setup):
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        j, store = journal_setup
        self._register_tool_call(j, "call_3", "browser_screenshot")
        cmd = Command(
            update={
                "artifacts": ["/mnt/user-data/outputs/shot.png"],
                "messages": [ToolMessage("Saved browser screenshot", tool_call_id="call_3")],
            }
        )
        j.on_tool_end(cmd, run_id=uuid4())
        j.record_delivery()
        await j.flush()

        events = await store.list_events("t1", "r1")
        content = next(e for e in events if e["event_type"] == "run.delivery")["content"]
        assert content["presented"] == 1
        assert content["by_tool"] == {"browser_screenshot": ["/mnt/user-data/outputs/shot.png"]}

    @pytest.mark.anyio
    async def test_duplicate_path_tool_pair_recorded_once(self, journal_setup):
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        j, store = journal_setup
        self._register_tool_call(j, "call_4", "present_files")
        for _ in range(2):
            j.on_tool_end(
                Command(
                    update={
                        "artifacts": ["/mnt/user-data/outputs/report.md"],
                        "messages": [ToolMessage("Successfully presented files", tool_call_id="call_4")],
                    }
                ),
                run_id=uuid4(),
            )
        j.record_delivery()
        await j.flush()

        events = await store.list_events("t1", "r1")
        content = next(e for e in events if e["event_type"] == "run.delivery")["content"]
        assert content["presented"] == 1
        assert content["paths"] == ["/mnt/user-data/outputs/report.md"]

    @pytest.mark.anyio
    async def test_unattributed_artifacts_counted_without_by_tool_entry(self, journal_setup):
        from langchain_core.messages import ToolMessage
        from langgraph.types import Command

        j, store = journal_setup
        # No _register_tool_call: attribution missing (e.g. tool_call names map miss).
        cmd = Command(
            update={
                "artifacts": ["/mnt/user-data/outputs/anon.txt"],
                "messages": [ToolMessage("ok", tool_call_id="call_unknown")],
            }
        )
        j.on_tool_end(cmd, run_id=uuid4())
        j.record_delivery()
        await j.flush()

        events = await store.list_events("t1", "r1")
        content = next(e for e in events if e["event_type"] == "run.delivery")["content"]
        assert content["presented"] == 1
        assert content["paths"] == ["/mnt/user-data/outputs/anon.txt"]
        assert content["by_tool"] == {}