From c85917a8128eeaf8efda686f822304db6d7739f6 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Fri, 28 Aug 2026 16:04:02 +0800 Subject: [PATCH] feat(core): support live component replacement (#505) * feat(core): support live component replacement * fix(core): serialize replacement with app lifecycle * fix(core): preserve state during component replacement * fix(core): roll back partial component startup * chore(deps): add 'web' extra to reme-ai and remove reme_studio from core dependencies - Updated reme-ai dependency to include 'as' and 'web' extras - Removed reme_studio from core dependency list to avoid duplication or unnecessary install * fix(core): preserve persisted state on startup failure --- reme/application.py | 188 ++++++++++-- reme/components/base_component.py | 56 +++- .../embedding_store/local_embedding_store.py | 3 +- .../file_catalog/base_file_catalog.py | 3 +- reme/components/file_graph/base_file_graph.py | 3 +- .../components/file_store/local_file_store.py | 9 +- .../keyword_index/base_keyword_index.py | 3 +- tests/unit/test_base_component.py | 68 +++++ tests/unit/test_embedded_consumer_compat.py | 284 ++++++++++++++++++ tests/unit/test_file_catalog.py | 24 ++ tests/unit/test_job.py | 2 + 11 files changed, 599 insertions(+), 44 deletions(-) diff --git a/reme/application.py b/reme/application.py index fb811a9d..b4b0cb69 100644 --- a/reme/application.py +++ b/reme/application.py @@ -2,9 +2,10 @@ import asyncio import heapq +from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from pathlib import Path -from typing import AsyncGenerator, TypeVar +from typing import Any, AsyncGenerator, TypeVar from . import __version__ from .components import ApplicationContext, BaseComponent @@ -26,6 +27,7 @@ class Application(BaseComponent): runtime = resolve_plugin_runtime(kwargs) self.context = ApplicationContext(registry=runtime.registry, **runtime.config) self._started_components: list[BaseComponent] = [] + self._component_mutation_lock = asyncio.Lock() self._setup_workspace_directories() logger = get_logger( @@ -133,11 +135,16 @@ class Application(BaseComponent): # ----- Dependency ordering ------------------------------------------ - def _topological_order(self) -> list[BaseComponent]: + def _topological_order( + self, + replacement: tuple[_NodeKey, BaseComponent] | None = None, + ) -> list[BaseComponent]: """Return components in dependency order via Kahn's algorithm; raise on missing dep or cycle.""" nodes: dict[_NodeKey, BaseComponent] = { (ctype, name): comp for ctype, group in self.context.components.items() for name, comp in group.items() } + if replacement is not None: + nodes[replacement[0]] = replacement[1] in_degree, dependents = self._build_dependency_graph(nodes) ready = [k for k, d in in_degree.items() if d == 0] @@ -179,22 +186,23 @@ class Application(BaseComponent): async def _start(self) -> None: """Start components, then jobs as base > stream > background > cron.""" - pool_size = self.config.thread_pool_max_workers - if pool_size > 0: - self.context.thread_pool = ThreadPoolExecutor(max_workers=pool_size) - self.logger.info(f"Thread pool created with max_workers={pool_size}") - try: - components = self._topological_order() - jobs = list(self.context.jobs.values()) - base_jobs = [j for j in jobs if not isinstance(j, (StreamJob, BackgroundJob))] - stream_jobs = [j for j in jobs if isinstance(j, StreamJob)] - background_jobs = [j for j in jobs if isinstance(j, BackgroundJob) and not isinstance(j, CronJob)] - cron_jobs = [j for j in jobs if isinstance(j, CronJob)] - for c in components + base_jobs + stream_jobs + background_jobs + cron_jobs: - await self._start_one(c) - except Exception: - await self._close() - raise + async with self._component_mutation_lock: + pool_size = self.config.thread_pool_max_workers + if pool_size > 0: + self.context.thread_pool = ThreadPoolExecutor(max_workers=pool_size) + self.logger.info(f"Thread pool created with max_workers={pool_size}") + try: + components = self._topological_order() + jobs = list(self.context.jobs.values()) + base_jobs = [j for j in jobs if not isinstance(j, (StreamJob, BackgroundJob))] + stream_jobs = [j for j in jobs if isinstance(j, StreamJob)] + background_jobs = [j for j in jobs if isinstance(j, BackgroundJob) and not isinstance(j, CronJob)] + cron_jobs = [j for j in jobs if isinstance(j, CronJob)] + for c in components + base_jobs + stream_jobs + background_jobs + cron_jobs: + await self._start_one(c) + except Exception: + await self._close_started_components() + raise async def _start_one(self, c: BaseComponent) -> None: """Start one component and record it for ordered shutdown.""" @@ -209,6 +217,11 @@ class Application(BaseComponent): async def _close(self) -> None: """Close in reverse start order so every peer outlives its dependents.""" + async with self._component_mutation_lock: + await self._close_started_components() + + async def _close_started_components(self) -> None: + """Close resources while the caller serializes component mutations.""" for c in reversed(self._started_components): try: await c.close() @@ -221,17 +234,136 @@ class Application(BaseComponent): async def update_component(self, component_enum: ComponentType, name: str, /, **kwargs) -> BaseComponent: """Update an existing component by type/name; never creates missing components.""" - component_type = component_type_name(component_enum) - group = self.context.components.get(component_type) - if not group or name not in group: - raise KeyError(f"Component '{name}' not found in {component_type}") + async with self._component_mutation_lock: + component_type = component_type_name(component_enum) + group = self.context.components.get(component_type) + if not group or name not in group: + raise KeyError(f"Component '{name}' not found in {component_type}") - component = group[name] - for key, value in kwargs.items(): - if not hasattr(component, key): - raise AttributeError(f"Component {component_type}:{name} has no attribute '{key}'") - setattr(component, key, value) - return component + component = group[name] + for key in kwargs: + if not hasattr(component, key): + raise AttributeError(f"Component {component_type}:{name} has no attribute '{key}'") + for key, value in kwargs.items(): + 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, + name: str, + /, + *, + config: ComponentConfig | Mapping[str, Any], + runtime_updates: Mapping[str, Any] | None = None, + ) -> BaseComponent: + """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 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) + group = self.context.components.get(component_type) + config_group = self.config.components.get(component_type) + if not group or name not in group or config_group is None: + raise KeyError(f"Component '{name}' not found in {component_type}") + + old_component = group[name] + replacement_config = ( + config.model_copy(deep=True) + if isinstance(config, ComponentConfig) + else ComponentConfig.model_validate(dict(config)) + ) + replacement = self._instantiate( + component_type, + replacement_config, + label=f"Component '{name}'", + expected_type=BaseComponent, + name=name, + ) + for key, value in (runtime_updates or {}).items(): + if not hasattr(replacement, key): + raise AttributeError( + f"Replacement {component_type}:{name} has no attribute '{key}'", + ) + setattr(replacement, key, value) + + node_key = (component_type, name) + replacement_order = Application._topological_order(self, replacement=(node_key, replacement)) + was_started = old_component.is_started + + consumers: list[tuple[BaseComponent, str]] = [] + candidates: list[BaseComponent] = [ + component for components in self.context.components.values() for component in components.values() + ] + candidates.extend(self.context.jobs.values()) + if self.context.service is not None: + candidates.append(self.context.service) + for consumer in candidates: + if consumer is old_component: + continue + for attr, dependency in consumer.dependency_bindings.items(): + if ( + dependency.ctype == component_type + and dependency.name == name + and consumer.__dict__.get(attr) is old_component + ): + 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() + await replacement.start() + + # 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 replacement_started_components is not None: + self._started_components = replacement_started_components + + if was_started: + try: + await old_component.close() + except Exception as exc: # The committed replacement remains authoritative. + self.logger.exception( + f"Failed to close replaced component {component_type}:{name}: {exc}", + ) + return replacement # ----- Job execution ------------------------------------------------- diff --git a/reme/components/base_component.py b/reme/components/base_component.py index efd45f6a..74ef4633 100644 --- a/reme/components/base_component.py +++ b/reme/components/base_component.py @@ -98,6 +98,10 @@ class BaseComponent(ComponentMixin, ABC): self._is_started: bool = False self._lock: asyncio.Lock = asyncio.Lock() + # Preserve bind() declarations after their Dependency placeholders are + # resolved. Application uses these specs to validate and rewire a live + # component replacement without guessing from arbitrary attributes. + self._binding_specs: dict[str, Dependency] = {} # Components created via bind() default_factory in standalone mode; # their lifecycle is owned by this component. self._owned: list["BaseComponent"] = [] @@ -138,13 +142,23 @@ class BaseComponent(ComponentMixin, ABC): @property def dependencies(self) -> list[Dependency]: - """All unresolved dependency placeholders on this instance.""" - return [v for v in self.__dict__.values() if isinstance(v, Dependency)] + """All declared dependencies, including bindings resolved at start.""" + bindings = dict(self._binding_specs) + bindings.update((attr, value) for attr, value in self.__dict__.items() if isinstance(value, Dependency)) + return list(bindings.values()) + + @property + def dependency_bindings(self) -> dict[str, Dependency]: + """Map dependency attributes to their stable bind specifications.""" + bindings = dict(self._binding_specs) + bindings.update((attr, value) for attr, value in self.__dict__.items() if isinstance(value, Dependency)) + return bindings async def _resolve_bindings(self) -> None: """Replace every ``Dependency`` attribute with its resolved target.""" for attr, dep in list(self.__dict__.items()): if isinstance(dep, Dependency): + self._binding_specs[attr] = dep self._resolve_one(attr, dep) def _resolve_one(self, attr: str, dep: Dependency) -> None: @@ -209,15 +223,41 @@ class BaseComponent(ComponentMixin, ABC): # ----- Lifecycle control -------------------------------------------- async def start(self) -> None: - """Start the component once: resolve deps → start owned → run _start.""" + """Start once, rolling back partial resources when any startup stage fails.""" async with self._lock: if self._is_started: return - await self._resolve_bindings() - for owned in self._owned: - await owned.start() - await self._start() - self._is_started = True + started_owned: list[BaseComponent] = [] + start_hook_entered = False + try: + await self._resolve_bindings() + for owned in self._owned: + await owned.start() + started_owned.append(owned) + start_hook_entered = True + await self._start() + self._is_started = True + except BaseException: + await self._rollback_start(start_hook_entered, started_owned) + raise + + async def _rollback_start( + self, + start_hook_entered: bool, + started_owned: list["BaseComponent"], + ) -> None: + """Best-effort cleanup that preserves the original startup error.""" + if start_hook_entered: + try: + await self._close() + except BaseException as exc: + self.logger.exception(f"Failed to roll back partially started component {self.name}: {exc}") + for owned in reversed(started_owned): + try: + await owned.close() + except BaseException as exc: + self.logger.exception(f"Failed to close owned component {owned.name} during startup rollback: {exc}") + self._is_started = False async def close(self) -> None: """Close the component once: run _close → close owned in reverse order.""" diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 5c9b557b..84d2c293 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -67,7 +67,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore): await self.load() async def _close(self) -> None: - await self.dump() + if self.is_started: + await self.dump() async def health_check(self, timeout: float | None = None) -> bool: timeout = self.health_check_timeout if timeout is None else timeout diff --git a/reme/components/file_catalog/base_file_catalog.py b/reme/components/file_catalog/base_file_catalog.py index 4b6612f5..96c36315 100644 --- a/reme/components/file_catalog/base_file_catalog.py +++ b/reme/components/file_catalog/base_file_catalog.py @@ -17,7 +17,8 @@ class BaseFileCatalog(BaseComponent): await self.load() async def _close(self) -> None: - await self.dump() + if self.is_started: + await self.dump() await super()._close() async def load(self) -> None: diff --git a/reme/components/file_graph/base_file_graph.py b/reme/components/file_graph/base_file_graph.py index c34d53e0..98a14e4d 100644 --- a/reme/components/file_graph/base_file_graph.py +++ b/reme/components/file_graph/base_file_graph.py @@ -26,7 +26,8 @@ class BaseFileGraph(BaseComponent): await self.load() async def _close(self) -> None: - await self.dump() + if self.is_started: + await self.dump() await super()._close() async def load(self) -> None: diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 757b7038..2b11ea27 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -105,10 +105,11 @@ class LocalFileStore(BaseFileStore): # Persist only this store's local state here so each component writes # exactly once during shutdown. Preserve the historical dump() hook for # third-party subclasses that override it to write additional state. - if type(self).dump is LocalFileStore.dump: - await self._dump_owned_state() - else: - await self.dump() + if self.is_started: + if type(self).dump is LocalFileStore.dump: + await self._dump_owned_state() + else: + await self.dump() self.file_chunks.clear() await super()._close() diff --git a/reme/components/keyword_index/base_keyword_index.py b/reme/components/keyword_index/base_keyword_index.py index 3ae551ff..feac9e50 100644 --- a/reme/components/keyword_index/base_keyword_index.py +++ b/reme/components/keyword_index/base_keyword_index.py @@ -24,7 +24,8 @@ class BaseKeywordIndex(BaseComponent): await self.load() async def _close(self) -> None: - await self.dump() + if self.is_started: + await self.dump() @property def document_ids(self) -> Set[str]: diff --git a/tests/unit/test_base_component.py b/tests/unit/test_base_component.py index e510f032..df7e8ead 100644 --- a/tests/unit/test_base_component.py +++ b/tests/unit/test_base_component.py @@ -3,6 +3,7 @@ # pylint: disable=protected-access,missing-function-docstring,missing-class-docstring,attribute-defined-outside-init import asyncio +import contextlib import os import tempfile @@ -122,6 +123,25 @@ def test_dependencies_lists_unresolved(): assert len(deps) == 2 +def test_dependency_bindings_survive_resolution(): + async def run(): + from reme.components.application_context import ApplicationContext + + target = DepTarget(name="real_index") + ctx = ApplicationContext() + ctx.components = {ComponentEnum.KEYWORD_INDEX: {"real_index": target}} + comp = StubComponent(app_context=ctx) + comp.dep = BaseComponent.bind("real_index", DepTarget) + + await comp.start() + + assert comp.dep is target + assert comp.dependencies[0].name == "real_index" + assert comp.dependency_bindings["dep"].ctype == "keyword_index" + + asyncio.run(run()) + + # -- lifecycle ---------------------------------------------------------------- @@ -186,6 +206,54 @@ def test_close_closes_owned_when_parent_close_fails(): asyncio.run(run()) +def test_start_failure_rolls_back_partial_resources(): + class PartialStartComponent(BaseComponent): + component_type = ComponentEnum.FILE_CHUNKER + + def __init__(self): + super().__init__() + self.resource_open = False + self.task = None + self.rollback_count = 0 + + async def _start(self): + self.resource_open = True + self.task = asyncio.create_task(asyncio.Event().wait()) + raise RuntimeError("partial startup failed") + + async def _close(self): + self.rollback_count += 1 + self.resource_open = False + if self.task is not None: + self.task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self.task + self.task = None + + async def run(): + comp = PartialStartComponent() + owned = StubComponent(name="owned") + comp.dep = BaseComponent.bind( + "owned", + StubComponent, + default_factory=lambda: owned, + ) + + with pytest.raises(RuntimeError, match="partial startup failed"): + await comp.start() + + assert comp.is_started is False + assert comp.resource_open is False + assert comp.task is None + assert comp.rollback_count == 1 + assert owned.is_started is False + assert owned.close_count == 1 + await comp.close() + assert comp.rollback_count == 1 + + asyncio.run(run()) + + # -- standalone resolution ---------------------------------------------------- diff --git a/tests/unit/test_embedded_consumer_compat.py b/tests/unit/test_embedded_consumer_compat.py index 2570d04c..40d12795 100644 --- a/tests/unit/test_embedded_consumer_compat.py +++ b/tests/unit/test_embedded_consumer_compat.py @@ -1,10 +1,16 @@ """Compatibility tests for applications that embed ReMe in-process.""" +# pylint: disable=protected-access + import asyncio +import pytest + 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: @@ -41,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))) @@ -78,3 +96,269 @@ def test_qwenpaw_style_config_keeps_in_process_application_api(tmp_path): assert app.is_started is False asyncio.run(exercise_api()) + + +def test_update_component_validates_all_fields_before_mutation(tmp_path): + """A rejected field update does not leave earlier attributes changed.""" + app = ReMe(**_qwenpaw_style_config(str(tmp_path))) + component = app.context.components[ComponentEnum.AS_LLM]["default"] + original_model = component.model + + async def exercise_api() -> None: + with pytest.raises(AttributeError, match="does_not_exist"): + await app.update_component( + "as_llm", + "default", + model=object(), + does_not_exist=True, + ) + + assert component.model is original_model + + asyncio.run(exercise_api()) + + +def test_replace_component_rebinds_dependents_and_reuses_runtime_model(tmp_path): + """A live backend replacement switches the wrapper and every bind target.""" + app = ReMe(**_qwenpaw_style_config(str(tmp_path))) + old_model = object() + verified_model = object() + + async def exercise_api() -> None: + old_component = await app.update_component("as_llm", "default", model=old_model) + await app.start() + wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] + + replacement = await app.replace_component( + "as_llm", + "default", + config={ + "backend": "dashscope", + "model": "consumer-injected", + "credential": {"api_key": ""}, + }, + runtime_updates={"model": verified_model}, + ) + + assert isinstance(replacement, DashScopeAsLLM) + assert replacement.model is verified_model + assert replacement.is_started is True + assert old_component.is_started is False + assert wrapper.as_llm is replacement + assert app.context.components[ComponentEnum.AS_LLM]["default"] is replacement + assert app.config.components[ComponentEnum.AS_LLM]["default"].backend == "dashscope" + assert old_component not in app._started_components + assert replacement in app._started_components + await app.close() + + asyncio.run(exercise_api()) + + +def test_replace_component_before_start_preserves_dependency_order(tmp_path): + """Unresolved bind placeholders continue to resolve during normal startup.""" + app = ReMe(**_qwenpaw_style_config(str(tmp_path))) + verified_model = object() + + async def exercise_api() -> None: + replacement = await app.replace_component( + "as_llm", + "default", + config={ + "backend": "dashscope", + "model": "consumer-injected", + "credential": {"api_key": ""}, + }, + runtime_updates={"model": verified_model}, + ) + wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] + + assert wrapper.dependencies[0].name == "default" + await app.start() + assert wrapper.as_llm is replacement + assert replacement.is_started is True + await app.close() + + 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))) + old_component = app.context.components[ComponentEnum.AS_LLM]["default"] + start_entered = asyncio.Event() + allow_start = asyncio.Event() + + async def blocking_start() -> None: + start_entered.set() + await allow_start.wait() + + old_component._start = blocking_start + + async def exercise_api() -> None: + start_task = asyncio.create_task(app.start()) + await start_entered.wait() + replace_task = asyncio.create_task( + app.replace_component( + "as_llm", + "default", + config={ + "backend": "dashscope", + "model": "consumer-injected", + "credential": {"api_key": ""}, + }, + runtime_updates={"model": object()}, + ), + ) + await asyncio.sleep(0) + + assert replace_task.done() is False + assert app.context.components[ComponentEnum.AS_LLM]["default"] is old_component + + allow_start.set() + await start_task + replacement = await replace_task + assert replacement.is_started is True + await app.close() + + asyncio.run(exercise_api()) + + +def test_replace_component_start_failure_keeps_old_generation(tmp_path): + """Construction/start failures do not expose a partially replaced graph.""" + + class BrokenAsLLM(BaseAsLLM): + """Backend whose startup deterministically fails for rollback tests.""" + + component_type = ComponentEnum.AS_LLM + + async def _start(self) -> None: + raise RuntimeError("replacement failed") + + app = ReMe(**_qwenpaw_style_config(str(tmp_path))) + app.context.registry.add("broken", BrokenAsLLM, owner="test") + + 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"] + + with pytest.raises(RuntimeError, match="replacement failed"): + await app.replace_component( + "as_llm", + "default", + config={ + "backend": "broken", + "model": "unused", + "credential": {}, + }, + ) + + assert app.context.components[ComponentEnum.AS_LLM]["default"] is old_component + assert app.config.components[ComponentEnum.AS_LLM]["default"].backend == "openai" + assert wrapper.as_llm is old_component + assert old_component.is_started is True + assert old_component in app._started_components + await app.close() + + asyncio.run(exercise_api()) diff --git a/tests/unit/test_file_catalog.py b/tests/unit/test_file_catalog.py index 2c2d04a8..06c1c90a 100644 --- a/tests/unit/test_file_catalog.py +++ b/tests/unit/test_file_catalog.py @@ -10,6 +10,7 @@ import pytest from reme.components.file_catalog import LocalFileCatalog from reme.schema import FileNode +from reme.utils.jsonl_zst import write_jsonl_zst class temp_chdir: @@ -150,6 +151,29 @@ def test_persistence_roundtrip(backend_cls): asyncio.run(run()) +def test_start_failure_does_not_overwrite_persisted_catalog(tmp_path): + """A partial load must not dump incomplete in-memory state during rollback.""" + + async def run(): + with temp_chdir(tmp_path): + catalog = LocalFileCatalog() + original_lines = [ + make_node("a.md").model_dump_json(), + "not valid json", + make_node("b.md").model_dump_json(), + ] + write_jsonl_zst(catalog._catalog_file, original_lines) + original_bytes = catalog._catalog_file.read_bytes() + + with pytest.raises(ValueError): + await catalog.start() + + assert catalog.is_started is False + assert catalog._catalog_file.read_bytes() == original_bytes + + asyncio.run(run()) + + if __name__ == "__main__": print("\n=== FileCatalog Tests ===") for backend in BACKENDS: diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 8686e503..76b3a53e 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -350,6 +350,7 @@ def test_application_starts_jobs_base_stream_background_cron(): }, thread_pool=None, ) + app._component_mutation_lock = asyncio.Lock() app._topological_order = lambda: [] async def start_one(component): @@ -390,6 +391,7 @@ def test_application_start_failure_propagates_and_closes_started_components(): thread_pool=None, ) app._started_components = [] + app._component_mutation_lock = asyncio.Lock() app._topological_order = lambda: [good, bad] app.logger = MagicMock()