fix(resume): re-register agent system prompt on sub-session resume — resumed sub-agents ran with no system prompt

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()

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)

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[<agent_name>].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."