feat(decision): allow hosts to inject owned JEV model instances (#652)

* feat(decision): support host-owned model injection

* fix(decision): finish model cleanup before propagating shutdown cancellation

docs/en/context_compaction.md

@@ -159,3 +159,20 @@ the entire gain cannot be attributed to Jev ranking itself. Agent continuation s
 
 <!-- Payment has 10 units (maximum Top 10); actual N is min(N, available units). -->
 [Complete experiment: design, results, expected answers and reproduction](https://github.com/jinliyl/ReMe/blob/74e1251ddcd212205788f21ccd6b5ebd82ade7c3/benchmark/compaction/README.md)
+
+## Host-owned decision models
+
+Embedded hosts can reuse their existing AgentScope JEV instance, just as they inject an LLM:
+
+```python
+# app contains an as_decision_model.compaction component with backend: agentscope.
+await app.update_component("as_decision_model", "compaction", model=host_jev)
+```
+
+The YAML `model` value remains the provider model name. The runtime `model` property accepts a callable instance
+and does not construct a client on access. Inject before the first request; ReMe validates and converts JSON
+questions normally, but does not require separate ReMe credentials or close the borrowed client. The host must
+keep the instance alive until ReMe has finished its calls, then close it itself. Close the component before
+replacing an internally constructed model or any model with active requests. Shutdown waits for active calls and
+owned-client cleanup even when cancelled, then propagates cancellation. Component shutdown detaches the instance;
+inject again before restarting a borrowed-model component. Standalone credential-based behavior is unchanged.

docs/zh/context_compaction.md

@@ -147,3 +147,18 @@ agent 继续执行的成功率尚未验证。
 
 <!-- 支付轨迹有 10 个证据块，最多 Top 10；实际 N 为 min(N, 可用证据块数)。 -->
 [完整实验：设计、结果、标准答案与复现](https://github.com/jinliyl/ReMe/blob/74e1251ddcd212205788f21ccd6b5ebd82ade7c3/benchmark/compaction/README_ZH.md)
+
+## 注入宿主管理的决策模型
+
+嵌入式宿主可以复用已有的 AgentScope JEV 实例，与注入 LLM 的方式一致：
+
+```python
+# app 已配置 backend: agentscope 的 as_decision_model.compaction 组件。
+await app.update_component("as_decision_model", "compaction", model=host_jev)
+```
+
+YAML 中的 `model` 仍表示供应商模型名称；运行时 `model` 属性接收可调用的模型实例，读取属性不会创建客户端。
+在首次请求之前注入。ReMe 继续校验并转换 JSON 问题，但不再要求额外的 ReMe 凭据，也不会关闭借用的客户端。
+宿主应等待 ReMe 调用结束后自行关闭实例。替换 ReMe 自建模型或存在进行中请求的实例之前，须先关闭组件。
+关闭过程即使被取消，也会等待活动请求和自建客户端清理完成，再传播取消信号。组件关闭时会解除实例引用；
+重新启动使用外部模型的组件前需要再次注入。独立运行时按凭据创建模型的行为保持不变。

reme/components/as_decision_model/agentscope.py

@@ -31,8 +31,8 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel):
     """Wrap AgentScope's TypeSafe decision model with lazy provider loading.
 
     ``questions`` accepts AgentScope question objects or JSON dictionaries with
-    a ``type`` discriminator (``binary``, ``choice``, or ``score``). Only this
-    component's configuration supplies credentials and transport options.
+    a ``type`` discriminator (``binary``, ``choice``, or ``score``). Hosts may inject a borrowed model
+    through ``model``; otherwise this component constructs and owns its client.
     """
 
     def __init__(
@@ -53,10 +53,31 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel):
             retry_delay=retry_delay,
         )
         self._model: "ClassifierModelBase | None" = None
+        self._owns_model = False
         self._active_calls = 0
         self._idle = asyncio.Event()
         self._idle.set()
 
+    @property
+    def model(self) -> "ClassifierModelBase | None":
+        """Return the current instance without constructing a provider client."""
+        return self._model
+
+    @model.setter
+    def model(self, value: "ClassifierModelBase | None") -> None:
+        """Borrow a host-owned instance, including via Application.update_component.
+
+        Close the component before replacing an internally owned client. An
+        active request must finish before any replacement or detachment.
+        """
+        if value is self._model:
+            return
+        if value is not None and not callable(value):
+            raise TypeError("An injected decision model must be callable.")
+        if self._active_calls or self._owns_model:
+            raise RuntimeError("Close the decision component before replacing its active or owned model.")
+        self._model = value
+
     @staticmethod
     def _load_decision_model():
         """Load the optional AgentScope adapter on the first decision request."""
@@ -102,6 +123,7 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel):
                         credential=credential_cls(**credential),
                         **self.options.model_dump(exclude_none=True),
                     )
