feat(repair): propose a branch after a step is repaired more than once (carry-over of openadapt-agent#76)

packages/agent/tests/test_repair_branch.py

@@ -0,0 +1,73 @@
+"""A second demonstration on one step also proposes a branch candidate."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from types import SimpleNamespace
+
+from openadapt_agent.repair import RepairAccessPolicy, RepairDesk
+
+
+class _Store:
+    def __init__(self) -> None:
+        self.added: list[object] = []
+
+    def add_candidate(self, candidate):
+        self.added.append(candidate)
+        return SimpleNamespace(candidate_id=f"cand-{len(self.added)}")
+
+
+class _AgentTeach:
+    def __init__(self, fail: bool = False) -> None:
+        self.calls: list[list[tuple[str, str]]] = []
+        self.fail = fail
+
+    def teach_branch_from_calling_agent(self, examples, out, *, bundle, provenance):
+        self.calls.append(list(examples))
+        if self.fail:
+            raise RuntimeError("no shared screen text")
+        return SimpleNamespace(candidate="branch-candidate")
+
+
+def _demo(tmp: Path, name: str) -> tuple[dict, Path]:
+    run_dir = tmp / "runs" / name
+    recording = tmp / "teach" / name / "recording"
+    run_dir.mkdir(parents=True)
+    recording.mkdir(parents=True)
+    record = {"workflow": "wf", "step_id": "s_save", "run_dir": str(run_dir), "bundle_dir": "b"}
+    return record, recording
+
+
+def test_first_demo_is_remembered_second_proposes_a_branch(tmp_path: Path) -> None:
+    desk = RepairDesk(tmp_path / ".repair", RepairAccessPolicy(None))
+    store, teach = _Store(), _AgentTeach()
+
+    record, rec = _demo(tmp_path, "r1")
+    assert desk._teach_branch(record, rec, tmp_path / "w1", store, teach, None) is None
+    assert teach.calls == []
+
+    record, rec = _demo(tmp_path, "r2")
+    assert desk._teach_branch(record, rec, tmp_path / "w2", store, teach, None) == "cand-1"
+    assert [Path(run).name for run, _ in teach.calls[0]] == ["r1", "r2"]
+    assert store.added == ["branch-candidate"]
+
+
+def test_other_step_and_refusal_add_no_branch(tmp_path: Path) -> None:
+    desk = RepairDesk(tmp_path / ".repair", RepairAccessPolicy(None))
+    store, teach = _Store(), _AgentTeach(fail=True)
+    record, rec = _demo(tmp_path, "r1")
+    desk._teach_branch(record, rec, tmp_path / "w1", store, teach, None)
+    other, rec2 = _demo(tmp_path, "r2")
+    other["step_id"] = "s_other"
+    assert desk._teach_branch(other, rec2, tmp_path / "w2", store, teach, None) is None
+    assert teach.calls == []
+    record3, rec3 = _demo(tmp_path, "r3")
+    assert desk._teach_branch(record3, rec3, tmp_path / "w3", store, teach, None) is None
+    assert store.added == []
+
+
+def test_older_flow_without_branch_support_is_skipped(tmp_path: Path) -> None:
+    desk = RepairDesk(tmp_path / ".repair", RepairAccessPolicy(None))
+    for name in ("r1", "r2"):
+        record, rec = _demo(tmp_path, name)
+        assert desk._teach_branch(record, rec, tmp_path / name, _Store(), object(), None) is None

packages/flow/openadapt_flow/agent/repair.py

@@ -713,7 +713,7 @@ def _teach(
             }
         store = RepairStore(self.store_for(record["workflow"]))
         candidate = store.add_candidate(result.candidate)
-        return {
+        response: dict[str, Any] = {
             "status": "pending_approval",
             "candidate_id": candidate.candidate_id,
             "what_happened": (
@@ -722,6 +722,68 @@ def _teach(
                 "then, this workflow still stops at this step."
             ),
         }
+        try:
+            branch_id = self._teach_branch(
+                record, recording, work, store, agent_teach, provenance
+            )
+        except Exception as exc:  # the single-step candidate already stands
+            _LOG.info("branch proposal skipped: %s", exc)
+            branch_id = None
+        if branch_id is not None:
+            response["branch_candidate_id"] = branch_id
+            response["what_happened"] += (
+                " This step has now been repaired more than once, so OpenAdapt "
+                "also proposed one branch that covers these cases. It needs "
+                "the same review and approval."
+            )
+        return response
+
+    def _demos_path(self, record: dict[str, Any]) -> Path:
+        key = _sha(f"{record['workflow']}\0{record['step_id']}")[:16]
+        return self.root / "demos" / f"{key}.json"
+
+    def _teach_branch(
+        self,
+        record: dict[str, Any],
+        recording: Path,
+        work: Path,
+        store: Any,
+        agent_teach: Any,
+        provenance: Any,
+    ) -> Optional[str]:
+        """Remember this demonstration; with two or more, propose a branch.
+
+        The branch is a separate repair candidate. If Flow refuses it, the
+        single-step candidate stands on its own.
+        """
+        path = self._demos_path(record)
+        demos = (_read_json(path) or {}).get("demos", [])
+        demos = [
+            d
+            for d in demos
+            if isinstance(d, dict) and d.get("run_dir") != record["run_dir"]
+        ]
+        demos.append({"run_dir": record["run_dir"], "recording": str(recording)})
+        _write_json(path, {"demos": demos})
+        usable = [
+            (d["run_dir"], d["recording"])
+            for d in demos
+            if Path(d["run_dir"]).is_dir() and Path(d["recording"]).is_dir()
+        ]
+        teach_branch = getattr(agent_teach, "teach_branch_from_calling_agent", None)
+        if len(usable) < 2 or teach_branch is None:
+            return None
+        try:
+            result = teach_branch(
+                usable,
+                work / "branch",
+                bundle=record["bundle_dir"],
+                provenance=provenance,
+            )
+        except Exception as exc:
+            _LOG.info("branch from several demonstrations refused: %s", exc)
+            return None
+        return str(store.add_candidate(result.candidate).candidate_id)
 
     def active_bundle(self, workflow: str, default: Path) -> Path:
         """The approved repaired bundle, when a person activated one."""