From a6509ff479e3c558e7953e72efb8aa7be1a62134 Mon Sep 17 00:00:00 2001
From: Chen Wang <springoflotus@gmail.com>
Date: Fri, 16 Oct 2026 13:59:30 +0800
Subject: [PATCH] fix(agent_wrapper): release Codex threads after each turn
 (#621)

* fix(agent_wrapper): release Codex threads after each turn

Loaded Codex threads and their per-thread MCP server subprocesses were
never unloaded from the app-server, growing file descriptors and child
processes monotonically until the wrapper was closed (fd exhaustion,
os error 24). Reply, streaming reply, and compact_session now
unsubscribe the thread in a turn-scoped finally block by default; the
new release_threads=False option restores the previous behavior.
Unsubscribe failures are logged and never break the turn, and threads
remain resumable after being unloaded.

Closes agentscope-ai/ReMe#609

* fix(agent_wrapper): keep tool-context guard on failed release and cover stream startup

A failed thread/unsubscribe left the thread loaded in the app-server
but dropped its tool-context entry, so a later resume with changed
tool_context_id or injected_job_kwargs would silently bypass the
constraint check. Keep the guard entry unless the unsubscribe request
succeeds.

reply_stream now releases the thread from an outer try/finally scoped
right after _open_thread, so threads are also unloaded when turn/start
raises or the caller is cancelled while awaiting it.
---
 .../agent_wrapper/codex_agent_wrapper.py      |  65 +++-
 tests/unit/test_codex_agent_wrapper.py        | 332 ++++++++++++++++++
 2 files changed, 381 insertions(+), 16 deletions(-)

diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py
index 882fbfbf..ce2f3549 100644
--- a/reme/components/agent_wrapper/codex_agent_wrapper.py
+++ b/reme/components/agent_wrapper/codex_agent_wrapper.py
@@ -96,6 +96,7 @@ class CodexAgentWrapper(BaseAgentWrapper):
         codex_bin: str | None = None,
         config_overrides: list[str] | tuple[str, ...] | None = None,
         experimental_api: bool = True,
+        release_threads: bool = True,
         **kwargs,
     ) -> None:
         if "launch_args_override" in kwargs:
@@ -112,6 +113,7 @@ class CodexAgentWrapper(BaseAgentWrapper):
             experimental_api=experimental_api,
         )
         self._codex: AsyncCodex | None = None
+        self._release_threads = release_threads
         self._turn_lock = asyncio.Lock()
         self._mcp_snapshot_path: Path | None = None
         self._thread_tool_contexts: dict[str, str] = {}
@@ -403,6 +405,25 @@ class CodexAgentWrapper(BaseAgentWrapper):
                     self._mcp_snapshot_path.unlink(missing_ok=True)
                     self._mcp_snapshot_path = None
 
+    async def _release_thread(self, codex: AsyncCodex, thread_id: str) -> None:
+        """Unload a finished thread from the app-server so its MCP children exit."""
+        try:
+            from openai_codex.generated.v2_all import ThreadUnsubscribeResponse
+
+            # The SDK exposes no public unsubscribe; request() is public but
+            # reached through the private client attribute on SDK 0.162.x.
+            await codex._client.request(  # pylint: disable=protected-access
+                "thread/unsubscribe",
+                {"threadId": thread_id},
+                response_model=ThreadUnsubscribeResponse,
+            )
+        except Exception as exc:  # pylint: disable=broad-exception-caught
+            self.logger.warning(f"Failed to unsubscribe Codex thread {thread_id}: {exc}")
+            return
+        # The thread is still loaded when unsubscribe fails; keep its tool
+        # context so a later resume cannot silently bypass changed constraints.
+        self._thread_tool_contexts.pop(thread_id, None)
+
     @staticmethod
     def _serialize(value: Any) -> Any:
         from pydantic_core import to_jsonable_python
