rune/tests/test_orchestration_conversation_state.py

"""Canonical ConversationState and TaskState (#18 / P0-02).

Covers the records (CS1), TaskStateRepository (CS2),
ConversationStateRepository (CS3), composition + send_task integration
(CS4), and the legacy operator.json migration + canary/docs/CI wiring
(CS5).
"""

from __future__ import annotations

import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

import pytest

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from rune import rune_operator  # noqa: E402
from rune.orchestration.atomic import AtomicStore  # noqa: E402
from rune.orchestration.composition import (  # noqa: E402
    AUTHORITY_KINDS,
    build_conversation_state,
    build_root,
    compose_conversation_state,
)
from rune.orchestration.conversation_state import (  # noqa: E402
    OPEN_TASK_STATUSES,
    ConversationStateRepository,
    IllegalTaskTransitionError,
    TaskStateRepository,
    migrate_operator_tasks,
)
from rune.orchestration.event_store import DurableEventStore  # noqa: E402
from rune.orchestration.records import (  # noqa: E402
    SCHEMA_VERSION,
    ConversationMessageRecord,
    TaskStateRecord,
    TASK_STATUSES,
)
from rune.tool_protocol import OperatorResponse, ToolEvent, ToolRequest  # noqa: E402


@pytest.fixture()
def repos(tmp_path: Path):
    """One event log + hot store with both repositories bound to them."""
    events = DurableEventStore(tmp_path / "events.jsonl")
    hot = AtomicStore(tmp_path / "hot")
    tasks = TaskStateRepository(events, hot)
    conv = ConversationStateRepository(events, hot, tasks=tasks)
    return SimpleNamespaceFixture(tasks=tasks, conv=conv, events=events, tmp_path=tmp_path)


class SimpleNamespaceFixture:
    def __init__(self, tasks, conv, events, tmp_path):
        self.tasks = tasks
        self.conv = conv
        self.events = events
        self.tmp_path = tmp_path

    def restart(self):
        """Simulate a process restart: fresh instances over the same disk."""
        hot = AtomicStore(self.tmp_path / "hot")
        tasks = TaskStateRepository(self.events, hot)
        conv = ConversationStateRepository(self.events, hot, tasks=tasks)
        return SimpleNamespaceFixture(tasks=tasks, conv=conv, events=self.events, tmp_path=self.tmp_path)


# ---------------------------------------------------------------- CS1


def test_task_state_record_round_trip_and_defaults():
    rec = TaskStateRecord(title="Fix the login bug", channel_id="task", assigned_model="qwen3:14b")
    assert rec.task_id.startswith("task_")
    assert rec.status == "created"
    assert rec.created_at and rec.updated_at  # non-empty ISO timestamps
    assert rec.parent_task_id is None
    assert rec.outcome_ref == ""
    assert rec.schema_version == SCHEMA_VERSION
    restored = TaskStateRecord.from_dict(rec.to_dict())
    assert restored == rec
    # dict boundary: JSON round trip preserves every field
    assert TaskStateRecord.from_dict(json.loads(json.dumps(rec.to_dict()))) == rec


def test_conversation_message_record_round_trip_and_defaults():
    msg = ConversationMessageRecord(
        role="user",
        content="hello rune",
        channel_id="task",
        task_id="task_abc123",
        token_estimate=2,
    )
    assert msg.msg_id.startswith("msg_")
    assert msg.ts  # non-empty ISO timestamp
    restored = ConversationMessageRecord.from_dict(msg.to_dict())
    assert restored == msg
    assert ConversationMessageRecord.from_dict(json.loads(json.dumps(msg.to_dict()))) == msg


def test_task_state_status_vocabulary():
    assert set(TASK_STATUSES) == {
        "created",
        "running",
        "awaiting_approval",
        "research_required",
        "completed",
        "failed",
        "cancelled",
        "superseded",
    }
    assert OPEN_TASK_STATUSES == {"created", "running", "awaiting_approval", "research_required"}
    assert TaskStateRecord().status in TASK_STATUSES


