From 18c177603434ccd8440355cbd7db044d4e61efdc Mon Sep 17 00:00:00 2001 From: fredheathcote007 Date: Fri, 6 Nov 2026 22:01:25 +0000 Subject: [PATCH 1/3] fix(as_llm): close owned OpenAI clients on component shutdown --- reme/components/as_llm/__init__.py | 25 ++++++ tests/unit/test_as_llm_lazy.py | 124 +++++++++++++++++++++++++++++ 2 files changed, 149 insertions(+) diff --git a/reme/components/as_llm/__init__.py b/reme/components/as_llm/__init__.py index 7ea4e1bf0..335801ddd 100644 --- a/reme/components/as_llm/__init__.py +++ b/reme/components/as_llm/__init__.py @@ -31,6 +31,7 @@ class BaseAsLLM(BaseComponent): def __init__(self, **kwargs) -> None: super().__init__(**kwargs) self._model: ChatModelBase | None = None + self._owns_model = False async def _start(self) -> None: """Keep service startup independent of provider credentials.""" @@ -46,6 +47,7 @@ def model(self) -> ChatModelBase: def model(self, value: ChatModelBase | None) -> None: """Preserve explicit model injection used by standalone consumers.""" self._model = value + self._owns_model = False def initialize_model(self) -> None: """Construct the configured provider once, without making a remote request.""" @@ -57,6 +59,7 @@ def initialize_model(self) -> None: params_dict = kwargs.pop("parameters", None) parameters = model_cls.Parameters(**params_dict) if params_dict else None self._model = model_cls(credential=credential, parameters=parameters, **kwargs) + self._owns_model = True @R.register("openai") @@ -65,6 +68,28 @@ class OpenAIAsLLM(BaseAsLLM): credential_cls = OpenAICredential + @BaseAsLLM.model.setter + def model(self, value: ChatModelBase | None) -> None: + """Reject replacement while an internally owned client is still live.""" + if value is self._model: + return + if self._owns_model: + raise RuntimeError("Close the OpenAI component before replacing its owned model.") + BaseAsLLM.model.fset(self, value) + + async def _close(self) -> None: + """Close an internally owned async client on component shutdown.""" + if not self._owns_model or self._model is None: + return + model = self._model + client = getattr(model, "client", None) + try: + if client is not None: + await client.close() + finally: + self._model = None + self._owns_model = False + @R.register("anthropic") class AnthropicAsLLM(BaseAsLLM): diff --git a/tests/unit/test_as_llm_lazy.py b/tests/unit/test_as_llm_lazy.py index 09c8a646a..5ae3a7516 100644 --- a/tests/unit/test_as_llm_lazy.py +++ b/tests/unit/test_as_llm_lazy.py @@ -98,3 +98,127 @@ def test_explicit_model_assignment_preserves_injection(): llm.model = model assert llm.model is model llm.credential_cls.assert_not_called() + + +@pytest.mark.asyncio +async def test_owned_openai_client_closes_after_failed_call(): + """A failed model operation must not leave client cleanup until loop GC.""" + from reme.components.as_llm import OpenAIAsLLM + + llm = OpenAIAsLLM(backend="openai", model="synthetic") + model_cls = Mock() + model_cls.return_value.client.close = AsyncMock() + model_cls.return_value.call = AsyncMock(side_effect=ConnectionError("synthetic")) + credential_cls = Mock() + credential_cls.return_value.get_chat_model_class.return_value = model_cls + llm.credential_cls = credential_cls + await llm.start() + try: + model = llm.model + with pytest.raises(ConnectionError, match="synthetic"): + await model.call("synthetic", []) + assert model is model_cls.return_value + finally: + await llm.close() + model_cls.return_value.client.close.assert_awaited_once_with() + assert llm._model is None + await llm.close() + model_cls.return_value.client.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_injected_openai_client_is_borrowed(): + """The component must not close externally supplied model clients.""" + from reme.components.as_llm import OpenAIAsLLM + + llm = OpenAIAsLLM(backend="openai", model="synthetic") + model = Mock() + model.client.close = AsyncMock() + llm.model = model + await llm.start() + await llm.close() + model.client.close.assert_not_awaited() + assert llm._model is model + + +@pytest.mark.asyncio +async def test_owned_openai_client_close_failure_clears_state(): + """A failed close is propagated once without leaving ownership stranded.""" + from reme.components.as_llm import OpenAIAsLLM + + llm = OpenAIAsLLM(backend="openai", model="synthetic") + model_cls = Mock() + model_cls.return_value.client.close = AsyncMock(side_effect=RuntimeError("close failed")) + credential_cls = Mock() + credential_cls.return_value.get_chat_model_class.return_value = model_cls + llm.credential_cls = credential_cls + await llm.start() + _ = llm.model + + with pytest.raises(RuntimeError, match="close failed"): + await llm.close() + + assert llm._model is None + assert not llm._owns_model + assert not llm.is_started + await llm.close() + model_cls.return_value.client.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_owned_openai_model_requires_close_before_replacement(): + """Replacement cannot abandon an internally owned client synchronously.""" + from reme.components.as_llm import OpenAIAsLLM + + llm = OpenAIAsLLM(backend="openai", model="synthetic") + model_cls = Mock() + model_cls.return_value.client.close = AsyncMock() + credential_cls = Mock() + credential_cls.return_value.get_chat_model_class.return_value = model_cls + llm.credential_cls = credential_cls + await llm.start() + owned_model = llm.model + borrowed_model = Mock() + + with pytest.raises(RuntimeError, match="Close the OpenAI component"): + llm.model = borrowed_model + assert llm._model is owned_model + + await llm.close() + llm.model = borrowed_model + assert llm.model is borrowed_model + await llm.start() + await llm.close() + model_cls.return_value.client.close.assert_awaited_once_with() + + +@pytest.mark.parametrize("injected", [False, True], ids=["owned", "borrowed"]) +def test_openai_start_failure_rolls_back_owned_client(injected): + """Startup failure closes a pre-start owned client but never a borrowed one.""" + from reme.components.as_llm import OpenAIAsLLM + + class FailingOpenAIAsLLM(OpenAIAsLLM): + async def _start(self): + raise RuntimeError("start failed") + + llm = FailingOpenAIAsLLM(backend="openai", model="synthetic") + model = Mock() + model.client.close = AsyncMock() + if injected: + llm.model = model + else: + model_cls = Mock(return_value=model) + credential_cls = Mock() + credential_cls.return_value.get_chat_model_class.return_value = model_cls + llm.credential_cls = credential_cls + assert llm.model is model + + async def run(): + with pytest.raises(RuntimeError, match="start failed"): + await llm.start() + assert not llm.is_started + assert llm._model is (model if injected else None) + assert model.client.close.await_count == (0 if injected else 1) + await llm.close() + + asyncio.run(run()) From 7bdd4f71ce3a0b5896cce485d738007348040b4e Mon Sep 17 00:00:00 2001 From: fredheathcote007 Date: Fri, 6 Nov 2026 22:17:14 +0000 Subject: [PATCH 2/3] fix(as_llm): harden owned-client cleanup across lifecycle edges --- reme/components/as_llm/__init__.py | 47 ++++++++++- tests/unit/test_as_llm_lazy.py | 121 ++++++++++++++++++++++++++++- 2 files changed, 164 insertions(+), 4 deletions(-) diff --git a/reme/components/as_llm/__init__.py b/reme/components/as_llm/__init__.py index 335801ddd..0bb9c0f3d 100644 --- a/reme/components/as_llm/__init__.py +++ b/reme/components/as_llm/__init__.py @@ -1,5 +1,7 @@ """LLM model wrappers for AgentScope.""" +import asyncio + from agentscope.credential import ( AnthropicCredential, CredentialBase, @@ -32,6 +34,31 @@ def __init__(self, **kwargs) -> None: super().__init__(**kwargs) self._model: ChatModelBase | None = None self._owns_model = False + self._closed = False + + async def start(self) -> None: + """Start the wrapper and clean up a lazy model if startup fails early.""" + try: + await super().start() + except BaseException: + if self._owns_model and self._model is not None: + try: + await self._close() + except BaseException: + self.logger.exception("Failed to clean up model after startup failure") + raise + self._closed = False + + async def close(self) -> None: + """Close started or pre-start resources while preserving base semantics elsewhere.""" + try: + async with self._lock: + if not self._is_started: + await self._close() + return + await super().close() + finally: + self._closed = True async def _start(self) -> None: """Keep service startup independent of provider credentials.""" @@ -53,6 +80,8 @@ def initialize_model(self) -> None: """Construct the configured provider once, without making a remote request.""" if self._model is not None: return + if self._closed: + raise RuntimeError("Start the LLM component before accessing its model after close.") kwargs = dict(self.kwargs) credential = self.credential_cls(**kwargs.pop("credential", {})) model_cls = credential.get_chat_model_class() @@ -78,12 +107,26 @@ def model(self, value: ChatModelBase | None) -> None: BaseAsLLM.model.fset(self, value) async def _close(self) -> None: - """Close an internally owned async client on component shutdown.""" + """Close an internally owned async client without losing cancellation-safe cleanup.""" if not self._owns_model or self._model is None: return + cleanup = asyncio.create_task(self._release_model()) + cancellation: asyncio.CancelledError | None = None + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + cleanup.result() + if cancellation is not None: + raise cancellation + + async def _release_model(self) -> None: + """Release the owned model and clear state even when provider close fails.""" model = self._model - client = getattr(model, "client", None) try: + client = getattr(model, "client", None) if client is not None: await client.close() finally: diff --git a/tests/unit/test_as_llm_lazy.py b/tests/unit/test_as_llm_lazy.py index 5ae3a7516..33907df78 100644 --- a/tests/unit/test_as_llm_lazy.py +++ b/tests/unit/test_as_llm_lazy.py @@ -139,6 +139,7 @@ async def test_injected_openai_client_is_borrowed(): await llm.close() model.client.close.assert_not_awaited() assert llm._model is model + assert llm.model is model @pytest.mark.asyncio @@ -163,6 +164,113 @@ async def test_owned_openai_client_close_failure_clears_state(): assert not llm.is_started await llm.close() model_cls.return_value.client.close.assert_awaited_once_with() + with pytest.raises(RuntimeError, match="after close"): + _ = llm.model + credential_cls.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_repeated_cancellation_finishes_owned_openai_client_close(): + """Repeated shutdown cancellation cannot detach the provider close task.""" + from reme.components.as_llm import OpenAIAsLLM + + llm = OpenAIAsLLM(backend="openai", model="synthetic") + model_cls = Mock() + model_cls.return_value.client.close = AsyncMock() + closing, release, finished = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def close_client(): + closing.set() + await release.wait() + finished.set() + + model_cls.return_value.client.close.side_effect = close_client + credential_cls = Mock() + credential_cls.return_value.get_chat_model_class.return_value = model_cls + llm.credential_cls = credential_cls + await llm.start() + model = llm.model + + close_task = asyncio.create_task(llm.close()) + await asyncio.wait_for(closing.wait(), 2) + for message in ("first cancellation", "second cancellation", "third cancellation"): + close_task.cancel(message) + await asyncio.sleep(0) + assert not close_task.done() + assert llm._model is model + assert llm._owns_model + with pytest.raises(RuntimeError, match="Close the OpenAI component"): + llm.model = Mock() + + release.set() + with pytest.raises(asyncio.CancelledError, match="first cancellation"): + await asyncio.wait_for(close_task, 2) + assert finished.is_set() + assert llm._model is None + assert not llm._owns_model + assert not llm.is_started + model_cls.return_value.client.close.assert_awaited_once_with() + assert not any( + not task.done() and getattr(task.get_coro(), "__qualname__", "") == "OpenAIAsLLM._release_model" + for task in asyncio.all_tasks() + ) + + await llm.close() + model_cls.return_value.client.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_pre_start_owned_openai_model_is_closed_and_restartable(): + """A model created before start is closed by an unstarted close and can restart.""" + from reme.components.as_llm import OpenAIAsLLM + + llm = OpenAIAsLLM(backend="openai", model="synthetic") + first_model, second_model = Mock(), Mock() + first_model.client.close = AsyncMock() + second_model.client.close = AsyncMock() + model_cls = Mock(side_effect=[first_model, second_model]) + credential_cls = Mock() + credential_cls.return_value.get_chat_model_class.return_value = model_cls + llm.credential_cls = credential_cls + + assert llm.model is first_model + assert not llm.is_started + await llm.close() + first_model.client.close.assert_awaited_once_with() + assert llm._model is None + + with pytest.raises(RuntimeError, match="after close"): + _ = llm.model + await llm.close() + first_model.client.close.assert_awaited_once_with() + + await llm.start() + assert llm.model is second_model + await llm.close() + second_model.client.close.assert_awaited_once_with() + assert model_cls.call_count == 2 + + +def test_pre_start_borrowed_openai_model_can_be_replaced(): + """Pre-start injection and replacement remain borrowed and unclosed.""" + from reme.components.as_llm import OpenAIAsLLM + + llm = OpenAIAsLLM(backend="openai", model="synthetic") + first_model, second_model = Mock(), Mock() + first_model.client.close = AsyncMock() + second_model.client.close = AsyncMock() + llm.model = first_model + llm.model = second_model + + async def run(): + await llm.close() + assert llm.model is second_model + await llm.start() + await llm.close() + + asyncio.run(run()) + first_model.client.close.assert_not_awaited() + second_model.client.close.assert_not_awaited() @pytest.mark.asyncio @@ -192,13 +300,21 @@ async def test_owned_openai_model_requires_close_before_replacement(): model_cls.return_value.client.close.assert_awaited_once_with() +@pytest.mark.parametrize("failure_stage", ["resolve", "hook"]) @pytest.mark.parametrize("injected", [False, True], ids=["owned", "borrowed"]) -def test_openai_start_failure_rolls_back_owned_client(injected): +def test_openai_start_failure_rolls_back_owned_client(injected, failure_stage): """Startup failure closes a pre-start owned client but never a borrowed one.""" from reme.components.as_llm import OpenAIAsLLM class FailingOpenAIAsLLM(OpenAIAsLLM): + async def _resolve_bindings(self): + if failure_stage == "resolve": + raise RuntimeError("resolve failed") + await super()._resolve_bindings() + async def _start(self): + if failure_stage == "resolve": + await super()._start() raise RuntimeError("start failed") llm = FailingOpenAIAsLLM(backend="openai", model="synthetic") @@ -214,7 +330,8 @@ async def _start(self): assert llm.model is model async def run(): - with pytest.raises(RuntimeError, match="start failed"): + expected = "resolve failed" if failure_stage == "resolve" else "start failed" + with pytest.raises(RuntimeError, match=expected): await llm.start() assert not llm.is_started assert llm._model is (model if injected else None) From 0606dc5f631d36a8987c80862593e29ee5377418 Mon Sep 17 00:00:00 2001 From: fredheathcote007 Date: Fri, 6 Nov 2026 22:28:36 +0000 Subject: [PATCH 3/3] test(as_llm): document startup failure fixture --- tests/unit/test_as_llm_lazy.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/test_as_llm_lazy.py b/tests/unit/test_as_llm_lazy.py index 33907df78..58240809c 100644 --- a/tests/unit/test_as_llm_lazy.py +++ b/tests/unit/test_as_llm_lazy.py @@ -307,6 +307,8 @@ def test_openai_start_failure_rolls_back_owned_client(injected, failure_stage): from reme.components.as_llm import OpenAIAsLLM class FailingOpenAIAsLLM(OpenAIAsLLM): + """Simulate failures at different stages of component startup.""" + async def _resolve_bindings(self): if failure_stage == "resolve": raise RuntimeError("resolve failed")