343 lines
12 KiB
Python
343 lines
12 KiB
Python
"""Tests for Host Agent task storage, timeline, retention, and collision avoidance (task 1.5)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
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]
|
|
|
|
|
|
def test_task_metadata_creation_is_idempotent_and_keeps_source_correlation(
|
|
tmp_path,
|
|
) -> None:
|
|
metadata_store = TaskMetadataStore(tmp_path / "tasks.sqlite3")
|
|
task = _make_task("task-source")
|
|
|
|
metadata_store.create_task(
|
|
task,
|
|
source_task_id="cloud-task-source",
|
|
source_attempt=2,
|
|
)
|
|
metadata_store.create_task(task)
|
|
|
|
row = metadata_store.get_task(task.id)
|
|
assert row is not None
|
|
assert row["source_task_id"] == "cloud-task-source"
|
|
assert row["source_attempt"] == 2
|
|
|
|
|
|
def test_existing_task_database_gains_source_correlation_columns(tmp_path) -> None:
|
|
db_path = tmp_path / "legacy.sqlite3"
|
|
with sqlite3.connect(db_path) as connection:
|
|
connection.execute(
|
|
"""
|
|
create table tasks (
|
|
id text primary key,
|
|
goal text not null,
|
|
device_id text not null,
|
|
status text not null,
|
|
created_at text not null,
|
|
updated_at text not null,
|
|
completed_at text,
|
|
failure_reason text
|
|
)
|
|
"""
|
|
)
|
|
|
|
metadata_store = TaskMetadataStore(db_path)
|
|
task = _make_task("task-migrated")
|
|
metadata_store.create_task(
|
|
task,
|
|
source_task_id="cloud-task-migrated",
|
|
source_attempt=1,
|
|
)
|
|
|
|
row = metadata_store.get_task(task.id)
|
|
assert row is not None
|
|
assert row["source_task_id"] == "cloud-task-migrated"
|
|
assert row["source_attempt"] == 1
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 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()
|