# ---------------------------------------------------------------- CS2


def test_task_state_create_mints_canonical_record(repos):
    rec = repos.tasks.create("Fix login", "task", model="qwen3:14b")
    assert rec.task_id.startswith("task_")
    assert rec.status == "created"
    assert rec.channel_id == "task"
    assert rec.assigned_model == "qwen3:14b"
    # the transition is durably logged
    assert "task_state.created" in [e["event_type"] for e in repos.events.replay()]
    # the hot checkpoint survives a restart
    reborn = repos.restart()
    assert reborn.tasks.get(rec.task_id).title == "Fix login"
    assert reborn.tasks.get_active("task").task_id == rec.task_id


def test_task_state_amend_keeps_identity_and_updates_title(repos):
    rec = repos.tasks.create("Fix login", "task")
    amended = repos.tasks.amend(rec.task_id, "Fix login and add tests")
    assert amended.task_id == rec.task_id  # same identity — follow-up, not a new task
    assert amended.title == "Fix login and add tests"
    assert repos.tasks.get(rec.task_id).title == "Fix login and add tests"
    assert "task_state.amended" in [e["event_type"] for e in repos.events.replay()]
    with pytest.raises(KeyError):
        repos.tasks.amend("task_missing", "nope")


def test_task_state_transition_follows_status_spine(repos):
    rec = repos.tasks.create("Ship it", "task")
    repos.tasks.transition(rec.task_id, "running")
    repos.tasks.transition(rec.task_id, "awaiting_approval")
    repos.tasks.transition(rec.task_id, "completed")
    assert repos.tasks.get(rec.task_id).status == "completed"
    assert repos.tasks.get_active("task") is None  # terminal status clears active
    with pytest.raises(IllegalTaskTransitionError):
        repos.tasks.transition(rec.task_id, "running")  # completed is terminal
    with pytest.raises(IllegalTaskTransitionError):
        repos.tasks.transition(rec.task_id, "bogus")  # unknown status
    with pytest.raises(KeyError):
        repos.tasks.transition("task_missing", "running")


def test_task_state_get_active_returns_open_task_and_clears_on_terminal(repos):
    assert repos.tasks.get_active("task") is None
    rec = repos.tasks.create("Work", "task")
    assert repos.tasks.get_active("task").task_id == rec.task_id
    repos.tasks.transition(rec.task_id, "failed")
    assert repos.tasks.get_active("task") is None
    # per-channel isolation
    other = repos.tasks.create("Other", "chat")
    assert repos.tasks.get_active("chat").task_id == other.task_id
    assert repos.tasks.get_active("task") is None


def test_task_state_get_returns_record_or_none(repos):
    assert repos.tasks.get("task_nope") is None
    rec = repos.tasks.create("Work", "task")
    assert repos.tasks.get(rec.task_id).task_id == rec.task_id


def test_task_state_list_open_filters_by_channel_and_status(repos):
    a = repos.tasks.create("A", "task")
    b = repos.tasks.create("B", "task")
    c = repos.tasks.create("C", "chat")
    repos.tasks.transition(b.task_id, "completed")
    assert [t.task_id for t in repos.tasks.list_open("task")] == [a.task_id]
    assert [t.task_id for t in repos.tasks.list_open("chat")] == [c.task_id]
    assert all(t.status in OPEN_TASK_STATUSES for t in repos.tasks.list_open("chat"))


# ---------------------------------------------------------------- CS3


def test_conversation_append_persists_message(repos):
    msg = ConversationMessageRecord(role="user", content="hi", channel_id="task", task_id="t1")
    stored = repos.conv.append(msg)
    assert stored.msg_id == msg.msg_id
    assert "conversation.message" in [e["event_type"] for e in repos.events.replay()]
    reborn = repos.restart()
    assert [m.msg_id for m in reborn.conv.window("task", n=10)] == [msg.msg_id]


