"""Contracts for the local adapters around the pure agent engine."""

from __future__ import annotations

import threading
import time
from types import SimpleNamespace

import pytest

import openai4s.agent.runtime as runtime
from openai4s.agent.actions import CodeCell, NativeToolBatch, NativeToolCall
from openai4s.agent.events import OutcomeProduced, ReplyReceived, RunStarted
from openai4s.agent.models import ExecutionOutcome, ModelReply, RunState
from openai4s.agent.runtime import (
    ChatModel,
    CompactionPolicy,
    LocalActionExecutor,
    TranscriptEventSink,
    TranscriptTurn,
    format_observation,
)
from openai4s.tools.native import ToolSpec


def _native_call(
    index: int,
    *,
    name: str = "lookup",
    wire_id: str | None | object = ...,
    arguments: dict | None = None,
    parse_error: str | None = None,
) -> NativeToolCall:
    call_id = f"call_{index}"
    actual_wire_id = call_id if wire_id is ... else wire_id
    return NativeToolCall(
        id=call_id,
        wire_id=actual_wire_id,
        name=name,
        ordinal=index,
        raw_arguments='{"query":"ATP"}',
        arguments={"query": "ATP"} if arguments is None else arguments,
        parse_error=parse_error,
        provider_meta={"provider": "test"},
    )


class FakeDispatcher:
    def __init__(self, last_output=None):
        self.last_output = last_output
        self.calls = []

    def __call__(self, method, args):
        self.calls.append((method, args))
        return {"ok": True}


class FakeKernel:
    def __init__(self, result=None, after_execute=None):
        self.result = result or {
            "stdout": "",
            "stderr": "",
            "error": None,
            "usage": {},
        }
        self.after_execute = after_execute
        self.calls = []

    def execute(self, code, origin=None):
        self.calls.append((code, origin))
        if self.after_execute is not None:
            self.after_execute()
        return self.result


def _executor(*, kernel=None, dispatcher=None, gate=None, execute_r=None):
    return LocalActionExecutor(
        kernel or FakeKernel(),
        dispatcher or FakeDispatcher(),
        gate or (lambda code, messages: None),
        execute_r or (lambda code: {"stdout": "R", "error": None}),
    )


def test_chat_model_passes_native_schemas_and_is_blocking_by_default():
    spec = ToolSpec(
        "lookup",
        "Look up a fact.",
        {"type": "object", "properties": {"query": {"type": "string"}}},
    )
    cfg = object()
    calls = []

    def fake_chat(messages, received_cfg, **kwargs):
        calls.append((messages, received_cfg, kwargs))
        return {"content": "done"}

    model = ChatModel(cfg, fake_chat, tools=[spec])
    source = [{"role": "user", "content": "look it up"}]
    result = model.complete(source, lambda delta: None)

    assert result == {"content": "done"}
    assert calls == [(source, cfg, {"tools": (spec,)})]
    assert "on_delta" not in calls[0][2]


def test_chat_model_refreshes_callable_session_catalog_each_turn():
    first = ToolSpec("first", "", {"type": "object", "properties": {}})
    second = ToolSpec("second", "", {"type": "object", "properties": {}})
    active = [first]
    seen = []

    def fake_chat(messages, cfg, **kwargs):
        del messages, cfg
        seen.append(tuple(tool.name for tool in kwargs["tools"]))
        return {"content": "done"}

    model = ChatModel(object(), fake_chat, tools=lambda: tuple(active))
    model.complete([], lambda _delta: None)
    active.append(second)
    model.complete([], lambda _delta: None)

    assert seen == [("first",), ("first", "second")]


def test_chat_model_cancel_releases_owner_and_quarantines_late_provider_output():
    entered = threading.Event()
    release = threading.Event()
    provider_done = threading.Event()
    shared_cancel = threading.Event()
    late_cancel_checks = []
    deltas = []
    abandoned = []
    accounted = threading.Event()
    result = {}

    def fake_chat(messages, cfg, **kwargs):
        del messages, cfg
        entered.set()
        assert release.wait(5), "test did not release the fake provider"
        late_cancel_checks.append(kwargs["should_cancel"]())
        kwargs["on_delta"]("must not escape after Stop")
        provider_done.set()
        return {
            "content": "late reply",
            "usage": {"prompt_tokens": 17, "completion_tokens": 3},
        }

    def account(reply):
        abandoned.append(reply["usage"])
        accounted.set()

    model = ChatModel(
        object(),
        fake_chat,
        stream=True,
        cancellation=SimpleNamespace(cancelled=shared_cancel.is_set),
        abandoned_reply=account,
    )
    owner = threading.Thread(
        target=lambda: result.setdefault("reply", model.complete([], deltas.append))
    )
    owner.start()
    assert entered.wait(2), "fake provider call did not start"

    shared_cancel.set()
    owner.join(1)
    assert not owner.is_alive(), "Stop waited for the blocked provider request"
    assert result["reply"]["finish_reason"] == "cancelled"

    # Admission of the next queued turn clears this shared Event. The old
    # request's private latch must remain cancelled, or it can ABA-revive and
    # stream its late response into the new turn.
    shared_cancel.clear()
    release.set()
    assert provider_done.wait(2), "detached fake provider did not finish"
    assert accounted.wait(2), "late provider usage was not accounted"
    assert late_cancel_checks == [True]
    assert deltas == []
    assert abandoned == [{"prompt_tokens": 17, "completion_tokens": 3}]


