fix(codex): retain tool constraints until threads unload (#630)

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:

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