From d3abc0b5b0c962ac4ad2a8de10feded6e53075d4 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Fri, 16 Oct 2026 16:46:02 +0800 Subject: [PATCH] fix(codex): retain tool constraints until threads unload (#630) --- .../agent_wrapper/codex_agent_wrapper.py | 35 +++++- tests/unit/test_codex_agent_wrapper.py | 107 ++++++++++++++++++ 2 files changed, 137 insertions(+), 5 deletions(-) diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index ce2f3549..733b55e0 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -307,6 +307,8 @@ class CodexAgentWrapper(BaseAgentWrapper): fork_session = bool(kwargs.get("fork_session", False)) if fork_session and not thread_id: raise ValueError("fork_session=True requires resume or session_id") + if self._release_threads and self._thread_tool_contexts: + await self._prune_thread_tool_contexts(codex) requested_tool_context = json.dumps( { "tool_context_id": str(kwargs.get("tool_context_id") or ""), @@ -405,8 +407,32 @@ class CodexAgentWrapper(BaseAgentWrapper): self._mcp_snapshot_path.unlink(missing_ok=True) self._mcp_snapshot_path = None + async def _prune_thread_tool_contexts(self, codex: AsyncCodex) -> None: + """Forget tool constraints only after the server confirms a thread is unloaded.""" + try: + from openai_codex.generated.v2_all import ThreadLoadedListResponse + + loaded_ids: set[str] = set() + cursor = None + while True: + page = await codex._client.request( # pylint: disable=protected-access + "thread/loaded/list", + {"cursor": cursor}, + response_model=ThreadLoadedListResponse, + ) + loaded_ids.update(page.data) + cursor = page.next_cursor + if not cursor: + break + except Exception as exc: # pylint: disable=broad-exception-caught + self.logger.warning(f"Failed to list loaded Codex threads: {exc}") + return + self._thread_tool_contexts = { + thread_id: context for thread_id, context in self._thread_tool_contexts.items() if thread_id in loaded_ids + } + async def _release_thread(self, codex: AsyncCodex, thread_id: str) -> None: - """Unload a finished thread from the app-server so its MCP children exit.""" + """Unsubscribe so the app-server can unload the thread once it is idle.""" try: from openai_codex.generated.v2_all import ThreadUnsubscribeResponse @@ -419,10 +445,9 @@ class CodexAgentWrapper(BaseAgentWrapper): ) 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) + # Unsubscribe is not an unload acknowledgement. In particular, a + # cancelled turn can still be running with its original MCP constraints. + # Keep the guard until thread/loaded/list confirms the thread is gone. @staticmethod def _serialize(value: Any) -> Any: diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index d1a4326b..c3072dce 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -1209,11 +1209,14 @@ class _RecordingClient: def __init__(self, fail: bool = False): self.requests = [] self.fail = fail + self.loaded_threads = ["thread-1"] 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") + if method == "thread/loaded/list": + return SimpleNamespace(data=self.loaded_threads, next_cursor=None) return SimpleNamespace(status="unsubscribed") @@ -1533,3 +1536,107 @@ async def test_reply_stream_releases_thread_when_turn_start_is_cancelled(tmp_pat _assert_unsubscribed(client) assert not wrapper._thread_tool_contexts # pylint: disable=protected-access + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("original", "changed"), + [ + ({"tool_context_id": "ctx-a"}, {"tool_context_id": "ctx-b"}), + ({"injected_job_kwargs": {"tenant": "a"}}, {"injected_job_kwargs": {"tenant": "b"}}), + ], +) +async def test_cancelled_reply_keeps_loaded_thread_constraints(tmp_path, monkeypatch, original, changed): + wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") + client = _RecordingClient() + entered = asyncio.Event() + + class FakeThread: + id = "thread-1" + calls = 0 + + async def run(self, _inputs, **_kwargs): + self.calls += 1 + if self.calls > 1: + return _TurnResult(final_response="old MCP constraints") + entered.set() + await asyncio.Event().wait() + + monkeypatch.setattr( + "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", + _fake_codex(client, FakeThread()), + ) + task = asyncio.create_task(wrapper.reply("answer", **original)) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + try: + _assert_unsubscribed(client) + with pytest.raises(ValueError, match="cannot change"): + await wrapper.reply("answer", session_id="thread-1", **changed) + finally: + await wrapper.close() + + +@pytest.mark.asyncio +async def test_unloaded_thread_can_resume_with_new_constraints(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()), + ) + try: + await wrapper.reply("answer", tool_context_id="ctx-a") + assert "thread-1" in wrapper._thread_tool_contexts + await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-a") + with pytest.raises(ValueError, match="cannot change"): + await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-b") + # Unsubscribe schedules delayed unload; only a later loaded-list query + # establishes that a resume will create a fresh runtime and MCP config. + client.loaded_threads.clear() + result = await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-b") + assert result["last_message"] == "done" + assert json.loads(wrapper._thread_tool_contexts["thread-1"])["tool_context_id"] == "ctx-b" + finally: + await wrapper.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_second_page", [False, True]) +async def test_prune_thread_constraints_reads_all_pages_before_removing_entries(tmp_path, fail_second_page): + from openai_codex.generated.v2_all import ThreadLoadedListResponse + + wrapper, _job = _wrapper(tmp_path) + wrapper._thread_tool_contexts = {"first": "a", "last": "b", "unloaded": "c"} + cursors = [] + + class Client: + async def request(self, method, params, *, response_model): + assert method == "thread/loaded/list" + assert response_model is ThreadLoadedListResponse + cursors.append(params["cursor"]) + if params["cursor"] is None: + return ThreadLoadedListResponse(data=["first"], nextCursor="page-2") + if fail_second_page: + raise RuntimeError("loaded-list unavailable") + return ThreadLoadedListResponse(data=["last"]) + + await wrapper._prune_thread_tool_contexts(SimpleNamespace(_client=Client())) + + assert cursors == [None, "page-2"] + expected = {"first": "a", "last": "b"} + if fail_second_page: + expected["unloaded"] = "c" + assert wrapper._thread_tool_contexts == expected