@@ -418,7 +439,11 @@ class CodexAgentWrapper(BaseAgentWrapper):
         async with self._turn_lock:
             codex = await self._get_codex()
             thread = await self._open_thread(codex, kwargs)
-            result = await thread.run(inputs, **self._turn_kwargs(kwargs))
+            try:
+                result = await thread.run(inputs, **self._turn_kwargs(kwargs))
+            finally:
+                if self._release_threads:
+                    await self._release_thread(codex, thread.id)
 
         final_response = result.final_response or ""
         response = {
@@ -445,7 +470,11 @@ class CodexAgentWrapper(BaseAgentWrapper):
         async with self._turn_lock:
             codex = await self._get_codex()
             thread = await codex.thread_resume(session_id)
-            await thread.compact()
+            try:
+                await thread.compact()
+            finally:
+                if self._release_threads:
+                    await self._release_thread(codex, thread.id)
 
     @classmethod
     # pylint: disable=too-many-return-statements
@@ -572,19 +601,23 @@ class CodexAgentWrapper(BaseAgentWrapper):
         async with self._turn_lock:
             codex = await self._get_codex()
             thread = await self._open_thread(codex, kwargs)
-            turn = await thread.turn(inputs, **self._turn_kwargs(kwargs))
-            stream = turn.stream()
-            completed = False
             try:
-                async for event in stream:
-                    if event.method == "turn/completed":
-                        completed = True
-                    for chunk in self._event_to_chunks(event, thread.id):
-                        yield chunk
+                turn = await thread.turn(inputs, **self._turn_kwargs(kwargs))
+                stream = turn.stream()
+                completed = False
+                try:
+                    async for event in stream:
+                        if event.method == "turn/completed":
+                            completed = True
+                        for chunk in self._event_to_chunks(event, thread.id):
+                            yield chunk
+                finally:
+                    if not completed:
+                        try:
+                            await turn.interrupt()
+                        except Exception as exc:  # pylint: disable=broad-exception-caught
+                            self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}")
+                    await stream.aclose()
             finally:
-                if not completed:
-                    try:
-                        await turn.interrupt()
-                    except Exception as exc:  # pylint: disable=broad-exception-caught
-                        self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}")
-                await stream.aclose()
+                if self._release_threads:
+                    await self._release_thread(codex, thread.id)
diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py
index c342a3b9..d1a4326b 100644
--- a/tests/unit/test_codex_agent_wrapper.py
+++ b/tests/unit/test_codex_agent_wrapper.py
@@ -653,6 +653,8 @@ async def test_compact_session_uses_native_thread_operation(tmp_path, monkeypatc
     observed = {}
 
     class FakeThread:
+        id = "thread-1"
+
         async def compact(self):
             observed["compacted"] = True
 
@@ -1201,3 +1203,333 @@ def test_default_config_provides_codex_oauth_wrapper(monkeypatch):
     assert "api_key" not in oauth
     assert codex["auth_mode"] == "api_key"
     assert codex["sandbox"] == "full-access"
+
+
+class _RecordingClient:
+    def __init__(self, fail: bool = False):
+        self.requests = []
+        self.fail = fail
+
+    async def request(self, method, params, *, response_model):
+        self.requests.append({"method": method, "params": params, "response_model": response_model})
+        if self.fail:
+            raise RuntimeError("unsubscribe failed")
+        return SimpleNamespace(status="unsubscribed")
+
+
+def _assert_unsubscribed(client, thread_id="thread-1"):
+    from openai_codex.generated.v2_all import ThreadUnsubscribeResponse
+
+    assert client.requests == [
+        {
+            "method": "thread/unsubscribe",
+            "params": {"threadId": thread_id},
+            "response_model": ThreadUnsubscribeResponse,
+        },
+    ]
+
+
+def _fake_codex(client, thread):
+    class FakeCodex:
+        def __init__(self, _config):
+            self._client = client
+
+        async def account(self):
+            return SimpleNamespace(account=SimpleNamespace())
+
+        async def close(self):
+            return None
+
+        async def thread_start(self, **_kwargs):
+            return thread
+
+        async def thread_resume(self, _thread_id, **_kwargs):
+            return thread
+
+    return FakeCodex
+
+
+@pytest.mark.asyncio
+async def test_reply_releases_thread_after_turn(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
+    client = _RecordingClient()
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def run(self, _inputs, **_kwargs):
+            return _TurnResult(final_response="done")
+
+    monkeypatch.setattr(
+        "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex",
+        _fake_codex(client, FakeThread()),
+    )
+
+    result = await wrapper.reply("answer")
+    await wrapper.close()
+
+    assert result["last_message"] == "done"
+    _assert_unsubscribed(client)
+    assert not wrapper._thread_tool_contexts  # pylint: disable=protected-access
+
+
+@pytest.mark.asyncio
+async def test_reply_releases_thread_when_turn_fails(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
+    client = _RecordingClient()
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def run(self, _inputs, **_kwargs):
+            raise ValueError("turn failed")
+
+    monkeypatch.setattr(
+        "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex",
+        _fake_codex(client, FakeThread()),
+    )
+
+    with pytest.raises(ValueError, match="turn failed"):
+        await wrapper.reply("answer")
+    await wrapper.close()
+
+    _assert_unsubscribed(client)
+    assert not wrapper._thread_tool_contexts  # pylint: disable=protected-access
+
+
+@pytest.mark.asyncio
+async def test_reply_ignores_unsubscribe_failure(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
+    client = _RecordingClient(fail=True)
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def run(self, _inputs, **_kwargs):
+            return _TurnResult(final_response="done")
+
+    monkeypatch.setattr(
+        "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex",
+        _fake_codex(client, FakeThread()),
+    )
+
+    result = await wrapper.reply("answer")
+
+    assert result["last_message"] == "done"
+    assert len(client.requests) == 1
+    # The thread is still loaded, so its tool-context guard must survive.
+    assert "thread-1" in wrapper._thread_tool_contexts  # pylint: disable=protected-access
+
+    await wrapper.close()
+
+
+@pytest.mark.asyncio
+async def test_failed_release_keeps_tool_context_guard_on_resume(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
+    client = _RecordingClient(fail=True)
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def run(self, _inputs, **_kwargs):
+            return _TurnResult(final_response="done")
+
+    monkeypatch.setattr(
+        "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex",
+        _fake_codex(client, FakeThread()),
+    )
+
+    result = await wrapper.reply("answer", tool_context_id="ctx-a")
+    assert result["session_id"] == "thread-1"
+    assert len(client.requests) == 1
+
+    with pytest.raises(ValueError, match="cannot change"):
+        await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-b")
+    await wrapper.close()
+
+
+@pytest.mark.asyncio
+async def test_release_threads_disabled_keeps_threads_loaded(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth", release_threads=False)
+    client = _RecordingClient()
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def run(self, _inputs, **_kwargs):
+            return _TurnResult(final_response="done")
+
+    monkeypatch.setattr(
+        "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex",
+        _fake_codex(client, FakeThread()),
+    )
+
+    result = await wrapper.reply("answer")
+
+    assert result["last_message"] == "done"
+    assert not client.requests
+    assert "thread-1" in wrapper._thread_tool_contexts  # pylint: disable=protected-access
+
+    await wrapper.close()
+
+
+@pytest.mark.asyncio
+async def test_reply_stream_releases_thread_on_completion(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
+    client = _RecordingClient()
+
+    class FakeTurn:
+        id = "turn-1"
+
+        async def stream(self):
+            yield SimpleNamespace(
+                method="turn/completed",
+                payload=SimpleNamespace(
+                    turn=SimpleNamespace(
+                        id="turn-1",
+                        status=SimpleNamespace(value="completed"),
+                        duration_ms=1,
+                        error=None,
+                    ),
+                ),
+            )
+
+        async def interrupt(self):
+            raise AssertionError("Completed turns must not be interrupted")
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def turn(self, _inputs, **_kwargs):
+            return FakeTurn()
+
+    monkeypatch.setattr(
+        "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex",
+        _fake_codex(client, FakeThread()),
+    )
+
+    chunks = [chunk async for chunk in wrapper.reply_stream("answer")]
+    await wrapper.close()
+
+    assert chunks[-1].chunk_type == ChunkEnum.REPLY_END
+    _assert_unsubscribed(client)
+    assert not wrapper._thread_tool_contexts  # pylint: disable=protected-access
+
+
+@pytest.mark.asyncio
+async def test_reply_stream_releases_thread_on_early_close(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
+    client = _RecordingClient()
+    interrupt_count = 0
+
+    class FakeTurn:
+        id = "turn-1"
+
+        async def stream(self):
+            yield SimpleNamespace(
+                method="turn/started",
+                payload=SimpleNamespace(turn=SimpleNamespace(id=self.id)),
+            )
+            await asyncio.Event().wait()
+
+        async def interrupt(self):
+            nonlocal interrupt_count
+            interrupt_count += 1
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def turn(self, _inputs, **_kwargs):
+            return FakeTurn()
+
+    monkeypatch.setattr(
+        "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex",
+        _fake_codex(client, FakeThread()),
+    )
+
+    stream = wrapper.reply_stream("answer")
+    first = await anext(stream)
+    assert first.chunk_type == ChunkEnum.REPLY_START
+    await stream.aclose()
+    await wrapper.close()
+
+    assert interrupt_count == 1
+    _assert_unsubscribed(client)
+    assert not wrapper._thread_tool_contexts  # pylint: disable=protected-access
+
+
+@pytest.mark.asyncio
+async def test_compact_session_releases_thread(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path)
+    client = _RecordingClient()
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def compact(self):
+            return None
+
+    async def start():
+        return None
+
+    async def resume(_thread_id):
+        return FakeThread()
+
+    async def get_codex():
+        return SimpleNamespace(_client=client, thread_resume=resume)
+
+    monkeypatch.setattr(wrapper, "start", start)
+    monkeypatch.setattr(wrapper, "_get_codex", get_codex)
+
+    await wrapper.compact_session("thread-1")
+
+    _assert_unsubscribed(client)
+    assert not wrapper._thread_tool_contexts  # pylint: disable=protected-access
+
+
+@pytest.mark.asyncio
+async def test_reply_stream_releases_thread_when_turn_start_fails(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
+    client = _RecordingClient()
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def turn(self, _inputs, **_kwargs):
+            raise RuntimeError("turn/start failed")
+
+    monkeypatch.setattr(
+        "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex",
+        _fake_codex(client, FakeThread()),
+    )
+
+    with pytest.raises(RuntimeError, match="turn/start failed"):
+        _ = [chunk async for chunk in wrapper.reply_stream("answer")]
+    await wrapper.close()
+
+    _assert_unsubscribed(client)
+    assert not wrapper._thread_tool_contexts  # pylint: disable=protected-access
+
+
+@pytest.mark.asyncio
+async def test_reply_stream_releases_thread_when_turn_start_is_cancelled(tmp_path, monkeypatch):
+    wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
+    client = _RecordingClient()
+
+    class FakeThread:
+        id = "thread-1"
+
+        async def turn(self, _inputs, **_kwargs):
+            raise asyncio.CancelledError("cancelled during turn/start")
+
+    monkeypatch.setattr(
+        "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex",
+        _fake_codex(client, FakeThread()),
+    )
+
+    with pytest.raises(asyncio.CancelledError):
+        _ = [chunk async for chunk in wrapper.reply_stream("answer")]
+    await wrapper.close()
+
+    _assert_unsubscribed(client)
+    assert not wrapper._thread_tool_contexts  # pylint: disable=protected-access