scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-15 09:11:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

fix(codex): claim compaction ownership before SDK side effects

Browse source 
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 <lumworn@gmail.com>
This commit is contained in:
Cursor Agent 2026-11-02 12:00:06 +00:00
parent b5d0bbbc20
commit ffda20c451
No known key found for this signature in database
2 changed files with 240 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

100
reme/components/agent_wrapper/codex_agent_wrapper.py
View file

@ -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

160
tests/unit/test_codex_compaction_ownership.py Normal file
View file

@ -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"}