diff --git a/reme/application.py b/reme/application.py index a1b584e8..63b41240 100644 --- a/reme/application.py +++ b/reme/application.py @@ -248,6 +248,26 @@ class Application(BaseComponent): setattr(component, key, value) return component + def _replacement_shutdown_order( + self, + old_component: BaseComponent, + replacement: BaseComponent, + replacement_order: list[BaseComponent], + ) -> list[BaseComponent]: + """Precompute shutdown tracking using identity, not component hashing.""" + started_ids = {id(component) for component in self._started_components} + started_ids.discard(id(old_component)) + started_ids.add(id(replacement)) + component_ids = { + id(component) for components in self.context.components.values() for component in components.values() + } + non_components = [ + component + for component in self._started_components + if component is not old_component and id(component) not in component_ids + ] + return [component for component in replacement_order if id(component) in started_ids] + non_components + async def replace_component( self, component_enum: ComponentType, @@ -257,18 +277,19 @@ class Application(BaseComponent): config: ComponentConfig | Mapping[str, Any], runtime_updates: Mapping[str, Any] | None = None, ) -> BaseComponent: - """Atomically replace an existing component and its bind() references. + """Replace an existing component and synchronously commit its references. ``config`` is the complete declarative configuration for the new component. ``runtime_updates`` injects non-serializable live objects, such as an already verified model, before the replacement is started. - The new component is constructed, validated, and started while the old - component remains visible. The context, dependent bindings, persisted - application config, and shutdown order are then switched without an - await boundary. A start failure leaves the old generation untouched. - Hosts must still quiesce calls that may retain component references - across this operation. + The old component is dumped before the replacement starts so compatible + ``start()`` / ``load()`` implementations see its latest state. The + context, dependent bindings, in-memory application config, and shutdown + order are then switched without an await boundary. A dump or start + failure leaves the old generation authoritative. Hosts must quiesce + calls that may retain component references across this operation and + migrate state explicitly when changing between incompatible backends. """ async with self._component_mutation_lock: component_type = component_type_name(component_enum) @@ -300,8 +321,6 @@ class Application(BaseComponent): node_key = (component_type, name) replacement_order = Application._topological_order(self, replacement=(node_key, replacement)) was_started = old_component.is_started - if was_started: - await replacement.start() consumers: list[tuple[BaseComponent, str]] = [] candidates: list[BaseComponent] = [ @@ -321,27 +340,30 @@ class Application(BaseComponent): ): consumers.append((consumer, attr)) + replacement_started_components = ( + self._replacement_shutdown_order(old_component, replacement, replacement_order) if was_started else None + ) + if was_started: + await old_component.dump() + try: + await replacement.start() + except BaseException: + try: + await replacement.close() + except BaseException as cleanup_error: + self.logger.exception( + f"Failed to clean up replacement {component_type}:{name}: {cleanup_error}", + ) + raise + # Commit the new generation synchronously so observers cannot see # a context with only some dependency references updated. group[name] = replacement config_group[name] = replacement_config for consumer, attr in consumers: consumer.__dict__[attr] = replacement - if was_started: - started = set(self._started_components) - started.discard(old_component) - started.add(replacement) - component_set = { - component for components in self.context.components.values() for component in components.values() - } - non_components = [ - component - for component in self._started_components - if component is not old_component and component not in component_set - ] - self._started_components = [ - component for component in replacement_order if component in started - ] + non_components + if replacement_started_components is not None: + self._started_components = replacement_started_components if was_started: try: diff --git a/tests/unit/test_embedded_consumer_compat.py b/tests/unit/test_embedded_consumer_compat.py index 4da967d0..40d12795 100644 --- a/tests/unit/test_embedded_consumer_compat.py +++ b/tests/unit/test_embedded_consumer_compat.py @@ -10,6 +10,7 @@ from reme import ReMe from reme.components.agent_wrapper import AsAgentWrapper from reme.components.as_llm import BaseAsLLM, DashScopeAsLLM from reme.enumeration import ComponentEnum +from reme.schema import FileNode def _qwenpaw_style_config(workspace_dir: str) -> dict: @@ -46,6 +47,18 @@ def _qwenpaw_style_config(workspace_dir: str) -> dict: } +def _file_graph_config(workspace_dir: str) -> dict: + """Return a minimal application with one persistent file graph.""" + return { + "workspace_dir": workspace_dir, + "enable_logo": False, + "log_to_console": False, + "log_to_file": False, + "service": {"backend": "http"}, + "components": {"file_graph": {"default": {"backend": "local"}}}, + } + + def test_qwenpaw_style_config_preserves_optional_defaults(tmp_path): """New application fields remain optional for existing embedded configs.""" app = ReMe(**_qwenpaw_style_config(str(tmp_path))) @@ -168,6 +181,107 @@ def test_replace_component_before_start_preserves_dependency_order(tmp_path): asyncio.run(exercise_api()) +def test_replace_component_flushes_persistent_state_before_start(tmp_path): + """A replacement loads runtime state that the old generation had not dumped.""" + config = _file_graph_config(str(tmp_path)) + app = ReMe(**config) + + async def exercise_api() -> None: + await app.start() + old_graph = app.context.components[ComponentEnum.FILE_GRAPH]["default"] + await old_graph.upsert_nodes([FileNode(path="memory.md", st_mtime=1.0)]) + + replacement = await app.replace_component( + "file_graph", + "default", + config={"backend": "local"}, + ) + + assert [node.path for node in await replacement.get_nodes()] == ["memory.md"] + await app.close() + + restored_app = ReMe(**config) + await restored_app.start() + restored_graph = restored_app.context.components[ComponentEnum.FILE_GRAPH]["default"] + assert [node.path for node in await restored_graph.get_nodes()] == ["memory.md"] + await restored_app.close() + + asyncio.run(exercise_api()) + + +def test_replace_component_accepts_unhashable_plugin_component(tmp_path): + """Shutdown-order calculation relies on identity for legal unhashable plugins.""" + + class UnhashableAsLLM(BaseAsLLM): + """Plugin component whose equality contract intentionally disables hashing.""" + + component_type = ComponentEnum.AS_LLM + __hash__ = None + + def __eq__(self, other) -> bool: + return self is other + + app = ReMe(**_qwenpaw_style_config(str(tmp_path))) + app.context.registry.add("unhashable", UnhashableAsLLM, owner="test") + + async def exercise_api() -> None: + await app.update_component("as_llm", "default", model=object()) + await app.start() + + replacement = await app.replace_component( + "as_llm", + "default", + config={"backend": "unhashable"}, + runtime_updates={"model": object()}, + ) + + assert any(component is replacement for component in app._started_components) + await app.close() + assert replacement.is_started is False + + asyncio.run(exercise_api()) + + +def test_replace_component_dump_failure_keeps_old_generation(tmp_path): + """A failed state flush prevents replacement startup and public mutation.""" + + class ObservedAsLLM(BaseAsLLM): + """Replacement backend that records whether startup was attempted.""" + + component_type = ComponentEnum.AS_LLM + start_calls = 0 + + async def _start(self) -> None: + type(self).start_calls += 1 + + app = ReMe(**_qwenpaw_style_config(str(tmp_path))) + app.context.registry.add("observed", ObservedAsLLM, owner="test") + + async def failing_dump() -> None: + raise RuntimeError("state flush failed") + + async def exercise_api() -> None: + old_component = await app.update_component("as_llm", "default", model=object()) + await app.start() + wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] + old_component.dump = failing_dump + + with pytest.raises(RuntimeError, match="state flush failed"): + await app.replace_component( + "as_llm", + "default", + config={"backend": "observed"}, + ) + + assert ObservedAsLLM.start_calls == 0 + assert app.context.components[ComponentEnum.AS_LLM]["default"] is old_component + assert wrapper.as_llm is old_component + assert old_component.is_started is True + await app.close() + + asyncio.run(exercise_api()) + + def test_replace_component_waits_for_application_start(tmp_path): """Replacement cannot race the dependency graph while startup is in progress.""" app = ReMe(**_qwenpaw_style_config(str(tmp_path)))