def test_a_call_cancelled_before_it_starts_never_reaches_the_provider():
    """Stop can land between ``Thread.start()`` and ``chat_fn``'s first line.

    The detached design exists because a urllib request already on the wire
    cannot be recalled. In this window nothing has been sent yet, so there is
    nothing to salvage: the request must not go out at all. It would be billed
    for a turn nobody is reading, and on the Web path ``chat_fn`` is
    ``_invoke_model_with_auto_budget``, which snapshots the LIVE Auto Mode run
    id at its own entry -- by then the run of the turn admitted *after* Stop.
    """

    calls: list[dict] = []
    probes = {"n": 0}

    def cancelled_after_spawn() -> bool:
        # False only for the pre-spawn check; true from the moment the thread
        # could be running, whichever side of the handoff asks first.
        probes["n"] += 1
        return probes["n"] > 1

    def fake_chat(messages, cfg, **kwargs):
        del messages, cfg
        calls.append(kwargs)
        return {"content": "must not be requested", "usage": {"prompt_tokens": 9}}

    accounted: list[object] = []
    model = ChatModel(
        object(),
        fake_chat,
        cancellation=SimpleNamespace(cancelled=cancelled_after_spawn),
        abandoned_reply=accounted.append,
    )

    reply = model.complete([], lambda _text: None)

    assert reply["finish_reason"] == "cancelled"
    assert calls == [], "a cancelled call still reached the provider"
    # Nothing was requested, so there is no usage to meter and no reply to
    # quarantine -- the accounting sink must stay untouched rather than
    # inventing a zero.
    assert accounted == []


def test_chat_model_projects_stream_deltas_on_the_owning_thread():
    shared_cancel = threading.Event()
    provider_threads = []
    callback_threads = []
    deltas = []

    def fake_chat(messages, cfg, **kwargs):
        del messages, cfg
        provider_threads.append(threading.get_ident())
        kwargs["on_delta"]("one")
        kwargs["on_delta"]("two")
        return {"content": "one two"}

    owner_thread = threading.get_ident()
    model = ChatModel(
        object(),
        fake_chat,
        stream=True,
        cancellation=SimpleNamespace(cancelled=shared_cancel.is_set),
    )

    def collect_delta(text):
        callback_threads.append(threading.get_ident())
        deltas.append(text)

    reply = model.complete([], collect_delta)

    assert reply["content"] == "one two"
    assert provider_threads and provider_threads != [owner_thread]
    assert callback_threads == [owner_thread, owner_thread]
    assert deltas == ["one", "two"]


def test_chat_model_bounds_detached_provider_calls(monkeypatch):
    monkeypatch.setattr(
        runtime,
        "_PROVIDER_CALL_BUDGET",
        runtime._DetachedCallBudget(1, per_scope_limit=4),
    )
    entered = threading.Event()
    release = threading.Event()
    shared_cancel = threading.Event()

    def blocked_chat(messages, cfg, **kwargs):
        del messages, cfg, kwargs
        entered.set()
        assert release.wait(5)
        return {"content": "late"}

    model = ChatModel(
        object(),
        blocked_chat,
        cancellation=SimpleNamespace(cancelled=shared_cancel.is_set),
    )
    owner = threading.Thread(target=lambda: model.complete([], lambda _text: None))
    owner.start()
    assert entered.wait(2)
    shared_cancel.set()
    owner.join(1)
    assert not owner.is_alive()
    assert runtime._PROVIDER_CALL_BUDGET.outstanding() == 1

    second = ChatModel(
        object(),
        blocked_chat,
        cancellation=SimpleNamespace(cancelled=lambda: False),
    )
    with pytest.raises(RuntimeError, match="cancelled model requests"):
        second.complete([], lambda _text: None)

    release.set()
    deadline = time.monotonic() + 2
    while runtime._PROVIDER_CALL_BUDGET.outstanding() and time.monotonic() < deadline:
        time.sleep(0.01)
    assert runtime._PROVIDER_CALL_BUDGET.outstanding() == 0


def test_chat_model_budget_never_charges_a_live_provider_call():
    """Only *detached* calls are bounded.

    Charging every cancellable request would bound healthy concurrency: each
    delegated child carries a cancellation, the fan-out cap of 48 is per node,
    and the session cap is 1000 -- so a nested fan-out routinely holds more
    live calls than the budget and would start refusing requests nobody
    cancelled.
    """

    budget = runtime._DetachedCallBudget(2, per_scope_limit=99)
    live = [budget.track() for _ in range(5)]
    for _ in range(10):
        budget.admit()  # no live call is charged, so admission stays open
    assert budget.outstanding() == 0

    live[0].detach()
    live[1].detach()
    assert budget.outstanding() == 2
    with pytest.raises(RuntimeError, match="cancelled model requests"):
        budget.admit()

    live[0].settle()
    budget.admit()
    for call in live:
        call.settle()
    assert budget.outstanding() == 0