+                    self._owns_model = True
                 except ImportError as exc:
                     raise ImportError(
                         "Decision model inference requires a compatible typesafe-sdk. "
@@ -120,10 +142,24 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel):
                 self._idle.set()
 
     async def _close(self) -> None:
+        """Finish draining and cleanup before propagating shutdown cancellation."""
+        cleanup = asyncio.create_task(self._release_model())
+        cancellation = None
+        while not cleanup.done():
+            try:
+                await asyncio.shield(cleanup)
+            except asyncio.CancelledError as exc:
+                # Keep the lifecycle lock until cleanup finishes, even after
+                # repeated cancellation. No cleanup task outlives close().
+                cancellation = exc
+        cleanup.result()
+        if cancellation is not None:
+            raise cancellation
+
+    async def _release_model(self) -> None:
         """Release the owned TypeSafe client and allow a later fresh start."""
-        try:
-            await self._idle.wait()
-        finally:
-            model, self._model = self._model, None
-            if model is not None:
-                await model.client.aclose()
+        await self._idle.wait()
+        model, owned = self._model, self._owns_model
+        self._model, self._owns_model = None, False
+        if model is not None and owned:
+            await model.client.aclose()

tests/unit/test_decision_model.py

@@ -238,6 +238,88 @@ def test_failed_or_cancelled_call_does_not_block_close(provider, failure):
     asyncio.run(run())
 
 
+@pytest.mark.parametrize("borrowed", [False, True])
+def test_cancelled_close_drains_requests_and_releases_model(provider, borrowed):
+    """Shutdown defers cancellation until active calls and model cleanup finish."""
+    _, model = provider
+
+    async def run():
+        started, release = asyncio.Event(), asyncio.Event()
+
+        async def reply(**_kwargs):
+            started.set()
+            await release.wait()
+
+        model.side_effect = reply
+        component = AgentScopeAsDecisionModel(credential={"api_key": "test-key"})
+        if borrowed:
+            component.model = model
+        call = asyncio.create_task(component(state="hello", questions={"q": {"type": "binary"}}))
+        await asyncio.wait_for(started.wait(), 2)
+        close = asyncio.create_task(component.close())
+        await asyncio.sleep(0)
+        try:
+            for _ in range(2):
+                close.cancel()
+                await asyncio.sleep(0)
+                assert not close.done()
+                assert component.is_started
+                assert component.model is model
+                model.client.aclose.assert_not_awaited()
+            release.set()
+            await asyncio.wait_for(call, 2)
+            with pytest.raises(asyncio.CancelledError):
+                await asyncio.wait_for(close, 2)
+            assert component.model is None
+            assert not component.is_started
+            assert not component._owns_model
+            await component.close()
+            assert model.client.aclose.await_count == (0 if borrowed else 1)
+        finally:
+            release.set()
+            await asyncio.gather(call, close, return_exceptions=True)
+
+    asyncio.run(run())
+
+
+def test_cancelled_close_finishes_owned_client_cleanup(provider):
+    """Repeated shutdown cancellation cannot interrupt the client's own close."""
+    _, model = provider
+
+    async def run():
+        closing, release, finished = asyncio.Event(), asyncio.Event(), asyncio.Event()
+
+        async def close_client():
+            closing.set()
+            await release.wait()
+            finished.set()
+
+        model.client.aclose.side_effect = close_client
+        component = AgentScopeAsDecisionModel(credential={"api_key": "test-key"})
+        await component(state="hello", questions={"q": {"type": "binary"}})
+        close = asyncio.create_task(component.close())
+        await asyncio.wait_for(closing.wait(), 2)
+        try:
+            for _ in range(2):
+                close.cancel()
+                await asyncio.sleep(0)
+                assert not close.done()
+                assert component.is_started
+            release.set()
+            with pytest.raises(asyncio.CancelledError):
+                await asyncio.wait_for(close, 2)
+            assert finished.is_set()
+            assert component.model is None
+            assert not component.is_started
+            await component.close()
+            model.client.aclose.assert_awaited_once()
+        finally:
+            release.set()
+            await asyncio.gather(close, return_exceptions=True)
+
+    asyncio.run(run())
+
+
 def test_client_close_error_clears_cached_model(provider):
     """A failed close cannot leave a stale client selected on the next start."""
     factory, model = provider