def test_conversation_window_returns_last_n_chronological(repos):
    for i in range(12):
        repos.conv.append(ConversationMessageRecord(role="user", content=f"m{i}", channel_id="task"))
    win = repos.conv.window("task", n=8)
    assert [m.content for m in win] == [f"m{i}" for i in range(4, 12)]
    assert repos.conv.window("task") == win  # default n=8
    repos.conv.append(ConversationMessageRecord(role="user", content="chat-only", channel_id="chat"))
    assert all(m.channel_id == "task" for m in repos.conv.window("task", n=100))
    assert repos.conv.window("task", n=0) == []


def test_conversation_messages_for_task_filters(repos):
    repos.conv.append(ConversationMessageRecord(role="user", content="a", channel_id="task", task_id="t1"))
    repos.conv.append(ConversationMessageRecord(role="assistant", content="b", channel_id="task", task_id="t1"))
    repos.conv.append(ConversationMessageRecord(role="user", content="c", channel_id="task", task_id="t2"))
    assert [m.content for m in repos.conv.messages_for_task("t1")] == ["a", "b"]
    assert repos.conv.messages_for_task("missing") == []


def test_conversation_latest_task_returns_most_recent(repos):
    first = repos.tasks.create("First", "task")
    assert repos.conv.latest_task("task").task_id == first.task_id
    second = repos.tasks.create("Second", "task")
    assert repos.conv.latest_task("task").task_id == second.task_id
    assert repos.conv.latest_task("chat") is None


# ---------------------------------------------------------------- CS4


def test_authority_kinds_declare_conversation_and_task_state():
    assert "conversation_state" in AUTHORITY_KINDS
    assert "task_state" in AUTHORITY_KINDS


def test_build_conversation_state_wires_single_authorities(tmp_path: Path):
    paths = {
        "events": tmp_path / "e.jsonl",
        "evidence": tmp_path / "ev",
        "state": tmp_path / "state",
    }
    root = build_conversation_state(paths)
    root.assert_single_authority()
    tasks = root.get("task_state")
    conv = root.get("conversation_state")
    assert isinstance(tasks, TaskStateRepository)
    assert isinstance(conv, ConversationStateRepository)
    # one shared durable log carries both ledgers
    rec = tasks.create("Work", "task")
    conv.append(ConversationMessageRecord(role="user", content="go", channel_id="task", task_id=rec.task_id))
    assert len(root.get("event_store").replay()) == 2
    # compose_conversation_state also wires into an existing root
    root2 = build_root({"events": tmp_path / "e2.jsonl", "evidence": tmp_path / "ev2"})
    auth = compose_conversation_state(root2, tmp_path / "state2")
    root2.assert_single_authority()
    assert auth.task_state is root2.get("task_state")
    assert auth.conversation is root2.get("conversation_state")
    assert auth.conversation.latest_task("task") is None


def _completed_write_event():
    from rune.tool_protocol import ToolEvent, ToolRequest

    return ToolEvent(
        kind="completed",
        request=ToolRequest(
            call_id="w1",
            name="write",
            arguments={"path": "/tmp/x", "content": "x"},
            source="model",
        ),
        result={"success": True, "result": "wrote /tmp/x"},
    )