def test_one_session_cannot_stack_cancelled_calls_against_a_stale_ledger(
    monkeypatch,
):
    """Stop releases the turn while the request keeps billing.

    ``enforce_llm_quota`` is check-then-call: it reads stored ``llm_*``
    counters and reserves nothing for a call in flight, so every Stop-and-
    resend passes a ledger that has not been charged for the request still
    running. Only this budget decides how far one session can outrun its own
    accounting, so the per-scope limit has to bite well before the process one.
    """

    monkeypatch.setattr(
        runtime,
        "_PROVIDER_CALL_BUDGET",
        runtime._DetachedCallBudget(128, per_scope_limit=2),
    )
    release = threading.Event()
    started = threading.Semaphore(0)

    def blocked_chat(messages, cfg, **kwargs):
        del messages, cfg, kwargs
        started.release()
        assert release.wait(5)
        return {"content": "late"}

    def stop_one(scope):
        shared = threading.Event()
        model = ChatModel(
            object(),
            blocked_chat,
            cancellation=SimpleNamespace(cancelled=shared.is_set),
            call_scope=scope,
        )
        owner = threading.Thread(target=lambda: model.complete([], lambda _t: None))
        owner.start()
        assert started.acquire(timeout=2)
        shared.set()
        owner.join(2)
        assert not owner.is_alive()

    try:
        stop_one("frame-a")
        stop_one("frame-a")
        # A third Stop on the same session would be a third concurrent billed
        # request against counters none of them have reached yet.
        with pytest.raises(RuntimeError, match="this session already has"):
            ChatModel(
                object(),
                blocked_chat,
                cancellation=SimpleNamespace(cancelled=lambda: False),
                call_scope="frame-a",
            ).complete([], lambda _t: None)

        # Another session is unaffected: the bound is per session, and the
        # process ceiling is far away.
        stop_one("frame-b")
        assert runtime._PROVIDER_CALL_BUDGET.outstanding("frame-a") == 2
        assert runtime._PROVIDER_CALL_BUDGET.outstanding("frame-b") == 1
    finally:
        release.set()

    deadline = time.monotonic() + 5
    while runtime._PROVIDER_CALL_BUDGET.outstanding() and time.monotonic() < deadline:
        time.sleep(0.01)
    assert runtime._PROVIDER_CALL_BUDGET.outstanding("frame-a") == 0


def test_native_batch_returns_one_canonical_tool_message_per_call(monkeypatch):
    dispatched = []

    def fake_execute(dispatcher, call):
        dispatched.append((dispatcher, call))
        return (f"result for {call['name']}", call["name"] == "lookup")

    monkeypatch.setattr(runtime, "execute_tool_call", fake_execute)
    dispatcher = FakeDispatcher(last_output={"stale": "must not submit"})
    executor = _executor(dispatcher=dispatcher)
    calls = (
        _native_call(0, wire_id=None),
        _native_call(1, name="delegate", arguments={"task": "check"}),
    )

    outcome = executor.execute(
        NativeToolBatch(calls), ModelReply(), RunState([{"role": "user"}])
    )

    assert list(outcome.history_messages) == [
        {
            "role": "tool",
            "tool_call_id": "call_0",
            "wire_id": None,
            "name": "lookup",
            "content": "result for lookup",
            "is_error": False,
        },
        {
            "role": "tool",
            "tool_call_id": "call_1",
            "wire_id": "call_1",
            "name": "delegate",
            "content": "result for delegate",
            "is_error": True,
        },
    ]
    assert all(message["role"] != "user" for message in outcome.history_messages)
    assert [call for _, call in dispatched] == [
        {"name": "lookup", "arguments": {"query": "ATP"}},
        {"name": "delegate", "arguments": {"task": "check"}},
    ]
    assert outcome.completion is None and outcome.stop_reason is None


def test_writing_native_call_uses_delegated_capture_hooks(monkeypatch):
    events = []

    class Hooks:
        def before_native(self, call):
            events.append(("before", call.id))
            return "snapshot-token"

        def after_native(self, call, token, result):
            events.append(("after", call.id, token, result))

    monkeypatch.setattr(
        runtime,
        "execute_tool_call",
        lambda dispatcher, call: ("wrote file", True),
    )
    executor = LocalActionExecutor(
        FakeKernel(),
        FakeDispatcher(),
        lambda code, messages: None,
        lambda code: {"stdout": "R", "error": None},
        cell_hooks=Hooks(),
    )
    call = _native_call(
        0,
        name="write_file",
        arguments={"path": "child.txt", "content": "child bytes"},
    )

    outcome = executor.execute(
        NativeToolBatch((call,)), ModelReply(), RunState([{"role": "user"}])
    )

    assert outcome.history_messages[0]["is_error"] is False
    assert events == [
        ("before", "call_0"),
        ("after", "call_0", "snapshot-token", ("wrote file", True)),
    ]


def test_native_parse_error_never_dispatches(monkeypatch):
    def unexpected_dispatch(*args):
        raise AssertionError(f"parse-error call was dispatched: {args!r}")

    monkeypatch.setattr(runtime, "execute_tool_call", unexpected_dispatch)
    malformed = NativeToolCall(
        id="bad_0",
        wire_id="wire_bad_0",
        name="lookup",
        ordinal=0,
        raw_arguments='{"query":',
        arguments=None,
        parse_error="invalid JSON",
        provider_meta={},
    )

    outcome = _executor().execute(
        NativeToolBatch((malformed,)), ModelReply(), RunState([])
    )

    assert outcome.history_messages[0]["tool_call_id"] == "bad_0"
    assert outcome.history_messages[0]["wire_id"] == "wire_bad_0"
    assert outcome.history_messages[0]["is_error"] is True
    assert "invalid JSON" in outcome.history_messages[0]["content"]