@@ -350,3 +432,80 @@ def test_adapter_uses_sdk_default_or_explicit_model(monkeypatch, model_name):
         model.client.aclose.assert_awaited_once()
 
     asyncio.run(run())
+
+
+def test_injected_model_is_borrowed_and_accepts_json_questions(provider, tmp_path):
+    """Application injection shares the host instance without owning its client."""
+    from reme.application import Application
+
+    factory, model = provider
+
+    async def run():
+        app = Application(
+            workspace_dir=str(tmp_path),
+            enable_logo=False,
+            log_to_console=False,
+            log_to_file=False,
+            service={"backend": "http"},
+            components={"as_decision_model": {"compaction": {"backend": "agentscope"}}},
+        )
+        component = await app.update_component("as_decision_model", "compaction", model=model)
+        assert component.model is model
+        async with app:
+            await component(state="hello", questions={"q": {"type": "binary"}})
+        factory.assert_not_called()
+        model.assert_awaited_once()
+        model.client.aclose.assert_not_awaited()
+        assert component.model is None
+
+    asyncio.run(run())
+
+
+def test_model_replacement_guards_owned_clients(provider):
+    """Replacing an owned client requires an explicit close to prevent leaks."""
+    _, model = provider
+
+    async def run():
+        component = AgentScopeAsDecisionModel(credential={"api_key": "test-key"})
+        await component(state="hello", questions={"q": {"type": "binary"}})
+        component.model = model
+        with pytest.raises(RuntimeError, match="Close"):
+            component.model = None
+        await component.close()
+        component.model = model
+        await component.close()
+        # The borrowed instance is never closed a second time.
+        model.client.aclose.assert_awaited_once()
+        with pytest.raises(TypeError, match="callable"):
+            component.model = "model-name"
+
+    asyncio.run(run())
+
+
+def test_borrowed_model_cannot_be_replaced_while_in_use(provider):
+    """Host ownership still obeys the component's in-flight request boundary."""
+    _, model = provider
+
+    async def run():
+        started, release = asyncio.Event(), asyncio.Event()
+
+        async def reply(**_kwargs):
+            started.set()
+            await release.wait()
+
+        model.side_effect = reply
+        component = AgentScopeAsDecisionModel()
+        component.model = model
+        task = asyncio.create_task(component(state="hello", questions={"q": {"type": "binary"}}))
+        await asyncio.wait_for(started.wait(), 2)
+        with pytest.raises(RuntimeError, match="Close"):
+            component.model = None
+        close = asyncio.create_task(component.close())
+        await asyncio.sleep(0)
+        assert not close.done()
+        release.set()
+        await task
+        await close
+        model.client.aclose.assert_not_awaited()
+
+    asyncio.run(run())