mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-17 04:37:00 +00:00
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.
This commit is contained in:
Chen Wang
GitHub
parent
72e243b7d6
commit
a6509ff479
2 changed files with 381 additions and 16 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue