diff --git a/reme/components/as_llm/__init__.py b/reme/components/as_llm/__init__.py index 7ea4e1bf0..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, @@ -31,6 +33,32 @@ class BaseAsLLM(BaseComponent): 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.""" @@ -46,17 +74,21 @@ 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.""" 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() 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 +97,42 @@ 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 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 + try: + client = getattr(model, "client", None) + 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..58240809c 100644 --- a/tests/unit/test_as_llm_lazy.py +++ b/tests/unit/test_as_llm_lazy.py @@ -98,3 +98,246 @@ 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 + 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() + 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 +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("failure_stage", ["resolve", "hook"]) +@pytest.mark.parametrize("injected", [False, True], ids=["owned", "borrowed"]) +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): + """Simulate failures at different stages of component startup.""" + + 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") + 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(): + 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) + assert model.client.close.await_count == (0 if injected else 1) + await llm.close() + + asyncio.run(run())