def test_native_limit_skips_dispatch_but_never_drops_tool_results(monkeypatch):
    dispatched = []

    def fake_execute(dispatcher, call):
        dispatched.append(call)
        return "ok", True

    monkeypatch.setattr(runtime, "execute_tool_call", fake_execute)
    calls = tuple(_native_call(index) for index in range(18))

    outcome = _executor().execute(NativeToolBatch(calls), ModelReply(), RunState([]))

    assert len(dispatched) == runtime.MAX_TOOL_CALLS_PER_TURN == 16
    assert len(outcome.history_messages) == len(calls) == 18
    assert [message["tool_call_id"] for message in outcome.history_messages] == [
        f"call_{index}" for index in range(18)
    ]
    assert all(message["role"] == "tool" for message in outcome.history_messages)
    assert all(message["is_error"] for message in outcome.history_messages[16:])
    assert all(
        "was not run" in message["content"] for message in outcome.history_messages[16:]
    )


def test_code_observation_notes_extra_cells_and_only_submit_sets_completion():
    dispatcher = FakeDispatcher()
    gate_calls = []
    first_kernel = FakeKernel(
        {"stdout": "42\n", "stderr": "", "error": None, "usage": {}}
    )
    executor = _executor(
        kernel=first_kernel,
        dispatcher=dispatcher,
        gate=lambda code, messages: gate_calls.append((code, messages)),
    )
    reply = ModelReply(
        content=(
            "```python\nprint(6 * 7)\n```\n" "```python\nraise AssertionError\n```"
        )
    )

    first = executor.execute(
        CodeCell("python", "print(6 * 7)\n"),
        reply,
        RunState([{"role": "user", "content": "compute"}]),
    )

    assert first.completion is None
    assert first.history_messages[0]["role"] == "user"
    assert "stdout:\n42" in first.observation
    assert runtime.MULTI_CELL_NOTE in first.observation
    assert first_kernel.calls == [("print(6 * 7)\n", "agent")]
    assert len(gate_calls) == 1

    submitted = {"output": {"answer": 42}, "completion_bullets": ["done"]}
    second_kernel = FakeKernel(
        after_execute=lambda: setattr(dispatcher, "last_output", submitted)
    )
    second_executor = _executor(kernel=second_kernel, dispatcher=dispatcher)
    second = second_executor.execute(
        CodeCell("python", "host.submit_output(...)"),
        ModelReply(content="```python\nhost.submit_output(...)\n```"),
        RunState([]),
    )
    assert second.completion is submitted
    assert second.stop_reason is None


def test_cli_revalidates_mid_cell_completion_after_capture_hooks():
    class RevalidatingDispatcher(FakeDispatcher):
        def __init__(self):
            super().__init__()
            self.revalidations = 0

        def revalidate_pending_completion(self):
            self.revalidations += 1
            if self.last_output is None:
                return None
            self.last_output = None
            return "verified source bytes changed after submission"

    dispatcher = RevalidatingDispatcher()
    submitted = {"output": {"answer": 42}, "completion_bullets": ["Computed it"]}
    kernel = FakeKernel(
        after_execute=lambda: setattr(dispatcher, "last_output", submitted)
    )

    outcome = _executor(kernel=kernel, dispatcher=dispatcher).execute(
        CodeCell("python", "host.submit_output(...); mutate_source()"),
        ModelReply(content="```python\nhost.submit_output(...)\n```"),
        RunState([]),
    )

    assert dispatcher.revalidations == 1
    assert outcome.completion is None
    assert "rejected after cell capture" in outcome.observation
    assert "source bytes changed" in outcome.observation


def test_code_observation_notes_an_incomplete_tail_after_the_executed_cell():
    kernel = FakeKernel({"stdout": "first\n", "stderr": "", "error": None, "usage": {}})
    executor = _executor(kernel=kernel, dispatcher=FakeDispatcher())
    reply = ModelReply(
        content=("```python\nprint('first')\n```\n" "```r\nresult <- unfinished(\n")
    )

    outcome = executor.execute(
        CodeCell("python", "print('first')\n"), reply, RunState([])
    )

    assert kernel.calls == [("print('first')\n", "agent")]
    assert runtime.MULTI_CELL_NOTE in outcome.observation


def test_none_action_keeps_legacy_tool_fallback_as_user_history(monkeypatch):
    calls = []

    def fake_run(dispatcher, parsed_calls, errors, *, on_result=None):
        calls.append((dispatcher, parsed_calls, errors))
        for parsed in parsed_calls:
            if on_result is not None:
                on_result(parsed, True)
        return "[Tool Results]\nlegacy result"

    monkeypatch.setattr(runtime, "run_tool_calls", fake_run)
    dispatcher = FakeDispatcher()
    executor = _executor(dispatcher=dispatcher)
    fenced = "```tool\n" '{"name":"list_dir","arguments":{"path":"."}}\n' "```"

    legacy = executor.execute(None, ModelReply(content=fenced), RunState([]))
    prose = executor.execute(None, ModelReply(content="plain prose"), RunState([]))

    assert len(calls) == 1
    assert calls[0][1] == [{"name": "list_dir", "arguments": {"path": "."}}]
    assert legacy.history_messages == (
        {"role": "user", "content": "[Tool Results]\nlegacy result"},
    )
    assert prose.observation == runtime.NO_CODE_NUDGE
    assert prose.history_messages == (
        {"role": "user", "content": runtime.NO_CODE_NUDGE},
    )