class SendTaskIntegrationTests(unittest.TestCase):
    def setUp(self) -> None:
        self.tempdir = tempfile.TemporaryDirectory()
        self.tmp = Path(self.tempdir.name)
        self.operator = rune_operator.RuneOperator(str(self.tmp / "operator.json"))

    def tearDown(self) -> None:
        self.tempdir.cleanup()

    @staticmethod
    def pending_response() -> OperatorResponse:
        request = ToolRequest(
            call_id="call-1",
            name="write_file",
            arguments={"path": "/tmp/a", "content": "new"},
            source="ollama-native",
        )
        return OperatorResponse(
            success=True,
            tool_events=[ToolEvent("pending_approval", request)],
            pending_approval=request,
            pending_state={"messages": []},
            model="qwen3:8b",
        )

    def test_send_task_amends_active_task_on_follow_up(self):
        with patch.object(self.operator, "_execute_on_model", return_value=self.pending_response()):
            first = self.operator.send_task("change the login flow", "qwen3:8b")

        active = self.operator._task_state.get_active("task")
        assert active is not None
        assert active.task_id == first["task"]["task_id"]
        assert active.status == "awaiting_approval"

        # Completion-truth contract: mutation tasks need mutating tool
        # evidence to terminate completed (2026-09-10).
        done = OperatorResponse(
            success=True,
            user_text="Done.",
            model="qwen3:8b",
            tool_events=[_completed_write_event()],
        )
        with patch.object(self.operator, "_execute_on_model", return_value=done):
            second = self.operator.send_task("yes continue with that", "qwen3:8b")

        # follow-up AMENDED the active task instead of minting a new one
        assert second["task"]["task_id"] == first["task"]["task_id"]
        assert self.operator._task_state.get(second["task"]["task_id"]).status == "completed"
        assert self.operator._task_state.get_active("task") is None
        messages = self.operator._conversation.messages_for_task(first["task"]["task_id"])
        assert [(m.role, m.content) for m in messages] == [
            ("user", "change the login flow"),
            ("user", "yes continue with that"),
            ("assistant", "Done."),
        ]

    def test_send_task_creates_new_task_with_parent_lineage_and_messages(self):
        done = OperatorResponse(
            success=True,
            user_text="All done.",
            model="qwen3:8b",
            tool_events=[_completed_write_event()],
        )
        with patch.object(self.operator, "_execute_on_model", return_value=done):
            first = self.operator.send_task("write the parser", "qwen3:8b")
            second = self.operator.send_task("now write its tests", "qwen3:8b")

        assert second["task"]["task_id"] != first["task"]["task_id"]
        second_rec = self.operator._task_state.get(second["task"]["task_id"])
        assert second_rec.parent_task_id == first["task"]["task_id"]
        assert second_rec.status == "completed"
        assert self.operator._task_state.get_active("task") is None
        window = self.operator._conversation.window("task", n=10)
        assert [(m.role, m.content) for m in window] == [
            ("user", "write the parser"),
            ("assistant", "All done."),
            ("user", "now write its tests"),
            ("assistant", "All done."),
        ]
        assert window[2].task_id == second["task"]["task_id"]


# ---------------------------------------------------------------- CS5


def test_migrate_operator_tasks_imports_legacy_ledger(repos):
    legacy = {
        "identity": "Rune",
        "current_task": "cc33dd44",
        "tasks": [
            {
                "task_id": "aa11bb22",
                "description": "hello",
                "status": "completed",
                "assigned_model": "qwen2.5:7b",
                "result": "hi there",
                "created_at": "2026-08-22T00:25:02+00:00",
                "completed_at": "2026-08-22T00:25:03+00:00",
            },
            {
                "task_id": "cc33dd44",
                "description": "running thing",
                "status": "running",
                "assigned_model": "m",
                "result": None,
                "created_at": "2026-08-22T01:00:00+00:00",
                "completed_at": None,
            },
            {
                "task_id": "ee55ff66",
                "description": "broken thing",
                "status": "failed",
                "assigned_model": "m",
                "result": "boom",
                "created_at": "2026-08-22T02:00:00+00:00",
                "completed_at": None,
            },
        ],
    }
    path = repos.tmp_path / "operator.json"
    path.write_text(json.dumps(legacy))

    report = migrate_operator_tasks(path, repos.tasks, repos.conv)
    assert report.imported == 3
    assert report.skipped == 0
    assert report.messages == 4  # 3 user + 1 assistant (completed result only)
    assert repos.tasks.get("aa11bb22").status == "completed"
    assert repos.tasks.get("cc33dd44").status == "running"
    assert repos.tasks.get("ee55ff66").status == "failed"
    assert repos.tasks.get("aa11bb22").channel_id == "task"
    # legacy current_task pointer is honored
    assert repos.tasks.get_active("task").task_id == "cc33dd44"
    # conversation history is reconstructed, nothing lost
    assert [(m.role, m.content) for m in repos.conv.messages_for_task("aa11bb22")] == [
        ("user", "hello"),
        ("assistant", "hi there"),
    ]
    assert [m.content for m in repos.conv.messages_for_task("ee55ff66")] == ["broken thing"]

    # idempotent: a second pass imports nothing
    report2 = migrate_operator_tasks(path, repos.tasks, repos.conv)
    assert report2.imported == 0
    assert report2.skipped == 3


