From 14bc8b2fa1294c4c3089d3eaf3718819fba5a73d Mon Sep 17 00:00:00 2001 From: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> Date: Fri, 28 Aug 2026 14:56:02 -0700 Subject: [PATCH] =?UTF-8?q?fix(resume):=20re-register=20agent=20system=20p?= =?UTF-8?q?rompt=20on=20sub-session=20resume=20=E2=80=94=20resumed=20sub-a?= =?UTF-8?q?gents=20ran=20with=20no=20system=20prompt?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause (session_spawner.py): 1. spawn_sub_session() registers the agent's system instruction as an in-memory FACTORY (session_spawner.py:766-802, context.set_system_prompt_factory) rather than a persisted message. context-simple's SimpleContextManager builds the system message into a per-request COPY inside get_messages_for_request() and never writes it into self.messages, so it is never present in transcript.jsonl. SessionStore._save_transcript also explicitly skips system/developer role messages when persisting, so no persisted transcript for a sub-session ever contains a system message, regardless of mode. 2. resume_sub_session() (session_spawner.py:923-1406) constructed AmplifierSession directly and restored the transcript (was :1296-1300) -- but never re-derived or re-registered the system instruction anywhere in that 483-line function. Every subsequent request on a resumed sub-session ran with system_msgs == []. 3. A live API probe confirmed omitting the system prompt on a chained request CLEARS the provider's server-held prompt rather than preserving it (provider-openai omits `instructions`, provider-anthropic omits `system`, when system_msgs is empty). Evidence: 8 of 92 captured sessions were delegate-resumed sub-sessions. ALL 8 lost their system instructions starting from the first post-resume request and never recovered -- 111 of 2,460 requests (4.5%) ran with no system prompt. Perfect correlation, zero counter-examples. The `instructions`/`system` key was ABSENT from the wire request, not present as an empty string. Fix (session_spawner.py, resume_sub_session, inserted before the transcript-restore block): - Recover the agent's system instruction from persisted metadata: metadata["agent_overlay"] (the exact agent_config dict spawn_sub_session saved at session_spawner.py:878), falling back to metadata["config"]["agents"][] for sessions saved before agent_overlay existed or with an empty inherit-as-is overlay. - Re-expand @-mentions using the just-restored mention_resolver / mention_deduplicator / session.working_dir capabilities, via the same amplifier_foundation.mentions.expand_mentions_in_instruction() helper the spawn path and the resumed-instruction path already use. This mirrors spawn's behavior exactly rather than injecting the raw, unexpanded body. - Register via context.set_system_prompt_factory() when supported (the same closure shape as the spawn path), with an add_message() fallback for context modules without factory support. - If no instruction is recoverable at all, log a loud warning naming the session and agent -- previously this failure mode was completely silent. Resume still proceeds (never raises). - The nested/grandchild resume capability (child_resume_capability, registered inside resume_sub_session) calls resume_sub_session() directly, so the fix applies transitively with no separate change needed. Also (session_runner.py, ~line 238): corrected a stale/misleading comment on the ROOT-session resume path. That "preserve fresh system prompt / re-inject if transcript lacks one" guard is dead code today: in factory mode, context.get_messages() only ever returns self.messages, which never contains a system-role message (the factory's output is injected ephemerally by get_messages_for_request(), never persisted) -- so `system_msgs` is always empty and the guard's re-injection branch never fires. The comment previously claimed this guard was the safety mechanism; it is not. Root-session resume is actually safe because create_session() (PreparedBundle, called earlier in the same flow) already re-registers a fresh set_system_prompt_factory() before this code runs, and context.set_messages() never touches that registered factory. Behavior is unchanged -- comment only. Cache note: resume re-expands @-mentions from disk at resume time. If a mentioned file changed since the original spawn, the reassembled instructions will legitimately differ from the spawn-time version -- a one-time cache miss, not a bug. Tests (tests/test_resume_system_prompt.py, new file, fixture style mirrors TestCapabilityRegistrationIntegration in test_session_spawner.py): - test_resume_reregisters_system_prompt_via_factory: factory-capable context receives a working factory whose output contains the persisted instruction. FAILS on unfixed code (fake_context.factory stays None). - test_resume_adds_system_message_when_no_factory_support: context without factory support gets a system-role add_message() call instead, ordered before the restored transcript history. - test_resume_falls_back_to_merged_config_agents_map: covers metadata saved before agent_overlay existed. - test_resume_with_no_recoverable_instruction_warns_but_succeeds: no instruction anywhere in metadata -> loud warning fires, resume still succeeds. Fail-before proof: all 4 new tests run against the pre-fix code and FAIL (4 failed in 0.05s). After the fix: all 4 PASS. Full suite: `uv run pytest -q` -> 1497 passed, 1 skipped, 13 deselected, 1 xfailed (baseline was ~1493 passed; +4 for the new tests, zero regressions). `uv run ruff check` on all three changed/added files: All checks passed. 🤖 Generated with [Amplifier](https://github.com/microsoft/amplifier) Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com> --- amplifier_app_cli/session_runner.py | 24 ++- amplifier_app_cli/session_spawner.py | 76 +++++++- tests/test_resume_system_prompt.py | 268 +++++++++++++++++++++++++++ 3 files changed, 364 insertions(+), 4 deletions(-) create mode 100644 tests/test_resume_system_prompt.py diff --git a/amplifier_app_cli/session_runner.py b/amplifier_app_cli/session_runner.py index 16cb6d2b..2468aae1 100644 --- a/amplifier_app_cli/session_runner.py +++ b/amplifier_app_cli/session_runner.py @@ -235,9 +235,27 @@ async def create_initialized_session( context = session.coordinator.get("context") if context and hasattr(context, "set_messages"): - # CRITICAL: create_session() already added a fresh system prompt. - # We need to preserve it because the transcript might have lost its system message - # during compaction (bug fixed in context-simple, but old sessions are affected). + # NOTE ON WHY ROOT RESUME IS SAFE FROM SYSTEM-PROMPT LOSS: + # create_session() (PreparedBundle, called before this function runs) + # already registered a system-prompt FACTORY via + # context.set_system_prompt_factory() -- it does not add a static + # system message to context.messages. get_messages_for_request() + # calls that factory fresh on every request, independent of + # self.messages, and context.set_messages() below only replaces + # self.messages -- it never touches the registered factory. So the + # factory keeps producing the system prompt on every subsequent + # request regardless of what this block does. + # + # The preserve/re-inject logic below is therefore DEAD CODE on the + # current factory-based path: context.get_messages() returns only + # self.messages, which never contains a system-role message in + # factory mode (the factory's output is injected ephemerally by + # get_messages_for_request(), never persisted) -- so + # `system_msgs` is always empty and `fresh_system_msg` stays None. + # It is kept only as a defensive fallback for a context module + # using the older add_message()-based static system message + # convention (pre-factory), where a real system message COULD + # live in self.messages and get lost by a transcript replace. fresh_system_msg = None if hasattr(context, "get_messages"): current_msgs = await context.get_messages() diff --git a/amplifier_app_cli/session_spawner.py b/amplifier_app_cli/session_spawner.py index 4d1b0fa5..2200b35a 100644 --- a/amplifier_app_cli/session_spawner.py +++ b/amplifier_app_cli/session_spawner.py @@ -1293,8 +1293,82 @@ async def child_resume_capability(sub_session_id: str, instruction: str) -> dict }, ) - # Restore transcript to context + # Re-register the agent's system prompt on resume. + # + # Mirrors the spawn path (see the "Inject agent's system instruction" + # block above, ~line 764). That block registers the system instruction + # via context.set_system_prompt_factory() rather than a persisted + # message: context-simple builds the system message into a per-request + # COPY and never writes it into self.messages, so it is never present in + # the saved transcript. SessionStore._save_transcript also explicitly + # skips system/developer role messages when persisting, so this holds + # even for a context module using the add_message() fallback below. + # + # Restoring the transcript alone (next block) therefore restores ZERO + # system-role messages -- every subsequent request on a resumed + # sub-session ran with no system prompt at all, and omitting it on a + # chained request CLEARS the provider's server-held prompt rather than + # preserving it. Recover the same instruction the original spawn used + # and re-register it through the same mechanism. context = child_session.coordinator.get("context") + agent_overlay = metadata.get("agent_overlay") or {} + resume_system_instruction = agent_overlay.get("instruction") or agent_overlay.get( + "system", {} + ).get("instruction") + if not resume_system_instruction: + # Fallback for metadata saved before agent_overlay existed, or an + # empty inherit-as-is overlay: recover the declaration from the + # merged config's own agents map, keyed by agent_name. + _resume_agents_cfg = merged_config.get("agents") or {} + _resume_agent_decl = _resume_agents_cfg.get(agent_name) or {} + resume_system_instruction = _resume_agent_decl.get( + "instruction" + ) or _resume_agent_decl.get("system", {}).get("instruction") + + if resume_system_instruction: + # Expand @-mentions exactly like the spawn path does, using the + # just-restored resolver/deduplicator/working_dir capabilities. + _resume_sys_resolver = child_session.coordinator.get_capability( + "mention_resolver" + ) + if _resume_sys_resolver is not None: + from amplifier_foundation.mentions import expand_mentions_in_instruction + + _resume_sys_dedup = child_session.coordinator.get_capability( + "mention_deduplicator" + ) + _resume_sys_wd = child_session.coordinator.get_capability( + "session.working_dir" + ) + _resume_sys_rel = Path(_resume_sys_wd) if _resume_sys_wd else Path.cwd() + resume_system_instruction = await expand_mentions_in_instruction( + resume_system_instruction, + resolver=_resume_sys_resolver, + deduplicator=_resume_sys_dedup, + relative_to=_resume_sys_rel, + ) + if context and hasattr(context, "set_system_prompt_factory"): + _resolved_resume_system_instruction = resume_system_instruction + + async def _resume_system_prompt_factory() -> str: + return _resolved_resume_system_instruction + + await context.set_system_prompt_factory(_resume_system_prompt_factory) + elif context and hasattr(context, "add_message"): + await context.add_message( + {"role": "system", "content": resume_system_instruction} + ) + else: + logger.warning( + "Sub-session %s (agent=%s): no system instruction recoverable from " + "persisted metadata (agent_overlay / config.agents) on resume. " + "This resumed session will run WITHOUT a system prompt for all " + "subsequent requests -- proceeding with resume anyway.", + sub_session_id, + agent_name, + ) + + # Restore transcript to context if context and hasattr(context, "add_message"): for message in transcript: await context.add_message(message) diff --git a/tests/test_resume_system_prompt.py b/tests/test_resume_system_prompt.py new file mode 100644 index 00000000..0c43dd35 --- /dev/null +++ b/tests/test_resume_system_prompt.py @@ -0,0 +1,268 @@ +"""Tests for system-prompt re-injection on sub-session resume. + +Bug: resumed sub-sessions ran with NO system prompt at all. + +The spawn path (spawn_sub_session) registers the agent's system instruction +via a *factory* (context.set_system_prompt_factory) rather than a persisted +message. context-simple builds the system message into a per-request COPY +and never writes it into self.messages -- so it is never present in +transcript.jsonl. SessionStore._save_transcript also explicitly skips +system/developer role messages when persisting, so this holds even for a +context module using the add_message() fallback. + +resume_sub_session only ever restored the transcript. It never re-derived +or re-registered the system instruction, so every resumed sub-session ran +every subsequent request with system_msgs == [] -- and a live API probe +confirmed omitting the system prompt on a chained request CLEARS the +provider's server-held prompt rather than preserving it (see PR for the +full evidence chain). + +These tests exercise the actual resume_sub_session() code path using the +same capability-registration-integration style as +TestCapabilityRegistrationIntegration in test_session_spawner.py: a fully +mocked AmplifierSession/coordinator plus a small fake context standing in +for context-simple, so we can assert directly on the system-prompt +registration contract without needing a real provider or context module. +""" + +from __future__ import annotations + +import logging +from unittest.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import patch + +import pytest +from amplifier_app_cli.session_spawner import resume_sub_session +from amplifier_app_cli.session_store import SessionStore + +pytestmark = pytest.mark.anyio + +SENTINEL_INSTRUCTION = "SENTINEL-INSTRUCTION-42: you are the test agent." + + +@pytest.fixture(scope="module") +def anyio_backend(): + """Configure anyio to use asyncio backend only.""" + return "asyncio" + + +class _FakeContextWithFactory: + """Stands in for context-simple: supports the factory-based system prompt.""" + + def __init__(self) -> None: + self.factory = None + self.messages: list[dict] = [] + + async def set_system_prompt_factory(self, factory) -> None: + self.factory = factory + + async def add_message(self, message: dict) -> None: + self.messages.append(message) + + async def get_messages(self) -> list[dict]: + return self.messages + + +class _FakeContextAddMessageOnly: + """Stands in for a context module with NO factory support (fallback path). + + Deliberately has no set_system_prompt_factory attribute at all, so + hasattr(context, "set_system_prompt_factory") is False, exactly like the + hasattr guard in resume_sub_session / spawn_sub_session. + """ + + def __init__(self) -> None: + self.messages: list[dict] = [] + + async def add_message(self, message: dict) -> None: + self.messages.append(message) + + async def get_messages(self) -> list[dict]: + return self.messages + + +def _base_metadata(session_id: str, **overrides) -> dict: + metadata = { + "session_id": session_id, + "parent_id": "parent-123", + "agent_name": "test-agent", + "config": { + "session": {"orchestrator": "loop-basic", "context": "context-simple"} + }, + "working_dir": "/test/project", + "self_delegation_depth": 0, + } + metadata.update(overrides) + return metadata + + +async def _resume_with_fake_context( + fake_context, session_id: str, instruction: str = "follow-up" +) -> None: + """Run resume_sub_session() against a fully mocked AmplifierSession. + + `fake_context` is wired in as the resumed session's "context" capability + (mirrors coordinator.get("context") in the production code). + """ + + def mock_get(name): + if name == "context": + return fake_context + return None + + mock_coordinator = MagicMock() + mock_coordinator.register_capability = MagicMock() + # mention_resolver intentionally returns None: the mention-expansion + # branch (both for the follow-up instruction AND the system instruction) + # is exercised elsewhere (test_session_spawner.py); returning None here + # keeps these tests focused on the factory/add_message wiring itself. + mock_coordinator.get_capability = MagicMock(return_value=None) + mock_coordinator.get = MagicMock(side_effect=mock_get) + mock_coordinator.mount = AsyncMock() + + mock_session = MagicMock() + mock_session.coordinator = mock_coordinator + mock_session.initialize = AsyncMock() + mock_session.execute = AsyncMock(return_value="response") + mock_session.cleanup = AsyncMock() + + with patch( + "amplifier_app_cli.session_spawner.AmplifierSession", + return_value=mock_session, + ): + with patch("amplifier_app_cli.ui.CLIApprovalSystem"): + with patch("amplifier_app_cli.ui.CLIDisplaySystem"): + with patch("amplifier_app_cli.paths.create_foundation_resolver"): + await resume_sub_session(session_id, instruction) + + +class TestResumeSystemPromptReinjection: + """resume_sub_session must re-register the agent's system prompt. + + Regression coverage for the silent system-prompt-loss bug: a resumed + sub-session ran every subsequent request with system_msgs == [] because + only the transcript (never the system instruction) was restored. + """ + + async def test_resume_reregisters_system_prompt_via_factory( + self, tmp_path, monkeypatch + ): + """Factory-capable context: resume must call set_system_prompt_factory + with a factory that reproduces the original agent instruction. + + FAILS BEFORE THE FIX: resume_sub_session never calls + set_system_prompt_factory at all, so fake_context.factory stays None. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + + store = SessionStore() + session_id = "test-resume-system-prompt-factory" + # Transcript deliberately has ZERO system-role messages -- this is + # exactly what SessionStore._save_transcript always produces (it + # skips system/developer messages), and what every real resumed + # session actually persists. + transcript = [{"role": "user", "content": "hi"}] + metadata = _base_metadata( + session_id, + agent_overlay={"instruction": SENTINEL_INSTRUCTION}, + ) + store.save(session_id, transcript, metadata) + + fake_context = _FakeContextWithFactory() + await _resume_with_fake_context(fake_context, session_id) + + assert fake_context.factory is not None, ( + "resume_sub_session must call set_system_prompt_factory() so " + "the resumed session has a system prompt on every subsequent " + "request -- this is the fix for the silent system-prompt-loss " + "bug (resumed sub-agents ran with no system prompt)." + ) + produced = await fake_context.factory() + assert SENTINEL_INSTRUCTION in produced + + async def test_resume_adds_system_message_when_no_factory_support( + self, tmp_path, monkeypatch + ): + """Fallback path: a context without factory support gets an + add_message() system-role message instead (mirrors the spawn path's + own hasattr-gated fallback). + """ + monkeypatch.setenv("HOME", str(tmp_path)) + + store = SessionStore() + session_id = "test-resume-system-prompt-fallback" + transcript = [{"role": "user", "content": "hi"}] + metadata = _base_metadata( + session_id, + agent_overlay={"instruction": SENTINEL_INSTRUCTION}, + ) + store.save(session_id, transcript, metadata) + + fake_context = _FakeContextAddMessageOnly() + await _resume_with_fake_context(fake_context, session_id) + + system_messages = [ + m for m in fake_context.messages if m.get("role") == "system" + ] + assert system_messages, ( + "resume_sub_session must add a system-role message when the " + "context module has no set_system_prompt_factory support." + ) + assert SENTINEL_INSTRUCTION in system_messages[0]["content"] + + # The system message must land before the restored transcript + # history, mirroring how a live conversation is structured (system + # message first, then user/assistant turns). + assert fake_context.messages[0]["role"] == "system" + + async def test_resume_falls_back_to_merged_config_agents_map( + self, tmp_path, monkeypatch + ): + """When agent_overlay carries no instruction (e.g. an empty + inherit-as-is overlay, or a session saved before agent_overlay + existed), fall back to config.agents[].instruction. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + + store = SessionStore() + session_id = "test-resume-system-prompt-config-fallback" + transcript: list[dict] = [] + metadata = _base_metadata(session_id, agent_overlay={}) + metadata["config"]["agents"] = { + "test-agent": {"instruction": SENTINEL_INSTRUCTION} + } + store.save(session_id, transcript, metadata) + + fake_context = _FakeContextWithFactory() + await _resume_with_fake_context(fake_context, session_id) + + assert fake_context.factory is not None + produced = await fake_context.factory() + assert SENTINEL_INSTRUCTION in produced + + async def test_resume_with_no_recoverable_instruction_warns_but_succeeds( + self, tmp_path, monkeypatch, caplog + ): + """No instruction anywhere in metadata -> loud warning, resume still + succeeds. Today this failure mode is completely silent; the fix + must surface it rather than leaving the resumed session mute. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + + store = SessionStore() + session_id = "test-resume-system-prompt-missing" + transcript: list[dict] = [] + metadata = _base_metadata(session_id) # no agent_overlay at all + store.save(session_id, transcript, metadata) + + fake_context = _FakeContextWithFactory() + with caplog.at_level(logging.WARNING): + await _resume_with_fake_context(fake_context, session_id) + + assert fake_context.factory is None + assert any( + "system" in record.getMessage().lower() + and "no" in record.getMessage().lower() + for record in caplog.records + ), "Missing system instruction on resume must be logged loudly, not silently swallowed."