def test_none_action_can_fall_back_to_code_when_native_tools_are_unavailable():
    executor = LocalActionExecutor(
        FakeKernel(),
        FakeDispatcher(),
        lambda code, messages: None,
        lambda code: {"stdout": "R", "error": None},
        prose_nudge=runtime.NO_NATIVE_COMPLETION_NUDGE,
    )

    outcome = executor.execute(None, ModelReply(content="plain prose"), RunState([]))

    assert outcome.observation == runtime.NO_NATIVE_COMPLETION_NUDGE
    assert "host.submit_output" in outcome.observation
    assert "finalize_response" not in outcome.observation


def test_incomplete_cell_is_not_executed_and_requests_a_full_replacement():
    kernel = FakeKernel()
    executor = _executor(kernel=kernel, dispatcher=FakeDispatcher())

    outcome = executor.execute(
        None,
        ModelReply(content="```python\nvalue = compute(\n"),
        RunState([]),
    )

    assert outcome.observation == runtime.INCOMPLETE_CELL_NUDGE
    assert "NOTHING" in outcome.observation
    assert kernel.calls == []


def test_compaction_expands_tail_to_keep_assistant_tool_group_atomic(monkeypatch):
    assistant = {
        "role": "assistant",
        "content": "",
        "tool_calls": [{"id": f"call_{index}"} for index in range(5)],
    }
    tools = [
        {"role": "tool", "tool_call_id": f"call_{index}", "content": "ok"}
        for index in range(5)
    ]
    messages = [
        {"role": "system", "content": "system"},
        {"role": "user", "content": "task"},
        {"role": "assistant", "content": "old"},
        {"role": "user", "content": "old observation"},
        assistant,
        *tools,
    ]
    captured = {}

    def fake_compact(
        messages,
        cfg,
        *,
        keep_recent,
        archive_dir,
        archive_metadata,
        large_output_chars,
        **_kwargs,
    ):
        captured.update(
            keep_recent=keep_recent,
            tail=messages[-keep_recent:],
            archive_dir=archive_dir,
            archive_metadata=archive_metadata,
            large_output_chars=large_output_chars,
        )
        return messages

    monkeypatch.setattr(runtime, "compact", fake_compact)
    # A one-token window with a zero trigger ratio forces compaction on any
    # non-empty history; the policy has a single trigger path.
    cfg = SimpleNamespace(
        compaction_dir="archive", context_window_tokens=1, compaction_trigger_ratio=0.0
    )

    prepared = CompactionPolicy(cfg).prepare(RunState(messages))

    assert prepared is messages
    assert captured["keep_recent"] == 6
    assert captured["tail"] == [assistant, *tools]
    assert captured["archive_dir"] == "archive"
    assert captured["archive_metadata"].active_kernel_generation is None
    assert captured["large_output_chars"] == runtime.DEFAULT_LARGE_OUTPUT_CHARS


def test_compaction_circuit_breaker_stops_repeated_low_yield_calls(monkeypatch):
    messages = [
        {"role": "system", "content": "system"},
        {"role": "user", "content": "task"},
        {"role": "assistant", "content": "old"},
        {"role": "user", "content": "old observation"},
        {"role": "assistant", "content": "recent"},
    ]
    attempts = []

    def no_yield(messages, cfg, **kwargs):
        attempts.append(kwargs)
        return messages

    monkeypatch.setattr(runtime, "compact", no_yield)
    policy = CompactionPolicy(
        SimpleNamespace(
            compaction_dir="archive",
            context_window_tokens=1,
            compaction_trigger_ratio=0.0,
        )
    )
    state = RunState(messages)

    policy.prepare(state)
    policy.prepare(state)
    policy.prepare(state)

    assert len(attempts) == 2
    assert policy.low_yield_streak == 2
    assert policy.circuit_open is True
    assert state.metadata["compaction_circuit_open"] is True


def test_executor_records_kernel_generation_change_for_safe_handoff():
    kernel = FakeKernel()
    kernel.generation = 9
    state = RunState(
        [{"role": "user", "content": "continue"}],
        metadata={"active_kernel_generation": 8},
    )

    _executor(kernel=kernel).execute(
        CodeCell("python", "print('ready')"),
        ModelReply(content="```python\nprint('ready')\n```"),
        state,
    )

    assert state.metadata["previous_kernel_generation"] == 8
    assert state.metadata["active_kernel_generation"] == 9
    assert state.metadata["kernel_restarted"] is True


