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."