From ffda20c45148f27a6cae7744c82d864762978743 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 2 Nov 2026 12:00:06 +0000 Subject: [PATCH] fix(codex): claim compaction ownership before SDK side effects Validate the workspace alias and settle the cross-workspace claim before thread resume or compact, so a losing request cannot compact. Drop a claim made for this call when the SDK rejects the operation, and leave an existing claim in place. Co-authored-by: celanwang --- .../agent_wrapper/codex_agent_wrapper.py | 100 ++++++++--- tests/unit/test_codex_compaction_ownership.py | 160 ++++++++++++++++++ 2 files changed, 240 insertions(+), 20 deletions(-) create mode 100644 tests/unit/test_codex_compaction_ownership.py diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index d90eb73e..7fd43ee4 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -456,17 +456,46 @@ class CodexAgentWrapper(BaseAgentWrapper): state.owned_thread_ids = ids return ids - async def _record_thread_ownership(self, alias: str, thread_id: str) -> None: + def _require_thread_alias(self, alias: str) -> str: + """Resolve an omitted alias and reject names that are not configured.""" workspaces = _context_workspaces(self.app_context) if not workspaces: - return - alias = self._normalize_thread_alias(alias) + return alias + resolved = self._normalize_thread_alias(alias) + if resolved not in workspaces: + configured = sorted(workspaces) + raise ValueError(f"Unknown workspace {resolved!r}; configured: {configured}") + return resolved + + @staticmethod + def _write_owned_thread_ids(path: Path, thread_ids: set[str]) -> None: + """Atomically replace one workspace's ownership file.""" + path.parent.mkdir(parents=True, exist_ok=True) + fd, raw_path = tempfile.mkstemp(prefix=".owned_threads-", suffix=".tmp", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + json.dump({"thread_ids": sorted(thread_ids)}, stream, indent=1) + os.replace(raw_path, path) + finally: + Path(raw_path).unlink(missing_ok=True) + + async def _record_thread_ownership(self, alias: str, thread_id: str) -> bool: + """Claim ``thread_id`` for ``alias``. + + Returns True when this call newly records the thread. An unknown alias, + or a thread already owned by another workspace, raises ValueError before + the claim is stored. + """ + workspaces = _context_workspaces(self.app_context) + if not workspaces: + return False + alias = self._require_thread_alias(alias) async with self.app_context.codex_thread_ownership_lock: for owner_alias in workspaces: if thread_id not in await self._owned_thread_ids(owner_alias): continue if owner_alias == alias: - return + return False raise ValueError( f"Codex thread {thread_id} belongs to workspace {owner_alias!r}; " f"it cannot be used from workspace {alias!r}", @@ -476,21 +505,40 @@ class CodexAgentWrapper(BaseAgentWrapper): ids = await self._owned_thread_ids(alias) path = self._thread_ownership_path(alias) if path is None: + if thread_id in ids: + return False ids.add(thread_id) - return + return True async with state.lock: if thread_id in ids: - return + return False updated = {*ids, thread_id} - path.parent.mkdir(parents=True, exist_ok=True) - fd, raw_path = tempfile.mkstemp(prefix=".owned_threads-", suffix=".tmp", dir=path.parent) - try: - with os.fdopen(fd, "w", encoding="utf-8") as stream: - json.dump({"thread_ids": sorted(updated)}, stream, indent=1) - os.replace(raw_path, path) - finally: - Path(raw_path).unlink(missing_ok=True) + self._write_owned_thread_ids(path, updated) ids.add(thread_id) + return True + + async def _forget_thread_ownership(self, alias: str, thread_id: str) -> None: + """Drop a claim reserved for an SDK operation that did not succeed.""" + workspaces = _context_workspaces(self.app_context) + if not workspaces: + return + alias = self._normalize_thread_alias(alias) + async with self.app_context.codex_thread_ownership_lock: + ids = await self._owned_thread_ids(alias) + if thread_id not in ids: + return + path = self._thread_ownership_path(alias) + if path is None: + ids.discard(thread_id) + return + state = self._codex_thread_state(alias) + async with state.lock: + if thread_id not in ids: + return + updated = set(ids) + updated.discard(thread_id) + self._write_owned_thread_ids(path, updated) + ids.discard(thread_id) async def _check_thread_ownership(self, thread_id: str, alias: str) -> None: """Reject resume/fork/compact of a thread owned by another workspace. @@ -648,17 +696,29 @@ class CodexAgentWrapper(BaseAgentWrapper): return response async def compact_session(self, session_id: str, workspace: str | None = None) -> None: - """Start native compaction of a Codex thread owned by the bound workspace.""" + """Compact a Codex thread after its workspace wins ownership. + + The alias is validated and the cross-workspace claim is settled before + any SDK call, so a losing request cannot compact. A claim made for this + call is removed when the SDK rejects the operation; an existing claim + is left in place. Thread release still runs once a thread was resumed. + """ await self.start() + alias = str(workspace or "") async with self._turn_lock: - await self._check_thread_ownership(session_id, str(workspace or "")) - codex = await self._get_codex() - thread = await codex.thread_resume(session_id) + newly_claimed = await self._record_thread_ownership(alias, session_id) + codex = None + thread = None try: + codex = await self._get_codex() + thread = await codex.thread_resume(session_id) await thread.compact() - await self._record_thread_ownership(str(workspace or ""), thread.id) + except BaseException: + if newly_claimed: + await self._forget_thread_ownership(alias, session_id) + raise finally: - if self._release_threads: + if thread is not None and codex is not None and self._release_threads: await self._release_thread(codex, thread.id) @classmethod diff --git a/tests/unit/test_codex_compaction_ownership.py b/tests/unit/test_codex_compaction_ownership.py new file mode 100644 index 00000000..47562045 --- /dev/null +++ b/tests/unit/test_codex_compaction_ownership.py @@ -0,0 +1,160 @@ +"""Compaction claims a Codex thread before any SDK side effect.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio +import json +from types import SimpleNamespace + +import pytest + +from reme.components import ApplicationContext +from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper + + +def _owned_on_disk(tmp_path, alias) -> set[str]: + path = tmp_path / alias / "mem_session" / "codex" / "owned_threads.json" + if not path.is_file(): + return set() + return set(json.loads(path.read_text(encoding="utf-8"))["thread_ids"]) + + +@pytest.mark.asyncio +async def test_compact_arbitrates_before_sdk_side_effects(tmp_path, monkeypatch): + """Compaction claims ownership before resume/compact, including under concurrency.""" + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + alpha = CodexAgentWrapper(name="codex-alpha", app_context=context, release_threads=False) + beta = CodexAgentWrapper(name="codex-beta", app_context=context, release_threads=False) + sdk_ops: list[str] = [] + entered_sdk = asyncio.Event() + release_sdk = asyncio.Event() + + class PausingThread: + """Resume pauses inside the SDK so a competing resume can try to claim.""" + + id = "shared-thread" + + async def compact(self): + sdk_ops.append("compact") + + async def alpha_resume(_thread_id): + sdk_ops.append("alpha-resume") + entered_sdk.set() + await release_sdk.wait() + return PausingThread() + + async def beta_resume(thread_id, **_kwargs): + sdk_ops.append("beta-resume") + return SimpleNamespace(id=thread_id) + + async def alpha_codex(): + return SimpleNamespace(thread_resume=alpha_resume) + + monkeypatch.setattr(alpha, "_get_codex", alpha_codex) + compact_task = asyncio.create_task(alpha.compact_session("shared-thread", workspace="alpha")) + await asyncio.wait_for(entered_sdk.wait(), timeout=2) + + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await beta._open_thread( + SimpleNamespace(thread_resume=beta_resume), + {"session_id": "shared-thread", "workspace": "beta"}, + ) + + release_sdk.set() + await asyncio.wait_for(compact_task, timeout=2) + assert sdk_ops == ["alpha-resume", "compact"] + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == {"shared-thread"} + assert "shared-thread" not in (context.workspaces["beta"].codex_thread_state.owned_thread_ids or set()) + assert _owned_on_disk(tmp_path, "alpha") == {"shared-thread"} + assert "shared-thread" not in _owned_on_disk(tmp_path, "beta") + + # An alias that was never configured fails before either SDK method. + unknown_ops: list[str] = [] + + async def unknown_resume(_thread_id): + unknown_ops.append("resume") + return PausingThread() + + async def unknown_codex(): + return SimpleNamespace(thread_resume=unknown_resume) + + monkeypatch.setattr(alpha, "_get_codex", unknown_codex) + with pytest.raises(ValueError, match="Unknown workspace 'typo'"): + await alpha.compact_session("unrecorded-thread", workspace="typo") + assert not unknown_ops + assert "unrecorded-thread" not in (context.workspaces["alpha"].codex_thread_state.owned_thread_ids or set()) + + # Once beta owns a thread, alpha's compaction does not reach the SDK. + await beta._open_thread( + SimpleNamespace(thread_resume=beta_resume), + {"session_id": "beta-thread", "workspace": "beta"}, + ) + loser_ops: list[str] = [] + + async def loser_resume(_thread_id): + loser_ops.append("resume") + return PausingThread() + + async def loser_codex(): + return SimpleNamespace(thread_resume=loser_resume) + + monkeypatch.setattr(alpha, "_get_codex", loser_codex) + with pytest.raises(ValueError, match="belongs to workspace 'beta'"): + await alpha.compact_session("beta-thread", workspace="alpha") + assert not loser_ops + assert context.workspaces["beta"].codex_thread_state.owned_thread_ids == {"beta-thread"} + + +@pytest.mark.asyncio +async def test_compact_drops_new_claim_when_sdk_rejects(tmp_path, monkeypatch): + """A failed resume or compact releases a claim this call just made.""" + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + alpha = CodexAgentWrapper(name="codex-alpha", app_context=context) + beta = CodexAgentWrapper(name="codex-beta", app_context=context, release_threads=False) + released: list[str] = [] + + class RecordingClient: + """Capture thread ids passed to the SDK unsubscribe request.""" + + async def request(self, method, params, *, response_model): + released.append(params["threadId"]) + assert method == "thread/unsubscribe" + assert response_model is not None + return SimpleNamespace() + + class RejectingThread: + """Resume succeeds and compaction is rejected by the SDK.""" + + def __init__(self, thread_id): + self.id = thread_id + + async def compact(self): + raise RuntimeError("sdk rejected") + + async def resume(thread_id): + if thread_id == "resume-fails": + raise RuntimeError("resume rejected") + return RejectingThread(thread_id) + + async def get_codex(): + return SimpleNamespace(_client=RecordingClient(), thread_resume=resume) + + monkeypatch.setattr(alpha, "_get_codex", get_codex) + + with pytest.raises(RuntimeError, match="resume rejected"): + await alpha.compact_session("resume-fails", workspace="alpha") + with pytest.raises(RuntimeError, match="sdk rejected"): + await alpha.compact_session("new-thread", workspace="alpha") + await alpha._record_thread_ownership("alpha", "kept-thread") + with pytest.raises(RuntimeError, match="sdk rejected"): + await alpha.compact_session("kept-thread", workspace="alpha") + + owned = context.workspaces["alpha"].codex_thread_state.owned_thread_ids or set() + assert "resume-fails" not in owned + assert "new-thread" not in owned + assert owned == {"kept-thread"} + assert _owned_on_disk(tmp_path, "alpha") == {"kept-thread"} + assert released == ["new-thread", "kept-thread"] + + await beta._record_thread_ownership("beta", "new-thread") + assert context.workspaces["beta"].codex_thread_state.owned_thread_ids == {"new-thread"}