From 55fd7b1f0c9538998d4a1b6eeee14c2b0d587200 Mon Sep 17 00:00:00 2001 From: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:38:24 -0700 Subject: [PATCH] fix: register mention_resolver capability before session.initialize() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit @namespace: skill sources (e.g. tool-skills' '@wayfinder:skills') could not resolve at module mount because the mention_resolver capability was only registered AFTER session.initialize() — modules deferred resolution to the first provider:request. Anything snapshotting module state between mount and first prompt saw an incomplete catalog: measured live, 19 skill slash-commands were unavailable cold in the CLI and never recovered (its registry freezes at startup). - create_session(): build + register mention_resolver/mention_deduplicator BEFORE initialize(), unconditionally (inputs are compose-time known). The later guarded system-prompt block reuses the registered instance. - spawn(): child sessions previously never got the capability at all — now registered before child initialize(), mirroring create_session(). - 5 new tests (capability visible during mount for both paths; unconditional; explicit call-order proof). Suite: 1696 passed (+5), ruff + pyright clean. Verified end-to-end in a DTU (CLI built from this branch): cold /skills lists the @-mention-sourced skill, cold /help lists its shortcut, cold dispatch works — all three red on the unfixed baseline. 🤖 Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- amplifier_foundation/bundle/_prepared.py | 94 +++++-- tests/test_cost_bridge_foundation.py | 16 ++ .../test_eager_mention_resolver_capability.py | 254 ++++++++++++++++++ 3 files changed, 343 insertions(+), 21 deletions(-) create mode 100644 tests/test_eager_mention_resolver_capability.py diff --git a/amplifier_foundation/bundle/_prepared.py b/amplifier_foundation/bundle/_prepared.py index a2c3b4f1..a9498bb5 100644 --- a/amplifier_foundation/bundle/_prepared.py +++ b/amplifier_foundation/bundle/_prepared.py @@ -577,6 +577,41 @@ async def create_session( "session.working_dir", str(effective_working_dir.resolve()) ) + # Register the mention resolver (and deduplicator) capabilities BEFORE + # initialize() so modules mounted during session.initialize() can resolve + # @namespace:... sources eagerly at mount time via + # get_capability("mention_resolver") instead of getting None. + # + # ROOT FIX (late skill-source resolution): previously this registration + # happened AFTER session.initialize() (guarded by "does the bundle have + # instruction/context content"), so any module mounted during initialize() + # -- e.g. tool-skills resolving an @namespace:skills source -- saw no + # resolver and had to defer resolution to the first provider:request. + # Anything that snapshots module state between mount and first prompt + # (e.g. a CLI slash-command registry) would then see an incomplete catalog. + # + # Registration is unconditional (cheap to construct) because bundle + # namespace resolution (_build_bundles_for_resolver) depends only on + # self.bundle.source_base_paths / self.bundle.name -- both already fully + # populated by bundle load/compose time, well before create_session() runs + # -- and does NOT depend on whether the bundle has inline instruction or + # context content. + from amplifier_foundation.mentions import BaseMentionResolver + from amplifier_foundation.mentions import ContentDeduplicator + + bundles_for_resolver = self._build_bundles_for_resolver(self.bundle) + # Use session_cwd for local @-mentions, fall back to bundle.base_path + resolver_base = session_cwd or self.bundle.base_path or Path.cwd() + initial_resolver = BaseMentionResolver( + bundles=bundles_for_resolver, + base_path=resolver_base, + ) + initial_deduplicator = ContentDeduplicator() + session.coordinator.register_capability("mention_resolver", initial_resolver) + session.coordinator.register_capability( + "mention_deduplicator", initial_deduplicator + ) + # Initialize the session (loads all modules) await session.initialize() @@ -630,27 +665,15 @@ async def create_session( lambda: [_MENTIONS_RESOLVED_EVENT], ) - from amplifier_foundation.mentions import BaseMentionResolver - from amplifier_foundation.mentions import ContentDeduplicator - - # Register resolver and deduplicator as capabilities for tools to use - # (e.g., filesystem tool's read_file can resolve @mention paths) - # Note: These are created once for capability registration, but the factory - # creates fresh instances each call for accurate file re-reading - bundles_for_resolver = self._build_bundles_for_resolver(self.bundle) - # Use session_cwd for local @-mentions, fall back to bundle.base_path - resolver_base = session_cwd or self.bundle.base_path or Path.cwd() - initial_resolver = BaseMentionResolver( - bundles=bundles_for_resolver, - base_path=resolver_base, - ) - initial_deduplicator = ContentDeduplicator() - session.coordinator.register_capability( - "mention_resolver", initial_resolver - ) - session.coordinator.register_capability( - "mention_deduplicator", initial_deduplicator - ) + # NOTE: The "mention_resolver" / "mention_deduplicator" capabilities + # are already registered above, BEFORE session.initialize() (see the + # eager-resolver block earlier in this method). They are reused here + # rather than rebuilt -- a single source of truth for the capability, + # no duplicate BaseMentionResolver/ContentDeduplicator construction, + # and no extra register_capability() replace() churn. The system + # prompt factory below builds its own fresh resolver/deduplicator + # instances per-call regardless (files may change mid-session), so + # nothing here depends on the registered instances directly. # Create and register the system prompt factory factory = self._create_system_prompt_factory( @@ -836,6 +859,35 @@ async def spawn( "session.working_dir", str(effective_child_cwd.resolve()) ) + # Register the mention resolver (and deduplicator) capabilities BEFORE + # initialize() -- same root fix as create_session(): modules mounted + # during child_session.initialize() (e.g. tool-skills resolving an + # @namespace:skills source) need a real resolver at mount time, not None. + # + # Evidence this was previously missing entirely (worse than merely late): + # unlike create_session(), spawn() never registered "mention_resolver" / + # "mention_deduplicator" on the child coordinator at any point -- before + # or after initialize(). Child sessions do not inherit capabilities from + # the parent coordinator automatically (only session.working_dir is + # explicitly copied above), so this was a real gap for spawned children, + # not just a timing issue. + from amplifier_foundation.mentions import BaseMentionResolver + from amplifier_foundation.mentions import ContentDeduplicator + + child_bundles_for_resolver = self._build_bundles_for_resolver( + effective_bundle + ) + child_session.coordinator.register_capability( + "mention_resolver", + BaseMentionResolver( + bundles=child_bundles_for_resolver, + base_path=effective_child_cwd, + ), + ) + child_session.coordinator.register_capability( + "mention_deduplicator", ContentDeduplicator() + ) + await child_session.initialize() # Register mentions:resolved on observability.events for child sessions. diff --git a/tests/test_cost_bridge_foundation.py b/tests/test_cost_bridge_foundation.py index f0510b23..d71751fe 100644 --- a/tests/test_cost_bridge_foundation.py +++ b/tests/test_cost_bridge_foundation.py @@ -71,6 +71,14 @@ async def test_spawn_calls_bridge_child_cost_with_parent(): bundle.base_path = None bundle.instruction = None bundle.context = None + # Root fix (eager mention_resolver registration): spawn() now calls + # _build_bundles_for_resolver(effective_bundle) unconditionally, before + # child_session.initialize(). That helper reads bundle.source_base_paths + # (dict) and bundle.name (str) -- give this minimal bundle mock real + # values for both so it behaves like a real Bundle instead of tripping + # dataclasses.replace() on a MagicMock. + bundle.source_base_paths = {} + bundle.name = "test-bundle" from amplifier_foundation.bundle._prepared import PreparedBundle @@ -135,6 +143,14 @@ async def test_spawn_does_not_call_bridge_without_parent(): bundle.base_path = None bundle.instruction = None bundle.context = None + # Root fix (eager mention_resolver registration): spawn() now calls + # _build_bundles_for_resolver(effective_bundle) unconditionally, before + # child_session.initialize(). That helper reads bundle.source_base_paths + # (dict) and bundle.name (str) -- give this minimal bundle mock real + # values for both so it behaves like a real Bundle instead of tripping + # dataclasses.replace() on a MagicMock. + bundle.source_base_paths = {} + bundle.name = "test-bundle" from amplifier_foundation.bundle._prepared import PreparedBundle diff --git a/tests/test_eager_mention_resolver_capability.py b/tests/test_eager_mention_resolver_capability.py new file mode 100644 index 00000000..5077794d --- /dev/null +++ b/tests/test_eager_mention_resolver_capability.py @@ -0,0 +1,254 @@ +"""Tests for eager `mention_resolver` capability registration (root fix). + +ROOT FIX (late skill-source resolution): PreparedBundle.create_session() and +PreparedBundle.spawn() now register the "mention_resolver" (and +"mention_deduplicator") capability BEFORE session.initialize() / +child_session.initialize() runs, so that any module mounted during +initialize() -- e.g. tool-skills resolving an ``@namespace:skills`` source -- +can call ``coordinator.get_capability("mention_resolver")`` at mount time and +get a real resolver instance instead of ``None``. + +Before this fix: +- create_session() registered the capability AFTER session.initialize(), and + only when the bundle had inline instruction/context/pending_context content. +- spawn() never registered the capability for the child session at all, at + any point. + +Contract asserted here (this is the NEW, deliberate contract -- see also +tests/test_mentions_resolved_event.py::TestObservabilityRegistration, which +covers the still-guarded ``mentions:resolved`` observability registration +that intentionally stays conditional): + +1. A module mounted during session.initialize()/child_session.initialize() + observes a non-None, real BaseMentionResolver instance. +2. Registration is unconditional -- it happens even for a bundle with no + instruction/context (namespace resolution doesn't depend on that content). +3. register_capability("mention_resolver", ...) is called strictly before + initialize() is awaited. +4. The same holds for spawn()'s child session. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from amplifier_foundation.bundle import Bundle +from amplifier_foundation.bundle._prepared import BundleModuleResolver +from amplifier_foundation.bundle._prepared import PreparedBundle +from amplifier_foundation.mentions import BaseMentionResolver + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class _FakeCoordinator: + """Coordinator stand-in with REAL capability storage. + + A bare MagicMock's get_capability() always returns the same canned value + regardless of what register_capability() was called with, which cannot + prove ordering/visibility. This fake actually stores what is registered + so a test can prove a module mounted mid-initialize() sees exactly what + was registered before initialize() was called. + """ + + def __init__(self) -> None: + self._capabilities: dict[str, Any] = {} + self.mount = AsyncMock(side_effect=self._record_mount) + self.register_contributor = MagicMock() + self.hooks = AsyncMock() + self.hooks.list_handlers = MagicMock(return_value={}) + self.hooks.register = MagicMock(return_value=MagicMock()) + self.mounted_modules: list[str] = [] + + async def _record_mount( + self, mount_point: str, module: Any, name: str | None = None + ) -> None: + self.mounted_modules.append(mount_point) + + def register_capability(self, name: str, value: Any) -> None: + self._capabilities[name] = value + + def get_capability(self, name: str) -> Any: + return self._capabilities.get(name) + + def get(self, mount_point: str, name: str | None = None) -> Any: + return None + + +class _FakeSession: + """AmplifierSession stand-in whose initialize() simulates a module mount + reading get_capability("mention_resolver") -- exactly what tool-skills + does when resolving an @namespace:skills mount-time source.""" + + def __init__(self) -> None: + self.coordinator = _FakeCoordinator() + self.observed_resolver_during_init: Any = "NOT_CAPTURED" + self.execute = AsyncMock(return_value="ok") + self.cleanup = AsyncMock() + self.session_id = "fake-session-id" + + async def initialize(self) -> None: + self.observed_resolver_during_init = self.coordinator.get_capability( + "mention_resolver" + ) + + +def _make_prepared(bundle: Bundle) -> PreparedBundle: + return PreparedBundle( + mount_plan={}, + bundle=bundle, + resolver=BundleModuleResolver(module_paths={}), + ) + + +# --------------------------------------------------------------------------- +# create_session() +# --------------------------------------------------------------------------- + + +class TestCreateSessionEagerResolver: + @pytest.mark.asyncio + async def test_module_mounted_during_initialize_sees_real_resolver(self) -> None: + """The core proof: a module mounted DURING initialize() sees a real + resolver, not None.""" + bundle = Bundle(name="test", instruction="Hello") + prepared = _make_prepared(bundle) + fake_session = _FakeSession() + + with patch("amplifier_core.AmplifierSession", return_value=fake_session): + await prepared.create_session() + + assert fake_session.observed_resolver_during_init is not None, ( + "mention_resolver capability was None during session.initialize() -- " + "modules mounted at this point (e.g. tool-skills) cannot resolve " + "@namespace:... sources eagerly." + ) + assert isinstance( + fake_session.observed_resolver_during_init, BaseMentionResolver + ) + + @pytest.mark.asyncio + async def test_registered_even_for_bundle_with_no_instruction_or_context( + self, + ) -> None: + """Registration is now UNCONDITIONAL: even a bundle with no + instruction/context/pending_context gets a real resolver at mount + time, because bundle namespace resolution does not depend on that + content.""" + bundle = Bundle(name="empty") + prepared = _make_prepared(bundle) + fake_session = _FakeSession() + + with patch("amplifier_core.AmplifierSession", return_value=fake_session): + await prepared.create_session() + + assert fake_session.observed_resolver_during_init is not None + assert isinstance( + fake_session.observed_resolver_during_init, BaseMentionResolver + ) + + @pytest.mark.asyncio + async def test_capability_registered_before_initialize_is_called(self) -> None: + """Explicit ordering proof: register_capability("mention_resolver", ...) + is called strictly before initialize() runs.""" + bundle = Bundle(name="test", instruction="Hi") + prepared = _make_prepared(bundle) + + call_order: list[str] = [] + + class _OrderedCoordinator(_FakeCoordinator): + def register_capability(self, name: str, value: Any) -> None: + if name == "mention_resolver": + call_order.append("register_mention_resolver") + super().register_capability(name, value) + + class _OrderedSession(_FakeSession): + def __init__(self) -> None: + super().__init__() + self.coordinator = _OrderedCoordinator() + + async def initialize(self) -> None: + call_order.append("initialize") + await super().initialize() + + fake_session = _OrderedSession() + with patch("amplifier_core.AmplifierSession", return_value=fake_session): + await prepared.create_session() + + assert "register_mention_resolver" in call_order, ( + "mention_resolver was never registered" + ) + assert "initialize" in call_order, "initialize() was never called" + assert call_order.index("register_mention_resolver") < call_order.index( + "initialize" + ), f"registration did not happen before initialize(): {call_order}" + + +# --------------------------------------------------------------------------- +# spawn() +# --------------------------------------------------------------------------- + + +class TestSpawnEagerResolver: + @pytest.mark.asyncio + async def test_child_module_mounted_during_initialize_sees_real_resolver( + self, + ) -> None: + """spawn()'s child session previously never registered mention_resolver + at all (a strictly worse gap than late registration). This proves the + child now has a real resolver visible during child_session.initialize(). + """ + parent_bundle = Bundle(name="parent") + child_bundle = Bundle(name="child", instruction="Do something") + prepared = _make_prepared(parent_bundle) + fake_child = _FakeSession() + + with patch("amplifier_core.AmplifierSession", return_value=fake_child): + await prepared.spawn(child_bundle, "Do something", compose=False) + + assert fake_child.observed_resolver_during_init is not None, ( + "spawn() did not register mention_resolver before " + "child_session.initialize()" + ) + assert isinstance( + fake_child.observed_resolver_during_init, BaseMentionResolver + ) + + @pytest.mark.asyncio + async def test_child_capability_registered_before_initialize_is_called( + self, + ) -> None: + parent_bundle = Bundle(name="parent") + child_bundle = Bundle(name="child", instruction="Do something") + prepared = _make_prepared(parent_bundle) + + call_order: list[str] = [] + + class _OrderedCoordinator(_FakeCoordinator): + def register_capability(self, name: str, value: Any) -> None: + if name == "mention_resolver": + call_order.append("register_mention_resolver") + super().register_capability(name, value) + + class _OrderedSession(_FakeSession): + def __init__(self) -> None: + super().__init__() + self.coordinator = _OrderedCoordinator() + + async def initialize(self) -> None: + call_order.append("initialize") + await super().initialize() + + fake_child = _OrderedSession() + with patch("amplifier_core.AmplifierSession", return_value=fake_child): + await prepared.spawn(child_bundle, "Do something", compose=False) + + assert "register_mention_resolver" in call_order + assert "initialize" in call_order + assert call_order.index("register_mention_resolver") < call_order.index( + "initialize" + ), f"registration did not happen before initialize(): {call_order}"