def test_operator_auto_migrates_legacy_state_on_init(tmp_path: Path):
    legacy = {
        "identity": "Rune",
        "tasks": [
            {
                "task_id": "zz99",
                "description": "old task",
                "status": "completed",
                "assigned_model": "m",
                "result": "done",
                "created_at": "2026-08-22T00:00:00+00:00",
                "completed_at": "2026-08-22T00:01:00+00:00",
            }
        ],
    }
    (tmp_path / "operator.json").write_text(json.dumps(legacy))
    operator = rune_operator.RuneOperator(str(tmp_path / "operator.json"))
    assert operator._task_state.get("zz99") is not None
    assert operator._task_state.get("zz99").status == "completed"
    assert operator._conversation.messages_for_task("zz99")[0].content == "old task"
    # module-level history reader works against the installed singleton
    with patch.object(rune_operator, "_operator", operator):
        assert rune_operator.get_conversation_window("task", n=8)[0]["content"] == "old task"


def test_canary_script_runs_clean():
    result = subprocess.run(
        [sys.executable, str(ROOT / "src/rune/canary_conversation_state.py")],
        capture_output=True,
        text=True,
        timeout=60,
    )
    assert result.returncode == 0, result.stderr
    assert "CONVERSATION STATE CANARY PASSED" in result.stdout


def test_docs_exist_and_ci_wiring():
    docs = (ROOT / "docs/conversation-state.md").read_text()
    assert "src/rune/orchestration/conversation_state.py" in docs
    assert "pytest" in docs
    ci = (ROOT / ".gitea/workflows/ci.yml").read_text()
    assert "canary_conversation_state.py" in ci


def test_self_edges_allow_idempotent_reassertion(repos):
    # P0-02 quality B1+B2 discharge: chained approvals re-enter
    # awaiting_approval; amend-while-running re-enters running. Self-edges
    # make both idempotent instead of crashing with IllegalTaskTransition.
    repo = repos.tasks
    task = repo.create("test task", channel_id="task", model="m")
    repo.transition(task.task_id, "running")
    # Re-assert running (B2: amend-while-running)
    repo.transition(task.task_id, "running")
    assert repo.get(task.task_id).status == "running"
    # Re-enter awaiting_approval (B1: chained approval)
    repo.transition(task.task_id, "awaiting_approval")
    repo.transition(task.task_id, "awaiting_approval")
    assert repo.get(task.task_id).status == "awaiting_approval"
    # Back to running then awaiting again (normal multi-tool loop)
    repo.transition(task.task_id, "running")
    repo.transition(task.task_id, "awaiting_approval")
    assert repo.get(task.task_id).status == "awaiting_approval"


def test_terminal_states_still_reject_all(repos):
    # Self-edges only apply to non-terminal states.
    import pytest
    from rune.orchestration.conversation_state import IllegalTaskTransitionError

    repo = repos.tasks
    task = repo.create("done task", channel_id="task", model="m")
    repo.transition(task.task_id, "completed")
    with pytest.raises(IllegalTaskTransitionError):
        repo.transition(task.task_id, "running")
    with pytest.raises(IllegalTaskTransitionError):
        repo.transition(task.task_id, "completed")