"""Tests for conversation fork-from-assistant-message.""" from __future__ import annotations from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from octop.infra.agents.threads.fork import ( _fork_title, find_assistant_fork_index, fork_dashboard_thread, write_checkpoint_messages, ) from octop.infra.errors import ErrorCode, OctopError from octop.infra.gateway.threads import ThreadRegistry def test_fork_title_suffix() -> None: assert _fork_title("My chat", "en") == "My chat (fork)" assert _fork_title("My chat", "zh") == "My chat（分叉）" assert _fork_title(None, "en") == "Forked chat" assert _fork_title("", "zh") == "分叉对话" def test_find_assistant_fork_index_by_suffix_count() -> None: messages = [ HumanMessage(content="one", id="h1"), AIMessage(content="a1", id="a1"), HumanMessage(content="two", id="h2"), AIMessage(content="a2", id="a2"), HumanMessage(content="three", id="h3"), AIMessage(content="a3", id="a3"), ] assert find_assistant_fork_index(messages, message_id="x", assistant_turns_from_end=2) == 3 assert find_assistant_fork_index(messages, message_id="x", assistant_turns_from_end=1) == 5 assert find_assistant_fork_index(messages, message_id="x", assistant_turns_from_end=3) == 1 def test_find_assistant_fork_index_skips_tool_call_shells_even_with_text() -> None: messages = [ HumanMessage(content="q", id="h1"), AIMessage( content="let me search", id="a-tool", tool_calls=[{"id": "c1", "name": "search", "args": {}}], ), ToolMessage(content="hit", id="t1", tool_call_id="c1"), AIMessage(content="final answer", id="a-final"), ] # Intermediate tool-call AI must not count as an answer turn. assert find_assistant_fork_index(messages, message_id="x", assistant_turns_from_end=1) == 3 assert find_assistant_fork_index(messages, assistant_turns_from_end=1) == 3 def test_find_assistant_fork_index_trusts_turns_when_content_mismatches() -> None: messages = [ HumanMessage(content="keep", id="h1"), AIMessage(content="first", id="a1"), HumanMessage(content="next", id="h2"), AIMessage(content="fork-me", id="a2"), ] # Client UUID + drifted content must not override the preferred locator. idx = find_assistant_fork_index( messages, message_id="client-uuid", content="other text", assistant_turns_from_end=1, ) assert idx == 3 def test_find_assistant_fork_index_by_content_fallback() -> None: messages = [ HumanMessage(content="alpha", id="h1"), AIMessage(content="reply-a", id="a1"), HumanMessage(content="beta", id="h2"), AIMessage(content="reply-b", id="a2"), ] assert find_assistant_fork_index(messages, message_id="missing", content="reply-b") == 3 def test_find_assistant_fork_index_missing_raises() -> None: messages = [HumanMessage(content="only user", id="h1")] with pytest.raises(OctopError) as exc: find_assistant_fork_index(messages, message_id="a1") assert exc.value.code is ErrorCode.NOT_FOUND @pytest.mark.asyncio async def test_write_checkpoint_messages_skips_empty() -> None: graph = SimpleNamespace(aupdate_state=AsyncMock()) harness = SimpleNamespace(graph=graph) await write_checkpoint_messages(harness, thread_id="thr_new", messages=[]) graph.aupdate_state.assert_not_called() @pytest.mark.asyncio async def test_fork_dashboard_thread_copies_through_assistant_and_rebinds() -> None: source_id = "thr_src" messages = [ HumanMessage(content="first", id="h1"), AIMessage(content="ok", id="a1"), ToolMessage(content="tool", id="t1", tool_call_id="c1"), HumanMessage(content="second", id="h2"), AIMessage(content="later", id="a2"), ] store: dict[str, list[Any]] = {source_id: list(messages)} graph = SimpleNamespace() async def aget_state(config: dict[str, Any]) -> SimpleNamespace: tid = config["configurable"]["thread_id"] return SimpleNamespace(values={"messages": list(store.get(tid, []))}) async def aupdate_state(config: dict[str, Any], values: dict[str, Any], **_k: Any) -> None: tid = config["configurable"]["thread_id"] store.setdefault(tid, []).extend(list(values.get("messages") or [])) graph.aget_state = aget_state graph.aupdate_state = aupdate_state harness = SimpleNamespace(graph=graph) source = MagicMock( thread_id=source_id, agent_id="agt_1", title="Original title", model_ref="openai/gpt-4o", reasoning_mode="auto", reasoning_effort=None, ) created: dict[str, Any] = {} async def rebind(**kwargs: Any) -> None: created["rebind"] = kwargs registry = MagicMock() registry.create_thread.return_value = "thr_fork" registry.get_thread.return_value = MagicMock( title="Original title", last_active=9, created_at=8, ) registry.rebind = rebind out = await fork_dashboard_thread( thread_registry=registry, harness=harness, source=source, user_id=7, message_id="a1", content="ok", assistant_turns_from_end=2, ) assert out["thread_id"] == "thr_fork" assert out["source_thread_id"] == source_id assert out["copied_messages"] == 2 assert [m.id for m in store["thr_fork"]] == ["h1", "a1"] registry.update_title.assert_called_once_with("thr_fork", "Original title (fork)") registry.touch_last_active.assert_not_called() registry.update_composer.assert_called_once() assert created["rebind"]["thread_id"] == "thr_fork" assert created["rebind"]["session_key"] == ThreadRegistry.dashboard_key( agent_id="agt_1", user_id=7 ) @pytest.mark.asyncio async def test_fork_dashboard_thread_from_latest_assistant() -> None: source_id = "thr_src" messages = [ HumanMessage(content="only", id="h1"), AIMessage(content="answer", id="a1"), ] store: dict[str, list[Any]] = {source_id: list(messages)} graph = SimpleNamespace() async def aget_state(config: dict[str, Any]) -> SimpleNamespace: tid = config["configurable"]["thread_id"] return SimpleNamespace(values={"messages": list(store.get(tid, []))}) async def aupdate_state(config: dict[str, Any], values: dict[str, Any], **_k: Any) -> None: tid = config["configurable"]["thread_id"] store.setdefault(tid, []).extend(list(values.get("messages") or [])) graph.aget_state = aget_state graph.aupdate_state = aupdate_state harness = SimpleNamespace(graph=graph) registry = MagicMock() registry.create_thread.return_value = "thr_fork" registry.get_thread.return_value = MagicMock(title="only", last_active=0, created_at=1) registry.rebind = AsyncMock() out = await fork_dashboard_thread( thread_registry=registry, harness=harness, source=MagicMock( thread_id=source_id, agent_id="agt_1", title="only", model_ref=None, reasoning_mode=None, reasoning_effort=None, ), user_id=1, assistant_turns_from_end=1, ) assert out["copied_messages"] == 2 assert [m.id for m in store["thr_fork"]] == ["h1", "a1"] registry.update_title.assert_called_once() @pytest.mark.asyncio async def test_fork_dashboard_thread_deletes_row_when_write_fails() -> None: source_id = "thr_src" messages = [ HumanMessage(content="first", id="h1"), AIMessage(content="ok", id="a1"), HumanMessage(content="second", id="h2"), AIMessage(content="later", id="a2"), ] graph = SimpleNamespace() async def aget_state(_config: dict[str, Any]) -> SimpleNamespace: return SimpleNamespace(values={"messages": list(messages)}) graph.aget_state = aget_state graph.aupdate_state = AsyncMock(side_effect=RuntimeError("checkpoint write failed")) harness = SimpleNamespace(graph=graph) registry = MagicMock() registry.create_thread.return_value = "thr_fork" registry.rebind = AsyncMock() with pytest.raises(RuntimeError, match="checkpoint write failed"): await fork_dashboard_thread( thread_registry=registry, harness=harness, source=MagicMock( thread_id=source_id, agent_id="agt_1", title="t", model_ref=None, reasoning_mode=None, reasoning_effort=None, ), user_id=1, message_id="a2", assistant_turns_from_end=1, ) registry.delete_thread.assert_called_once_with("thr_fork")