OpenOPC/tests/test_stale_checkpoint_selfheal.py

"""Regression tests: orphaned task-wait checkpoints must never swallow user messages.

Reproduces the production failure where a company-mode review left a pending
``task_user_input`` checkpoint behind (the runtime carried the flow forward via
approval-card grants and a fresh review attempt, never replying through the
checkpoint). The user's next chat message in the primary session was captured
by that orphan row, routed into the deprecated multi-agent resume path with an
empty task list, and answered with an empty string.
"""

from __future__ import annotations

import tempfile
import unittest
import uuid
from pathlib import Path
from unittest.mock import AsyncMock

from opc.core.models import (
    DelegationWorkItem,
    ExecutionCheckpoint,
    ExecutionMode,
    Phase,
    Task,
    TaskStatus,
)
from opc.engine import OPCEngine


class StaleCheckpointSelfHealTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self) -> None:
        self._tmp = tempfile.TemporaryDirectory()
        self.engine = OPCEngine(opc_home=Path(self._tmp.name), project_id="default")
        await self.engine.initialize()
        self.store = self.engine.store

    async def asyncTearDown(self) -> None:
        await self.store.close()
        self._tmp.cleanup()

    async def _save_task(self, status: TaskStatus, session_id: str) -> Task:
        task = Task(
            title="Review #1: competitive analysis",
            session_id=session_id,
            project_id="default",
            status=status,
        )
        await self.store.save_task(task)
        return task

    async def _save_wait_checkpoint(self, task: Task, session_id: str, **payload_extra) -> ExecutionCheckpoint:
        checkpoint = ExecutionCheckpoint(
            project_id="default",
            session_id=session_id,
            checkpoint_type="task_user_input",
            task_id=task.id,
            payload={
                "task_id": task.id,
                "session_id": session_id,
                "execution_mode": ExecutionMode.COMPANY_MODE.value,
                "task_ids": [task.id],
                "prompt": "Tool execution blocked by autonomy policy",
                "review_level": "human",
                **payload_extra,
            },
        )
        await self.store.save_execution_checkpoint(checkpoint)
        return checkpoint

    async def _checkpoint_status(self, checkpoint_id: str) -> str:
        rows = await self.store.get_execution_checkpoints(project_id="default")
        for row in rows:
            if row.checkpoint_id == checkpoint_id:
                return str(row.status or "")
        return "<not found>"

    async def test_orphan_checkpoint_is_lazily_resolved_and_not_returned(self) -> None:
        session_id = str(uuid.uuid4())
        task = await self._save_task(TaskStatus.DONE, session_id)
        checkpoint = await self._save_wait_checkpoint(task, session_id)

        found = await self.engine.get_latest_pending_checkpoint_for_session(session_id)

        self.assertIsNone(found)
        self.assertEqual(await self._checkpoint_status(checkpoint.checkpoint_id), "stale")

    async def test_waiting_checkpoint_is_still_returned(self) -> None:
        session_id = str(uuid.uuid4())
        task = await self._save_task(TaskStatus.AWAITING_HUMAN, session_id)
        checkpoint = await self._save_wait_checkpoint(task, session_id)

        found = await self.engine.get_latest_pending_checkpoint_for_session(session_id)

        self.assertIsNotNone(found)
        self.assertEqual(found.checkpoint_id, checkpoint.checkpoint_id)
        self.assertEqual(await self._checkpoint_status(checkpoint.checkpoint_id), "pending")

    async def test_user_message_falls_through_to_normal_processing(self) -> None:
        """A plain follow-up question must not be consumed by an orphan checkpoint."""
        session_id = str(uuid.uuid4())
        task = await self._save_task(TaskStatus.DONE, session_id)
        checkpoint = await self._save_wait_checkpoint(task, session_id)

        result = await self.engine._maybe_resume_checkpoint(
            "你的交付文件在哪里？", session_id=session_id
        )

        self.assertIsNone(result)  # None → caller processes it as a fresh message
        self.assertEqual(await self._checkpoint_status(checkpoint.checkpoint_id), "stale")

    async def test_explicit_reply_to_dead_checkpoint_reports_inactive(self) -> None:
        session_id = str(uuid.uuid4())
        task = await self._save_task(TaskStatus.DONE, session_id)
        checkpoint = await self._save_wait_checkpoint(task, session_id)

        result = await self.engine._maybe_resume_checkpoint(
            "approve",
            session_id=session_id,
            reply_metadata={
                "response_to_checkpoint_id": checkpoint.checkpoint_id,
                "response_to_checkpoint_type": "task_user_input",
            },
        )

        self.assertEqual(result, "This request is no longer active.")
        self.assertEqual(await self._checkpoint_status(checkpoint.checkpoint_id), "stale")

    async def _link_work_item(self, task: Task, phase: Phase) -> None:
        item = DelegationWorkItem(
            work_item_id=f"review::{uuid.uuid4()}::v1",
            title="Review #1",
            phase=phase,
        )
        await self.store.save_delegation_work_item(item)
        await self.store.link_work_item_runtime_task(item.work_item_id, task.id)

    async def test_checkpoint_with_closed_work_item_is_stale_even_if_task_not_terminal(self) -> None:
        """Production case: the swallowed resume left the task at PENDING, but the
        review attempt's work item had long been approved — the checkpoint is dead."""
        session_id = str(uuid.uuid4())
        task = await self._save_task(TaskStatus.PENDING, session_id)
        await self._link_work_item(task, Phase.APPROVED)
        checkpoint = await self._save_wait_checkpoint(task, session_id)

        found = await self.engine.get_latest_pending_checkpoint_for_session(session_id)

        self.assertIsNone(found)
        self.assertEqual(await self._checkpoint_status(checkpoint.checkpoint_id), "stale")

    async def test_checkpoint_with_open_work_item_is_kept(self) -> None:
        session_id = str(uuid.uuid4())
        task = await self._save_task(TaskStatus.RUNNING, session_id)
        await self._link_work_item(task, Phase.RUNNING)
        checkpoint = await self._save_wait_checkpoint(task, session_id)

        found = await self.engine.get_latest_pending_checkpoint_for_session(session_id)

        self.assertIsNotNone(found)
        self.assertEqual(found.checkpoint_id, checkpoint.checkpoint_id)
        self.assertEqual(await self._checkpoint_status(checkpoint.checkpoint_id), "pending")

    async def test_task_settling_supersedes_pending_wait_checkpoints(self) -> None:
        session_id = str(uuid.uuid4())
        task = await self._save_task(TaskStatus.AWAITING_HUMAN, session_id)
        checkpoint = await self._save_wait_checkpoint(task, session_id)

        await self.engine._supersede_stale_task_wait_checkpoints(task.id, reason="test settle")

        self.assertEqual(await self._checkpoint_status(checkpoint.checkpoint_id), "superseded")

    async def test_resume_with_unresolvable_siblings_never_returns_empty(self) -> None:
        """Sibling ids that are work-item ids (not task UUIDs) must not empty the task list."""
        session_id = str(uuid.uuid4())
        task = await self._save_task(TaskStatus.AWAITING_HUMAN, session_id)
        checkpoint = await self._save_wait_checkpoint(
            task,
            session_id,
            task_ids=["review::467f36ff::v1"],  # work-item id, unresolvable as a task
            company_work_item_plan=None,
        )

        self.engine._execute_single_agent = AsyncMock(return_value="resumed reply")
        self.engine._execute_multi_agent = AsyncMock(return_value="")
        self.engine._execute_company_mode = AsyncMock(return_value="")

        result = await self.engine._resume_task_checkpoint(checkpoint, "继续")

        self.assertEqual(result, "resumed reply")
        self.engine._execute_multi_agent.assert_not_awaited()
        self.engine._execute_company_mode.assert_not_awaited()
        (called_tasks, _agent), _ = self.engine._execute_single_agent.await_args
        self.assertEqual([t.id for t in called_tasks], [task.id])
        self.assertEqual(await self._checkpoint_status(checkpoint.checkpoint_id), "resolved")

    async def test_resume_with_plan_routes_to_company_mode(self) -> None:
        session_id = str(uuid.uuid4())
        task = await self._save_task(TaskStatus.AWAITING_HUMAN, session_id)
        checkpoint = await self._save_wait_checkpoint(
            task,
            session_id,
            company_work_item_plan={"profile": "corporate"},
        )

        self.engine._execute_company_mode = AsyncMock(return_value="company resumed")
        self.engine._execute_multi_agent = AsyncMock(return_value="")

        result = await self.engine._resume_task_checkpoint(checkpoint, "继续")

        self.assertEqual(result, "company resumed")
        self.engine._execute_multi_agent.assert_not_awaited()


if __name__ == "__main__":
    unittest.main()