fix(train): plan policy updates outside the policy-root tree lock

openviking/session/train/domain.py

@@ -64,12 +64,17 @@ class PolicySet:
     request_context: Any | None = field(default=None, repr=False, compare=False)
 
     @asynccontextmanager
-    async def lock(self):
+    async def lock(self, *, timeout_secs: float = 300.0):
         """Acquire a tree lock for the whole policy root directory.
 
-        Policy updates serialize on this lock so concurrent realtime/batch
-        training jobs plan and apply against a freshly reloaded policy set.
-        ``timeout=None`` means wait indefinitely until the lock is available.
+        Policy updates serialize on this lock.  Keep the critical section
+        short and local (reload + apply): long-running work such as remote
+        LLM planning must run before acquiring the lock, otherwise concurrent
+        writers of the same root are blocked until they hit the acquisition
+        timeout.
+
+        ``timeout_secs`` bounds how long acquisition waits for the current
+        holder before failing; it defaults to 300 seconds.
         """
 
         if self.viking_fs is None:
@@ -81,7 +86,9 @@ async def lock(self):
             raise RuntimeError("PolicySet.viking_fs must provide _uri_to_path for locking")
 
         path = uri_to_path(self.root_uri, ctx=self.request_context)
-        lease = await self.viking_fs._async_agfs.pathlock_acquire_tree(path, timeout_secs=300.0)
+        lease = await self.viking_fs._async_agfs.pathlock_acquire_tree(
+            path, timeout_secs=timeout_secs
+        )
         try:
             yield lease
         finally:

openviking/session/train/engine.py

@@ -23,6 +23,11 @@
     SemanticGradient,
 )
 
+# Bounded number of optimistic re-plans when the policy set changed while the
+# optimizer was planning.  Keeps the (potentially minutes-long) LLM replanning
+# from livelocking against sustained concurrent writers.
+_PLAN_APPLY_MAX_REPLANS = 2
+
 
 @dataclass(slots=True)
 class PolicyTrainingEngine:
@@ -84,17 +89,53 @@ async def plan_and_apply(
         policy_set: ExperienceSet,
         ctx: PipelineContext,
     ) -> tuple[PolicyUpdatePlan, PolicyApplyResult]:
-        async with policy_set.lock() as transaction_handle:
-            latest_policy_set = await policy_set.reload()
+        """Plan against a lock-free snapshot, then apply in a short critical section.
+
+        ``policy_optimizer.plan`` may block on a remote LLM for minutes.  Holding
+        the policy-root tree lock across that call starves every concurrent
+        writer of the same root until the lock acquisition timeout fails their
+        commits, so planning runs outside the lock.  The lock is taken only to
+        re-check that the snapshot the plan was built against is still the
+        latest one before applying.  When another writer committed meanwhile,
+        the plan is rebuilt against the fresh snapshot (bounded by
+        ``_PLAN_APPLY_MAX_REPLANS``); once the bound is exhausted the last plan
+        is applied under the lock with a conflict marker, and the updater's
+        per-item base-content guard reports stale targets as errors instead of
+        silently overwriting concurrent changes.
+        """
+
+        replans = 0
+        while True:
+            planned_snapshot = await policy_set.reload()
             plan = await self.policy_optimizer.plan(
                 gradients,
-                latest_policy_set,
+                planned_snapshot,
                 ctx.optimization_context,
             )
