agentic-mobile-control/tests/test_host_agent_task_storage.py

"""Tests for Host Agent task storage, timeline, retention, and collision avoidance (task 1.5)."""

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from pathlib import Path

from core.models import Task
from host_agent.config import HostAgentConfig
from host_agent.retention import prune_task_history
from storage.artifact_store import ArtifactStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline
from tests.fakes import PNG_10X20


def _make_task(
    task_id: str,
    *,
    goal: str = "test goal",
    device_id: str = "iphone-1",
    created_at: datetime | None = None,
    status: str = "created",
) -> Task:
    ts = created_at or datetime(2026, 7, 1, tzinfo=UTC)
    return Task(
        id=task_id,
        goal=goal,
        device_id=device_id,
        status=status,  # type: ignore[arg-type]
        created_at=ts,
        updated_at=ts,
    )


# --------------------------------------------------------------------------- #
# Persist during execution
# --------------------------------------------------------------------------- #


def test_step_transitions_persist_during_execution(tmp_path) -> None:
    metadata_store = TaskMetadataStore(tmp_path / "tasks.sqlite3")
    timeline = Timeline(ArtifactStore(tmp_path / "history"))

    task = _make_task("task-persist")
    metadata_store.create_task(task)

    # Simulate the execution loop writing step transitions.
    metadata_store.update_task(task.id, status="running")
    timeline.append(
        task_id=task.id,
        scene={"screen": {"width": 1, "height": 1}, "elements": []},
        prompt="step 1",
        tool_call={"action": "tap"},
        result={"ok": True},
        screenshot=PNG_10X20,
    )
    timeline.append(
        task_id=task.id,
        scene={"screen": {"width": 1, "height": 1}, "elements": []},
        prompt="step 2",
        tool_call={"action": "swipe"},
        result={"ok": True},
        screenshot=PNG_10X20,
    )
    metadata_store.update_task(task.id, status="completed", completed=True)

    row = metadata_store.get_task(task.id)
    assert row is not None
    assert row["status"] == "completed"
    assert row["completed_at"] is not None

    records = timeline.read(task.id)
    assert len(records) == 2
    assert [r["index"] for r in records] == [1, 2]


# --------------------------------------------------------------------------- #
# Queryable after completion (in-memory Task discarded)
# --------------------------------------------------------------------------- #


def test_task_history_queryable_after_in_memory_task_discarded(tmp_path) -> None:
    db_path = tmp_path / "tasks.sqlite3"
    history_dir = tmp_path / "history"

    metadata_store = TaskMetadataStore(db_path)
    timeline = Timeline(ArtifactStore(history_dir))

    task = _make_task("task-survive")
    metadata_store.create_task(task)
    metadata_store.update_task(task.id, status="running")
    timeline.append(
        task_id=task.id,
        scene={"screen": {"width": 1, "height": 1}, "elements": []},
        prompt="do something",
        tool_call={"action": "tap"},
        result={"ok": True},
        screenshot=PNG_10X20,
    )
    metadata_store.update_task(task.id, status="completed", completed=True)

    # Drop ALL in-memory references.
    del metadata_store
    del timeline
    del task

    # Re-open from the same paths.
    reopened_store = TaskMetadataStore(db_path)
    reopened_timeline = Timeline(ArtifactStore(history_dir))

    row = reopened_store.get_task("task-survive")
    assert row is not None
    assert row["status"] == "completed"
    assert row["goal"] == "test goal"

    records = reopened_timeline.read("task-survive")
    assert len(records) == 1
    assert records[0]["tool_call"] == {"action": "tap"}


# --------------------------------------------------------------------------- #
# Retention prunes tasks beyond configured threshold
# --------------------------------------------------------------------------- #


def _seed_n_tasks(
    store: TaskMetadataStore,
    n: int,
    *,
    base: datetime = datetime(2026, 7, 1, tzinfo=UTC),
) -> list[str]:
    ids: list[str] = []
    for i in range(n):
        task = _make_task(f"task-{i:02d}", created_at=base + timedelta(hours=i))
        store.create_task(task)
        ids.append(task.id)
    return ids


def test_retention_prunes_by_max_count(tmp_path) -> None:
    metadata_store = TaskMetadataStore(tmp_path / "tasks.sqlite3")
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    ids = _seed_n_tasks(metadata_store, 10)

    config = HostAgentConfig(
        control_plane_url="https://control.example",
        task_retention_max_count=5,
        task_retention_max_age_days=365,  # large so count is the binding constraint
    )
    pruned = prune_task_history(
        metadata_store,
        timeline,
        config=config,
        now=datetime(2026, 7, 1, tzinfo=UTC) + timedelta(days=1),
    )
    assert pruned == 5
    remaining = metadata_store.list_task_ids()
    assert len(remaining) == 5
    # list_tasks is ordered by created_at desc, so the 5 newest survive.
    expected_survivors = set(ids[5:])
    assert set(remaining) == expected_survivors


