diff --git a/reme/application.py b/reme/application.py index fb811a9d..9746fb32 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] @@ -221,17 +228,123 @@ 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 + + async def replace_component( + self, + component_enum: ComponentType, + name: str, + /, + *, + config: ComponentConfig | Mapping[str, Any], + runtime_updates: Mapping[str, Any] | None = None, + ) -> BaseComponent: + """Atomically replace an existing component and its bind() 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. + """ + 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 = self._topological_order((node_key, replacement)) + was_started = old_component.is_started + if was_started: + await replacement.start() + + 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)) + + # 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 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..303af29a 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: diff --git a/tests/unit/test_base_component.py b/tests/unit/test_base_component.py index e510f032..5636ab74 100644 --- a/tests/unit/test_base_component.py +++ b/tests/unit/test_base_component.py @@ -122,6 +122,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 ---------------------------------------------------------------- diff --git a/tests/unit/test_embedded_consumer_compat.py b/tests/unit/test_embedded_consumer_compat.py index 2570d04c..c5cadf2e 100644 --- a/tests/unit/test_embedded_consumer_compat.py +++ b/tests/unit/test_embedded_consumer_compat.py @@ -1,9 +1,14 @@ """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 @@ -78,3 +83,126 @@ 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_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())