fix(core): preserve state during component replacement

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:

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