def test_transcript_sink_projects_only_reply_and_observed_outcome_events():
    transcript = []
    logs = []
    sink = TranscriptEventSink(transcript, log=lambda *parts: logs.append(parts))
    reply = ModelReply(content="working")

    sink.emit(RunStarted(max_turns=2, history_size=1))
    sink.emit(ReplyReceived(reply, turn=0))
    sink.emit(OutcomeProduced(ExecutionOutcome(observation="observed"), turn=0))
    sink.emit(OutcomeProduced(ExecutionOutcome(), turn=1))

    assert transcript == [
        TranscriptTurn("assistant", "working"),
        TranscriptTurn("observation", "observed"),
    ]
    assert len(logs) == 2


def test_format_observation_preserves_stable_error_and_usage_protocol():
    observation = format_observation(
        {
            "stdout": "value\n",
            "stderr": "warning\n",
            "error": "boom\n",
            "trace": {"error_lineno": 3},
            "usage": {"wall_s": 1.0, "cpu_s": 0.5, "peak_rss_kb": 64},
        }
    )

    assert observation == (
        "[Observation]\n"
        "stdout:\nvalue\n"
        "stderr:\nwarning\n"
        "ERROR (cell line 3):\nboom\n"
        "[system] The cell stopped at the first exception. Statements after "
        "that line did not run, and their variables/files must not be assumed "
        "to exist. Repair with one complete cell beginning before the failed "
        "dependency; never send only a continuation fragment.\n"
        "[usage wall=1.0s cpu=0.5s rss=64kb]"
    )


def test_format_observation_reminds_models_that_host_is_injected():
    observation = format_observation(
        {"error": "ModuleNotFoundError: No module named 'host'"}
    )

    assert "never `import host`" in observation


def test_a_reply_delivered_before_cancellation_keeps_its_usage(monkeypatch):
    """Cancellation that lands after ChatModel returned a real reply.

    ``ChatModel`` owns cancellation end to end: whatever it delivers as real
    has passed its own final check, and a late Stop is observed by the engine
    after the reply was recorded with its usage. There is no second wrapper
    left to re-check the raw signal and swap the reply for the no-op without
    ``abandon()`` -- which is the one exit that metered nothing.
    """
    from openai4s.agent.engine import AgentEngine
    from openai4s.agent.events import ReplyReceived
    from openai4s.agent.models import ExecutionOutcome

    usage = {"prompt_tokens": 17, "completion_tokens": 3}
    stop = threading.Event()
    abandoned = []
    model = ChatModel(
        object(),
        lambda messages, cfg, **kwargs: {"content": "real reply", "usage": dict(usage)},
        cancellation=SimpleNamespace(cancelled=stop.is_set),
        abandoned_reply=abandoned.append,
    )

    class StopRightAfterDelivery:
        def complete(self, messages, on_delta):
            reply = model.complete(messages, on_delta)
            stop.set()  # Stop lands the instant the reply is delivered
            return reply

    received = []

    class Events:
        def emit(self, event):
            if isinstance(event, ReplyReceived):
                received.append(event.reply)

    class Executor:
        def execute(self, action, reply, state):
            return ExecutionOutcome()

    engine = AgentEngine(
        StopRightAfterDelivery(),
        Executor(),
        event_sink=Events(),
        cancellation=SimpleNamespace(cancelled=stop.is_set),
    )
    result = engine.run([{"role": "user", "content": "go"}])

    assert result.stop_reason == "cancelled"
    assert [reply.usage for reply in received] == [usage]
    assert abandoned == []  # delivered as real: metered once, on the normal path


def test_agent_composes_chat_model_without_a_second_cancellation_wrapper():
    import ast
    import inspect

    import openai4s.agent.loop as loop_mod

    tree = ast.parse(inspect.getsource(loop_mod))
    assert not [
        node.name
        for node in ast.walk(tree)
        if isinstance(node, ast.ClassDef) and "Cancellation" in node.name
    ], "loop.py wraps ChatModel in a second cancellation check again"
    chat_model_calls = [
        node
        for node in ast.walk(tree)
        if isinstance(node, ast.Call)
        and isinstance(node.func, ast.Name)
        and node.func.id == "ChatModel"
    ]
    assert chat_model_calls
    assert all(
        {"cancellation", "abandoned_reply"}
        <= {keyword.arg for keyword in call.keywords}
        for call in chat_model_calls
    )


def test_delegated_children_are_bounded_by_the_process_ceiling_only():
    """A child's detached calls must not fill the session's accounting slots.

    Keyed on the session root, one Stop on a fan-out -- or a parent cell's
    ordinary ``host.stop_child`` on a few in-flight siblings -- filled the
    root's four slots and the parent's very next model call was refused.
    """
    from openai4s.agent.loop import Agent

    store = SimpleNamespace(resolve_frame_scope=lambda fid: {"root_frame_id": "root-1"})
    root = Agent(use_skills=False, allow_delegate=False, frame_id="branch-1")
    child = Agent(
        use_skills=False, allow_delegate=False, frame_id="child-1", delegate_depth=1
    )
    # The scope is resolved through the dispatcher's Store at call time.
    root.dispatcher = SimpleNamespace(store=store)
    child.dispatcher = SimpleNamespace(store=store)
    assert root._provider_call_scope() == "root-1"
    assert child._provider_call_scope() is None


def test_late_accounting_survives_a_sink_that_raises_a_base_exception():
    accounting = runtime._LateAccounting()
    first_done = threading.Event()
    second = []
    second_done = threading.Event()

    def exploding(reply):
        del reply
        first_done.set()
        raise KeyboardInterrupt("injected")

    def recording(reply):
        second.append(reply)
        second_done.set()

    accounting.submit(exploding, {"usage": {}})
    assert first_done.wait(2)
    accounting.submit(recording, {"usage": {"input_tokens": 1}})
    assert second_done.wait(2), "the late-accounting worker died with its first sink"
    assert second == [{"usage": {"input_tokens": 1}}]


def test_call_context_is_captured_on_the_owning_thread_before_the_handoff():
    """The provider thread must not read the session's identity itself.

    On the Web path ``chat_fn`` used to snapshot the live Auto Mode run at its
    own entry -- on the detached thread, possibly after Stop released the turn
    and the next queued turn installed a new run. ``call_context`` is captured
    by the owner before the thread starts and handed over as a value.
    """
    live = {"run": "run-of-this-turn"}
    entered = threading.Event()
    release = threading.Event()
    seen = []

    def fake_chat(messages, cfg, **kwargs):
        del messages, cfg
        entered.set()
        assert release.wait(5)
        seen.append(kwargs.get("call_context"))
        return {"content": "ok"}

    model = ChatModel(
        object(),
        fake_chat,
        cancellation=SimpleNamespace(cancelled=lambda: False),
        call_context=lambda: {"run_id": live["run"]},
    )
    owner = threading.Thread(target=lambda: model.complete([], lambda _t: None))
    owner.start()
    assert entered.wait(2)
    live["run"] = "run-installed-after-stop"  # the next turn moved in
    release.set()
    owner.join(2)
    assert seen == [{"run_id": "run-of-this-turn"}]


def test_call_context_is_absent_when_no_provider_asked_for_one():
    seen = []

    def fake_chat(messages, cfg, **kwargs):
        del messages, cfg
        seen.append("call_context" in kwargs)
        return {"content": "ok"}

    ChatModel(object(), fake_chat).complete([], lambda _t: None)
    ChatModel(
        object(), fake_chat, cancellation=SimpleNamespace(cancelled=lambda: False)
    ).complete([], lambda _t: None)
    assert seen == [False, False]


def test_the_cancel_probe_carries_the_mid_stream_policy():
    """Metered sessions drain an abandoned stream; the rest abort it."""
    seen = []

    def fake_chat(messages, cfg, **kwargs):
        del messages, cfg
        seen.append(kwargs["should_cancel"].abort_stream)
        return {"content": "ok"}

    probe = SimpleNamespace(cancelled=lambda: False)
    ChatModel(object(), fake_chat, cancellation=probe).complete([], lambda _t: None)
    ChatModel(
        object(), fake_chat, cancellation=probe, drain_cancelled_stream=True
    ).complete([], lambda _t: None)
    assert seen == [True, False]


def test_a_delegated_child_of_a_metered_session_drains_too():
    from openai4s.agent.loop import Agent

    owners = {"root-1": {"user_id": "member-1", "project_id": "p"}}
    store = SimpleNamespace(
        resolve_frame_scope=lambda fid: {"root_frame_id": "root-1"},
        team=SimpleNamespace(session_owner=lambda sid: owners.get(sid)),
    )
    child = Agent(
        use_skills=False, allow_delegate=False, frame_id="child-1", delegate_depth=1
    )
    child.dispatcher = SimpleNamespace(store=store)
    assert child._session_is_metered() is True
    owners.clear()
    assert child._session_is_metered() is False
    child.dispatcher = None  # no Store at all: unmetered, never an error
    assert child._session_is_metered() is False


# --- R Cell attempts name the R kernel generation ---------------------------
#
# The CLI bound a durable generation only to Python Cell attempts, although
# ``Agent._execute_r`` registers the R worker with the same recorder right
# before running the Cell. An R attempt therefore named no generation. It must
# name the R row, but only when a kernel produced the result: a spawn failure,
# a dead worker or a cancellation returns a synthesized dict without
# ``stdout``, and the recorder's R row may then still name the previous, dead
# worker — binding it would be provenance that is wrong rather than absent.


class _GenerationAttemptStore:
    def __init__(self) -> None:
        self.allocated = 0
        self.bound: list[tuple[str, str]] = []
        self.finished: list[tuple[str, str]] = []

    def allocate_execution_attempt(self, *, group_id, producing_cell_id):
        self.allocated += 1
        return {"attempt_id": f"attempt-{self.allocated}"}

    def mark_execution_attempt_started(self, attempt_id):
        pass

    def mark_execution_attempt_response(self, attempt_id):
        pass

    def mark_execution_attempt_capture(self, attempt_id):
        pass

    def bind_execution_attempt_generation(self, attempt_id, generation_id):
        self.bound.append((attempt_id, generation_id))

    def finish_execution_attempt(self, attempt_id, *, terminal_state):
        self.finished.append((attempt_id, terminal_state))


class _Recorder:
    """The KernelGenerationRecorder surface the executor reads."""

    def __init__(self, open_rows: dict[str, str]) -> None:
        self.open_rows = dict(open_rows)

    def observe(self, kernel, *, language="python"):
        return self.open_rows.get(language)

    def current(self, language="python"):
        return self.open_rows.get(language)