def test_retention_prunes_by_max_age(tmp_path) -> None:
    metadata_store = TaskMetadataStore(tmp_path / "tasks.sqlite3")
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    base = datetime(2026, 7, 1, tzinfo=UTC)
    _seed_n_tasks(metadata_store, 10, base=base)

    config = HostAgentConfig(
        control_plane_url="https://control.example",
        task_retention_max_count=100,  # large so age is the binding constraint
        task_retention_max_age_days=3,
    )
    # "now" is 5 days after base; tasks 0..2 are within 3 days of now-5d... actually:
    # tasks created at base+0h .. base+9h.  now = base + 5 days.
    # cutoff = now - 3 days = base + 2 days.  Tasks with created_at >= base+2d survive.
    # Only tasks created at base+0h..base+9h -- all are < base+2d, so all pruned.
    # Let's use a now that's closer.
    now = base + timedelta(hours=5)
    pruned = prune_task_history(metadata_store, timeline, config=config, now=now)
    # cutoff = now - 3 days. All tasks are within 3 days. None pruned by age.
    # Actually max_count=100 so nothing pruned at all.
    # Let's redo: use max_age_days to actually prune older tasks.
    assert pruned == 0  # all tasks are younger than 3 days

    # Now use a stricter age that prunes some.
    config_strict = HostAgentConfig(
        control_plane_url="https://control.example",
        task_retention_max_count=100,
        task_retention_max_age_days=1,
    )
    # cutoff = now - 1 day = base + 5h - 24h = base - 19h.  Still all tasks are after that.
    # So we need a much later "now".
    late_now = base + timedelta(days=10)
    pruned2 = prune_task_history(
        metadata_store, timeline, config=config_strict, now=late_now
    )
    # All tasks are ~10 days old, max_age=1 day. All pruned.
    assert pruned2 == 10
    assert metadata_store.list_task_ids() == []


def test_retention_whichever_keeps_fewer_wins(tmp_path) -> None:
    """The more restrictive bound between max_count and max_age wins."""
    metadata_store = TaskMetadataStore(tmp_path / "tasks.sqlite3")
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    base = datetime(2026, 7, 1, tzinfo=UTC)
    _seed_n_tasks(metadata_store, 10, base=base)

    # max_count=8 would keep 8, max_age=2 days from base+3d would keep tasks
    # created within 2 days of base+3d = base+1d..base+9h. Only tasks at base+0h..base+9h.
    # cutoff = base+3d - 2d = base+1d = base+24h.  Tasks created at base+0h..base+9h
    # are all < base+24h, so keep_by_age = {} (0 tasks). keep_by_count = newest 8.
    # len(keep_by_age)=0 <= len(keep_by_count)=8, so keep_by_age wins -> 0 kept, all pruned.
    config = HostAgentConfig(
        control_plane_url="https://control.example",
        task_retention_max_count=8,
        task_retention_max_age_days=2,
    )
    now = base + timedelta(days=3)
    pruned = prune_task_history(metadata_store, timeline, config=config, now=now)
    assert pruned == 10
    assert metadata_store.list_task_ids() == []


def test_retention_deletes_timeline_alongside_metadata(tmp_path) -> None:
    metadata_store = TaskMetadataStore(tmp_path / "tasks.sqlite3")
    timeline = Timeline(ArtifactStore(tmp_path / "history"))

    # Create a task with timeline records.
    task = _make_task("task-to-prune")
    metadata_store.create_task(task)
    timeline.append(
        task_id=task.id,
        scene={"screen": {"width": 1, "height": 1}, "elements": []},
        prompt="step",
        tool_call={"action": "tap"},
        result={"ok": True},
        screenshot=PNG_10X20,
    )
    assert len(timeline.read(task.id)) == 1

    config = HostAgentConfig(
        control_plane_url="https://control.example",
        task_retention_max_count=0,
        task_retention_max_age_days=365,
    )
    pruned = prune_task_history(
        metadata_store,
        timeline,
        config=config,
        now=datetime(2026, 7, 1, tzinfo=UTC),
    )
    assert pruned == 1
    assert metadata_store.get_task(task.id) is None
    assert timeline.read(task.id) == []


# --------------------------------------------------------------------------- #
# No collision between Host Agent and Runtime default paths
# --------------------------------------------------------------------------- #


def test_host_agent_default_db_path_differs_from_runtime_default() -> None:
    host_agent_config = HostAgentConfig(control_plane_url="https://control.example")
    # Use posix paths for cross-platform comparison.
    assert host_agent_config.task_progress_db_path.as_posix() == (
        "host_agent_data/task_progress.sqlite3"
    )

    # Runtime's TaskMetadataStore default.
    runtime_default = TaskMetadataStore.__init__.__defaults__[0]
    assert Path(runtime_default).as_posix() == "tasks/tasks.sqlite3"

    assert host_agent_config.task_progress_db_path != Path(runtime_default)


def test_two_metadata_stores_at_different_paths_are_isolated(tmp_path) -> None:
    store_a = TaskMetadataStore(tmp_path / "a.sqlite3")
    store_b = TaskMetadataStore(tmp_path / "b.sqlite3")

    task_a = _make_task("task-in-a")
    store_a.create_task(task_a)

    assert store_a.get_task("task-in-a") is not None
    assert store_b.get_task("task-in-a") is None
    assert "task-in-a" not in store_b.list_task_ids()
    assert "task-in-a" in store_a.list_task_ids()