diff --git a/AGENTS.md b/AGENTS.md index c8ca0916..90f22893 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -136,6 +136,10 @@ JSON output paths must restore the console's configured backing stream, not a dynamic temporary stdout capture. Run-command diagnostics before headless execution belong on stderr so JSON stdout remains one parseable payload. +Both chat and headless entry points use shared pre-turn transcript repair and +persist before dispatch. A failed repair checkpoint must prevent execution; +a successful one must not suppress saving a subsequent failed turn. +Test this through actual `execute_single`, not only extracted interactive helpers. ## Fork event ownership and cumulative cost diff --git a/amplifier_app_cli/main.py b/amplifier_app_cli/main.py index d82a3d08..1ccc431f 100644 --- a/amplifier_app_cli/main.py +++ b/amplifier_app_cli/main.py @@ -80,7 +80,11 @@ from .provider_diagnostics import format_model_line from .provider_diagnostics import invoke_list_models from .provider_diagnostics import test_provider_connectivity -from .session_runner import SessionConfig, create_initialized_session +from .session_runner import ( + SessionConfig, + create_initialized_session, + repair_transcript_if_needed, +) from .session_store import SessionStore from .stdout_offload import patch_stdout_offloaded as patch_stdout from .ui.dashboard_renderer import DashboardRenderer @@ -4002,57 +4006,8 @@ async def _save_session(): # Helper to detect and repair broken transcripts before each turn async def _repair_transcript_if_needed(): - """Pre-turn transcript repair. - - Detects and fixes orphaned tool calls, ordering violations, and - incomplete assistant turns left by interrupted operations (Ctrl+C, - SIGKILL, OOM, MCP transport failures). - - Uses the same foundation diagnosis library as resume-time repair - (session_runner.py), but operates on live in-memory context messages - rather than on-disk transcript files. Runs once per turn; the scan - is a pure in-memory walk (<10 ms for typical sessions). - """ - context = session.coordinator.get("context") - if not context or not hasattr(context, "get_messages"): - return - - try: - messages = await context.get_messages() - if not messages: - return - - from amplifier_foundation.session import ( - diagnose_transcript, - repair_transcript, - ) - - diagnosis = diagnose_transcript(messages) - if diagnosis["status"] != "broken": - return - - failure_modes = diagnosis.get("failure_modes", []) - orphan_ids = diagnosis.get("orphaned_tool_ids", []) - - # Repair and update context in-place - repaired = repair_transcript(messages, diagnosis) - if hasattr(context, "set_messages"): - await context.set_messages(repaired) - - # Persist immediately so the fix survives further interruptions - await _save_session() - - logger.warning( - "Pre-turn transcript repair: %s (orphaned tool calls: %s).", - ", ".join(failure_modes), - ", ".join(orphan_ids) if orphan_ids else "none", - ) - except ImportError: - # Foundation not available (non-standard setup) — skip repair - pass - except Exception as e: - # Repair must never block the session — log and continue - logger.debug("Pre-turn transcript repair failed: %s", e) + """Repair live history and checkpoint through the interactive save path.""" + await repair_transcript_if_needed(session, _save_session) # Helper to execute a prompt with Ctrl+C handling async def _execute_with_interrupt(prompt_text: str) -> bool: @@ -4854,6 +4809,12 @@ def _goal_sigint_handler(signum, frame): _original_sigint_handler = signal.signal(signal.SIGINT, _goal_sigint_handler) try: + await repair_transcript_if_needed(session, _persist_session) + if handoff.requested.is_set(): + return + # A successful repair checkpoint must not suppress saving messages + # added by a subsequent failed turn. Failed repair saves never get here. + session_save_attempted = False execution_started = True response = await session.execute(prompt) finally: diff --git a/amplifier_app_cli/session_runner.py b/amplifier_app_cli/session_runner.py index 2b1aa1a7..91955cf8 100644 --- a/amplifier_app_cli/session_runner.py +++ b/amplifier_app_cli/session_runner.py @@ -30,6 +30,7 @@ import sys import uuid from collections import Counter +from collections.abc import Awaitable, Callable from dataclasses import dataclass from dataclasses import field from pathlib import Path @@ -51,6 +52,55 @@ logger = logging.getLogger(__name__) +async def repair_transcript_if_needed( + session: AmplifierSession, save_session: Callable[[], Awaitable[Any]] +) -> None: + """Apply Foundation's best-effort repair, then persist before dispatch. + + Unsafe associations may remain broken; Foundation owns that decision. + A failed checkpoint must propagate rather than dispatch an unsaved repair. + """ + try: + context = session.coordinator.get("context") + if ( + not context + or not hasattr(context, "get_messages") + or not hasattr(context, "set_messages") + ): + return + + messages = await context.get_messages() + if not messages: + return + + from amplifier_foundation.session import ( + diagnose_transcript, + repair_transcript, + ) + + diagnosis = diagnose_transcript(messages) + if diagnosis["status"] != "broken": + return + + failure_modes = diagnosis.get("failure_modes", []) + orphan_ids = diagnosis.get("orphaned_tool_ids", []) + repaired = repair_transcript(messages, diagnosis) + await context.set_messages(repaired) + except Exception as exc: + # Foundation availability, diagnosis and context mutation are best-effort. + logger.debug("Pre-turn transcript repair failed: %s", exc) + return + + # Deliberately outside the best-effort boundary, including ImportError and + # cancellation: never dispatch a repair whose persistence failed. + await save_session() + logger.warning( + "Pre-turn transcript repair: %s (orphaned tool calls: %s).", + ", ".join(failure_modes), + ", ".join(orphan_ids) if orphan_ids else "none", + ) + + @dataclass class SessionConfig: """All parameters needed to create and initialize a session. @@ -299,9 +349,9 @@ async def create_initialized_session( ) # Step 7: Restore transcript (resume only) - # NOTE: Transcript repair (orphaned tool calls, ordering violations) is handled - # by _repair_transcript_if_needed() in main.py, which runs before every LLM call. - # This covers both resume and mid-session Ctrl+C cases in a single code path. + # Initialization only restores history. Chat and headless host entry points + # call shared pre-turn repair and persist it before execution, once their + # save callbacks exist (including after mid-session interruptions). if config.is_resume and config.initial_transcript: transcript_to_restore = config.initial_transcript diff --git a/tests/test_pre_turn_repair.py b/tests/test_pre_turn_repair.py index c8f238fe..946c3a43 100644 --- a/tests/test_pre_turn_repair.py +++ b/tests/test_pre_turn_repair.py @@ -10,13 +10,16 @@ import importlib import json import logging +import signal from copy import deepcopy from datetime import UTC, datetime +from io import StringIO from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, Mock import pytest +from rich.console import Console from amplifier_foundation.session import ( diagnose_transcript, @@ -318,6 +321,7 @@ def _installed_interactive_helpers(**bindings): "UTC": UTC, "Path": Path, "logger": logging.getLogger(__name__), + "repair_transcript_if_needed": main.repair_transcript_if_needed, **bindings, } # Trusted installed CLI definitions only; never transcript or user input. @@ -452,3 +456,441 @@ async def test_pre_turn_receipts_survive_native_save_and_resume( json.loads(line) for line in backup.read_text(encoding="utf-8").splitlines() ] == persisted assert messages == original + + +def _headless_messages(case): + messages = ( + _healthy_messages()[1:] + if case == "healthy" + else _ordering_violation_messages()[1:] + + [{"role": "assistant", "content": "Observed result recorded"}] + ) + # Persist a canonical assistant payload: native serialization omits None. + messages[1]["content"] = "" + receipt = next(m for m in messages if m["role"] == "tool") + receipt.update( + caller="fixture-caller", + cost={"amount": "0.125", "currency": "USD"}, + extra={"receipt": ["fixture", 7]}, + metadata={"source": "observed", "timestamp": "2026-01-01T00:00:00+00:00"}, + ) + if case == "parallel": + missing = deepcopy(messages[1]["tool_calls"][0]) + missing["id"] = "tc_missing" + messages[1]["tool_calls"].append(missing) + return messages + + +def _without_seq(message): + """Ignore only Simple's documented sequence restamping, not receipt fields.""" + result = deepcopy(message) + if "metadata" in result: + result["metadata"].pop("_seq", None) + if not result["metadata"]: + result.pop("metadata") + return result + + +def _assert_headless_repair(messages, original, case): + assert diagnose_transcript(messages)["status"] == "healthy" + results = [m for m in messages if m["role"] == "tool"] + observed = [m for m in results if m["tool_call_id"] == "tc_1"] + assert len(observed) == 1 + assert _without_seq(observed[0]) == next(m for m in original if m["role"] == "tool") + assert observed[0]["content"] != SYNTHETIC_TOOL_RESULT_CONTENT + assert {m["tool_call_id"] for m in results} == ( + {"tc_1", "tc_missing"} if case == "parallel" else {"tc_1"} + ) + if case == "parallel": + assert results[1]["tool_call_id"] == "tc_missing" + assert results[1]["content"] == SYNTHETIC_TOOL_RESULT_CONTENT + call_index = next(i for i, m in enumerate(messages) if m.get("tool_calls")) + assert messages[call_index + 1 : call_index + 1 + len(results)] == results + assert [_without_seq(m) for m in messages if m["role"] == "user"] == [ + m for m in original if m["role"] == "user" + ] + stripped = [_without_seq(m) for m in messages] + assert all(m in stripped for m in original) + + +@pytest.fixture +def headless_resume(isolated_home, tmp_path, monkeypatch): + """Real initializer, Simple context, Foundation and native store; no provider. + + Bundle/session construction, first-run setup and writer acquisition are + doubles. The normal native-persistence opt-out avoids taking host locks. + """ + from amplifier_app_cli import session_handoff, session_runner, shared_root_state + from amplifier_app_cli.session_store import SessionStore + from amplifier_module_context_simple import SimpleContextManager + + main = importlib.import_module("amplifier_app_cli.main") + monkeypatch.delenv("AMPLIFIER_HOME", raising=False) + monkeypatch.setenv("AMPLIFIER_CONTEXT_INTELLIGENCE_BASE_PATH", str(tmp_path / "ci")) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr( + "amplifier_app_cli.commands.init.check_first_run", lambda: False + ) + monkeypatch.setattr( + shared_root_state, "_shared_root_platform_supported", lambda: False + ) + monkeypatch.setattr( + shared_root_state, "warn_windows_native_persistence", lambda: None + ) + private_console = Console(file=StringIO()) + monkeypatch.setattr(main, "console", private_console) + close_tty = Mock() + monkeypatch.setattr(main, "close_dedicated_tty_input", close_tty) + store = SessionStore() + monkeypatch.setattr(main, "SessionStore", lambda: store) + monkeypatch.setattr(session_runner, "SessionStore", lambda: store) + + def make(case="late", *, root=None): + original = _headless_messages(case) + store.save( + "headless-resume", + original, + { + "session_id": "headless-resume", + "created": "2026-01-01T00:00:00+00:00", + "name": "Receipt fixture", + "description": "Preserve observed evidence", + "fixture_metadata": {"keep": [1, 2]}, + }, + ) + seed, seed_metadata = store.load("headless-resume") + save = Mock(wraps=store.save) + monkeypatch.setattr(store, "save", save) + context = SimpleContextManager( + max_tokens=100_000, compaction_notice_enabled=False + ) + hooks = MagicMock() + hooks.emit = AsyncMock() + coordinator = MagicMock(session_state={}) + coordinator.get.side_effect = lambda key: { + "context": context, + "hooks": hooks, + "providers": {}, + }.get(key) + coordinator.get_capability.return_value = None + session = SimpleNamespace( + session_id="headless-resume", + config={}, + coordinator=coordinator, + execute=AsyncMock(return_value="fixture response"), + cleanup=AsyncMock(), + ) + + async def create_bundle_session(**kwargs): + assert await context.get_messages() == [] + session.config = kwargs["config"].config + return session + + monkeypatch.setattr( + session_runner, "_create_bundle_session", create_bundle_session + ) + if root is not None: + monkeypatch.setattr( + shared_root_state, "_shared_root_platform_supported", lambda: True + ) + root.read.return_value = (seed, seed_metadata) + monkeypatch.setattr( + session_handoff, "acquire_root", AsyncMock(return_value=root) + ) + # No real writer registration for this deliberately refusing fake writer. + monkeypatch.setattr(session_handoff.CLIHandoff, "start", AsyncMock()) + + async def invoke(output_format="json"): + await main.execute_single( + prompt="Continue the fixture", + config={}, + search_paths=[], + verbose=False, + session_id=session.session_id, + bundle_name="fixture-bundle", + initial_transcript=deepcopy(seed), + output_format=output_format, + ) + + return SimpleNamespace( + context=context, + session=session, + store=store, + save=save, + original=deepcopy(original), + seed=seed, + seed_metadata=seed_metadata, + invoke=invoke, + close_tty=close_tty, + console=private_console, + ) + + return make + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", ["healthy", "late", "parallel"]) +@pytest.mark.parametrize("output_format", ["json", "json-trace", "text"]) +async def test_execute_single_repairs_and_checkpoints_before_dispatch( + headless_resume, capsys, case, output_format +): + """The actual headless boundary must see repaired memory AND disk on entry.""" + env = headless_resume(case) + before_execute = [] + original_sigint = signal.getsignal(signal.SIGINT) + + async def execute(prompt): + messages = await env.context.get_messages() + _assert_headless_repair(messages, env.original, case) + assert env.save.call_count == (0 if case == "healthy" else 1) + persisted, metadata = env.store.load(env.session.session_id) + assert [_without_seq(m) for m in persisted] == [ + _without_seq(m) for m in messages + ] + assert all(metadata[k] == v for k, v in env.seed_metadata.items()) + before_execute.extend(deepcopy(messages)) + await env.context.add_message({"role": "user", "content": prompt}) + await env.context.add_message( + {"role": "assistant", "content": "fixture response"} + ) + return "fixture response" + + env.session.execute.side_effect = execute + await env.invoke(output_format) + assert signal.getsignal(signal.SIGINT) == original_sigint + env.session.execute.assert_awaited_once_with("Continue the fixture") + env.session.cleanup.assert_awaited_once() + env.close_tty.assert_called_once() + assert env.save.call_count == (1 if case == "healthy" else 2) + persisted, metadata = env.store.load(env.session.session_id) + assert persisted == await env.context.get_messages() + assert all(metadata[k] == v for k, v in env.seed_metadata.items()) + assert metadata["bundle"] == "fixture-bundle" + backup = env.store.base_dir / env.session.session_id / "transcript.jsonl.backup" + backup_messages = [ + json.loads(line) for line in backup.read_text(encoding="utf-8").splitlines() + ] + assert [_without_seq(m) for m in backup_messages] == [ + _without_seq(m) for m in before_execute + ] + output = capsys.readouterr() + if output_format != "text": + payload = json.loads(output.out) # Exactly one object; no repair diagnostics. + assert payload["status"] == "success" + assert payload["response"] == "fixture response" + assert env.console._file.getvalue() == "" + if output_format == "json-trace": + assert payload["execution_trace"] == [] + else: + assert "fixture response" in env.console._file.getvalue() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("writer", ["native", "refusing-root"]) +async def test_execute_single_failed_repair_checkpoint_never_dispatches_or_retries( + headless_resume, capsys, writer +): + root = None + if writer == "refusing-root": + root = SimpleNamespace( + read=Mock(), + checkpoint=Mock(side_effect=RuntimeError("writer refused")), + release=Mock(), + ) + env = headless_resume(root=root) + if root is None: + env.save.side_effect = OSError("checkpoint failed") + original_sigint = signal.getsignal(signal.SIGINT) + with pytest.raises(SystemExit) as exc: + await env.invoke() + assert exc.value.code == 1 + assert signal.getsignal(signal.SIGINT) == original_sigint + env.session.execute.assert_not_awaited() + env.session.cleanup.assert_awaited_once() + env.close_tty.assert_called_once() + assert env.store.load(env.session.session_id) == (env.seed, env.seed_metadata) + if root is None: + assert env.save.call_count == 1 + else: + root.checkpoint.assert_called_once() + root.release.assert_called_once() + env.save.assert_not_called() + payload = json.loads(capsys.readouterr().out) + assert payload["status"] == "error" + assert payload["error"] == ( + "checkpoint failed" if root is None else "writer refused" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError]) +async def test_execute_single_repair_checkpoint_does_not_suppress_failed_turn_save( + headless_resume, capsys, failure +): + env = headless_resume() + appended = [] + + async def execute(prompt): + assert env.save.call_count == 1 + messages = await env.context.get_messages() + _assert_headless_repair(messages, env.original, "late") + assert env.store.load(env.session.session_id)[0] == messages + await env.context.add_message({"role": "user", "content": prompt}) + await env.context.add_message( + {"role": "assistant", "content": "partial response"} + ) + appended.extend(deepcopy(await env.context.get_messages())) + raise failure("turn interrupted") + + async def cleanup(): + assert env.save.call_count == 2 + assert env.store.load(env.session.session_id)[0] == appended + + env.session.execute.side_effect = execute + env.session.cleanup.side_effect = cleanup + original_sigint = signal.getsignal(signal.SIGINT) + expected = SystemExit if failure is RuntimeError else asyncio.CancelledError + with pytest.raises(expected): + await env.invoke() + assert signal.getsignal(signal.SIGINT) == original_sigint + env.session.cleanup.assert_awaited_once() + env.close_tty.assert_called_once() + assert env.save.call_count == 2 + assert ( + env.store.load(env.session.session_id)[0][-1]["content"] == "partial response" + ) + output = capsys.readouterr().out + if failure is RuntimeError: + assert json.loads(output)["status"] == "error" + else: + assert output == "" + + +@pytest.mark.asyncio +async def test_execute_single_rechecks_handoff_after_repair_checkpoint( + headless_resume, monkeypatch +): + from amplifier_app_cli.session_handoff import CLIHandoff + + env = headless_resume() + controls = [] + + async def start(control): + controls.append(control) + + async def finish(control, save, after_cleanup): + # Only the ownership protocol is faked, not repair, persistence or dispatch. + assert control.requested.is_set() + await save() + await control.initialized.cleanup() + await after_cleanup() + + monkeypatch.setattr(CLIHandoff, "start", start) + monkeypatch.setattr(CLIHandoff, "finish", finish) + real_get_messages = env.context.get_messages + + async def get_messages(): + # Count only after initialization: scan first, then _persist_session. + if controls: + get_messages.calls += 1 + if get_messages.calls == 2: + controls[0].requested.set() + return await real_get_messages() + + get_messages.calls = 0 + monkeypatch.setattr(env.context, "get_messages", get_messages) + await env.invoke() + env.session.execute.assert_not_awaited() + assert env.save.call_count == 1 + env.session.cleanup.assert_awaited_once() + env.close_tty.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [ImportError, RuntimeError, asyncio.CancelledError]) +async def test_shared_repair_checkpoint_errors_propagate(failure): + from amplifier_app_cli.session_runner import repair_transcript_if_needed + + context = SimpleNamespace( + get_messages=AsyncMock(return_value=_ordering_violation_messages()), + set_messages=AsyncMock(), + ) + session = SimpleNamespace(coordinator=SimpleNamespace(get=lambda _: context)) + save = AsyncMock(side_effect=failure("checkpoint failed")) + with pytest.raises(failure, match="checkpoint failed"): + await repair_transcript_if_needed(session, save) + context.set_messages.assert_awaited_once() + save.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_shared_repair_leaves_ambiguous_receipts_broken_without_discarding_evidence(): + from amplifier_app_cli.session_runner import repair_transcript_if_needed + from amplifier_module_context_simple import SimpleContextManager + + original = _headless_messages("late") + duplicate = deepcopy(next(m for m in original if m["role"] == "tool")) + duplicate["content"] = "A different observed result" + original.append(duplicate) + assert diagnose_transcript(original)["status"] == "broken" + context = SimpleContextManager(max_tokens=100_000, compaction_notice_enabled=False) + await context.set_messages(original) + session = SimpleNamespace(coordinator=SimpleNamespace(get=lambda _: context)) + save = AsyncMock() + await repair_transcript_if_needed(session, save) + messages = await context.get_messages() + assert [_without_seq(m) for m in messages] == original + assert diagnose_transcript(messages)["status"] == "broken" + save.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ["diagnose", "foundation", "get", "set", "lookup"]) +async def test_shared_repair_scan_and_context_failures_are_best_effort( + monkeypatch, stage +): + from amplifier_app_cli.session_runner import repair_transcript_if_needed + + context = SimpleNamespace( + get_messages=AsyncMock(return_value=_ordering_violation_messages()), + set_messages=AsyncMock(), + ) + lookup = Mock(return_value=context) + if stage == "diagnose": + monkeypatch.setattr( + "amplifier_foundation.session.diagnose_transcript", + Mock(side_effect=RuntimeError("diagnosis failed")), + ) + elif stage == "foundation": + monkeypatch.delattr("amplifier_foundation.session.repair_transcript") + elif stage in ("get", "set"): + getattr(context, f"{stage}_messages").side_effect = RuntimeError( + "context failed" + ) + else: + lookup.side_effect = RuntimeError("lookup failed") + session = SimpleNamespace(coordinator=SimpleNamespace(get=lookup)) + save = AsyncMock() + await repair_transcript_if_needed(session, save) + save.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("state", ["absent", "no-get", "no-set", "empty", "healthy"]) +async def test_shared_repair_unavailable_or_healthy_context_does_not_save(state): + from amplifier_app_cli.session_runner import repair_transcript_if_needed + + context = SimpleNamespace( + get_messages=AsyncMock( + return_value=[] if state == "empty" else _healthy_messages() + ), + set_messages=AsyncMock(), + ) + if state == "absent": + context = None + elif state in ("no-get", "no-set"): + delattr(context, "get_messages" if state == "no-get" else "set_messages") + session = SimpleNamespace(coordinator=SimpleNamespace(get=lambda _: context)) + save = AsyncMock() + await repair_transcript_if_needed(session, save) + save.assert_not_awaited()