def _generation_executor(execute_r, recorder, kernel=None):
    store = _GenerationAttemptStore()
    executor = LocalActionExecutor(
        kernel or FakeKernel(),
        FakeDispatcher(),
        lambda code, messages: None,
        execute_r,
        action_ledger=SimpleNamespace(current_group_id="group-1", store=store),
        generation_recorder=recorder,
    )
    return executor, store


def _run_r(executor, state=None):
    executor.execute(
        CodeCell("r", "print(1 + 1)\n"),
        ModelReply(content="```r\nprint(1 + 1)\n```"),
        state if state is not None else RunState([]),
    )


def test_cli_r_cell_attempt_binds_the_r_kernel_generation():
    executor, store = _generation_executor(
        lambda code, **kwargs: {"stdout": "[1] 2\n", "stderr": "", "error": None},
        _Recorder({"r": "gen-r-live", "python": "gen-py"}),
    )

    _run_r(executor)

    assert store.bound == [("attempt-1", "gen-r-live")]
    assert store.finished == [("attempt-1", "completed")]


def test_cli_r_cell_error_result_from_the_kernel_still_names_its_generation():
    executor, store = _generation_executor(
        lambda code, **kwargs: {
            "stdout": "",
            "stderr": "",
            "error": "Error in stop('boom'): boom",
        },
        _Recorder({"r": "gen-r-live"}),
    )

    _run_r(executor)

    assert store.bound == [("attempt-1", "gen-r-live")]
    assert store.finished == [("attempt-1", "failed")]


@pytest.mark.parametrize(
    "synthesized",
    [
        {"error": "R kernel unavailable: Rscript not found"},
        {"error": "R kernel failed: kernel worker exited unexpectedly"},
        {"error": "Interrupted", "interrupted": True},
    ],
    ids=["spawn-failed", "worker-died", "cancelled-before-run"],
)
def test_cli_r_soft_error_never_binds_a_stale_generation(synthesized):
    executor, store = _generation_executor(
        lambda code, **kwargs: dict(synthesized),
        # The recorder's open R row still names the worker that died before.
        _Recorder({"r": "gen-r-dead"}),
    )

    _run_r(executor)

    assert store.bound == []
    assert len(store.finished) == 1


def test_cli_r_and_python_attempts_name_their_own_generations():
    class GenerationKernel(FakeKernel):
        generation = 3

        def execute(self, code, origin=None, cell_id=None):
            return super().execute(code, origin)

    executor, store = _generation_executor(
        lambda code, **kwargs: {"stdout": "[1] 2\n", "error": None},
        _Recorder({"r": "gen-r", "python": "gen-py"}),
        kernel=GenerationKernel(),
    )
    state = RunState([])

    _run_r(executor, state)
    executor.execute(
        CodeCell("python", "print(2 + 2)\n"),
        ModelReply(content="```python\nprint(2 + 2)\n```"),
        state,
    )

    assert [generation for _attempt, generation in store.bound] == ["gen-r", "gen-py"]


def test_cli_r_attempt_without_a_recorder_binds_nothing():
    executor, store = _generation_executor(
        lambda code, **kwargs: {"stdout": "[1] 2\n", "error": None},
        None,
    )

    _run_r(executor)

    assert store.bound == []


def _real_r_available() -> bool:
    from openai4s.kernel.r_kernel import resolve_r_interpreter

    return resolve_r_interpreter() is not None


@pytest.mark.skipif(
    not _real_r_available(), reason="no Rscript resolvable on this machine"
)
def test_cli_run_r_cell_attempt_names_the_durable_r_generation(monkeypatch, tmp_path):
    """End to end: a real R worker and a real Python worker in one CLI run."""
    import openai4s.agent.loop as loop_mod
    from openai4s.agent import Agent
    from openai4s.config import get_config
    from openai4s.store import get_store

    replies = [
        "```r\nprint(1 + 1)\n```",
        "```python\nprint(2 + 2)\n```",
        "```python\nhost.submit_output({'summary': 'ok'}, ['Printed 2 and 4'])\n```",
    ]

    def chat(messages, cfg, **kwargs):
        del messages, cfg, kwargs
        return {"content": replies.pop(0), "usage": {}}

    monkeypatch.setattr(loop_mod, "chat", chat)
    monkeypatch.chdir(tmp_path)
    agent = Agent(cfg=get_config(), use_skills=False, allow_delegate=False)
    result = agent.run("Run one R cell, then one Python cell")
    assert result["stop_reason"] == "submitted"

    store = get_store(get_config().db_path)
    r_generation = store.latest_kernel_generation(agent.frame_id, "r")
    py_generation = store.latest_kernel_generation(agent.frame_id, "python")
    assert r_generation is not None and py_generation is not None
    attempts = sorted(
        store.list_execution_attempts(root_frame_id=agent.frame_id),
        key=lambda attempt: attempt["allocated_at"],
    )
    assert len(attempts) == 3
    assert (
        attempts[0]["generation_id"] == r_generation["generation_id"]
    ), "the R cell's execution attempt names no kernel generation"
    assert attempts[1]["generation_id"] == py_generation["generation_id"]
    assert r_generation["generation_id"] != py_generation["generation_id"]