TaskTool child remains running and task mapping is lost after Agent Server process restart

"""Deterministic process-boundary reproducer for TaskTool child restart state.

This is reproduction equipment, not a recovery implementation.  The controller
starts a real ConversationService in a child process, waits until a native
TaskTool delegate is blocked inside a fake (zero-network) LLM, kills that process,
then starts a fresh ConversationService and performs exactly one native parent
run.  No task ID or child result is injected into the resumed parent.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import shutil
import subprocess
import sys
import time
from collections.abc import Sequence
from pathlib import Path
from typing import Any
from uuid import UUID

from pydantic import PrivateAttr

from openhands.agent_server.conversation_service import ConversationService
from openhands.agent_server.models import StartConversationRequest
from openhands.sdk import LLM, Agent, Tool
from openhands.sdk.conversation.state import ConversationExecutionStatus
from openhands.sdk.event import Event, ObservationEvent
from openhands.sdk.llm import Message, MessageToolCall, TextContent
from openhands.sdk.llm.llm_response import LLMResponse
from openhands.sdk.llm.streaming import TokenCallbackType
from openhands.sdk.subagent.registry import _reset_registry_for_tests, register_agent
from openhands.sdk.testing import TestLLM
from openhands.sdk.tool.tool import ToolDefinition
from openhands.sdk.workspace import LocalWorkspace
from openhands.tools.task import TaskToolSet
from openhands.tools.task.definition import TaskObservation


def write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(value, indent=2, sort_keys=True, default=str))
    tmp.replace(path)


def read_json(path: Path) -> Any:
    return json.loads(path.read_text())


def text_message(text: str) -> Message:
    return Message(role="assistant", content=[TextContent(text=text)])


def task_call() -> Message:
    return Message(
        role="assistant",
        content=[TextContent(text="Delegating to the worker.")],
        tool_calls=[
            MessageToolCall(
                id="restart-repro-task-call",
                name="task",
                arguments=json.dumps(
                    {
                        "prompt": "Remain active until the server process is restarted.",
                        "subagent_type": "blocking_worker",
                        "description": "restart lifecycle probe",
                    }
                ),
                origin="completion",
            )
        ],
    )


class BlockingTestLLM(TestLLM):
    """A deterministic fake LLM that records entry and never returns."""

    _marker: str = PrivateAttr(default="")

    def __init__(self, *, marker: str, **data: Any) -> None:
        super().__init__(**data)
        self._marker = marker

    def completion(
        self,
        messages: list[Message],
        tools: Sequence[ToolDefinition] | None = None,
        add_security_risk_prediction: bool = False,
        on_token: TokenCallbackType | None = None,
        call_context: Any = None,
        **kwargs: Any,
    ) -> LLMResponse:
        write_json(
            Path(self._marker),
            {
                "pid": os.getpid(),
                "entered_at_unix": time.time(),
                "message_count": len(messages),
                "tool_count": len(tools or []),
            },
        )
        while True:
            time.sleep(60)


def placeholder_llm(usage_id: str) -> LLM:
    return LLM(
        usage_id=usage_id,
        model="openai/gpt-4o",
    )


def register_blocking_worker(marker: Path) -> None:
    _reset_registry_for_tests()
    worker_llm = BlockingTestLLM(
        marker=str(marker),
        model="deterministic-blocking-test-model",
        usage_id="restart-repro-child",
        scripted_responses=[],
    )
    register_agent(
        name="blocking_worker",
        factory_func=lambda inherited_llm: Agent(llm=worker_llm, tools=[]),
        description="Deterministic worker used only by the restart reproducer",
    )


def event_summary(conversation_dir: Path) -> list[dict[str, Any]]:
    result: list[dict[str, Any]] = []
    events_dir = conversation_dir / "events"
    if not events_dir.exists():
        return result
    for path in sorted(
        events_dir.glob("event-*.json"),
        key=lambda p: int(p.name.split("-")[1]),
    ):
        event = Event.model_validate_json(path.read_text())
        row: dict[str, Any] = {
            "file": path.name,
            "type": type(event).__name__,
            "id": str(event.id),
            "parent_id": str(event.parent_id) if event.parent_id else None,
        }
        for name in ("tool_name", "tool_call_id", "action_id", "error"):
            value = getattr(event, name, None)
            if value is not None:
                row[name] = str(value)
        if isinstance(event, ObservationEvent) and isinstance(
            event.observation, TaskObservation
        ):
            row["task_observation"] = event.observation.model_dump(mode="json")
        result.append(row)
    return result


def locate_parent(conversations_dir: Path) -> tuple[Path, dict[str, Any]]:
    parents = [p for p in conversations_dir.iterdir() if (p / "meta.json").exists()]
    if len(parents) != 1:
        raise RuntimeError(f"expected one parent conversation, found {parents}")
    parent_dir = parents[0]
    return parent_dir, read_json(parent_dir / "base_state.json")


def locate_child(parent_dir: Path) -> tuple[Path, dict[str, Any]]:
    children = [
        p
        for p in (parent_dir / "subagents").iterdir()
        if (p / "base_state.json").exists()
    ]
    if len(children) != 1:
        raise RuntimeError(f"expected one child conversation, found {children}")
    child_dir = children[0]
    return child_dir, read_json(child_dir / "base_state.json")


def manager_snapshot(conversation: Any) -> dict[str, Any]:
    try:
        task_tool = conversation.agent.tools_map["task"]
        manager = task_tool.executor._manager
        return {
            "available": True,
            "task_ids": sorted(manager._tasks),
            "tasks": {
                task_id: {
                    "status": str(task.status),
                    "conversation_id": str(task.conversation_id),
                    "has_live_conversation": task.conversation is not None,
                    "result": task.result,
                    "error": task.error,
                }
                for task_id, task in manager._tasks.items()
            },
            "persistence_dir": str(manager._persistence_dir),
        }
    except Exception as exc:
        return {"available": False, "error": repr(exc)}


async def phase1(evidence: Path) -> None:
    conversations = evidence / "state" / "conversations"
    workspace = evidence / "workspace"
    workspace.mkdir(parents=True, exist_ok=True)
    marker = evidence / "child_llm_entered.json"
    register_blocking_worker(marker)

    parent_llm = TestLLM.from_messages(
        [task_call(), text_message("This response must not be reached before restart.")],
        model="deterministic-parent-test-model",
        usage_id="restart-repro-parent-before",
    )
    request = StartConversationRequest(
        agent=Agent(
            llm=placeholder_llm("restart-repro-placeholder"),
            tools=[Tool(name=TaskToolSet.name)],
        ),
        workspace=LocalWorkspace(working_dir=str(workspace)),
        autotitle=False,
    )

    # Do not use a context manager: the controller deliberately kills this
    # process, bypassing graceful EventService/TaskManager cleanup.
    service = ConversationService(conversations_dir=conversations)
    await service.__aenter__()
    info, _ = await service.start_conversation(request)
    event_service = await service.get_event_service(info.id)
    assert event_service is not None
    conversation = event_service.get_conversation()
    conversation.switch_llm(parent_llm)
    await event_service.send_message(
        Message(role="user", content=[TextContent(text="Run the delegated probe.")]),
        run=False,
    )
    write_json(
        evidence / "phase1_started.json",
        {"pid": os.getpid(), "parent_conversation_id": str(info.id)},
    )
    await event_service.run()

    # The child writes its marker only after native TaskTool construction and
    # entry into its fake LLM.  Snapshot the live in-memory registry before the
    # controller kills this process.
    while not marker.exists():
        await asyncio.sleep(0.01)
    write_json(
        evidence / "phase1_runtime_snapshot.json",
        {
            "parent_execution_status": conversation.state.execution_status.value,
            "manager": manager_snapshot(conversation),
        },
    )
    while True:
        await asyncio.sleep(60)


async def phase2(evidence: Path) -> None:
    conversations = evidence / "state" / "conversations"
    workspace = evidence / "workspace"
    marker = evidence / "phase2-unused-child-marker.json"
    register_blocking_worker(marker)

    parent_dir, parent_base = locate_parent(conversations)
    parent_id = UUID(parent_base["id"])
    async with ConversationService(conversations_dir=conversations) as service:
        recovered = await service.get_event_service(parent_id)
        assert recovered is not None
        conversation = recovered.get_conversation()
        # ConversationService.start() has already performed its stock crash
        # recovery. Install a deterministic parent response, as upstream tests do.
        conversation.switch_llm(
            TestLLM.from_messages(
                [text_message("Parent completed one native post-restart run.")],
                model="deterministic-parent-test-model",
                usage_id="restart-repro-parent-after",
            )
        )
        conversation._ensure_agent_ready()
        manager_before = manager_snapshot(conversation)
        status_before = conversation.state.execution_status.value
        events_before = event_summary(parent_dir)

        # The sole native resume action in this reproducer.
        await recovered.run()
        run_task = recovered._run_task
        if run_task is not None:
            await asyncio.wait_for(asyncio.shield(run_task), timeout=30)

        parent_dir, parent_after = locate_parent(conversations)
        child_dir, child_after = locate_child(parent_dir)
        write_json(
            evidence / "phase2_result.json",
            {
                "pid": os.getpid(),
                "native_resume_calls": 1,
                "parent_id": str(parent_id),
                "parent_status_before_resume": status_before,
                "parent_status_after_resume": parent_after["execution_status"],
                "manager_before_resume": manager_before,
                "manager_after_resume": manager_snapshot(conversation),
                "events_before_resume": events_before,
                "events_after_resume": event_summary(parent_dir),
                "child_id": child_after["id"],
                "child_status_after_resume": child_after["execution_status"],
                "child_observability_metadata": child_after.get(
                    "observability_metadata"
                ),
                "child_events_after_resume": event_summary(child_dir),
                "child_result_observations": [
                    row
                    for row in event_summary(child_dir)
                    if row["type"] in {"ObservationEvent", "AgentErrorEvent"}
                ],
            },
        )


def controller(evidence: Path) -> None:
    evidence.mkdir(parents=True, exist_ok=True)
    if any(evidence.iterdir()):
        raise RuntimeError(f"refusing to reuse non-empty evidence dir: {evidence}")

    phase1_log = (evidence / "phase1.log").open("w", encoding="utf-8")
    proc = subprocess.Popen(
        [sys.executable, __file__, "phase1", "--evidence", str(evidence)],
        stdout=phase1_log,
        stderr=subprocess.STDOUT,
        text=True,
    )
    try:
        marker = evidence / "phase1_runtime_snapshot.json"
        deadline = time.monotonic() + 60
        while time.monotonic() < deadline and not marker.exists():
            if proc.poll() is not None:
                raise RuntimeError(f"phase1 exited early with code {proc.returncode}")
            time.sleep(0.05)
        if not marker.exists():
            raise TimeoutError("active child registry snapshot was not written")

        parent_dir, parent_before = locate_parent(evidence / "state" / "conversations")
        child_dir, child_before = locate_child(parent_dir)
        write_json(
            evidence / "before_restart.json",
            {
                "phase1_pid": proc.pid,
                "parent_id": parent_before["id"],
                "parent_execution_status": parent_before["execution_status"],
                "task_registry": read_json(marker)["manager"],
                "parent_events": event_summary(parent_dir),
                "child_id": child_before["id"],
                "child_execution_status": child_before["execution_status"],
                "child_observability_metadata": child_before.get(
                    "observability_metadata"
                ),
                "child_events": event_summary(child_dir),
            },
        )
        proc.kill()
        proc.wait(timeout=10)
        write_json(
            evidence / "restart.json",
            {
                "method": "SIGKILL of the process owning ConversationService",
                "phase1_returncode": proc.returncode,
                "killed_at_unix": time.time(),
            },
        )
    finally:
        if proc.poll() is None:
            proc.kill()
            proc.wait(timeout=10)
        phase1_log.close()

    phase2_log = (evidence / "phase2.log").open("w", encoding="utf-8")
    resumed = subprocess.run(
        [sys.executable, __file__, "phase2", "--evidence", str(evidence)],
        stdout=phase2_log,
        stderr=subprocess.STDOUT,
        text=True,
        timeout=60,
        check=False,
    )
    phase2_log.close()
    write_json(evidence / "phase2_exit.json", {"returncode": resumed.returncode})
    if resumed.returncode != 0:
        raise RuntimeError(f"phase2 failed with code {resumed.returncode}")

    before = read_json(evidence / "before_restart.json")
    after = read_json(evidence / "phase2_result.json")
    write_json(
        evidence / "summary.json",
        {
            "before_restart": {
                "parent_status": before["parent_execution_status"],
                "child_status": before["child_execution_status"],
                "parent_task_observation_count": sum(
                    1 for row in before["parent_events"] if "task_observation" in row
                ),
            },
            "after_restart_before_resume": {
                "parent_status": after["parent_status_before_resume"],
                "manager_task_ids": after["manager_before_resume"].get("task_ids"),
                "parent_has_generic_recovery_error": any(
                    row["type"] == "AgentErrorEvent"
                    for row in after["events_before_resume"]
                ),
            },
            "after_one_native_resume": {
                "parent_status": after["parent_status_after_resume"],
                "child_status": after["child_status_after_resume"],
                "manager_task_ids": after["manager_after_resume"].get("task_ids"),
                "parent_task_observation_count": sum(
                    1
                    for row in after["events_after_resume"]
                    if "task_observation" in row
                ),
                "child_result_observation_count": len(
                    after["child_result_observations"]
                ),
                "native_resume_calls": after["native_resume_calls"],
            },
        },
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", choices=("controller", "phase1", "phase2"))
    parser.add_argument("--evidence", required=True, type=Path)
    args = parser.parse_args()
    if args.mode == "controller":
        controller(args.evidence)
    elif args.mode == "phase1":
        asyncio.run(phase1(args.evidence))
    else:
        asyncio.run(phase2(args.evidence))


if __name__ == "__main__":
    main()