-            apply_result = await self.policy_updater.apply(
-                plan,
-                latest_policy_set,
-                ctx.apply_context or latest_policy_set.request_context,
-                transaction_handle=transaction_handle,
-            )
-        return plan, apply_result
+            async with policy_set.lock() as transaction_handle:
+                latest_policy_set = await policy_set.reload()
+                snapshot_current = _policy_sets_match(latest_policy_set, planned_snapshot)
+                if snapshot_current or replans >= _PLAN_APPLY_MAX_REPLANS:
+                    if not snapshot_current:
+                        plan.metadata = {
+                            **dict(plan.metadata or {}),
+                            "optimistic_apply_after_conflict": True,
+                            "replan_attempts": replans + 1,
+                        }
+                    apply_result = await self.policy_updater.apply(
+                        plan,
+                        latest_policy_set,
+                        ctx.apply_context or latest_policy_set.request_context,
+                        transaction_handle=transaction_handle,
+                    )
+                    return plan, apply_result
+            # The policy set moved on while planning: the lock was released on
+            # exiting the context manager, so re-plan outside it against the
+            # fresh snapshot.
+            replans += 1
+
+
+def _policy_sets_match(current: ExperienceSet, planned: ExperienceSet) -> bool:
+    """Whether ``planned`` is still the latest committed view of the policy root."""
+
+    return current.root_uri == planned.root_uri and list(current.policies) == list(planned.policies)

tests/session/train/test_fakes.py

@@ -3,7 +3,9 @@
 
 from __future__ import annotations
 
+import asyncio
 import threading
+import time
 from types import SimpleNamespace
 from typing import Any
 
@@ -217,6 +219,77 @@ def fake_request_context(account_id: str = "default", user_id: str = "u"):
     )
 
 
+class LockingAGFS:
+    """Async AGFS fake implementing tree path locks with real mutual exclusion.
+
+    ``PolicySet.lock`` acquires tree locks through
+    ``viking_fs._async_agfs.pathlock_acquire_tree``; this fake enforces actual
+    exclusion between concurrent acquirers of the same path, honors the
+    ``timeout_secs`` deadline, and records every acquire/release plus the hold
+    interval of each lease so tests can assert which code ran while the lock
+    was held.
+
+    ``lock_budget_secs`` caps how long any single acquisition may wait for a
+    holder regardless of the requested timeout; exceeding it raises
+    ``TimeoutError``.  It exists to keep tests fast when simulating the
+    production 300s lock timeout.
+    """
+
+    def __init__(self, *, lock_budget_secs: float | None = None):
+        self.lock_budget_secs = lock_budget_secs
+        self._locks: dict[str, asyncio.Lock] = {}
+        self.acquire_calls: list[dict[str, Any]] = []
+        self.release_calls: list[dict[str, Any]] = []
+        self.holds: list[dict[str, float]] = []
+
+    def _lock_for(self, path: str) -> asyncio.Lock:
+        if path not in self._locks:
+            self._locks[path] = asyncio.Lock()
+        return self._locks[path]
+
+    def is_locked(self, path: str) -> bool:
+        lock = self._locks.get(path)
+        return lock is not None and lock.locked()
+
+    async def pathlock_acquire_tree(
+        self,
+        path: str,
+        timeout_secs: float = 0.0,
+        owner_lease_ref: Any = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        del owner_lease_ref, kwargs
+        self.acquire_calls.append({"path": path, "timeout_secs": timeout_secs})
+        deadline: float | None = None
+        if timeout_secs and timeout_secs > 0:
+            deadline = timeout_secs
+        if self.lock_budget_secs is not None:
+            deadline = min(deadline, self.lock_budget_secs) if deadline else self.lock_budget_secs
+        lock = self._lock_for(path)
+        try:
+            await asyncio.wait_for(lock.acquire(), deadline)
+        except asyncio.TimeoutError:
+            raise TimeoutError(
+                f"tree lock acquisition timed out after {timeout_secs}s: {path}"
+            ) from None
+        lease = {"path": path, "acquired_at": time.monotonic()}
+        return lease
+
+    async def pathlock_release(self, lease: dict[str, Any], **kwargs: Any) -> None:
+        del kwargs
+        release = dict(lease)
+        release["released_at"] = time.monotonic()
+        self.release_calls.append(release)
+        self.holds.append(
+            {
+                "path": release["path"],
+                "start": release["acquired_at"],
+                "end": release["released_at"],
+            }
+        )
+        self._lock_for(release["path"]).release()
+
+
 def _norm_path(path: str) -> str:
     if not path:
         return "/"