fix(as_llm): close owned OpenAI clients safely

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):

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())