nanobot/tests/webui/test_forking.py

from __future__ import annotations

import asyncio
import threading
from types import SimpleNamespace
from unittest.mock import ANY, AsyncMock, MagicMock

import pytest

import nanobot.webui.forking as forking
from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY


def test_create_fork_rebuilds_missing_transcript_and_saves_clean_title(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    forked = SimpleNamespace(messages=[{"role": "user", "content": "hello"}], metadata={})
    manager = MagicMock()
    manager.fork_session_before_user_index.return_value = forked
    rebuild = MagicMock()
    marker = MagicMock()
    monkeypatch.setattr(forking.uuid, "uuid4", lambda: "fork-id")
    monkeypatch.setattr(forking, "fork_transcript_before_user_index", lambda *_args: False)
    monkeypatch.setattr(forking, "write_session_messages_as_transcript", rebuild)
    monkeypatch.setattr(forking, "append_fork_marker", marker)

    result = forking.create_webui_chat_fork(
        manager,
        source_chat_id="source",
        before_user_index=2,
        title="  Useful fork  ",
    )

    assert result == ("fork-id", "websocket:fork-id")
    manager.fork_session_before_user_index.assert_called_once_with(
        "websocket:source",
        "websocket:fork-id",
        2,
    )
    rebuild.assert_called_once_with("websocket:fork-id", forked.messages)
    marker.assert_called_once_with("websocket:fork-id")
    assert forked.metadata[WEBUI_TITLE_METADATA_KEY] == "Useful fork"
    manager.save.assert_called_once_with(forked, fsync=True)


def test_create_fork_rolls_back_session_and_transcript_together(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    manager = MagicMock()
    manager.fork_session_before_user_index.return_value = SimpleNamespace(
        messages=[],
        metadata={},
    )
    delete_transcript = MagicMock()
    monkeypatch.setattr(forking.uuid, "uuid4", lambda: "failed-fork")
    monkeypatch.setattr(
        forking,
        "fork_transcript_before_user_index",
        MagicMock(side_effect=OSError("disk full")),
    )
    monkeypatch.setattr(forking, "delete_webui_transcript", delete_transcript)

    with pytest.raises(OSError, match="disk full"):
        forking.create_webui_chat_fork(
            manager,
            source_chat_id="source",
            before_user_index=1,
        )

    delete_transcript.assert_called_once_with("websocket:failed-fork")
    manager.delete_session.assert_called_once_with("websocket:failed-fork")


def test_create_fork_stops_before_transcript_work_when_source_is_invalid(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    manager = MagicMock()
    manager.fork_session_before_user_index.return_value = None
    fork_transcript = MagicMock()
    monkeypatch.setattr(forking, "fork_transcript_before_user_index", fork_transcript)

    result = forking.create_webui_chat_fork(
        manager,
        source_chat_id="missing",
        before_user_index=1,
    )

    assert result is None
    fork_transcript.assert_not_called()


@pytest.mark.parametrize(
    ("envelope", "detail"),
    [
        ({"source_chat_id": "bad/id", "before_user_index": 0}, "invalid source_chat_id"),
        ({"source_chat_id": "source", "before_user_index": True}, "invalid before_user_index"),
        ({"source_chat_id": "source", "before_user_index": -1}, "invalid before_user_index"),
    ],
)
@pytest.mark.asyncio
async def test_fork_handler_rejects_invalid_protocol_input(
    envelope: dict[str, object],
    detail: str,
) -> None:
    connection = object()
    channel = SimpleNamespace(
        send_webui_protocol_error=AsyncMock(),
        gateway=SimpleNamespace(session_manager=MagicMock()),
    )

    await forking.handle_webui_fork_chat(channel, connection, envelope)

    channel.send_webui_protocol_error.assert_awaited_once_with(connection, detail)


@pytest.mark.asyncio
async def test_fork_handler_reports_unavailable_session_manager() -> None:
    connection = object()
    channel = SimpleNamespace(
        send_webui_protocol_error=AsyncMock(),
        gateway=SimpleNamespace(session_manager=None),
    )

    await forking.handle_webui_fork_chat(
        channel,
        connection,
        {"source_chat_id": "source", "before_user_index": 0},
    )

    channel.send_webui_protocol_error.assert_awaited_once_with(
        connection,
        "session_manager_unavailable",
    )


@pytest.mark.asyncio
async def test_fork_handler_maps_invalid_source_and_internal_failure_to_stable_errors(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    connection = object()
    channel = SimpleNamespace(
        send_webui_protocol_error=AsyncMock(),
        gateway=SimpleNamespace(session_manager=MagicMock()),
        logger=SimpleNamespace(warning=MagicMock()),
    )
    envelope = {"source_chat_id": "source", "before_user_index": 0}
    monkeypatch.setattr(forking, "create_webui_chat_fork", lambda *_args, **_kwargs: None)

    await forking.handle_webui_fork_chat(channel, connection, envelope)
    channel.send_webui_protocol_error.assert_awaited_once_with(
        connection,
        "invalid fork source or index",
    )

    channel.send_webui_protocol_error.reset_mock()
    monkeypatch.setattr(
        forking,
        "create_webui_chat_fork",
        MagicMock(side_effect=RuntimeError("broken transcript")),
    )
    await forking.handle_webui_fork_chat(channel, connection, envelope)

    channel.logger.warning.assert_called_once_with("fork_chat failed: {}", ANY)
    channel.send_webui_protocol_error.assert_awaited_once_with(connection, "fork_chat_failed")


@pytest.mark.asyncio
async def test_fork_cancellation_waits_for_creation_and_client_attachment(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    started = threading.Event()
    release = threading.Event()
    connection = object()

    def blocked_create(*_args, **_kwargs) -> tuple[str, str]:
        started.set()
        assert release.wait(timeout=1)
        return "fork-id", "websocket:fork-id"

    channel = SimpleNamespace(
        send_webui_protocol_error=AsyncMock(),
        attach_webui_fork=AsyncMock(),
        gateway=SimpleNamespace(session_manager=object()),
        logger=SimpleNamespace(warning=MagicMock()),
    )
    monkeypatch.setattr(forking, "create_webui_chat_fork", blocked_create)
    task = asyncio.create_task(
        forking.handle_webui_fork_chat(
            channel,
            connection,
            {"source_chat_id": "source", "before_user_index": 0},
        )
    )
    assert await asyncio.to_thread(started.wait, 1)
    try:
        task.cancel()
        await asyncio.sleep(0)
        assert not task.done()
        channel.attach_webui_fork.assert_not_awaited()
    finally:
        release.set()

    with pytest.raises(asyncio.CancelledError):
        await asyncio.wait_for(task, timeout=1)

    channel.attach_webui_fork.assert_awaited_once_with(
        connection,
        fork_id="fork-id",
        fork_key="websocket:fork-id",
    )
    channel.send_webui_protocol_error.assert_not_awaited()