From 0bdbcd91a8d6024a7bd938687bcb427178be8c0e Mon Sep 17 00:00:00 2001 From: Mateusz Konopelski Date: Tue, 14 Jul 2026 14:15:26 +0200 Subject: [PATCH 1/2] fix(tui): give retry in-app feedback via a shared retry core MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pressing retry on a failed generation suspended the full-screen TUI to the raw terminal (print + Enter prompt), which reads as a crash. The retry decision logic was also duplicated between the CLI handler and the MCP tool. Extract one print-free core, services.retry.retry_generation_core, returning a sum type (AlreadyRunning / PendingRejectedBeforeCodegen / PendingNotFailed / Queued / BackendError) — every field required, no Optionals. It operates on an already-resolved id; resolution and the "no session" case stay at each surface where they genuinely differ. - TUI: _retry_flow calls the core directly and renders in-app (toasts for queued/already-running, MessageScreen for pending/backend errors); no more suspend-to-terminal. actions.do_retry removed. - CLI (cmd_retry_generation) and MCP (retry_generation) become thin callers of the core, keeping their own print/exit-code and chat_json rendering. CLI now also guards the pending state and returns a clean exit 1 on backend error instead of crashing. - refresh_status now re-arms the poll timer when a run leaves a terminal state (failed -> retry -> pending), fixing a latent staleness bug. Tests: new test_retry_core.py and test_server_retry.py; CLI/TUI tests updated. make unit-tests: 673 passed. --- mcp_server/cli.py | 50 ++++++++------ mcp_server/server.py | 89 ++++++++++++------------- mcp_server/services/retry.py | 81 +++++++++++++++++++++++ mcp_server/tests/test_cli.py | 67 +++++++++++++++++-- mcp_server/tests/test_retry_core.py | 81 +++++++++++++++++++++++ mcp_server/tests/test_server_retry.py | 94 +++++++++++++++++++++++++++ mcp_server/tests/test_tui_actions.py | 12 +--- mcp_server/tests/test_tui_app.py | 82 ++++++++++++++++++----- mcp_server/tui/actions.py | 14 ++-- mcp_server/tui/app.py | 41 +++++++++++- 10 files changed, 502 insertions(+), 109 deletions(-) create mode 100644 mcp_server/services/retry.py create mode 100644 mcp_server/tests/test_retry_core.py create mode 100644 mcp_server/tests/test_server_retry.py diff --git a/mcp_server/cli.py b/mcp_server/cli.py index d64deec..e6b0119 100644 --- a/mcp_server/cli.py +++ b/mcp_server/cli.py @@ -259,8 +259,14 @@ async def cmd_check_status(args: argparse.Namespace) -> int: async def cmd_retry_generation(args: argparse.Namespace) -> int: """retry-generation: retry a failed generation.""" from services.session import set_project_root, resolve_generation_id - from services.specflow_backend import call_backend_endpoint - from services.tool_helpers import check_status_safe, is_generation_in_progress + from services.retry import ( + retry_generation_core, + AlreadyRunning, + PendingRejectedBeforeCodegen, + PendingNotFailed, + Queued, + BackendError, + ) root = resolve_root(args.root_path) print(f"Using project root: {root}") @@ -271,23 +277,29 @@ async def cmd_retry_generation(args: argparse.Namespace) -> int: print("No previous generation found. Run `specflow run-generation` to start one.") return 0 - status_data = await check_status_safe(generation_id) - if is_generation_in_progress(status_data): - print( - "ERROR: A generation is already running. Wait for it to finish before retrying.", - file=sys.stderr, - ) - return 1 - - response_text = await call_backend_endpoint( - endpoint=f"/api/v1/generation-sessions/{generation_id}/retry", - method="POST", - timeout_seconds=30, - ) - data = json.loads(response_text) - print(json.dumps(data, indent=2)) - print("\nRetry queued. Generation will resume from the last checkpoint on the same workspaces.") - return 0 + match await retry_generation_core(generation_id): + case AlreadyRunning(): + print( + "ERROR: A generation is already running. Wait for it to finish before retrying.", + file=sys.stderr, + ) + return 1 + case PendingRejectedBeforeCodegen(error=error): + print( + f"Last run was rejected before codegen: {error}\n" + "Fix files locally and run `specflow run-generation` — not retry." + ) + return 0 + case PendingNotFailed(): + print("Session is pending but has not failed. Use `specflow run-generation`, not retry.") + return 0 + case Queued(backend_data=backend_data): + print(json.dumps(backend_data, indent=2)) + print("\nRetry queued. Generation will resume from the last checkpoint on the same workspaces.") + return 0 + case BackendError(error=error): + print(f"ERROR: Couldn't retry: {error}", file=sys.stderr) + return 1 async def cmd_download_outputs(args: argparse.Namespace) -> int: diff --git a/mcp_server/server.py b/mcp_server/server.py index 649d503..1cf2e2e 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -57,6 +57,14 @@ resolve_workspace_count, ) from services.user_response import brief_sentences, chat_json +from services.retry import ( + AlreadyRunning, + BackendError, + PendingNotFailed, + PendingRejectedBeforeCodegen, + Queued, + retry_generation_core, +) BACKEND_URL = os.getenv("BACKEND_URL", "http://127.0.0.1:8000") logger.info("MCP Server starting with Backend URL: %s", BACKEND_URL) @@ -617,58 +625,47 @@ async def retry_generation( if sp.is_absolute(): set_project_root(sp.parent) - try: - generation_id = resolve_generation_id(generation_id) - if not generation_id: - return chat_json( - "No previous generation found. Run `run_generation` to start one.", - ) + generation_id = resolve_generation_id(generation_id) + if not generation_id: + return chat_json( + "No previous generation found. Run `run_generation` to start one.", + ) - status_data = await check_status_safe(generation_id) - if is_generation_in_progress(status_data): + match await retry_generation_core(generation_id): + case AlreadyRunning(): return _rejection_chat_payload({ "error": "A generation is already running. Wait for it to finish before retrying.", "code": PrecheckRejectionCode.GENERATION_ALREADY_RUNNING, "generation_id": generation_id, }) - - if status_data: - st = (status_data.get("status") or "").lower() - if st == GenerationStatus.PENDING and (status_data.get("error") or "").strip(): - return chat_json( - brief_sentences( - f"Last run was rejected before codegen: {status_data.get('error')} " - "Fix files locally and call `run_generation` — not `retry_generation`." - ), - details=status_data, - generation_id=generation_id, - ) - if st == GenerationStatus.PENDING: - return chat_json( - "Session is pending but has not failed. Use `run_generation`, not `retry_generation`.", - details=status_data, - generation_id=generation_id, - ) - - response_text = await call_backend_endpoint( - endpoint=f"/api/v1/generation-sessions/{generation_id}/retry", - method="POST", - timeout_seconds=30, - ) - backend_data = json.loads(response_text) - return chat_json( - "Retry queued. Generation will resume from the last checkpoint on the same workspaces.", - details=backend_data, - generation_id=generation_id, - status=backend_data.get("status"), - retry_count=backend_data.get("retry_count"), - ) - - except Exception as e: - return chat_json( - "Couldn't retry. The SpecFlow server may be unreachable.", - details={"error": str(e)}, - ) + case PendingRejectedBeforeCodegen(status_data=status_data, error=error): + return chat_json( + brief_sentences( + f"Last run was rejected before codegen: {error} " + "Fix files locally and call `run_generation` — not `retry_generation`." + ), + details=status_data, + generation_id=generation_id, + ) + case PendingNotFailed(status_data=status_data): + return chat_json( + "Session is pending but has not failed. Use `run_generation`, not `retry_generation`.", + details=status_data, + generation_id=generation_id, + ) + case Queued(backend_data=backend_data): + return chat_json( + "Retry queued. Generation will resume from the last checkpoint on the same workspaces.", + details=backend_data, + generation_id=generation_id, + status=backend_data.get("status"), + retry_count=backend_data.get("retry_count"), + ) + case BackendError(error=error): + return chat_json( + "Couldn't retry. The SpecFlow server may be unreachable.", + details={"error": error}, + ) def main(): diff --git a/mcp_server/services/retry.py b/mcp_server/services/retry.py new file mode 100644 index 0000000..c26ba4d --- /dev/null +++ b/mcp_server/services/retry.py @@ -0,0 +1,81 @@ +"""Shared retry decision logic — one source of truth for CLI, MCP, and TUI. + +The core operates on an already-resolved generation id and returns a sum type: +exactly one outcome dataclass, each carrying only the fields it always has. +Callers own resolution ("which session?") and rendering (print / chat_json / +toast) — those genuinely differ per surface. The core owns the status→retry +decision and is non-raising for every expected outcome. +""" +from __future__ import annotations + +import json +from dataclasses import dataclass + +from schemas.generation_workflow_enums import GenerationStatus +from services.specflow_backend import call_backend_endpoint +from services.tool_helpers import check_status_safe, is_generation_in_progress + + +@dataclass(frozen=True) +class AlreadyRunning: + """Session is running/initializing — retry is not allowed.""" + + +@dataclass(frozen=True) +class PendingRejectedBeforeCodegen: + """Session is pending with a rejection error — user must fix files, not retry.""" + status_data: dict + error: str + + +@dataclass(frozen=True) +class PendingNotFailed: + """Session is pending but not failed — nothing to retry yet.""" + status_data: dict + + +@dataclass(frozen=True) +class Queued: + """Retry accepted — backend reset the session and re-fired the workflow.""" + backend_data: dict + + +@dataclass(frozen=True) +class BackendError: + """The retry POST failed (server unreachable / HTTP error).""" + error: str + + +RetryOutcome = ( + AlreadyRunning + | PendingRejectedBeforeCodegen + | PendingNotFailed + | Queued + | BackendError +) + + +async def retry_generation_core(generation_id: str) -> RetryOutcome: + """Decide and perform a retry for an already-resolved generation id.""" + status_data = await check_status_safe(generation_id) + + if is_generation_in_progress(status_data): + return AlreadyRunning() + + if status_data: + st = (status_data.get("status") or "").lower() + if st == GenerationStatus.PENDING: + error = (status_data.get("error") or "").strip() + if error: + return PendingRejectedBeforeCodegen(status_data=status_data, error=error) + return PendingNotFailed(status_data=status_data) + + try: + text = await call_backend_endpoint( + endpoint=f"/api/v1/generation-sessions/{generation_id}/retry", + method="POST", + timeout_seconds=30, + ) + return Queued(backend_data=json.loads(text)) + except Exception as e: # noqa: BLE001 - surfaced as a structured outcome, not raised + return BackendError(error=str(e)) diff --git a/mcp_server/tests/test_cli.py b/mcp_server/tests/test_cli.py index b64e39a..9e29e36 100644 --- a/mcp_server/tests/test_cli.py +++ b/mcp_server/tests/test_cli.py @@ -285,12 +285,13 @@ async def test_calls_retry_endpoint(self, tmp_project): args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) - # check_status_safe returns failed → proceed to retry POST + # check_status_safe returns failed → proceed to retry POST. Patched at + # services.retry.* where retry_generation_core binds these primitives. with patch( - "services.tool_helpers.check_status_safe", + "services.retry.check_status_safe", new_callable=AsyncMock, ) as mock_status, patch( - "services.specflow_backend.call_backend_endpoint", + "services.retry.call_backend_endpoint", new_callable=AsyncMock, ) as mock_ep: mock_status.return_value = {"status": "failed"} @@ -311,7 +312,7 @@ async def test_blocks_when_already_running(self, tmp_project, capsys): args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) with patch( - "services.tool_helpers.check_status_safe", + "services.retry.check_status_safe", new_callable=AsyncMock, ) as mock_cs: mock_cs.return_value = {"status": "running"} @@ -319,6 +320,64 @@ async def test_blocks_when_already_running(self, tmp_project, capsys): assert code == 1 + @pytest.mark.asyncio + async def test_no_session_returns_zero_without_backend_call(self, tmp_project): + from cli import cmd_retry_generation + + # No write_session → resolve_generation_id returns None. + args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) + + with patch( + "services.retry.check_status_safe", + new_callable=AsyncMock, + ) as mock_status: + code = await cmd_retry_generation(args) + + assert code == 0 + assert mock_status.call_count == 0 # core never called + + @pytest.mark.asyncio + async def test_pending_not_failed_returns_zero_without_post(self, tmp_project): + from cli import cmd_retry_generation + from services.session import write_session + write_session("gen-pending", tmp_project) + + args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) + + with patch( + "services.retry.check_status_safe", + new_callable=AsyncMock, + ) as mock_status, patch( + "services.retry.call_backend_endpoint", + new_callable=AsyncMock, + ) as mock_ep: + mock_status.return_value = {"status": "pending"} + code = await cmd_retry_generation(args) + + assert code == 0 + assert mock_ep.call_count == 0 # pending short-circuits before POST + + @pytest.mark.asyncio + async def test_backend_error_returns_one(self, tmp_project): + from cli import cmd_retry_generation + from services.session import write_session + write_session("gen-err", tmp_project) + + args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) + + with patch( + "services.retry.check_status_safe", + new_callable=AsyncMock, + ) as mock_status, patch( + "services.retry.call_backend_endpoint", + new_callable=AsyncMock, + ) as mock_ep: + mock_status.return_value = {"status": "failed"} + mock_ep.side_effect = Exception("Backend returned HTTP 500") + code = await cmd_retry_generation(args) + + assert code == 1 + # --------------------------------------------------------------------------- # cmd_download_outputs diff --git a/mcp_server/tests/test_retry_core.py b/mcp_server/tests/test_retry_core.py new file mode 100644 index 0000000..21b6953 --- /dev/null +++ b/mcp_server/tests/test_retry_core.py @@ -0,0 +1,81 @@ +"""Tests for the shared retry core (services/retry.py). + +One case per outcome variant. Patches are applied at ``services.retry.*`` because +the core binds ``check_status_safe`` / ``call_backend_endpoint`` at import time. +""" +import json +from unittest.mock import AsyncMock, patch + +import pytest + +from services.retry import ( + AlreadyRunning, + BackendError, + PendingNotFailed, + PendingRejectedBeforeCodegen, + Queued, + retry_generation_core, +) + +GID = "est-abc123" + + +def _patches(status_return, post=None): + status = patch("services.retry.check_status_safe", new=AsyncMock(return_value=status_return)) + post_mock = AsyncMock(return_value=post) if not isinstance(post, Exception) else AsyncMock(side_effect=post) + endpoint = patch("services.retry.call_backend_endpoint", new=post_mock) + return status, endpoint, post_mock + + +@pytest.mark.asyncio +async def test_already_running(): + status, endpoint, post = _patches({"status": "running"}) + with status as st, endpoint: + outcome = await retry_generation_core(GID) + assert isinstance(outcome, AlreadyRunning) + assert st.await_count == 1 + assert post.await_count == 0 # no POST when already running + + +@pytest.mark.asyncio +async def test_pending_rejected_before_codegen(): + status, endpoint, post = _patches({"status": "pending", "error": "Missing required files"}) + with status, endpoint: + outcome = await retry_generation_core(GID) + assert isinstance(outcome, PendingRejectedBeforeCodegen) + assert outcome.error == "Missing required files" + assert outcome.status_data["status"] == "pending" + assert post.await_count == 0 + + +@pytest.mark.asyncio +async def test_pending_not_failed(): + status, endpoint, post = _patches({"status": "pending"}) + with status, endpoint: + outcome = await retry_generation_core(GID) + assert isinstance(outcome, PendingNotFailed) + assert outcome.status_data["status"] == "pending" + assert post.await_count == 0 + + +@pytest.mark.asyncio +async def test_queued(): + status, endpoint, post = _patches( + {"status": "failed"}, post=json.dumps({"status": "pending", "retry_count": 1}) + ) + with status as st, endpoint: + outcome = await retry_generation_core(GID) + assert isinstance(outcome, Queued) + assert outcome.backend_data == {"status": "pending", "retry_count": 1} + assert st.await_count == 1 + assert post.await_count == 1 # exactly one POST + + +@pytest.mark.asyncio +async def test_backend_error(): + status, endpoint, post = _patches({"status": "failed"}, post=Exception("Backend returned HTTP 500")) + with status, endpoint: + outcome = await retry_generation_core(GID) + assert isinstance(outcome, BackendError) + assert "500" in outcome.error + assert post.await_count == 1 diff --git a/mcp_server/tests/test_server_retry.py b/mcp_server/tests/test_server_retry.py new file mode 100644 index 0000000..65828d0 --- /dev/null +++ b/mcp_server/tests/test_server_retry.py @@ -0,0 +1,94 @@ +"""Characterization + parity tests for the MCP ``retry_generation`` tool. + +Locks the ``chat_json`` shapes for every outcome so the core refactor preserves +them. `retry_generation` is a plain async function; `ctx` is only consumed by +`apply_project_root_from_context`, which we stub. Backend primitives are patched +at both `server.*` (resolution/guard) and `services.retry.*` (the core's POST + +status read) so these tests hold before and after the refactor. +""" +import json +from contextlib import ExitStack +from unittest.mock import AsyncMock, patch + +import pytest + +import server + +GID = "est-abc123" + + +async def _run(status_return, *, resolved=GID, post=None): + """Invoke retry_generation with all IO stubbed; return (payload dict, post mock).""" + post_mock = ( + AsyncMock(side_effect=post) if isinstance(post, Exception) + else AsyncMock(return_value=post) + ) + status_mock = AsyncMock(return_value=status_return) + with ExitStack() as stack: + stack.enter_context(patch("server.apply_project_root_from_context", new=AsyncMock())) + stack.enter_context(patch("server.resolve_generation_id", return_value=resolved)) + stack.enter_context(patch("server.check_status_safe", new=status_mock)) + stack.enter_context(patch("services.retry.check_status_safe", new=status_mock)) + stack.enter_context(patch("services.retry.call_backend_endpoint", new=post_mock)) + stack.enter_context(patch("server.call_backend_endpoint", new=post_mock)) + payload = await server.retry_generation(generation_id=None, ctx=object()) + return json.loads(payload), post_mock + + +@pytest.mark.asyncio +async def test_no_session(): + with ( + patch("server.apply_project_root_from_context", new=AsyncMock()), + patch("server.resolve_generation_id", return_value=None), + ): + payload = await server.retry_generation(generation_id=None, ctx=object()) + data = json.loads(payload) + assert "No previous generation found" in data["message"] + + +@pytest.mark.asyncio +async def test_already_running(): + data, post = await _run({"status": "running"}) + assert "already running" in data["message"].lower() + assert data["code"] == "GENERATION_ALREADY_RUNNING" + assert data["generation_id"] == GID + assert post.await_count == 0 + + +@pytest.mark.asyncio +async def test_pending_rejected_before_codegen(): + data, post = await _run({"status": "pending", "error": "Missing required files"}) + assert "rejected before codegen" in data["message"].lower() + assert data["details"]["error"] == "Missing required files" + assert data["generation_id"] == GID + assert post.await_count == 0 + + +@pytest.mark.asyncio +async def test_pending_not_failed(): + data, post = await _run({"status": "pending"}) + assert "pending but has not failed" in data["message"].lower() + assert data["details"]["status"] == "pending" + assert data["generation_id"] == GID + assert post.await_count == 0 + + +@pytest.mark.asyncio +async def test_queued(): + data, post = await _run( + {"status": "failed"}, post=json.dumps({"status": "pending", "retry_count": 1}) + ) + assert "retry queued" in data["message"].lower() + assert data["status"] == "pending" + assert data["retry_count"] == 1 + assert data["details"] == {"status": "pending", "retry_count": 1} + assert data["generation_id"] == GID + assert post.await_count == 1 + + +@pytest.mark.asyncio +async def test_backend_error(): + data, post = await _run({"status": "failed"}, post=Exception("Backend returned HTTP 500")) + assert "couldn't retry" in data["message"].lower() + assert "500" in data["details"]["error"] + assert post.await_count == 1 diff --git a/mcp_server/tests/test_tui_actions.py b/mcp_server/tests/test_tui_actions.py index bedcb05..495fead 100644 --- a/mcp_server/tests/test_tui_actions.py +++ b/mcp_server/tests/test_tui_actions.py @@ -3,9 +3,10 @@ The wrappers must delegate to the existing cli.cmd_* handlers (single source of truth for guards/precheck/backend calls), passing the namespace fields each handler reads. We patch the handlers and assert delegation + key arguments. +(Retry no longer goes through here — it calls the retry core directly; see +test_tui_app.TestDashboardActionFlows.) """ -from pathlib import Path from unittest.mock import AsyncMock, patch import pytest @@ -13,15 +14,6 @@ from tui import actions -@pytest.mark.asyncio -async def test_do_retry_delegates_with_root(): - with patch("cli.cmd_retry_generation", new=AsyncMock(return_value=0)) as m: - rc = await actions.do_retry(Path("/proj")) - assert rc == 0 - ns = m.await_args.args[0] - assert ns.root_path == "/proj" - - @pytest.mark.asyncio async def test_do_clear_set_forces_yes(): with patch("cli.cmd_clear_workspace", new=AsyncMock(return_value=0)) as m: diff --git a/mcp_server/tests/test_tui_app.py b/mcp_server/tests/test_tui_app.py index 07aabfb..ae0d8b7 100644 --- a/mcp_server/tests/test_tui_app.py +++ b/mcp_server/tests/test_tui_app.py @@ -1117,15 +1117,15 @@ async def test_cancel_stops_countdown_timer(self): class TestDashboardActionFlows: - """Retry/clear worker flows: confirmation gating and clear eligibility. + """Retry/clear worker flows: confirmation gating and in-app feedback. - ``_run_suspended`` (which suspends the app and blocks on ``input()``) is - patched out so we exercise only the decision logic; ``actions.do_*`` are - patched to assert which CLI handler the flow ultimately reaches. + Retry now calls ``retry_generation_core`` directly and renders results via + ``notify``/``MessageScreen`` — it must never suspend the app. We patch the + core to force each outcome and assert the flow's response. """ @pytest.mark.asyncio - async def test_retry_runs_action_when_confirmed(self): + async def test_retry_calls_core_and_refreshes_without_suspending(self): a, b, c = _gate_ready() with a, b, c, patch("tui.app.poll_once", new=AsyncMock(return_value=_running_payload())): app = tui_app.SpecFlowTUI(root=Path("/tmp/x"), generation_id="gen_x", poll_interval=999) @@ -1134,19 +1134,46 @@ async def test_retry_runs_action_when_confirmed(self): screen = app.screen with ( patch.object(app, "push_screen_wait", new=AsyncMock(return_value=True)), - patch.object(screen, "_run_suspended", new=AsyncMock()) as run_susp, - # do_retry is async; force a sync mock so the flow's - # ``do_retry(root, generation_id)`` yields a sentinel, not a live coroutine. patch( - "tui.app.actions.do_retry", new=MagicMock(return_value="retry-coro") - ) as do_retry, + "tui.app.retry_generation_core", + new=AsyncMock(return_value=tui_app.Queued(backend_data={"status": "pending"})), + ) as core, + patch.object(screen, "refresh_status", new=AsyncMock()) as refresh, + patch.object(screen, "notify") as notify, + patch.object(app, "suspend") as suspend, ): await screen._retry_flow() - do_retry.assert_called_once_with(app.root, "gen_x") - run_susp.assert_awaited_once_with("retry-coro") + core.assert_awaited_once_with("gen_x") + refresh.assert_awaited_once() + suspend.assert_not_called() # never drops to the terminal + assert any("queued" in str(c.args[0]).lower() for c in notify.call_args_list) @pytest.mark.asyncio - async def test_retry_skips_action_when_cancelled(self): + async def test_retry_backend_error_shows_message_screen(self): + a, b, c = _gate_ready() + with a, b, c, patch("tui.app.poll_once", new=AsyncMock(return_value=_running_payload())): + app = tui_app.SpecFlowTUI(root=Path("/tmp/x"), generation_id="gen_x", poll_interval=999) + async with app.run_test() as pilot: + await pilot.pause() + screen = app.screen + psw = AsyncMock(return_value=True) + with ( + patch.object(app, "push_screen_wait", new=psw), + patch( + "tui.app.retry_generation_core", + new=AsyncMock(return_value=tui_app.BackendError(error="unreachable")), + ), + patch.object(screen, "refresh_status", new=AsyncMock()), + patch.object(app, "suspend") as suspend, + ): + await screen._retry_flow() + suspend.assert_not_called() + # push_screen_wait called twice: ConfirmScreen, then the error MessageScreen. + assert psw.await_count == 2 + assert isinstance(psw.await_args_list[1].args[0], tui_app.MessageScreen) + + @pytest.mark.asyncio + async def test_retry_skips_core_when_cancelled(self): a, b, c = _gate_ready() with a, b, c, patch("tui.app.poll_once", new=AsyncMock(return_value=_running_payload())): app = tui_app.SpecFlowTUI(root=Path("/tmp/x"), generation_id="gen_x", poll_interval=999) @@ -1155,12 +1182,33 @@ async def test_retry_skips_action_when_cancelled(self): screen = app.screen with ( patch.object(app, "push_screen_wait", new=AsyncMock(return_value=False)), - patch.object(screen, "_run_suspended", new=AsyncMock()) as run_susp, - patch("tui.app.actions.do_retry") as do_retry, + patch("tui.app.retry_generation_core", new=AsyncMock()) as core, + patch.object(screen, "refresh_status", new=AsyncMock()) as refresh, ): await screen._retry_flow() - do_retry.assert_not_called() - run_susp.assert_not_awaited() + core.assert_not_awaited() + refresh.assert_not_awaited() + + @pytest.mark.asyncio + async def test_poll_timer_self_heals_after_terminal(self): + a, b, c = _gate_ready() + with a, b, c, patch("tui.app.poll_once", new=AsyncMock(return_value=_running_payload())): + app = tui_app.SpecFlowTUI(root=Path("/tmp/x"), generation_id="gen_x", poll_interval=999) + async with app.run_test() as pilot: + await pilot.pause() + screen = app.screen + # Terminal payload stops the poll timer. + with patch("tui.app.poll_once", new=AsyncMock(return_value={"status": "failed"})): + await screen.refresh_status() + assert screen._poll_timer is None + # A subsequent non-terminal payload re-arms exactly one timer... + with patch("tui.app.poll_once", new=AsyncMock(return_value={"status": "pending"})): + await screen.refresh_status() + assert screen._poll_timer is not None + armed = screen._poll_timer + # ...and refreshing again while non-terminal does not create a second. + await screen.refresh_status() + assert screen._poll_timer is armed @pytest.mark.asyncio async def test_clear_runs_for_eligible_set_when_confirmed(self): diff --git a/mcp_server/tui/actions.py b/mcp_server/tui/actions.py index 556a32c..a858685 100644 --- a/mcp_server/tui/actions.py +++ b/mcp_server/tui/actions.py @@ -1,19 +1,20 @@ """In-app actions — thin wrappers over the existing CLI command handlers. These deliberately call the same ``cmd_*`` coroutines the standalone -subcommands use (``cli.cmd_retry_generation`` etc.), so every guard, capacity +subcommands use (``cli.cmd_clear_workspace`` etc.), so every guard, capacity message, precheck, and backend call stays in exactly one place. The TUI never re-implements those flows; it suspends its screen (see ``app.py``) and runs the real handler, which prints its familiar output to the terminal. Each wrapper builds the ``argparse.Namespace`` the handler expects and returns its integer exit code. ``--yes`` is forced on for the workspace clear because -the TUI gathers confirmation through its own dialog before calling. +the TUI gathers confirmation through its own dialog before calling. (Retry does +not go through here — it calls ``services.retry.retry_generation_core`` directly +so it can render feedback in-app instead of suspending.) """ from __future__ import annotations -from pathlib import Path from types import SimpleNamespace import cli @@ -44,13 +45,6 @@ def _ns(**overrides: object) -> SimpleNamespace: return SimpleNamespace(**base) -async def do_retry(root: Path, generation_id: str | None = None) -> int: - """Retry the current generation (reuses ``cmd_retry_generation`` guards).""" - return await cli.cmd_retry_generation( - _ns(root_path=str(root), generation_id=generation_id) - ) - - async def do_clear_set(set_number: int) -> int: """Clear all members of a workspace set (confirmation handled by the TUI).""" return await cli.cmd_clear_workspace(_ns(set=set_number, yes=True)) diff --git a/mcp_server/tui/app.py b/mcp_server/tui/app.py index 243c00e..41e8d67 100644 --- a/mcp_server/tui/app.py +++ b/mcp_server/tui/app.py @@ -63,6 +63,14 @@ from cli import resolve_backend_config from services import local_env, validate_models from services.llm_tiers import LLM_TIER_KEYS +from services.retry import ( + AlreadyRunning, + BackendError, + PendingNotFailed, + PendingRejectedBeforeCodegen, + Queued, + retry_generation_core, +) from services.session import resolve_generation_id, set_project_root from services.specflow_backend import call_backend_endpoint_bytes from tui import actions, activity, mcp_clients, onboarding, render @@ -535,11 +543,15 @@ async def refresh_status(self) -> None: build_dashboard(payload, self.app.root, self._generation_id, self._selected_ws_id) ) self.refresh_bindings() - # A finished run will not change again — stop polling it. + # A finished run will not change again — stop polling it. Conversely, a + # run that left a terminal state (e.g. failed → retry → pending) must + # resume polling, so re-arm the timer if it was previously stopped. if payload and (payload.get("status") or "").lower() in TERMINAL_STATUSES: if self._poll_timer is not None: self._poll_timer.stop() self._poll_timer = None + elif self._poll_timer is None: + self._poll_timer = self.set_interval(self.app.poll_interval, self.refresh_status) def check_action(self, action: str, parameters: tuple[object, ...]) -> bool | None: if action == "retry": @@ -627,8 +639,31 @@ async def _retry_flow(self) -> None: countdown=0, ) ) - if ok: - await self._run_suspended(actions.do_retry(self.app.root, self._generation_id)) + if not ok: + return + self.notify("Retrying generation…", severity="information") + match await retry_generation_core(self._generation_id): + case Queued(): + self.notify( + "Retry queued. Resuming from the last checkpoint.", + severity="information", + ) + case AlreadyRunning(): + self.notify("A generation is already running.", severity="warning") + case PendingRejectedBeforeCodegen(error=error): + await self.app.push_screen_wait(MessageScreen( + "Retry", + f"Last run was rejected before codegen: {error}\n" + "Fix files locally and start a new generation — retry won't help.", + )) + case PendingNotFailed(): + await self.app.push_screen_wait(MessageScreen( + "Retry", + "Session is pending but has not failed — nothing to retry yet.", + )) + case BackendError(error=error): + await self.app.push_screen_wait(MessageScreen("Retry failed", error)) + await self.refresh_status() def action_clear(self) -> None: self.run_worker(self._clear_flow(), exclusive=True) From e3057f2f853ffb783575b648dd20558ec53bd01f Mon Sep 17 00:00:00 2001 From: Mateusz Konopelski Date: Tue, 14 Jul 2026 15:35:57 +0200 Subject: [PATCH 2/2] refactor(retry): defer eligibility to the backend, collapse the retry core MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Following the cancel flow (PR #48): retry validation belongs to the backend state machine (reset_for_retry accepts FAILED or stuck-PENDING), so the client-side status pre-check was a second validator — and it had drifted, refusing PENDING that the backend accepts. Drop the pre-check and the outcome sum type. services.retry.retry_generation_core is now a one-call helper: POST /retry, return the parsed dict, raise on failure. Each surface POSTs and renders success vs the backend's error message in its own style, mirroring _cancel_flow: - TUI _retry_flow: toast on success, MessageScreen with the backend reason on error. - CLI: print + exit 0/1. - MCP: chat_json; drops the client-derived `code`/pending-guidance fields in favour of the backend's message (single source of truth). No client status round-trip before the POST. Tests updated (core, CLI, MCP, TUI). make unit-tests: 671 passed. --- mcp_server/cli.py | 40 +++--------- mcp_server/server.py | 59 +++++------------- mcp_server/services/retry.py | 88 ++++++--------------------- mcp_server/tests/test_cli.py | 66 +++----------------- mcp_server/tests/test_retry_core.py | 84 ++++++------------------- mcp_server/tests/test_server_retry.py | 72 +++++++--------------- mcp_server/tests/test_tui_app.py | 7 ++- mcp_server/tui/app.py | 36 +++-------- 8 files changed, 101 insertions(+), 351 deletions(-) diff --git a/mcp_server/cli.py b/mcp_server/cli.py index e6b0119..ba83f70 100644 --- a/mcp_server/cli.py +++ b/mcp_server/cli.py @@ -259,14 +259,7 @@ async def cmd_check_status(args: argparse.Namespace) -> int: async def cmd_retry_generation(args: argparse.Namespace) -> int: """retry-generation: retry a failed generation.""" from services.session import set_project_root, resolve_generation_id - from services.retry import ( - retry_generation_core, - AlreadyRunning, - PendingRejectedBeforeCodegen, - PendingNotFailed, - Queued, - BackendError, - ) + from services.retry import retry_generation_core root = resolve_root(args.root_path) print(f"Using project root: {root}") @@ -277,29 +270,14 @@ async def cmd_retry_generation(args: argparse.Namespace) -> int: print("No previous generation found. Run `specflow run-generation` to start one.") return 0 - match await retry_generation_core(generation_id): - case AlreadyRunning(): - print( - "ERROR: A generation is already running. Wait for it to finish before retrying.", - file=sys.stderr, - ) - return 1 - case PendingRejectedBeforeCodegen(error=error): - print( - f"Last run was rejected before codegen: {error}\n" - "Fix files locally and run `specflow run-generation` — not retry." - ) - return 0 - case PendingNotFailed(): - print("Session is pending but has not failed. Use `specflow run-generation`, not retry.") - return 0 - case Queued(backend_data=backend_data): - print(json.dumps(backend_data, indent=2)) - print("\nRetry queued. Generation will resume from the last checkpoint on the same workspaces.") - return 0 - case BackendError(error=error): - print(f"ERROR: Couldn't retry: {error}", file=sys.stderr) - return 1 + try: + backend_data = await retry_generation_core(generation_id) + except Exception as exc: # noqa: BLE001 - backend validates state; surface its reason + print(f"ERROR: Couldn't retry: {exc}", file=sys.stderr) + return 1 + print(json.dumps(backend_data, indent=2)) + print("\nRetry queued. Generation will resume from the last checkpoint on the same workspaces.") + return 0 async def cmd_download_outputs(args: argparse.Namespace) -> int: diff --git a/mcp_server/server.py b/mcp_server/server.py index 1cf2e2e..1688fe1 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -57,14 +57,7 @@ resolve_workspace_count, ) from services.user_response import brief_sentences, chat_json -from services.retry import ( - AlreadyRunning, - BackendError, - PendingNotFailed, - PendingRejectedBeforeCodegen, - Queued, - retry_generation_core, -) +from services.retry import retry_generation_core BACKEND_URL = os.getenv("BACKEND_URL", "http://127.0.0.1:8000") logger.info("MCP Server starting with Backend URL: %s", BACKEND_URL) @@ -631,41 +624,21 @@ async def retry_generation( "No previous generation found. Run `run_generation` to start one.", ) - match await retry_generation_core(generation_id): - case AlreadyRunning(): - return _rejection_chat_payload({ - "error": "A generation is already running. Wait for it to finish before retrying.", - "code": PrecheckRejectionCode.GENERATION_ALREADY_RUNNING, - "generation_id": generation_id, - }) - case PendingRejectedBeforeCodegen(status_data=status_data, error=error): - return chat_json( - brief_sentences( - f"Last run was rejected before codegen: {error} " - "Fix files locally and call `run_generation` — not `retry_generation`." - ), - details=status_data, - generation_id=generation_id, - ) - case PendingNotFailed(status_data=status_data): - return chat_json( - "Session is pending but has not failed. Use `run_generation`, not `retry_generation`.", - details=status_data, - generation_id=generation_id, - ) - case Queued(backend_data=backend_data): - return chat_json( - "Retry queued. Generation will resume from the last checkpoint on the same workspaces.", - details=backend_data, - generation_id=generation_id, - status=backend_data.get("status"), - retry_count=backend_data.get("retry_count"), - ) - case BackendError(error=error): - return chat_json( - "Couldn't retry. The SpecFlow server may be unreachable.", - details={"error": error}, - ) + try: + backend_data = await retry_generation_core(generation_id) + except Exception as e: # noqa: BLE001 - backend validates state; surface its reason + return chat_json( + brief_sentences(f"Couldn't retry: {e}"), + details={"error": str(e)}, + generation_id=generation_id, + ) + return chat_json( + "Retry queued. Generation will resume from the last checkpoint on the same workspaces.", + details=backend_data, + generation_id=generation_id, + status=backend_data.get("status"), + retry_count=backend_data.get("retry_count"), + ) def main(): diff --git a/mcp_server/services/retry.py b/mcp_server/services/retry.py index c26ba4d..294239e 100644 --- a/mcp_server/services/retry.py +++ b/mcp_server/services/retry.py @@ -1,81 +1,27 @@ -"""Shared retry decision logic — one source of truth for CLI, MCP, and TUI. +"""Shared retry call — one definition of the retry endpoint for CLI, MCP, and TUI. -The core operates on an already-resolved generation id and returns a sum type: -exactly one outcome dataclass, each carrying only the fields it always has. -Callers own resolution ("which session?") and rendering (print / chat_json / -toast) — those genuinely differ per surface. The core owns the status→retry -decision and is non-raising for every expected outcome. +Retry eligibility is validated by the backend state machine (``reset_for_retry`` +accepts FAILED or stuck-PENDING); there is no client-side status pre-check to +drift from it. Callers POST here and render success (the returned dict) vs the +backend's error message (``str(exc)``) in their own style — mirroring the cancel +flow, which likewise defers to the backend. """ from __future__ import annotations import json -from dataclasses import dataclass -from schemas.generation_workflow_enums import GenerationStatus from services.specflow_backend import call_backend_endpoint -from services.tool_helpers import check_status_safe, is_generation_in_progress -@dataclass(frozen=True) -class AlreadyRunning: - """Session is running/initializing — retry is not allowed.""" +async def retry_generation_core(generation_id: str) -> dict: + """POST a retry for an already-resolved id; return the parsed backend response. - -@dataclass(frozen=True) -class PendingRejectedBeforeCodegen: - """Session is pending with a rejection error — user must fix files, not retry.""" - status_data: dict - error: str - - -@dataclass(frozen=True) -class PendingNotFailed: - """Session is pending but not failed — nothing to retry yet.""" - status_data: dict - - -@dataclass(frozen=True) -class Queued: - """Retry accepted — backend reset the session and re-fired the workflow.""" - backend_data: dict - - -@dataclass(frozen=True) -class BackendError: - """The retry POST failed (server unreachable / HTTP error).""" - error: str - - -RetryOutcome = ( - AlreadyRunning - | PendingRejectedBeforeCodegen - | PendingNotFailed - | Queued - | BackendError -) - - -async def retry_generation_core(generation_id: str) -> RetryOutcome: - """Decide and perform a retry for an already-resolved generation id.""" - status_data = await check_status_safe(generation_id) - - if is_generation_in_progress(status_data): - return AlreadyRunning() - - if status_data: - st = (status_data.get("status") or "").lower() - if st == GenerationStatus.PENDING: - error = (status_data.get("error") or "").strip() - if error: - return PendingRejectedBeforeCodegen(status_data=status_data, error=error) - return PendingNotFailed(status_data=status_data) - - try: - text = await call_backend_endpoint( - endpoint=f"/api/v1/generation-sessions/{generation_id}/retry", - method="POST", - timeout_seconds=30, - ) - return Queued(backend_data=json.loads(text)) - except Exception as e: # noqa: BLE001 - surfaced as a structured outcome, not raised - return BackendError(error=str(e)) + Raises on any failure — the backend validates state and returns the reason, + so callers surface ``str(exc)`` to the user. + """ + text = await call_backend_endpoint( + endpoint=f"/api/v1/generation-sessions/{generation_id}/retry", + method="POST", + timeout_seconds=30, + ) + return json.loads(text) diff --git a/mcp_server/tests/test_cli.py b/mcp_server/tests/test_cli.py index 9e29e36..84e635d 100644 --- a/mcp_server/tests/test_cli.py +++ b/mcp_server/tests/test_cli.py @@ -285,37 +285,31 @@ async def test_calls_retry_endpoint(self, tmp_project): args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) - # check_status_safe returns failed → proceed to retry POST. Patched at - # services.retry.* where retry_generation_core binds these primitives. + # POST succeeds → exit 0. Patched at services.retry.* where the helper binds it. with patch( - "services.retry.check_status_safe", - new_callable=AsyncMock, - ) as mock_status, patch( "services.retry.call_backend_endpoint", new_callable=AsyncMock, ) as mock_ep: - mock_status.return_value = {"status": "failed"} - mock_ep.return_value = json.dumps({"status": "running", "retry_count": 1}) + mock_ep.return_value = json.dumps({"status": "pending", "retry_count": 1}) code = await cmd_retry_generation(args) assert code == 0 - # Exactly one backend call: the retry POST - assert mock_ep.call_count == 1 - assert mock_status.call_count == 1 + assert mock_ep.call_count == 1 # single POST; backend validates state @pytest.mark.asyncio - async def test_blocks_when_already_running(self, tmp_project, capsys): + async def test_backend_rejection_returns_one(self, tmp_project): from cli import cmd_retry_generation from services.session import write_session write_session("gen-running", tmp_project) args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) + # Wrong-state / unreachable both raise from the helper → exit 1. with patch( - "services.retry.check_status_safe", + "services.retry.call_backend_endpoint", new_callable=AsyncMock, - ) as mock_cs: - mock_cs.return_value = {"status": "running"} + ) as mock_ep: + mock_ep.side_effect = Exception("Backend returned HTTP 400: Cannot retry in 'running' state") code = await cmd_retry_generation(args) assert code == 1 @@ -328,55 +322,13 @@ async def test_no_session_returns_zero_without_backend_call(self, tmp_project): args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) with patch( - "services.retry.check_status_safe", - new_callable=AsyncMock, - ) as mock_status: - code = await cmd_retry_generation(args) - - assert code == 0 - assert mock_status.call_count == 0 # core never called - - @pytest.mark.asyncio - async def test_pending_not_failed_returns_zero_without_post(self, tmp_project): - from cli import cmd_retry_generation - from services.session import write_session - write_session("gen-pending", tmp_project) - - args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) - - with patch( - "services.retry.check_status_safe", - new_callable=AsyncMock, - ) as mock_status, patch( "services.retry.call_backend_endpoint", new_callable=AsyncMock, ) as mock_ep: - mock_status.return_value = {"status": "pending"} code = await cmd_retry_generation(args) assert code == 0 - assert mock_ep.call_count == 0 # pending short-circuits before POST - - @pytest.mark.asyncio - async def test_backend_error_returns_one(self, tmp_project): - from cli import cmd_retry_generation - from services.session import write_session - write_session("gen-err", tmp_project) - - args = SimpleNamespace(root_path=str(tmp_project), command="retry-generation", generation_id=None) - - with patch( - "services.retry.check_status_safe", - new_callable=AsyncMock, - ) as mock_status, patch( - "services.retry.call_backend_endpoint", - new_callable=AsyncMock, - ) as mock_ep: - mock_status.return_value = {"status": "failed"} - mock_ep.side_effect = Exception("Backend returned HTTP 500") - code = await cmd_retry_generation(args) - - assert code == 1 + assert mock_ep.call_count == 0 # helper never called # --------------------------------------------------------------------------- diff --git a/mcp_server/tests/test_retry_core.py b/mcp_server/tests/test_retry_core.py index 21b6953..ed386d3 100644 --- a/mcp_server/tests/test_retry_core.py +++ b/mcp_server/tests/test_retry_core.py @@ -1,81 +1,33 @@ -"""Tests for the shared retry core (services/retry.py). +"""Tests for the shared retry helper (services/retry.py). -One case per outcome variant. Patches are applied at ``services.retry.*`` because -the core binds ``check_status_safe`` / ``call_backend_endpoint`` at import time. +The helper POSTs the retry and returns the parsed backend dict, raising on +failure (the backend validates eligibility). Patch ``call_backend_endpoint`` at +``services.retry.*`` where the helper binds it. """ import json from unittest.mock import AsyncMock, patch import pytest -from services.retry import ( - AlreadyRunning, - BackendError, - PendingNotFailed, - PendingRejectedBeforeCodegen, - Queued, - retry_generation_core, -) +from services.retry import retry_generation_core GID = "est-abc123" -def _patches(status_return, post=None): - status = patch("services.retry.check_status_safe", new=AsyncMock(return_value=status_return)) - post_mock = AsyncMock(return_value=post) if not isinstance(post, Exception) else AsyncMock(side_effect=post) - endpoint = patch("services.retry.call_backend_endpoint", new=post_mock) - return status, endpoint, post_mock - - @pytest.mark.asyncio -async def test_already_running(): - status, endpoint, post = _patches({"status": "running"}) - with status as st, endpoint: - outcome = await retry_generation_core(GID) - assert isinstance(outcome, AlreadyRunning) - assert st.await_count == 1 - assert post.await_count == 0 # no POST when already running - - -@pytest.mark.asyncio -async def test_pending_rejected_before_codegen(): - status, endpoint, post = _patches({"status": "pending", "error": "Missing required files"}) - with status, endpoint: - outcome = await retry_generation_core(GID) - assert isinstance(outcome, PendingRejectedBeforeCodegen) - assert outcome.error == "Missing required files" - assert outcome.status_data["status"] == "pending" - assert post.await_count == 0 - - -@pytest.mark.asyncio -async def test_pending_not_failed(): - status, endpoint, post = _patches({"status": "pending"}) - with status, endpoint: - outcome = await retry_generation_core(GID) - assert isinstance(outcome, PendingNotFailed) - assert outcome.status_data["status"] == "pending" - assert post.await_count == 0 - - -@pytest.mark.asyncio -async def test_queued(): - status, endpoint, post = _patches( - {"status": "failed"}, post=json.dumps({"status": "pending", "retry_count": 1}) - ) - with status as st, endpoint: - outcome = await retry_generation_core(GID) - assert isinstance(outcome, Queued) - assert outcome.backend_data == {"status": "pending", "retry_count": 1} - assert st.await_count == 1 - assert post.await_count == 1 # exactly one POST +async def test_returns_backend_dict_on_success(): + post = AsyncMock(return_value=json.dumps({"status": "pending", "retry_count": 1})) + with patch("services.retry.call_backend_endpoint", new=post): + data = await retry_generation_core(GID) + assert data == {"status": "pending", "retry_count": 1} + assert post.await_count == 1 + assert GID in post.await_args.kwargs["endpoint"] + assert post.await_args.kwargs["method"] == "POST" @pytest.mark.asyncio -async def test_backend_error(): - status, endpoint, post = _patches({"status": "failed"}, post=Exception("Backend returned HTTP 500")) - with status, endpoint: - outcome = await retry_generation_core(GID) - assert isinstance(outcome, BackendError) - assert "500" in outcome.error - assert post.await_count == 1 +async def test_raises_on_backend_failure(): + post = AsyncMock(side_effect=Exception("Backend returned HTTP 400: Cannot retry in 'running' state")) + with patch("services.retry.call_backend_endpoint", new=post): + with pytest.raises(Exception, match="running"): + await retry_generation_core(GID) diff --git a/mcp_server/tests/test_server_retry.py b/mcp_server/tests/test_server_retry.py index 65828d0..aefb408 100644 --- a/mcp_server/tests/test_server_retry.py +++ b/mcp_server/tests/test_server_retry.py @@ -1,13 +1,11 @@ -"""Characterization + parity tests for the MCP ``retry_generation`` tool. +"""Tests for the MCP ``retry_generation`` tool after deferring to the backend. -Locks the ``chat_json`` shapes for every outcome so the core refactor preserves -them. `retry_generation` is a plain async function; `ctx` is only consumed by -`apply_project_root_from_context`, which we stub. Backend primitives are patched -at both `server.*` (resolution/guard) and `services.retry.*` (the core's POST + -status read) so these tests hold before and after the refactor. +The tool resolves the id, POSTs via the shared helper, and renders success vs +the backend's error message. `ctx` is only consumed by +`apply_project_root_from_context`, which we stub. The POST is patched at +`services.retry.*` where the helper binds it. """ import json -from contextlib import ExitStack from unittest.mock import AsyncMock, patch import pytest @@ -17,20 +15,17 @@ GID = "est-abc123" -async def _run(status_return, *, resolved=GID, post=None): - """Invoke retry_generation with all IO stubbed; return (payload dict, post mock).""" +async def _run(*, resolved=GID, post=None): + """Invoke retry_generation with IO stubbed; return (payload dict, post mock).""" post_mock = ( AsyncMock(side_effect=post) if isinstance(post, Exception) else AsyncMock(return_value=post) ) - status_mock = AsyncMock(return_value=status_return) - with ExitStack() as stack: - stack.enter_context(patch("server.apply_project_root_from_context", new=AsyncMock())) - stack.enter_context(patch("server.resolve_generation_id", return_value=resolved)) - stack.enter_context(patch("server.check_status_safe", new=status_mock)) - stack.enter_context(patch("services.retry.check_status_safe", new=status_mock)) - stack.enter_context(patch("services.retry.call_backend_endpoint", new=post_mock)) - stack.enter_context(patch("server.call_backend_endpoint", new=post_mock)) + with ( + patch("server.apply_project_root_from_context", new=AsyncMock()), + patch("server.resolve_generation_id", return_value=resolved), + patch("services.retry.call_backend_endpoint", new=post_mock), + ): payload = await server.retry_generation(generation_id=None, ctx=object()) return json.loads(payload), post_mock @@ -46,38 +41,9 @@ async def test_no_session(): assert "No previous generation found" in data["message"] -@pytest.mark.asyncio -async def test_already_running(): - data, post = await _run({"status": "running"}) - assert "already running" in data["message"].lower() - assert data["code"] == "GENERATION_ALREADY_RUNNING" - assert data["generation_id"] == GID - assert post.await_count == 0 - - -@pytest.mark.asyncio -async def test_pending_rejected_before_codegen(): - data, post = await _run({"status": "pending", "error": "Missing required files"}) - assert "rejected before codegen" in data["message"].lower() - assert data["details"]["error"] == "Missing required files" - assert data["generation_id"] == GID - assert post.await_count == 0 - - -@pytest.mark.asyncio -async def test_pending_not_failed(): - data, post = await _run({"status": "pending"}) - assert "pending but has not failed" in data["message"].lower() - assert data["details"]["status"] == "pending" - assert data["generation_id"] == GID - assert post.await_count == 0 - - @pytest.mark.asyncio async def test_queued(): - data, post = await _run( - {"status": "failed"}, post=json.dumps({"status": "pending", "retry_count": 1}) - ) + data, post = await _run(post=json.dumps({"status": "pending", "retry_count": 1})) assert "retry queued" in data["message"].lower() assert data["status"] == "pending" assert data["retry_count"] == 1 @@ -87,8 +53,12 @@ async def test_queued(): @pytest.mark.asyncio -async def test_backend_error(): - data, post = await _run({"status": "failed"}, post=Exception("Backend returned HTTP 500")) +async def test_backend_rejection_or_error_is_surfaced(): + # Covers wrong-state rejections (HTTP 400) and unreachable backend alike — + # both raise from the helper and render the backend's reason. + data, post = await _run( + post=Exception("Backend returned HTTP 400: Cannot retry in 'running' state") + ) assert "couldn't retry" in data["message"].lower() - assert "500" in data["details"]["error"] - assert post.await_count == 1 + assert "running" in data["details"]["error"] + assert data["generation_id"] == GID diff --git a/mcp_server/tests/test_tui_app.py b/mcp_server/tests/test_tui_app.py index 4be308f..1859c3d 100644 --- a/mcp_server/tests/test_tui_app.py +++ b/mcp_server/tests/test_tui_app.py @@ -1209,7 +1209,7 @@ async def test_retry_calls_core_and_refreshes_without_suspending(self): patch.object(app, "push_screen_wait", new=AsyncMock(return_value=True)), patch( "tui.app.retry_generation_core", - new=AsyncMock(return_value=tui_app.Queued(backend_data={"status": "pending"})), + new=AsyncMock(return_value={"status": "pending", "retry_count": 1}), ) as core, patch.object(screen, "refresh_status", new=AsyncMock()) as refresh, patch.object(screen, "notify") as notify, @@ -1234,9 +1234,9 @@ async def test_retry_backend_error_shows_message_screen(self): patch.object(app, "push_screen_wait", new=psw), patch( "tui.app.retry_generation_core", - new=AsyncMock(return_value=tui_app.BackendError(error="unreachable")), + new=AsyncMock(side_effect=Exception("Backend returned HTTP 400: Cannot retry")), ), - patch.object(screen, "refresh_status", new=AsyncMock()), + patch.object(screen, "refresh_status", new=AsyncMock()) as refresh, patch.object(app, "suspend") as suspend, ): await screen._retry_flow() @@ -1244,6 +1244,7 @@ async def test_retry_backend_error_shows_message_screen(self): # push_screen_wait called twice: ConfirmScreen, then the error MessageScreen. assert psw.await_count == 2 assert isinstance(psw.await_args_list[1].args[0], tui_app.MessageScreen) + refresh.assert_not_awaited() # error returns before the refresh @pytest.mark.asyncio async def test_retry_skips_core_when_cancelled(self): diff --git a/mcp_server/tui/app.py b/mcp_server/tui/app.py index dad33b5..bf808f6 100644 --- a/mcp_server/tui/app.py +++ b/mcp_server/tui/app.py @@ -63,14 +63,7 @@ from cli import resolve_backend_config from services import local_env, validate_models from services.llm_tiers import LLM_TIER_KEYS -from services.retry import ( - AlreadyRunning, - BackendError, - PendingNotFailed, - PendingRejectedBeforeCodegen, - Queued, - retry_generation_core, -) +from services.retry import retry_generation_core from services.session import resolve_generation_id, set_project_root from services.specflow_backend import call_backend_endpoint, call_backend_endpoint_bytes from tui import actions, activity, mcp_clients, onboarding, render @@ -646,27 +639,12 @@ async def _retry_flow(self) -> None: if not ok: return self.notify("Retrying generation…", severity="information") - match await retry_generation_core(self._generation_id): - case Queued(): - self.notify( - "Retry queued. Resuming from the last checkpoint.", - severity="information", - ) - case AlreadyRunning(): - self.notify("A generation is already running.", severity="warning") - case PendingRejectedBeforeCodegen(error=error): - await self.app.push_screen_wait(MessageScreen( - "Retry", - f"Last run was rejected before codegen: {error}\n" - "Fix files locally and start a new generation — retry won't help.", - )) - case PendingNotFailed(): - await self.app.push_screen_wait(MessageScreen( - "Retry", - "Session is pending but has not failed — nothing to retry yet.", - )) - case BackendError(error=error): - await self.app.push_screen_wait(MessageScreen("Retry failed", error)) + try: + await retry_generation_core(self._generation_id) + except Exception as exc: # noqa: BLE001 - surface the backend's reason to the user + await self.app.push_screen_wait(MessageScreen("Retry failed", str(exc))) + return + self.notify("Retry queued. Resuming from the last checkpoint.", severity="information") await self.refresh_status() def action_cancel(self) -> None: