From 0cf407ff973eb21dc7f21001ea4354fc347afbbc Mon Sep 17 00:00:00 2001 From: Artsiom Kozak Date: Tue, 14 Jul 2026 08:52:52 +0200 Subject: [PATCH 1/5] feat(generation): support user cancellation of a running session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cancelling a generation previously just called fail_generation_session, which sent an "❌ Run Failed" notification, marked the session FAILED (indistinguishable from a real failure), and never stopped the running workflow task or its agent subprocesses. Introduce a distinct terminal CANCELLED status and make cancellation actually stop in-flight work through two mechanisms: - A process-local generation_id → asyncio.Task registry so the owning pod calls task.cancel() for immediate teardown. The task registers itself via asyncio.current_task(), so no plumbing is needed at the create_task sites (initial run, retry, and boot recovery). - Cooperative raise_if_cancelled() checks (throttled DB reads) in the orchestrator step loop, execute_all_phases, the parallel executors, and the P10Y poll loop, raising GenerationCancelledError as a cross-pod-safe fallback. The cancel endpoint sets CANCELLED first, then cancels the task, so the injected CancelledError and any GenerationCancelledError unwind against an already-terminal state and neither re-fails nor notifies. The state machine cancel() mirrors interrupted_by_shutdown(): notification-free, sets cancelled_by_user/cancelled_at markers, ends the API-key slot, and retains workspaces (Commandment II) — generated code is preserved and reclaimed by scheduled_wipe. agent_query now aclose()s the SDK query stream so the Claude Code CLI subprocess is torn down on cancel. A CANCELLED session is terminal: boot recovery and the stuck detectors already exclude it, and retry now refuses it with a clear message. --- backend/app/api/v1/generation_sessions.py | 68 ++++++++++++------ .../app/schemas/generation_workflow_enums.py | 5 ++ backend/app/services/claude_code.py | 15 ++++ backend/app/services/generation_session.py | 32 ++++++++- .../app/services/generation_session_retry.py | 6 ++ .../app/services/generation_task_registry.py | 56 +++++++++++++++ .../services/generation_workflow_runner.py | 15 ++++ backend/app/services/p10y/p10y_lib.py | 14 +++- backend/app/services/parallel_executor.py | 23 ++++++- backend/app/state/__init__.py | 11 ++- .../app/state/api_key_session_concurrency.py | 1 + backend/app/state/cancellation.py | 51 ++++++++++++++ backend/app/state/exceptions.py | 15 ++++ .../state/generation_session_state_machine.py | 43 +++++++++++- backend/app/state/transitions.py | 11 +++ backend/app/state/workflow_orchestrator.py | 16 ++++- .../multi_workspace_estimation_p10y.py | 13 ++++ backend/test/api/test_generation_sessions.py | 59 +++++++++++----- .../test_shutdown_interrupted_recovery.py | 16 +++++ .../test/services/test_generation_session.py | 37 ++++++++++ .../services/test_generation_session_retry.py | 15 ++++ .../services/test_generation_task_registry.py | 69 +++++++++++++++++++ .../test_generation_workflow_runner.py | 26 +++++++ backend/test/state/test_cancellation.py | 59 ++++++++++++++++ .../test_generation_session_state_machine.py | 50 ++++++++++++++ .../test/state/test_workflow_orchestrator.py | 31 +++++++++ backend/test/test_parallel_execution.py | 21 ++++++ 27 files changed, 731 insertions(+), 47 deletions(-) create mode 100644 backend/app/services/generation_task_registry.py create mode 100644 backend/app/state/cancellation.py create mode 100644 backend/test/services/test_generation_task_registry.py create mode 100644 backend/test/state/test_cancellation.py diff --git a/backend/app/api/v1/generation_sessions.py b/backend/app/api/v1/generation_sessions.py index 9b210da..27e2c87 100644 --- a/backend/app/api/v1/generation_sessions.py +++ b/backend/app/api/v1/generation_sessions.py @@ -77,6 +77,7 @@ from app.services.artifact_store import ARTIFACTS_BASE, ArtifactStore from app.services.agent_stream_broker import get_agent_stream_broker from app.services.mcp_prune import resolve_enabled_mcps_and_set_telemetry +from app.services import generation_task_registry from app.schemas.model_token_usage import ModelTokenUsage from app.schemas.workflow_usage_metrics import ( aggregate_model_usage_by_workspace, @@ -853,11 +854,17 @@ async def cancel_generation_session( generation_session_service: GenerationSessionService = Depends(get_generation_session_service) ): """ - Cancel a running generation session. + Cancel a running generation session (user-initiated). + + Transitions the session to the terminal CANCELLED state and stops the in-flight + workflow: on the pod that owns the run, the workflow task is cancelled immediately; + on other pods the workflow stops cooperatively at the next checkpoint. No failure + notification is sent (the user requested the stop). Generated code is preserved — + workspaces stay ALLOCATED (Commandment II) and are reclaimed by scheduled_wipe. A + CANCELLED session is never resumed on restart and cannot be retried. - Marks the session as failed with reason "cancelled" and releases workspaces. Can only cancel sessions in 'pending', 'initializing', or 'running' state. - + Example: ``` DELETE /api/v1/generation-sessions/est-abc123 @@ -866,37 +873,54 @@ async def cancel_generation_session( try: # Get current session document session_doc = await generation_session_service.get_generation_session_status(generation_id) - + current_status = session_doc["status"] - - # Can only cancel if not already completed/failed - if current_status in (GenerationStatus.COMPLETED.value, GenerationStatus.FAILED.value): + + # Can only cancel if not already terminal. + if current_status in ( + GenerationStatus.COMPLETED.value, + GenerationStatus.FAILED.value, + GenerationStatus.CANCELLED.value, + ): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Cannot cancel generation session in '{current_status}' state" ) - - # Mark as failed with cancellation reason - await generation_session_service.fail_generation_session( - generation_id=generation_id, - error="Cancelled by user" - ) - + + # Ordering matters: set CANCELLED FIRST (state machine + slot release, no notify), + # THEN tear down the running task. Because the state is already terminal when the + # injected CancelledError / cooperative GenerationCancelledError unwind, neither + # re-fails the session nor emits an error notification. + await generation_session_service.cancel_generation_session(generation_id) + + cancelled_locally = generation_task_registry.cancel_task(generation_id) + return CancelGenerationSessionResponse( generation_id=generation_id, - status="failed", - message="Generation session cancelled successfully" + status=GenerationStatus.CANCELLED.value, + message=( + "Generation session cancelled" + if cancelled_locally + else "Generation session cancelled; in-flight work will stop shortly" + ), ) - + except GenerationSessionNotFoundError as e: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=str(e) ) - + + except InvalidGenerationSessionStateError as e: + # Race: session became terminal between the status read and the transition. + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(e) + ) + except HTTPException: raise - + except Exception as e: logger.error(f"Failed to cancel generation session {generation_id}: {e}", exc_info=True) raise HTTPException( @@ -1066,6 +1090,7 @@ async def _run_generation_session_workflow( normalized_src_dir = generation_session_params.get("src_dir", "src") try: + generation_task_registry.register_current_task(generation_id) async with session.task_slot(generation_id=generation_id): logger.info("Starting app generation + P10Y workflow for %s", generation_id) agent_logger.info("Starting app generation + P10Y workflow for %s", generation_id) @@ -1099,9 +1124,10 @@ async def _run_generation_session_workflow( except Exception as e: # Single exit path — routes ContractRejection → reject (PENDING) and everything - # else → fail (FAILED). + # else → fail (FAILED). GenerationCancelledError is handled as a silent no-op there. await _handle_workflow_exception(e, generation_id, generation_session_service) finally: + generation_task_registry.deregister_task(generation_id) TelemetryContext.set_agent_query_totals_handler(None) @@ -1562,6 +1588,8 @@ async def download_generation_session_outputs( if not outputs_archived and not emergency_archived: est_status = session_doc.get("status") workspace_ids = session_doc.get("workspace_ids", []) + # CANCELLED is intentionally excluded: a user-cancelled run is not emergency-archived + # for download. Its workspaces stay ALLOCATED and are reclaimed by scheduled_wipe. recoverable_statuses = {GenerationStatus.FAILED.value, GenerationStatus.RUNNING.value} if est_status in recoverable_statuses and workspace_ids: diff --git a/backend/app/schemas/generation_workflow_enums.py b/backend/app/schemas/generation_workflow_enums.py index ec7fdc8..79a8d2b 100644 --- a/backend/app/schemas/generation_workflow_enums.py +++ b/backend/app/schemas/generation_workflow_enums.py @@ -17,14 +17,19 @@ class GenerationStatus(str, Enum): Lifecycle: PENDING → INITIALIZING → RUNNING → COMPLETED | FAILED + (any of PENDING | INITIALIZING | RUNNING) → CANCELLED (user-initiated) Spec analysis and planning run locally in the IDE before ``run_generation``. + + CANCELLED is terminal and, unlike FAILED, is never resumed by boot recovery + and never retryable — the user deliberately stopped the run. """ PENDING = "pending" INITIALIZING = "initializing" RUNNING = "running" COMPLETED = "completed" FAILED = "failed" + CANCELLED = "cancelled" class GenerationCheckpoint(str, Enum): diff --git a/backend/app/services/claude_code.py b/backend/app/services/claude_code.py index 845eccb..884c39c 100644 --- a/backend/app/services/claude_code.py +++ b/backend/app/services/claude_code.py @@ -1,5 +1,6 @@ import asyncio from collections.abc import AsyncIterator +import contextlib from datetime import datetime, timezone import logging import os @@ -65,6 +66,7 @@ ) from app.agents_sandboxing.claude_env_vars import build_redacted_env_overlay from app.database.factory import get_database +from app.state.cancellation import raise_if_cancelled from app.state.db_adapter import StateMachineDBAdapter from app.state.transitions import TriggeredBy from app.state.workspace_models import set_workspace_model_override @@ -895,6 +897,14 @@ async def agent_query( lf_tracer.finalize_pending() lf_tracer.end_generation_on_error("exception") raise + finally: + # Guarantee the SDK query generator — and its Claude Code CLI subprocess — is + # torn down on every exit path, including task.cancel() (CancelledError) during a + # user cancellation, where the stream would otherwise linger until GC. aclose() + # on an already-exhausted stream is a harmless no-op. + if hasattr(query_stream, "aclose"): + with contextlib.suppress(Exception): + await query_stream.aclose() if messages: logger.info(f"Agent returned {len(messages)} messages.") @@ -1189,6 +1199,11 @@ async def execute_all_phases( consecutive_errors = 0 for phase_num in range(phase_start, phase_end + 1): + # Cooperative cancellation: stop between phases if the user cancelled + # (cross-pod-safe fallback to the local task.cancel()). + if db_adapter is not None and request.generation_id: + await raise_if_cancelled(db_adapter, request.generation_id) + phase_info = planning_data.phases[phase_num - 1] if phase_info.applicable_agent_mcps is None: diff --git a/backend/app/services/generation_session.py b/backend/app/services/generation_session.py index 16984fb..613470e 100644 --- a/backend/app/services/generation_session.py +++ b/backend/app/services/generation_session.py @@ -721,7 +721,37 @@ async def fail_generation_session( except Exception as e: # Don't fail the session failure handling if notification fails logger.error(f"Failed to send Slack notification for generation session {generation_id}: {e}", exc_info=True) - + + async def cancel_generation_session( + self, + generation_id: str, + *, + triggered_by: str = TriggeredBy.CANCEL, + ) -> None: + """ + Mark the generation session CANCELLED (user-initiated). Notification-free. + + Deliberately does NOT route through ``fail_generation_session`` — a user + cancellation must not emit the "❌ Run Failed" Slack/email notification, and + must land in the distinct terminal CANCELLED state (never resumed, never + retried). Workspaces stay ALLOCATED (Commandment II); the state machine is the + single writer of status. + + Args: + generation_id: The session to cancel + triggered_by: TriggeredBy constant (defaults to CANCEL) + + Raises: + GenerationSessionNotFoundError: Session doesn't exist + InvalidGenerationSessionStateError: Session is already terminal + """ + session_doc = self._db.get(COL_GENERATION_SESSIONS, generation_id) + if not session_doc: + raise GenerationSessionNotFoundError(f"Generation session {generation_id} not found") + + await self._esm.cancel(generation_id, triggered_by=triggered_by) + logger.info("Generation session %s cancelled by user", generation_id) + async def get_generation_session_status(self, generation_id: str) -> Dict[str, Any]: """ Get current generation session status. diff --git a/backend/app/services/generation_session_retry.py b/backend/app/services/generation_session_retry.py index b783f8e..b12b1f6 100644 --- a/backend/app/services/generation_session_retry.py +++ b/backend/app/services/generation_session_retry.py @@ -90,6 +90,12 @@ async def retry_generation_session( raise GenerationSessionRetryError(f"Generation session {generation_id} not found") current_status = doc["status"] + if current_status == GenerationStatus.CANCELLED.value: + raise InvalidRetryStateError( + "Cannot retry a cancelled generation session. Cancellation is a " + "deliberate, terminal user action — create a new session from the same " + "spec to start again." + ) if current_status not in ( GenerationStatus.FAILED.value, GenerationStatus.PENDING.value, diff --git a/backend/app/services/generation_task_registry.py b/backend/app/services/generation_task_registry.py new file mode 100644 index 0000000..1ab9dfa --- /dev/null +++ b/backend/app/services/generation_task_registry.py @@ -0,0 +1,56 @@ +""" +Process-local registry mapping generation_id → the asyncio.Task running its workflow. + +Used by the cancel endpoint to immediately tear down an in-flight run on the pod that +owns it (``task.cancel()`` injects ``CancelledError`` at the current ``await``). This is +the fast path; the cooperative :func:`app.state.cancellation.raise_if_cancelled` checks +are the cross-pod-safe fallback for cancels that land on a different pod. + +The workflow task registers itself via ``asyncio.current_task()`` from inside its own +body, so there is no plumbing at the ``asyncio.create_task(...)`` call sites — the initial +run and the shared rerun path (manual retry + boot recovery) are all covered. +""" +import asyncio +import logging + +logger = logging.getLogger(__name__) + +_ACTIVE_TASKS: dict[str, asyncio.Task] = {} + + +def register_current_task(generation_id: str) -> None: + """Register the currently running task as the workflow task for ``generation_id``.""" + task = asyncio.current_task() + if task is not None: + _ACTIVE_TASKS[generation_id] = task + + +def deregister_task(generation_id: str) -> None: + """Remove the registry entry for ``generation_id`` — only if it is *this* task. + + Guarding on identity avoids a re-fired run (retry/recovery) that has already + registered its own task being clobbered by the previous task's ``finally`` cleanup. + """ + current = asyncio.current_task() + if _ACTIVE_TASKS.get(generation_id) is current: + _ACTIVE_TASKS.pop(generation_id, None) + + +def get_task(generation_id: str) -> asyncio.Task | None: + """Return the registered workflow task for ``generation_id``, or None.""" + return _ACTIVE_TASKS.get(generation_id) + + +def cancel_task(generation_id: str) -> bool: + """Cancel the registered workflow task if present and not already done. + + Returns True if a live task was found and ``cancel()`` was requested (immediate + local teardown), False otherwise (run owned by another pod, already finished, or + never registered — the cooperative status check handles those). + """ + task = _ACTIVE_TASKS.get(generation_id) + if task is not None and not task.done(): + task.cancel() + logger.info("Requested cancellation of local workflow task for %s", generation_id) + return True + return False diff --git a/backend/app/services/generation_workflow_runner.py b/backend/app/services/generation_workflow_runner.py index 0f22aba..a6e5c85 100644 --- a/backend/app/services/generation_workflow_runner.py +++ b/backend/app/services/generation_workflow_runner.py @@ -21,8 +21,10 @@ from app.services.claude_code import clear_workspace_caches from app.services.contract_validator import ContractRejection from app.services.generation_session import GenerationSessionService +from app.services import generation_task_registry from app.state.api_key_session_concurrency import OperationKind, SessionBeginOutcome from app.state.exceptions import ( + GenerationCancelledError, InvalidGenerationSessionStateError as SMInvalidGenerationSessionStateError, ) from app.state.transitions import TriggeredBy @@ -79,8 +81,18 @@ async def _handle_workflow_exception( released cleanly) so the user can fix their uploaded files and retry. Any other exception is a real failure → ``fail_generation_session``. + A ``GenerationCancelledError`` (raised by cooperative cancellation checks) is NOT a + failure: the session has already been transitioned to CANCELLED by the cancel endpoint, + so this is a silent no-op — no fail(), no reject(), no notification (the user knows). + Centralizing here is the structural fix — entry points can no longer diverge when a new exception type is introduced. """ + if isinstance(exc, GenerationCancelledError): + logger.info( + "Generation session %s cancelled by user — workflow stopped (no-op)", generation_id + ) + return + if isinstance(exc, ContractRejection): logger.info( "Generation session %s rejected by contract validator: %s", generation_id, exc @@ -239,6 +251,7 @@ async def rerun_generation_session( } try: + generation_task_registry.register_current_task(generation_id) TelemetryContext.set_user_context( user_email=user_email, generation_id=generation_id, @@ -294,6 +307,8 @@ async def rerun_generation_session( except Exception as e: # Single exit path — ContractRejection → reject (PENDING), else → fail (FAILED). + # GenerationCancelledError is handled as a silent no-op there. await _handle_workflow_exception(e, generation_id, generation_session_service) finally: + generation_task_registry.deregister_task(generation_id) TelemetryContext.set_agent_query_totals_handler(None) diff --git a/backend/app/services/p10y/p10y_lib.py b/backend/app/services/p10y/p10y_lib.py index 1fbbca7..4cf8637 100644 --- a/backend/app/services/p10y/p10y_lib.py +++ b/backend/app/services/p10y/p10y_lib.py @@ -8,6 +8,7 @@ from typing import Any, Dict, List, Optional from app.services.p10y.p10y_api_client import P10YInternalAPIClient +from app.state.cancellation import raise_if_cancelled # Commits whose first line starts with this prefix are excluded from P10Y / component breakdown # (e.g. user-provided initial seed: SKIP_initial_user_source, SKIP_generation_baseline). @@ -370,16 +371,20 @@ async def trigger_and_poll_p10y_metrics( organisation_id: int, workspace_name: str, logger: logging.Logger, + db_adapter=None, + generation_id: Optional[str] = None, ) -> None: """ Trigger P10Y metrics calculation and poll until processing is complete. - + Args: client: P10Y API client repository_id: P10Y repository ID organisation_id: P10Y organisation ID workspace_name: Workspace name for logging logger: Logger instance + db_adapter: Optional state-machine DB adapter for cooperative cancellation checks + generation_id: Optional generation id for cooperative cancellation checks """ # Trigger P10Y metrics calculation logger.info(f"Triggering P10Y metrics for repository {repository_id}") @@ -396,8 +401,13 @@ async def trigger_and_poll_p10y_metrics( poll_limit = 10 while not all_commits_processed and poll_counter < poll_limit: + # Cooperative cancellation: stop polling if the user cancelled + # (cross-pod-safe fallback to the local task.cancel()). + if db_adapter is not None and generation_id: + await raise_if_cancelled(db_adapter, generation_id) + logger.info(f"Polling attempt {poll_counter + 1} / {poll_limit}") - + commit_stats_polled = await client.get_commit_stats( organisation_id=organisation_id, repository_ids=[repository_id], diff --git a/backend/app/services/parallel_executor.py b/backend/app/services/parallel_executor.py index ba9328b..12648b8 100644 --- a/backend/app/services/parallel_executor.py +++ b/backend/app/services/parallel_executor.py @@ -12,6 +12,7 @@ from app.services.claude_code import agent_query from app.services.model_routing import classify_error from app.services.workspace_manager import WorkspaceManager +from app.state.exceptions import GenerationCancelledError @dataclass @@ -86,7 +87,13 @@ async def execute_parallel( tasks.append(task) results = await asyncio.gather(*tasks, return_exceptions=True) - + + # return_exceptions=True absorbs a cooperative cancellation into the results + # list — surface it so the entire generation aborts (session already CANCELLED). + for r in results: + if isinstance(r, GenerationCancelledError): + raise r + # Count successes and failures successes = sum(1 for r in results if isinstance(r, ParallelAgentResult) and r.success) failures = len(results) - successes @@ -146,7 +153,12 @@ async def _execute_single( success=True, duration_ms=duration_ms ) - + + except GenerationCancelledError: + # User cancellation must abort the whole fan-out, not degrade to a + # per-workspace failure result. Re-raise so execute_parallel propagates it. + raise + except Exception as e: duration_ms = (time.time() - start_time) * 1000 error_msg = str(e) @@ -363,7 +375,12 @@ async def process_workspace(workspace: WorkspaceSettings) -> ParallelGenerationR warning="Commit files not found - workspace may not have completed code generation", duration_ms=duration_ms ) - + + except GenerationCancelledError: + # User cancellation aborts the whole estimation fan-out (gather uses + # return_exceptions=False, so this propagates out immediately). + raise + except Exception as e: # Handle other errors duration_ms = (time.time() - start_time) * 1000 diff --git a/backend/app/state/__init__.py b/backend/app/state/__init__.py index fd287a7..9e804f7 100644 --- a/backend/app/state/__init__.py +++ b/backend/app/state/__init__.py @@ -8,11 +8,18 @@ WorkflowOrchestrator InvalidGenerationSessionStateError InvalidWorkspaceStateError + GenerationCancelledError + raise_if_cancelled """ from .generation_session_state_machine import GenerationSessionStateMachine from .workspace_state_machine import WorkspaceStateMachine from .workflow_orchestrator import WorkflowOrchestrator -from .exceptions import InvalidGenerationSessionStateError, InvalidWorkspaceStateError +from .exceptions import ( + InvalidGenerationSessionStateError, + InvalidWorkspaceStateError, + GenerationCancelledError, +) +from .cancellation import raise_if_cancelled __all__ = [ "GenerationSessionStateMachine", @@ -20,4 +27,6 @@ "WorkflowOrchestrator", "InvalidGenerationSessionStateError", "InvalidWorkspaceStateError", + "GenerationCancelledError", + "raise_if_cancelled", ] diff --git a/backend/app/state/api_key_session_concurrency.py b/backend/app/state/api_key_session_concurrency.py index e911ef1..fe0c348 100644 --- a/backend/app/state/api_key_session_concurrency.py +++ b/backend/app/state/api_key_session_concurrency.py @@ -51,6 +51,7 @@ class SessionEndReason(StrEnum): STUCK_DETECTED = "stuck_detected" FORCE_RELEASED = "force_released" BEGIN_ROLLBACK = "begin_rollback" + CANCELLED = "cancelled" _TTL_MINUTES: dict[OperationKind, int] = { diff --git a/backend/app/state/cancellation.py b/backend/app/state/cancellation.py new file mode 100644 index 0000000..23f7097 --- /dev/null +++ b/backend/app/state/cancellation.py @@ -0,0 +1,51 @@ +""" +Cooperative cancellation for in-flight generation workflows. + +Single source of truth for "has this generation been cancelled?". The user-facing +cancel path (``DELETE`` endpoint → ``GenerationSessionService.cancel_generation_session``) +transitions the session to CANCELLED and, on the owning pod, calls ``task.cancel()`` for +immediate teardown. This module is the cross-pod-safe fallback: long-running workflow / +phase / poll loops call :func:`raise_if_cancelled` at coarse boundaries so a cancel that +lands on a different pod (or a task that is not locally registered) still stops the run. + +Status is read through the async DB adapter and compared via ``parse_status`` — no second +status validator is introduced (Commandment: single source of truth). Reads are throttled +per generation so the helper stays cheap when called inside tight loops. +""" +import time + +from app.schemas.generation_workflow_enums import GenerationStatus, parse_status +from app.state.exceptions import GenerationCancelledError + +# generation_id -> last monotonic timestamp we hit the DB for this session. +_last_check: dict[str, float] = {} + +# Minimum seconds between DB reads for the same generation inside a hot loop. +_DEFAULT_MIN_INTERVAL_S = 2.0 + + +async def raise_if_cancelled( + db_adapter, + generation_id: str, + *, + min_interval_s: float = _DEFAULT_MIN_INTERVAL_S, +) -> None: + """Raise :class:`GenerationCancelledError` if the session is CANCELLED. + + Throttled: if called again for the same ``generation_id`` within ``min_interval_s`` + seconds, returns immediately without touching the DB. ``db_adapter`` is the async + state-machine DB adapter (``service.db_adapter`` / ``orchestrator._db``). + """ + now = time.monotonic() + if now - _last_check.get(generation_id, 0.0) < min_interval_s: + return + _last_check[generation_id] = now + + doc = await db_adapter.get_generation_session(generation_id) + if doc and parse_status(doc.get("status")) == GenerationStatus.CANCELLED: + raise GenerationCancelledError(generation_id) + + +def reset_cancellation_check(generation_id: str) -> None: + """Drop the throttle timestamp for a generation (test/cleanup helper).""" + _last_check.pop(generation_id, None) diff --git a/backend/app/state/exceptions.py b/backend/app/state/exceptions.py index f854ccb..f71595b 100644 --- a/backend/app/state/exceptions.py +++ b/backend/app/state/exceptions.py @@ -61,6 +61,21 @@ class CleanupTargetError(Exception): +class GenerationCancelledError(Exception): + """ + Raised by cooperative cancellation checks (``raise_if_cancelled``) when the + session has been transitioned to CANCELLED while its workflow is still running. + + Deliberately an ``Exception`` (not ``BaseException``) so it propagates through the + existing ``except Exception`` layers up to the shared workflow exception handler, + where it is intercepted and routed to a silent no-op (no fail(), no notification — + the session is already terminal). + """ + def __init__(self, generation_id: str): + self.generation_id = generation_id + super().__init__(f"Generation {generation_id} was cancelled by user") + + class MaxRetriesExceededError(Exception): """ Raised by reset_for_retry when retry_count >= max_retries. diff --git a/backend/app/state/generation_session_state_machine.py b/backend/app/state/generation_session_state_machine.py index 52eb23b..45d127e 100644 --- a/backend/app/state/generation_session_state_machine.py +++ b/backend/app/state/generation_session_state_machine.py @@ -223,6 +223,38 @@ async def interrupted_by_shutdown( ) return update + async def cancel( + self, generation_id: str, triggered_by: str, metadata: dict | None = None + ) -> dict: + """ + PENDING | INITIALIZING | RUNNING → CANCELLED, on explicit user request. + + Mirrors interrupted_by_shutdown()/stuck_detected() but is a deliberate, + terminal, user-initiated stop: notification-free (the user knows — no error + Slack/email), sets the queryable ``cancelled_by_user`` marker + ``cancelled_at`` + so boot-recovery / retry / analytics can distinguish it from a genuine failure. + + Side-effects: sets cancelled_at, cancelled_by_user=True; ends the API-key session + slot. Does NOT set failed_at. Workspaces remain ALLOCATED (Invariant 2) — the + generated code is preserved and reclaimed by scheduled_wipe; a CANCELLED session + is never resumed and never retried. + """ + extra_fields = { + "cancelled_at": datetime.now(timezone.utc), + "cancelled_by_user": True, + } + update = await self._transition( + generation_id, "cancel", + triggered_by=triggered_by, + metadata=metadata or {}, + extra_fields=extra_fields, + ) + await self._api_sessions.end( + generation_id=generation_id, + reason=SessionEndReason.CANCELLED, + ) + return update + async def _record_phase_progress_for_field( self, generation_id: str, @@ -669,8 +701,15 @@ def _require_running(self, doc: dict, action: str, generation_id: str) -> None: allowed_from=[GenerationStatus.RUNNING], ) - _TERMINAL_STATUS_VALUES = {GenerationStatus.COMPLETED.value, GenerationStatus.FAILED.value} - _NON_TERMINAL_STATUSES = [s for s in GenerationStatus if s not in {GenerationStatus.COMPLETED, GenerationStatus.FAILED}] + _TERMINAL_STATUS_VALUES = { + GenerationStatus.COMPLETED.value, + GenerationStatus.FAILED.value, + GenerationStatus.CANCELLED.value, + } + _NON_TERMINAL_STATUSES = [ + s for s in GenerationStatus + if s not in {GenerationStatus.COMPLETED, GenerationStatus.FAILED, GenerationStatus.CANCELLED} + ] def _require_non_terminal(self, doc: Optional[dict], action: str, generation_id: str) -> None: """Guard: raises if generation is missing or in a terminal state (COMPLETED/FAILED).""" diff --git a/backend/app/state/transitions.py b/backend/app/state/transitions.py index abd127b..a27895f 100644 --- a/backend/app/state/transitions.py +++ b/backend/app/state/transitions.py @@ -36,6 +36,7 @@ class TriggeredBy: MANUAL_RETRY = "api:manual_retry" RUN_GENERATION = "api:run_generation" RESEND_EMAIL = "api:resend_email" + CANCEL = "api:cancel_generation_session" FORCE_RELEASE = "admin:force_release" VALIDATE_CONTRACT = "api:validate_contract" @@ -115,6 +116,16 @@ def orchestrator_step(step_name: str) -> str: to_state=GenerationStatus.FAILED, required_fields=("triggered_by",), ), + "cancel": GenerationSessionTransition( + name="cancel", + from_states=frozenset([ + GenerationStatus.PENDING, + GenerationStatus.INITIALIZING, + GenerationStatus.RUNNING, + ]), + to_state=GenerationStatus.CANCELLED, + required_fields=("triggered_by",), + ), "reset_for_retry": GenerationSessionTransition( name="reset_for_retry", from_states=frozenset([GenerationStatus.FAILED]), diff --git a/backend/app/state/workflow_orchestrator.py b/backend/app/state/workflow_orchestrator.py index 8ef33ba..0783c22 100644 --- a/backend/app/state/workflow_orchestrator.py +++ b/backend/app/state/workflow_orchestrator.py @@ -24,7 +24,8 @@ parse_checkpoint, ) from app.services.contract_validator import ContractRejection -from app.state.exceptions import InvalidGenerationSessionStateError +from app.state.cancellation import raise_if_cancelled +from app.state.exceptions import GenerationCancelledError, InvalidGenerationSessionStateError logger = logging.getLogger(__name__) @@ -204,6 +205,10 @@ async def run( ) continue + # Cooperative cancellation: stop before starting the next step if the user + # cancelled (cross-pod-safe fallback to the local task.cancel()). + await raise_if_cancelled(self._db, generation_id) + logger.info("generation %s running step '%s'", generation_id, step.name) try: await impl() @@ -215,6 +220,15 @@ async def run( generation_id, step.name, ) raise + except GenerationCancelledError: + # User cancellation — session is already CANCELLED. Propagate without + # fail() (mirrors the ContractRejection passthrough); the shared workflow + # exception handler treats it as a silent no-op. + logger.info( + "generation %s step '%s' observed cancellation — propagating without fail()", + generation_id, step.name, + ) + raise except Exception as exc: logger.error( "generation %s step '%s' failed: %s", diff --git a/backend/app/workflows/multi_workspace_estimation_p10y.py b/backend/app/workflows/multi_workspace_estimation_p10y.py index e0d19fa..904b599 100644 --- a/backend/app/workflows/multi_workspace_estimation_p10y.py +++ b/backend/app/workflows/multi_workspace_estimation_p10y.py @@ -72,6 +72,7 @@ is_skip_mode_enabled, ) from app.services.workflow_stats import workflow_metrics +from app.state.cancellation import raise_if_cancelled from app.state.workspace_models import load_workspace_model_overrides, resolve_workspace_model from app.services.workspace_manager import WorkspaceManager from app.state.db_adapter import COL_GENERATION_SESSIONS, StateMachineDBAdapter @@ -83,6 +84,8 @@ async def _process_single_workspace( client: P10YInternalAPIClient, settings: Settings, logger: logging.Logger, + db_adapter=None, + generation_id: Optional[str] = None, ) -> WorkspaceEstimation | None: """ Process estimation for a single workspace. @@ -141,6 +144,8 @@ async def _process_single_workspace( organisation_id=settings.P10Y_ORGANISATION_ID, workspace_name=workspace.name, logger=logger, + db_adapter=db_adapter, + generation_id=generation_id, ) # Fetch and filter commit stats @@ -511,6 +516,12 @@ async def multi_workspace_estimation_p10y_workflow( # Create workspace manager (needed for parallel executor) workspace_manager = WorkspaceManager(settings, logger) + # Cooperative cancellation: skip the whole estimation phase if the user cancelled + # (cross-pod-safe fallback to the local task.cancel()). Code is already archived by + # this point (STEEL XI), so aborting here is safe. + if db_adapter is not None and request.generation_id: + await raise_if_cancelled(db_adapter, request.generation_id) + # Process workspaces in parallel (data collection only, no AI agents) logger.info("Executing estimations in parallel...") parallel_results = await execute_generation_parallel( @@ -522,6 +533,8 @@ async def multi_workspace_estimation_p10y_workflow( "client": client, "settings": settings, "logger": logger, + "db_adapter": db_adapter, + "generation_id": request.generation_id, }, logger=logger, ) diff --git a/backend/test/api/test_generation_sessions.py b/backend/test/api/test_generation_sessions.py index dd72d76..59b32ef 100644 --- a/backend/test/api/test_generation_sessions.py +++ b/backend/test/api/test_generation_sessions.py @@ -799,9 +799,10 @@ class TestCancelGenerationSession: def test_cancel_generation_success( self, client, mock_generation_session_service ): - """Test successful cancellation of running generation session.""" + """Cancellation transitions to CANCELLED (not FAILED), stops the local task, + and never routes through the error-notifying fail path.""" generation_id = "est-cancel-123" - + running_doc = { "generation_id": generation_id, "status": "running", @@ -819,29 +820,53 @@ def test_cancel_generation_success( "error": None, "user_email": "test@example.com" } - + mock_generation_session_service.get_generation_session_status = AsyncMock( return_value=running_doc ) + mock_generation_session_service.cancel_generation_session = AsyncMock() mock_generation_session_service.fail_generation_session = AsyncMock() - - response = client.delete( - f"/api/v1/generation-sessions/{generation_id}", - headers={"X-API-Key": "test-key"} - ) - + + with patch( + "app.services.generation_task_registry.cancel_task", return_value=True + ) as mock_cancel_task: + response = client.delete( + f"/api/v1/generation-sessions/{generation_id}", + headers={"X-API-Key": "test-key"} + ) + assert response.status_code == 200 data = response.json() - + assert data["generation_id"] == generation_id - assert data["status"] == "failed" - assert "cancelled" in data["message"].lower() - - mock_generation_session_service.fail_generation_session.assert_called_once_with( - generation_id=generation_id, - error="Cancelled by user" + assert data["status"] == "cancelled" + + mock_generation_session_service.cancel_generation_session.assert_called_once_with( + generation_id ) - + # The error-notifying failure path must NOT be used for a user cancellation. + mock_generation_session_service.fail_generation_session.assert_not_called() + mock_cancel_task.assert_called_once_with(generation_id) + + def test_cancel_generation_already_cancelled( + self, client, mock_generation_session_service + ): + """Cannot cancel a session that is already CANCELLED (terminal).""" + generation_id = "est-already-cancelled" + mock_generation_session_service.get_generation_session_status = AsyncMock( + return_value={"generation_id": generation_id, "status": "cancelled"} + ) + mock_generation_session_service.cancel_generation_session = AsyncMock() + + response = client.delete( + f"/api/v1/generation-sessions/{generation_id}", + headers={"X-API-Key": "test-key"} + ) + + assert response.status_code == 400 + assert "Cannot cancel" in response.json()["detail"] + mock_generation_session_service.cancel_generation_session.assert_not_called() + def test_cancel_generation_already_completed( self, client, mock_generation_session_service ): diff --git a/backend/test/jobs/test_shutdown_interrupted_recovery.py b/backend/test/jobs/test_shutdown_interrupted_recovery.py index b54c12e..615991b 100644 --- a/backend/test/jobs/test_shutdown_interrupted_recovery.py +++ b/backend/test/jobs/test_shutdown_interrupted_recovery.py @@ -119,6 +119,22 @@ async def test_ignores_marked_session_outside_window(db, gss, _patch_spawn): assert db.get_generation_session_data("gen-old")["status"] == GenerationStatus.FAILED +@pytest.mark.asyncio +async def test_ignores_cancelled_session(db, gss, _patch_spawn): + """A user-cancelled session must never be resumed on boot recovery.""" + rerun, _ = _patch_spawn + db.seed_generation_session("gen-cancelled", { + "status": GenerationStatus.CANCELLED, + "cancelled_by_user": True, + "cancelled_at": datetime.now(timezone.utc), + "state_history": [], + }) + recovered = await recover_interrupted_sessions(db, gss) + assert recovered == [] + rerun.assert_not_called() + assert db.get_generation_session_data("gen-cancelled")["status"] == GenerationStatus.CANCELLED + + @pytest.mark.asyncio async def test_non_failed_marked_session_ignored(db, gss, _patch_spawn): """A PENDING session carrying a stale marker must not be re-fired (only FAILED).""" diff --git a/backend/test/services/test_generation_session.py b/backend/test/services/test_generation_session.py index 016adb2..e573966 100644 --- a/backend/test/services/test_generation_session.py +++ b/backend/test/services/test_generation_session.py @@ -409,6 +409,43 @@ async def test_fail_generation_session(self, generation_session_service, sample_ ws = generation_session_service._db.get("workspaces", ws_id) assert ws["status"] == "allocated" + @pytest.mark.asyncio + async def test_cancel_generation_session_is_terminal_and_silent( + self, generation_session_service, sample_workspaces + ): + """cancel_generation_session → CANCELLED with markers, workspaces retained, no notify.""" + from unittest.mock import patch + + est_id = await generation_session_service.create_generation_session( + user_email="user-123@example.com", + parameters={"spec_file": "spec.md"} + ) + workspace_ids = await generation_session_service.start_generation_session(est_id) + + with patch("app.services.generation_session.notifications") as mock_notifications: + await generation_session_service.cancel_generation_session(est_id) + mock_notifications.notify.assert_not_called() + + est = generation_session_service._db.get(COL_GENERATION_SESSIONS, est_id) + assert est["status"] == "cancelled" + assert est["cancelled_by_user"] is True + assert est.get("cancelled_at") is not None + assert est.get("failed_at") is None + + # Commandment II: workspaces stay ALLOCATED — generated code preserved. + for ws_id in workspace_ids: + ws = generation_session_service._db.get("workspaces", ws_id) + assert ws["status"] == "allocated" + + @pytest.mark.asyncio + async def test_cancel_generation_session_not_found_raises( + self, generation_session_service + ): + from app.services.generation_session import GenerationSessionNotFoundError + + with pytest.raises(GenerationSessionNotFoundError): + await generation_session_service.cancel_generation_session("est-missing") + class TestProgressTracking: """Test progress tracking for retry continuity.""" diff --git a/backend/test/services/test_generation_session_retry.py b/backend/test/services/test_generation_session_retry.py index dbadffa..d488c8d 100644 --- a/backend/test/services/test_generation_session_retry.py +++ b/backend/test/services/test_generation_session_retry.py @@ -234,6 +234,21 @@ async def test_stuck_pending_retry_succeeds_without_incrementing_count( # retry_count must NOT be incremented again assert est["retry_count"] == 1 + @pytest.mark.asyncio + async def test_cannot_retry_cancelled(self, retry_service, sample_workspaces): + """A user-cancelled session is terminal and must never be retried.""" + db = retry_service._db + db.set(COL_GENERATION_SESSIONS, "est-cancelled", { + "status": GenerationStatus.CANCELLED.value, + "workspace_ids": [], + "code_archived": False, + "retry_count": 0, + "state_history": [], + }) + + with pytest.raises(InvalidRetryStateError, match="cancelled"): + await retry_service.retry_generation_session("est-cancelled") + @pytest.mark.asyncio async def test_cannot_retry_initializing(self, retry_service, sample_workspaces): """Cannot retry generation that is actively INITIALIZING (allocation in progress).""" diff --git a/backend/test/services/test_generation_task_registry.py b/backend/test/services/test_generation_task_registry.py new file mode 100644 index 0000000..73f54f6 --- /dev/null +++ b/backend/test/services/test_generation_task_registry.py @@ -0,0 +1,69 @@ +""" +Tests for the process-local generation task registry. +""" +import asyncio + +import pytest + +from app.services import generation_task_registry as reg + + +@pytest.fixture(autouse=True) +def _clean_registry(): + reg._ACTIVE_TASKS.clear() + yield + reg._ACTIVE_TASKS.clear() + + +@pytest.mark.asyncio +async def test_register_and_get_current_task(): + async def body(): + reg.register_current_task("est-1") + assert reg.get_task("est-1") is asyncio.current_task() + + await asyncio.create_task(body()) + + +@pytest.mark.asyncio +async def test_cancel_task_cancels_running_task(): + started = asyncio.Event() + + async def body(): + reg.register_current_task("est-1") + started.set() + await asyncio.sleep(3600) # would hang without cancellation + + task = asyncio.create_task(body()) + await started.wait() + assert reg.cancel_task("est-1") is True + with pytest.raises(asyncio.CancelledError): + await task + + +def test_cancel_task_returns_false_when_absent(): + assert reg.cancel_task("nope") is False + + +@pytest.mark.asyncio +async def test_deregister_only_removes_own_task(): + # A stale entry (a *different* task object) must not be clobbered by another + # task's deregister — models a re-fired run that already re-registered itself. + sentinel = object() + reg._ACTIVE_TASKS["est-1"] = sentinel # type: ignore[assignment] + + async def body(): + reg.deregister_task("est-1") # current task != sentinel → no-op + + await asyncio.create_task(body()) + assert reg._ACTIVE_TASKS.get("est-1") is sentinel + + +@pytest.mark.asyncio +async def test_deregister_removes_own_task(): + async def body(): + reg.register_current_task("est-1") + assert reg.get_task("est-1") is not None + reg.deregister_task("est-1") + assert reg.get_task("est-1") is None + + await asyncio.create_task(body()) diff --git a/backend/test/services/test_generation_workflow_runner.py b/backend/test/services/test_generation_workflow_runner.py index 85784ad..95a55ae 100644 --- a/backend/test/services/test_generation_workflow_runner.py +++ b/backend/test/services/test_generation_workflow_runner.py @@ -153,3 +153,29 @@ async def test_missing_parameters_fail_routed_not_raised(): sessions.try_begin_for_generation.assert_not_awaited() gss.start_generation_session.assert_not_called() gss.fail_generation_session.assert_awaited() + + +@pytest.mark.asyncio +async def test_handle_workflow_exception_cancelled_is_silent_noop(): + """GenerationCancelledError must NOT fail() or reject() — session already CANCELLED.""" + from app.services.generation_workflow_runner import _handle_workflow_exception + from app.state.exceptions import GenerationCancelledError + + gss = AsyncMock() + await _handle_workflow_exception(GenerationCancelledError("gen-1"), "gen-1", gss) + + gss.fail_generation_session.assert_not_called() + gss.reject_generation_session.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancellation_during_rerun_not_failed(): + """A cooperative cancellation raised mid-workflow is a silent no-op, never fail().""" + from app.state.exceptions import GenerationCancelledError + + gss, sessions = _make_gss(parameters={"spec_path": "specs/", "outputs_dir": "docs"}) + gss.start_generation_session.side_effect = GenerationCancelledError("gen-cancel") + + await rerun_generation_session("gen-cancel", generation_session_service=gss) + + gss.fail_generation_session.assert_not_called() diff --git a/backend/test/state/test_cancellation.py b/backend/test/state/test_cancellation.py new file mode 100644 index 0000000..395c447 --- /dev/null +++ b/backend/test/state/test_cancellation.py @@ -0,0 +1,59 @@ +""" +Tests for the cooperative-cancellation helper (app.state.cancellation). +""" +import pytest + +from app.schemas.generation_workflow_enums import GenerationStatus +from app.state.cancellation import raise_if_cancelled, reset_cancellation_check +from app.state.exceptions import GenerationCancelledError + + +class FakeAdapter: + """Minimal async DB adapter exposing get_generation_session + a call counter.""" + def __init__(self, status=None): + self.status = status + self.calls = 0 + + async def get_generation_session(self, generation_id: str) -> dict: + self.calls += 1 + return {"status": self.status} if self.status is not None else {} + + +@pytest.fixture(autouse=True) +def _reset(): + reset_cancellation_check("est-1") + yield + reset_cancellation_check("est-1") + + +@pytest.mark.asyncio +async def test_raises_when_cancelled(): + adapter = FakeAdapter(status=GenerationStatus.CANCELLED.value) + with pytest.raises(GenerationCancelledError): + await raise_if_cancelled(adapter, "est-1", min_interval_s=0.0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [ + GenerationStatus.RUNNING.value, + GenerationStatus.PENDING.value, + GenerationStatus.COMPLETED.value, +]) +async def test_noop_when_not_cancelled(status): + adapter = FakeAdapter(status=status) + await raise_if_cancelled(adapter, "est-1", min_interval_s=0.0) # does not raise + + +@pytest.mark.asyncio +async def test_noop_when_missing_doc(): + adapter = FakeAdapter(status=None) + await raise_if_cancelled(adapter, "est-1", min_interval_s=0.0) + + +@pytest.mark.asyncio +async def test_throttles_repeated_reads(): + adapter = FakeAdapter(status=GenerationStatus.RUNNING.value) + # First call hits the DB; the immediate second call is throttled (no DB read). + await raise_if_cancelled(adapter, "est-1", min_interval_s=1000.0) + await raise_if_cancelled(adapter, "est-1", min_interval_s=1000.0) + assert adapter.calls == 1 diff --git a/backend/test/state/test_generation_session_state_machine.py b/backend/test/state/test_generation_session_state_machine.py index ba5afe1..effea28 100644 --- a/backend/test/state/test_generation_session_state_machine.py +++ b/backend/test/state/test_generation_session_state_machine.py @@ -350,6 +350,56 @@ async def test_rejects_terminal_states(self, db, terminal): await esm.interrupted_by_shutdown("est-1", triggered_by="system:server_shutdown") +class TestCancel: + @pytest.mark.asyncio + @pytest.mark.parametrize("start_status", [ + GenerationStatus.RUNNING, + GenerationStatus.INITIALIZING, + GenerationStatus.PENDING, + ]) + async def test_cancel_sets_cancelled_with_markers(self, db, start_status): + esm = make_esm(db) + db.seed_generation_session("est-1", {"status": start_status, "state_history": []}) + await esm.cancel("est-1", triggered_by=TriggeredBy.CANCEL) + data = db.get_generation_session_data("est-1") + assert data["status"] == GenerationStatus.CANCELLED + assert data["cancelled_by_user"] is True + assert "cancelled_at" in data + # A user cancellation is NOT a failure — no failed_at, no error. + assert "failed_at" not in data + assert "error" not in data + entry = data["state_history"][-1] + assert entry["status"] == GenerationStatus.CANCELLED + assert entry["triggered_by"] == TriggeredBy.CANCEL + + @pytest.mark.asyncio + async def test_cancel_does_not_touch_workspaces(self, db): + """Commandment II — workspaces stay ALLOCATED, generated code preserved.""" + workspace_sm = make_workspace_sm_mock() + esm = make_esm(db, workspace_sm=workspace_sm) + db.seed_generation_session("est-1", { + "status": GenerationStatus.RUNNING, + "workspace_ids": ["ws-1", "ws-2"], + "state_history": [], + }) + await esm.cancel("est-1", triggered_by=TriggeredBy.CANCEL) + workspace_sm.archive_and_release.assert_not_called() + workspace_sm.allocation_rollback.assert_not_called() + assert db.get_generation_session_data("est-1")["workspace_ids"] == ["ws-1", "ws-2"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("terminal", [ + GenerationStatus.FAILED, + GenerationStatus.COMPLETED, + GenerationStatus.CANCELLED, + ]) + async def test_cancel_rejects_terminal_states(self, db, terminal): + esm = make_esm(db) + db.seed_generation_session("est-1", {"status": terminal, "state_history": []}) + with pytest.raises(InvalidGenerationSessionStateError): + await esm.cancel("est-1", triggered_by=TriggeredBy.CANCEL) + + class TestResetForRetry: @pytest.mark.asyncio async def test_reset_from_failed_to_pending(self, db): diff --git a/backend/test/state/test_workflow_orchestrator.py b/backend/test/state/test_workflow_orchestrator.py index c18557c..24644b8 100644 --- a/backend/test/state/test_workflow_orchestrator.py +++ b/backend/test/state/test_workflow_orchestrator.py @@ -170,6 +170,37 @@ async def step(): esm.complete.assert_called_once() +class TestCooperativeCancellation: + @pytest.mark.asyncio + async def test_cancelled_session_stops_without_fail(self, db): + """A CANCELLED session makes the step loop raise GenerationCancelledError and + must NOT call fail() (session is already terminal).""" + from app.state.cancellation import reset_cancellation_check + from app.state.exceptions import GenerationCancelledError + + reset_cancellation_check("est-cancel-orch") + db.seed_generation_session("est-cancel-orch", { + "status": GenerationStatus.CANCELLED, + "checkpoint": None, + }) + esm = AsyncMock(spec=GenerationSessionStateMachine) + esm.advance_checkpoint = AsyncMock(return_value={}) + esm.fail = AsyncMock(return_value={}) + esm.complete = AsyncMock(return_value={}) + orch = WorkflowOrchestrator(db, esm) + + ran = [] + async def step(): + ran.append("validate_contract") + + with pytest.raises(GenerationCancelledError): + await orch.run("est-cancel-orch", {"validate_contract": step}, triggered_by="test") + + # Cooperative check runs before the step impl and before fail(). + esm.fail.assert_not_called() + assert ran == [] + + class TestFailure: @pytest.mark.asyncio async def test_step_failure_calls_fail_with_step_name(self, db): diff --git a/backend/test/test_parallel_execution.py b/backend/test/test_parallel_execution.py index 8f41ffa..57f3589 100644 --- a/backend/test/test_parallel_execution.py +++ b/backend/test/test_parallel_execution.py @@ -212,6 +212,27 @@ async def mock_agent_function(workspace, manager, log): assert results[2].success is True assert "Simulated failure" in results[1].error + @pytest.mark.asyncio + async def test_execute_parallel_reraises_cancellation(self, mock_settings, workspaces): + """A cooperative cancellation must abort the whole fan-out, not degrade to a + per-workspace failure result absorbed by gather(return_exceptions=True).""" + from app.state.exceptions import GenerationCancelledError + + logger = Mock() + workspace_manager = WorkspaceManager(mock_settings, logger) + executor = ParallelAgentExecutor(workspace_manager, logger) + + async def mock_agent_function(workspace, manager, log): + if workspace.provider == "anthropic": + raise GenerationCancelledError("gen-1") + return AgentResult(result="ok", session_id="s") + + with pytest.raises(GenerationCancelledError): + await executor.execute_parallel( + workspaces=workspaces, + agent_fn=mock_agent_function, + ) + @pytest.mark.asyncio async def test_execute_parallel_workflow(self, mock_settings, workspaces): """Test parallel workflow execution.""" From 71b84cc27b802e7270f08afbd0134f9baf7e2264 Mon Sep 17 00:00:00 2001 From: Artsiom Kozak Date: Tue, 14 Jul 2026 09:23:21 +0200 Subject: [PATCH 2/5] fix(cancellation): always run the first cooperative check raise_if_cancelled used 0.0 as the "never checked" throttle sentinel. time.monotonic() has no fixed epoch, so on a freshly-booted host (e.g. a CI runner) the clock can be smaller than min_interval_s, making now - 0.0 < min_interval_s true and throttling the very first read. Use a None sentinel so the first check for a generation always reads. --- backend/app/state/cancellation.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/backend/app/state/cancellation.py b/backend/app/state/cancellation.py index 23f7097..64a9593 100644 --- a/backend/app/state/cancellation.py +++ b/backend/app/state/cancellation.py @@ -32,12 +32,17 @@ async def raise_if_cancelled( ) -> None: """Raise :class:`GenerationCancelledError` if the session is CANCELLED. - Throttled: if called again for the same ``generation_id`` within ``min_interval_s`` - seconds, returns immediately without touching the DB. ``db_adapter`` is the async - state-machine DB adapter (``service.db_adapter`` / ``orchestrator._db``). + The first check for a generation always reads; thereafter, if called again for the + same ``generation_id`` within ``min_interval_s`` seconds, it returns immediately + without touching the DB. (``time.monotonic()`` has no fixed epoch, so a "never + checked" sentinel must be distinct from any real timestamp — not 0.0 — otherwise a + small monotonic value on a freshly-booted host would throttle the very first read.) + ``db_adapter`` is the async state-machine DB adapter (``service.db_adapter`` / + ``orchestrator._db``). """ now = time.monotonic() - if now - _last_check.get(generation_id, 0.0) < min_interval_s: + last = _last_check.get(generation_id) + if last is not None and now - last < min_interval_s: return _last_check[generation_id] = now From 35c17e496a8160b93485692638dc360e6955b24e Mon Sep 17 00:00:00 2001 From: Artsiom Kozak Date: Tue, 14 Jul 2026 09:46:03 +0200 Subject: [PATCH 3/5] refactor(cancellation): make task registry an injected instance Replace the module-level global dict + functions with a GenerationTaskRegistry class. One instance is created at app startup and stored on app.state; the run/retry/cancel endpoints receive it via a get_generation_task_registry dependency, and it is threaded into the workflow task bodies (and boot recovery) as a parameter. This removes mutable import-time global state, aligns with the codebase's dependency-injection and OOP conventions, and lets tests inject the registry via dependency_overrides instead of reaching into module internals. Behavior is unchanged: same process-local fast cancel backed by the cooperative status check. --- backend/app/api/v1/generation_sessions.py | 25 ++++-- backend/app/core/app_lifecycle.py | 14 +++- .../app/jobs/shutdown_interrupted_recovery.py | 6 +- .../app/services/generation_task_registry.py | 84 ++++++++++--------- .../services/generation_workflow_runner.py | 9 +- backend/test/api/test_generation_sessions.py | 34 +++++--- .../services/test_generation_task_registry.py | 20 ++--- 7 files changed, 117 insertions(+), 75 deletions(-) diff --git a/backend/app/api/v1/generation_sessions.py b/backend/app/api/v1/generation_sessions.py index 27e2c87..03cb041 100644 --- a/backend/app/api/v1/generation_sessions.py +++ b/backend/app/api/v1/generation_sessions.py @@ -77,7 +77,7 @@ from app.services.artifact_store import ARTIFACTS_BASE, ArtifactStore from app.services.agent_stream_broker import get_agent_stream_broker from app.services.mcp_prune import resolve_enabled_mcps_and_set_telemetry -from app.services import generation_task_registry +from app.services.generation_task_registry import GenerationTaskRegistry from app.schemas.model_token_usage import ModelTokenUsage from app.schemas.workflow_usage_metrics import ( aggregate_model_usage_by_workspace, @@ -250,6 +250,11 @@ def get_generation_session_retry_service( return GenerationSessionRetryService(db, workspace_pool) +def get_generation_task_registry(request: Request) -> GenerationTaskRegistry: + """Provide the process-local task registry created once at app startup.""" + return request.app.state.generation_task_registry + + # ============================================================ # Endpoints # ============================================================ @@ -776,7 +781,8 @@ async def retry_generation_session( _: None = Depends(require_generation_session_owner_or_admin), retry_service: GenerationSessionRetryService = Depends(get_generation_session_retry_service), generation_session_service: GenerationSessionService = Depends(get_generation_session_service), - workspace_pool: WorkspacePoolService = Depends(get_workspace_pool) + workspace_pool: WorkspacePoolService = Depends(get_workspace_pool), + task_registry: GenerationTaskRegistry = Depends(get_generation_task_registry), ): """ Retry a failed generation session. @@ -816,6 +822,7 @@ async def retry_generation_session( generation_id, generation_session_service=generation_session_service, user_email=getattr(request.state, "user_email", None), + task_registry=task_registry, ) ) @@ -851,7 +858,8 @@ async def cancel_generation_session( generation_id: str, request: Request, _: None = Depends(require_generation_session_owner), - generation_session_service: GenerationSessionService = Depends(get_generation_session_service) + generation_session_service: GenerationSessionService = Depends(get_generation_session_service), + task_registry: GenerationTaskRegistry = Depends(get_generation_task_registry), ): """ Cancel a running generation session (user-initiated). @@ -893,7 +901,7 @@ async def cancel_generation_session( # re-fails the session nor emits an error notification. await generation_session_service.cancel_generation_session(generation_id) - cancelled_locally = generation_task_registry.cancel_task(generation_id) + cancelled_locally = task_registry.cancel_task(generation_id) return CancelGenerationSessionResponse( generation_id=generation_id, @@ -1076,6 +1084,7 @@ async def _run_generation_session_workflow( spec_path: str, session, generation_session_service, + task_registry: GenerationTaskRegistry | None = None, ) -> None: """Background task: run app generation + P10Y estimation, release session slot when done.""" agent_logger = create_agent_logger( @@ -1090,7 +1099,8 @@ async def _run_generation_session_workflow( normalized_src_dir = generation_session_params.get("src_dir", "src") try: - generation_task_registry.register_current_task(generation_id) + if task_registry is not None: + task_registry.register_current_task(generation_id) async with session.task_slot(generation_id=generation_id): logger.info("Starting app generation + P10Y workflow for %s", generation_id) agent_logger.info("Starting app generation + P10Y workflow for %s", generation_id) @@ -1127,7 +1137,8 @@ async def _run_generation_session_workflow( # else → fail (FAILED). GenerationCancelledError is handled as a silent no-op there. await _handle_workflow_exception(e, generation_id, generation_session_service) finally: - generation_task_registry.deregister_task(generation_id) + if task_registry is not None: + task_registry.deregister_task(generation_id) TelemetryContext.set_agent_query_totals_handler(None) @@ -1322,6 +1333,7 @@ async def run_generation_session( workspace_count: Optional[int] = Form(None, ge=1, le=3), workspace_pool: WorkspacePoolService = Depends(get_workspace_pool), generation_session_service: GenerationSessionService = Depends(get_generation_session_service), + task_registry: GenerationTaskRegistry = Depends(get_generation_task_registry), _owner: None = Depends(require_generation_session_owner_form), ): """ @@ -1515,6 +1527,7 @@ async def run_generation_session( spec_path=spec_path, session=session, generation_session_service=generation_session_service, + task_registry=task_registry, )) return RunGenerationSessionResponse( diff --git a/backend/app/core/app_lifecycle.py b/backend/app/core/app_lifecycle.py index aa34912..199f527 100644 --- a/backend/app/core/app_lifecycle.py +++ b/backend/app/core/app_lifecycle.py @@ -34,6 +34,7 @@ from app.jobs.stuck_initializing_detector import detect_stuck_initializing from app.jobs.stuck_running_detector import detect_stuck_running from app.services.generation_session import GenerationSessionService +from app.services.generation_task_registry import GenerationTaskRegistry from app.services.langfuse import tracer from app.services.openrouter_pricing import ensure_openrouter_pricing_cache from app.services.shutdown_handler import mark_active_sessions_interrupted @@ -128,13 +129,16 @@ def build_background_jobs(db, workspace_pool) -> list: ] -def start_boot_recovery(db, raw_db, workspace_pool) -> asyncio.Task: +def start_boot_recovery(db, raw_db, workspace_pool, task_registry: GenerationTaskRegistry) -> asyncio.Task: """One-shot task: auto-retry sessions interrupted by a previous shutdown (pod eviction). Non-blocking so it never delays the port-open. See docs/backend/graceful-shutdown-recovery.md. + Recovered runs register with ``task_registry`` so they remain cancellable. """ generation_session_service = GenerationSessionService(raw_db, workspace_pool) - return asyncio.create_task(recover_interrupted_sessions(db, generation_session_service)) + return asyncio.create_task( + recover_interrupted_sessions(db, generation_session_service, task_registry=task_registry) + ) async def run_shutdown_session_handling(db) -> None: @@ -169,9 +173,13 @@ async def lifespan(app: FastAPI): workspace_pool = WorkspacePoolService(raw_db) job_tasks = build_background_jobs(db, workspace_pool) + # Process-local registry of running workflow tasks, shared by the run/retry/cancel + # endpoints (via dependency injection) and boot recovery. + app.state.generation_task_registry = GenerationTaskRegistry() + tracer.init() - recovery_task = start_boot_recovery(db, raw_db, workspace_pool) + recovery_task = start_boot_recovery(db, raw_db, workspace_pool, app.state.generation_task_registry) # CRITICAL: Yield immediately so FastAPI starts listening (this is what Cloud Run checks). yield diff --git a/backend/app/jobs/shutdown_interrupted_recovery.py b/backend/app/jobs/shutdown_interrupted_recovery.py index 13f3fef..e6c5dd9 100644 --- a/backend/app/jobs/shutdown_interrupted_recovery.py +++ b/backend/app/jobs/shutdown_interrupted_recovery.py @@ -10,6 +10,7 @@ from app.core.notifications import notifications from app.schemas.generation_workflow_enums import GenerationStatus, WorkspaceStatus from app.services.generation_session import GenerationSessionService +from app.services.generation_task_registry import GenerationTaskRegistry from app.services.generation_workflow_runner import rerun_generation_session from app.state import GenerationSessionStateMachine from app.state.exceptions import ( @@ -62,11 +63,13 @@ async def recover_interrupted_sessions( generation_session_service: GenerationSessionService, *, window_minutes: int | None = None, + task_registry: GenerationTaskRegistry | None = None, ) -> list[str]: """Auto-retry shutdown-interrupted sessions found within the recovery window. One-shot (not a loop). Returns the generation IDs it re-fired. Each candidate is - independently guarded; one failure never blocks the others. + independently guarded; one failure never blocks the others. ``task_registry`` (when + provided) is threaded into each re-fired run so it stays cancellable. """ if not settings.AUTO_RECOVER_INTERRUPTED_SESSIONS: logger.info("shutdown_interrupted_recovery: disabled by config — skipping") @@ -132,6 +135,7 @@ async def recover_interrupted_sessions( gid, generation_session_service=generation_session_service, user_email=session.get("user_email"), + task_registry=task_registry, ) ) recovered.append(gid) diff --git a/backend/app/services/generation_task_registry.py b/backend/app/services/generation_task_registry.py index 1ab9dfa..e8e31c2 100644 --- a/backend/app/services/generation_task_registry.py +++ b/backend/app/services/generation_task_registry.py @@ -6,51 +6,53 @@ the fast path; the cooperative :func:`app.state.cancellation.raise_if_cancelled` checks are the cross-pod-safe fallback for cancels that land on a different pod. -The workflow task registers itself via ``asyncio.current_task()`` from inside its own -body, so there is no plumbing at the ``asyncio.create_task(...)`` call sites — the initial -run and the shared rerun path (manual retry + boot recovery) are all covered. +One instance per process is created at app startup and stored on ``app.state`` (see +``app_lifecycle``); endpoints receive it via dependency injection and it is threaded into +the workflow task bodies. The workflow task registers itself via ``asyncio.current_task()`` +from inside its own body, so callers only pass the registry — they don't manage handles. """ import asyncio import logging logger = logging.getLogger(__name__) -_ACTIVE_TASKS: dict[str, asyncio.Task] = {} - -def register_current_task(generation_id: str) -> None: - """Register the currently running task as the workflow task for ``generation_id``.""" - task = asyncio.current_task() - if task is not None: - _ACTIVE_TASKS[generation_id] = task - - -def deregister_task(generation_id: str) -> None: - """Remove the registry entry for ``generation_id`` — only if it is *this* task. - - Guarding on identity avoids a re-fired run (retry/recovery) that has already - registered its own task being clobbered by the previous task's ``finally`` cleanup. - """ - current = asyncio.current_task() - if _ACTIVE_TASKS.get(generation_id) is current: - _ACTIVE_TASKS.pop(generation_id, None) - - -def get_task(generation_id: str) -> asyncio.Task | None: - """Return the registered workflow task for ``generation_id``, or None.""" - return _ACTIVE_TASKS.get(generation_id) - - -def cancel_task(generation_id: str) -> bool: - """Cancel the registered workflow task if present and not already done. - - Returns True if a live task was found and ``cancel()`` was requested (immediate - local teardown), False otherwise (run owned by another pod, already finished, or - never registered — the cooperative status check handles those). - """ - task = _ACTIVE_TASKS.get(generation_id) - if task is not None and not task.done(): - task.cancel() - logger.info("Requested cancellation of local workflow task for %s", generation_id) - return True - return False +class GenerationTaskRegistry: + """In-memory, single-loop registry of active generation workflow tasks.""" + + def __init__(self) -> None: + self._tasks: dict[str, asyncio.Task] = {} + + def register_current_task(self, generation_id: str) -> None: + """Register the currently running task as the workflow task for ``generation_id``.""" + task = asyncio.current_task() + if task is not None: + self._tasks[generation_id] = task + + def deregister_task(self, generation_id: str) -> None: + """Remove the registry entry for ``generation_id`` — only if it is *this* task. + + Guarding on identity avoids a re-fired run (retry/recovery) that has already + registered its own task being clobbered by the previous task's ``finally`` cleanup. + """ + current = asyncio.current_task() + if self._tasks.get(generation_id) is current: + self._tasks.pop(generation_id, None) + + def get_task(self, generation_id: str) -> asyncio.Task | None: + """Return the registered workflow task for ``generation_id``, or None.""" + return self._tasks.get(generation_id) + + def cancel_task(self, generation_id: str) -> bool: + """Cancel the registered workflow task if present and not already done. + + Returns True if a live task was found and ``cancel()`` was requested (immediate + local teardown), False otherwise (run owned by another pod, already finished, or + never registered — the cooperative status check handles those). + """ + task = self._tasks.get(generation_id) + if task is not None and not task.done(): + task.cancel() + logger.info("Requested cancellation of local workflow task for %s", generation_id) + return True + return False diff --git a/backend/app/services/generation_workflow_runner.py b/backend/app/services/generation_workflow_runner.py index a6e5c85..0016645 100644 --- a/backend/app/services/generation_workflow_runner.py +++ b/backend/app/services/generation_workflow_runner.py @@ -21,7 +21,7 @@ from app.services.claude_code import clear_workspace_caches from app.services.contract_validator import ContractRejection from app.services.generation_session import GenerationSessionService -from app.services import generation_task_registry +from app.services.generation_task_registry import GenerationTaskRegistry from app.state.api_key_session_concurrency import OperationKind, SessionBeginOutcome from app.state.exceptions import ( GenerationCancelledError, @@ -218,6 +218,7 @@ async def rerun_generation_session( *, generation_session_service: GenerationSessionService, user_email: str | None = None, + task_registry: GenerationTaskRegistry | None = None, ) -> None: """Re-fire a generation session from its persisted parameters. @@ -251,7 +252,8 @@ async def rerun_generation_session( } try: - generation_task_registry.register_current_task(generation_id) + if task_registry is not None: + task_registry.register_current_task(generation_id) TelemetryContext.set_user_context( user_email=user_email, generation_id=generation_id, @@ -310,5 +312,6 @@ async def rerun_generation_session( # GenerationCancelledError is handled as a silent no-op there. await _handle_workflow_exception(e, generation_id, generation_session_service) finally: - generation_task_registry.deregister_task(generation_id) + if task_registry is not None: + task_registry.deregister_task(generation_id) TelemetryContext.set_agent_query_totals_handler(None) diff --git a/backend/test/api/test_generation_sessions.py b/backend/test/api/test_generation_sessions.py index 59b32ef..28921e8 100644 --- a/backend/test/api/test_generation_sessions.py +++ b/backend/test/api/test_generation_sessions.py @@ -16,6 +16,7 @@ _handle_workflow_exception, get_generation_session_retry_service, get_generation_session_service, + get_generation_task_registry, get_workspace_pool, list_generation_sessions, router as generation_sessions_router, @@ -38,7 +39,17 @@ def mock_retry_service(): @pytest.fixture -def app(test_app, mock_workspace_pool, mock_generation_session_service, mock_retry_service): +def mock_task_registry(): + """Mock GenerationTaskRegistry injected in place of the app.state singleton.""" + registry = Mock() + registry.register_current_task = Mock() + registry.deregister_task = Mock() + registry.cancel_task = Mock(return_value=True) + return registry + + +@pytest.fixture +def app(test_app, mock_workspace_pool, mock_generation_session_service, mock_retry_service, mock_task_registry): """Create FastAPI app with generation_sessions router and mocked dependencies.""" # Include router test_app.include_router(generation_sessions_router, prefix="/api/v1") @@ -52,7 +63,10 @@ def override_get_generation_session_service(): def override_get_generation_session_retry_service(): return mock_retry_service - + + def override_get_generation_task_registry(): + return mock_task_registry + # Override authorization dependencies to skip ownership check in tests # Tests will use db fixture to create sessions with matching user_email from app.api.dependencies import ( @@ -67,6 +81,7 @@ def override_require_generation_session_owner(): test_app.dependency_overrides[get_workspace_pool] = override_get_workspace_pool test_app.dependency_overrides[get_generation_session_service] = override_get_generation_session_service test_app.dependency_overrides[get_generation_session_retry_service] = override_get_generation_session_retry_service + test_app.dependency_overrides[get_generation_task_registry] = override_get_generation_task_registry test_app.dependency_overrides[require_generation_session_owner] = override_require_generation_session_owner test_app.dependency_overrides[require_generation_session_owner_form] = override_require_generation_session_owner test_app.dependency_overrides[require_generation_session_owner_or_admin] = override_require_generation_session_owner @@ -797,7 +812,7 @@ class TestCancelGenerationSession: """Tests for DELETE /generation-sessions/{generation_id} endpoint.""" def test_cancel_generation_success( - self, client, mock_generation_session_service + self, client, mock_generation_session_service, mock_task_registry ): """Cancellation transitions to CANCELLED (not FAILED), stops the local task, and never routes through the error-notifying fail path.""" @@ -827,13 +842,10 @@ def test_cancel_generation_success( mock_generation_session_service.cancel_generation_session = AsyncMock() mock_generation_session_service.fail_generation_session = AsyncMock() - with patch( - "app.services.generation_task_registry.cancel_task", return_value=True - ) as mock_cancel_task: - response = client.delete( - f"/api/v1/generation-sessions/{generation_id}", - headers={"X-API-Key": "test-key"} - ) + response = client.delete( + f"/api/v1/generation-sessions/{generation_id}", + headers={"X-API-Key": "test-key"} + ) assert response.status_code == 200 data = response.json() @@ -846,7 +858,7 @@ def test_cancel_generation_success( ) # The error-notifying failure path must NOT be used for a user cancellation. mock_generation_session_service.fail_generation_session.assert_not_called() - mock_cancel_task.assert_called_once_with(generation_id) + mock_task_registry.cancel_task.assert_called_once_with(generation_id) def test_cancel_generation_already_cancelled( self, client, mock_generation_session_service diff --git a/backend/test/services/test_generation_task_registry.py b/backend/test/services/test_generation_task_registry.py index 73f54f6..86b7cfd 100644 --- a/backend/test/services/test_generation_task_registry.py +++ b/backend/test/services/test_generation_task_registry.py @@ -5,18 +5,13 @@ import pytest -from app.services import generation_task_registry as reg - - -@pytest.fixture(autouse=True) -def _clean_registry(): - reg._ACTIVE_TASKS.clear() - yield - reg._ACTIVE_TASKS.clear() +from app.services.generation_task_registry import GenerationTaskRegistry @pytest.mark.asyncio async def test_register_and_get_current_task(): + reg = GenerationTaskRegistry() + async def body(): reg.register_current_task("est-1") assert reg.get_task("est-1") is asyncio.current_task() @@ -26,6 +21,7 @@ async def body(): @pytest.mark.asyncio async def test_cancel_task_cancels_running_task(): + reg = GenerationTaskRegistry() started = asyncio.Event() async def body(): @@ -41,6 +37,7 @@ async def body(): def test_cancel_task_returns_false_when_absent(): + reg = GenerationTaskRegistry() assert reg.cancel_task("nope") is False @@ -48,18 +45,21 @@ def test_cancel_task_returns_false_when_absent(): async def test_deregister_only_removes_own_task(): # A stale entry (a *different* task object) must not be clobbered by another # task's deregister — models a re-fired run that already re-registered itself. + reg = GenerationTaskRegistry() sentinel = object() - reg._ACTIVE_TASKS["est-1"] = sentinel # type: ignore[assignment] + reg._tasks["est-1"] = sentinel # type: ignore[assignment] async def body(): reg.deregister_task("est-1") # current task != sentinel → no-op await asyncio.create_task(body()) - assert reg._ACTIVE_TASKS.get("est-1") is sentinel + assert reg.get_task("est-1") is sentinel @pytest.mark.asyncio async def test_deregister_removes_own_task(): + reg = GenerationTaskRegistry() + async def body(): reg.register_current_task("est-1") assert reg.get_task("est-1") is not None From f6f7bc550d4bbe1480f5aaa7ae360adde78775d6 Mon Sep 17 00:00:00 2001 From: Artsiom Kozak Date: Tue, 14 Jul 2026 11:11:29 +0200 Subject: [PATCH 4/5] feat(tui): let users cancel a running generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the backend cancellation (PR #48) into the Textual TUI. Detail screen (DashboardScreen): add an `x` → cancel binding, shown only while the run is active (check_action gates it to pending/initializing/ running). Confirming pops a ConfirmScreen, then calls the existing DELETE /api/v1/generation-sessions/{id} for that specific session and refreshes in place — no terminal suspend, since cancellation is silent. Sessions list: a cancelled row now renders "⊘ CANCELLED BY USER". Recognize the status everywhere it matters: add a CANCELLED pill and put "cancelled" in TERMINAL_STATUSES (so the dashboard stops polling and the header pill renders), keep the cli_service mirror in sync, announce a cancel as a neutral "Generation cancelled" desktop milestone instead of a failure, and add CANCELLED to the mcp_server GenerationStatus mirror for client/backend parity. --- .../schemas/generation_workflow_enums.py | 1 + mcp_server/services/cli_service.py | 2 +- .../tests/test_generation_workflow_enums.py | 5 +- mcp_server/tests/test_tui_app.py | 73 +++++++++++++++++++ mcp_server/tests/test_tui_poller.py | 12 +++ mcp_server/tui/app.py | 49 ++++++++++++- mcp_server/tui/constants.py | 3 +- mcp_server/tui/poller.py | 5 ++ 8 files changed, 145 insertions(+), 5 deletions(-) diff --git a/mcp_server/schemas/generation_workflow_enums.py b/mcp_server/schemas/generation_workflow_enums.py index 305b2d8..948f701 100644 --- a/mcp_server/schemas/generation_workflow_enums.py +++ b/mcp_server/schemas/generation_workflow_enums.py @@ -14,6 +14,7 @@ class GenerationStatus(str, Enum): RUNNING = "running" COMPLETED = "completed" FAILED = "failed" + CANCELLED = "cancelled" class GenerationCheckpoint(str, Enum): diff --git a/mcp_server/services/cli_service.py b/mcp_server/services/cli_service.py index 5361170..6fd4556 100644 --- a/mcp_server/services/cli_service.py +++ b/mcp_server/services/cli_service.py @@ -29,7 +29,7 @@ logger = logging.getLogger(__name__) -_TERMINAL_STATUSES = frozenset({"completed", "failed"}) +_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"}) _NO_SUITABLE_IMPLEMENTATION = "no suitable implementation" _MACOS_BUNDLE_ID = "com.griddynamics.specflow" _MACOS_BUNDLE_NAME = "SpecFlow" diff --git a/mcp_server/tests/test_generation_workflow_enums.py b/mcp_server/tests/test_generation_workflow_enums.py index 57fdb7d..30d60ab 100644 --- a/mcp_server/tests/test_generation_workflow_enums.py +++ b/mcp_server/tests/test_generation_workflow_enums.py @@ -27,8 +27,11 @@ def test_completed(self): def test_failed(self): assert GenerationStatus.FAILED == "failed" + def test_cancelled(self): + assert GenerationStatus.CANCELLED == "cancelled" + def test_no_extra_values(self): - expected = {"pending", "initializing", "running", "completed", "failed"} + expected = {"pending", "initializing", "running", "completed", "failed", "cancelled"} assert {s.value for s in GenerationStatus} == expected diff --git a/mcp_server/tests/test_tui_app.py b/mcp_server/tests/test_tui_app.py index 7b70857..0e258f4 100644 --- a/mcp_server/tests/test_tui_app.py +++ b/mcp_server/tests/test_tui_app.py @@ -254,6 +254,18 @@ def test_unknown_checkpoint_falls_back_to_key(self): assert "custom_checkpoint" in label + def test_cancelled_shows_by_user(self): + label = tui_app._session_label( + { + "generation_id": "est-d67dcac6cbe5", + "status": "cancelled", + "checkpoint": "generation_done", + } + ) + + assert "CANCELLED BY USER" in label + assert "⊘" in label + class TestRunTuiNonTty: @pytest.mark.asyncio @@ -832,6 +844,67 @@ async def test_q_exits_sessions_screen(self): assert not app.is_running +class TestDashboardCancel: + def test_check_action_gates_cancel_by_status(self): + screen = tui_app.DashboardScreen("gen_x") + for active in ("pending", "initializing", "running"): + screen._payload = {"status": active} + assert screen.check_action("cancel", ()) is True + for terminal in ("completed", "failed", "cancelled"): + screen._payload = {"status": terminal} + assert screen.check_action("cancel", ()) is None + + @pytest.mark.asyncio + async def test_cancel_flow_confirms_then_calls_delete(self): + a, b, c = _gate_ready() + delete_mock = AsyncMock(return_value='{"status": "cancelled"}') + with ( + a, + b, + c, + patch("tui.app.poll_once", new=AsyncMock(return_value=_running_payload())), + patch("tui.app.call_backend_endpoint", new=delete_mock), + ): + 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() + assert isinstance(app.screen, tui_app.DashboardScreen) + await pilot.press("x") + await pilot.pause() + # Confirmation dialog appears before any backend call. + assert isinstance(app.screen, tui_app.ConfirmScreen) + delete_mock.assert_not_awaited() + await pilot.click("#confirm") # confirm + await pilot.pause() + + delete_mock.assert_awaited_once() + kwargs = delete_mock.await_args.kwargs + assert kwargs["endpoint"] == "/api/v1/generation-sessions/gen_x" + assert kwargs["method"] == "DELETE" + + @pytest.mark.asyncio + async def test_cancel_flow_declined_does_not_call_delete(self): + a, b, c = _gate_ready() + delete_mock = AsyncMock() + with ( + a, + b, + c, + patch("tui.app.poll_once", new=AsyncMock(return_value=_running_payload())), + patch("tui.app.call_backend_endpoint", new=delete_mock), + ): + 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() + await pilot.press("x") + await pilot.pause() + assert isinstance(app.screen, tui_app.ConfirmScreen) + await pilot.press("escape") # decline + await pilot.pause() + + delete_mock.assert_not_awaited() + + class TestAppNotifications: @pytest.mark.asyncio async def test_notifies_for_active_sessions_independent_of_screen(self): diff --git a/mcp_server/tests/test_tui_poller.py b/mcp_server/tests/test_tui_poller.py index 71f501d..c57cb47 100644 --- a/mcp_server/tests/test_tui_poller.py +++ b/mcp_server/tests/test_tui_poller.py @@ -47,6 +47,18 @@ def test_failed_fires(self): out = tracker.process({"status": "failed", "checkpoint": "generation_started"}) assert any("failed" in m.message for m in out) + def test_cancelled_fires_neutrally_once(self): + # A user cancellation is announced as "cancelled", never as a failure. + tracker = MilestoneTracker("gen_abc123def456") + tracker.process({"status": "running", "checkpoint": "generation_started"}) + out = tracker.process({"status": "cancelled", "checkpoint": "generation_started"}) + assert len(out) == 1 + assert "cancelled" in out[0].title.lower() + assert "cancelled" in out[0].message + assert not any("failed" in m.message for m in out) + # Terminal — fires exactly once. + assert tracker.process({"status": "cancelled", "checkpoint": "generation_started"}) == [] + def test_workspace_phase_progress_fires_once_per_change(self): tracker = MilestoneTracker("gen_abc123def456") tracker.process( diff --git a/mcp_server/tui/app.py b/mcp_server/tui/app.py index c909ad5..4f84f90 100644 --- a/mcp_server/tui/app.py +++ b/mcp_server/tui/app.py @@ -64,7 +64,7 @@ from services import local_env, validate_models from services.llm_tiers import LLM_TIER_KEYS from services.session import resolve_generation_id, set_project_root -from services.specflow_backend import call_backend_endpoint_bytes +from services.specflow_backend import call_backend_endpoint, call_backend_endpoint_bytes from tui import actions, activity, mcp_clients, onboarding, render from tui.config import ( EDITABLE_KEYS, @@ -466,6 +466,7 @@ class DashboardScreen(_SpecFlowScreen): BINDINGS = [ Binding("r", "retry", "retry"), + Binding("x", "cancel", "cancel"), Binding("w", "clear", "clear ws"), Binding("s", "settings", "settings"), Binding("b", "sessions", "sessions"), @@ -545,6 +546,9 @@ def check_action(self, action: str, parameters: tuple[object, ...]) -> bool | No if action == "retry": status = ((self._payload or {}).get("status") or "").lower() return True if status == "failed" else None + if action == "cancel": + status = ((self._payload or {}).get("status") or "").lower() + return True if status in {"pending", "initializing", "running"} else None if action == "open_report": return True if render.estimate_panel(self._payload) is not None else None return True @@ -630,6 +634,45 @@ async def _retry_flow(self) -> None: if ok: await self._run_suspended(actions.do_retry(self.app.root)) + def action_cancel(self) -> None: + self.run_worker(self._cancel_flow(), exclusive=True) + + async def _cancel_flow(self) -> None: + """Cancel this specific run via the backend DELETE, then refresh in place. + + Targets ``self._generation_id`` directly (the dashboard may show any session + opened from the list). No terminal suspend — cancellation produces no output. + """ + ok = await self.app.push_screen_wait( + ConfirmScreen( + "Cancel this generation? It stops all agents immediately and " + "cannot be resumed (the workspace is kept).", + countdown=0, + ) + ) + if not ok: + return + try: + await call_backend_endpoint( + endpoint=f"/api/v1/generation-sessions/{self._generation_id}", + method="DELETE", + timeout_seconds=30, + ) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 400: + self.notify( + "Generation already finished — nothing to cancel.", + severity="information", + ) + else: + self.notify(f"Couldn't cancel: {exc}", severity="warning") + return + except Exception as exc: # noqa: BLE001 - surface any failure to the user + self.notify(f"Couldn't cancel: {exc}", severity="warning") + return + self.notify("Generation cancelled.", severity="information") + await self.refresh_status() + def action_clear(self) -> None: self.run_worker(self._clear_flow(), exclusive=True) @@ -777,7 +820,9 @@ def _session_label(s: dict) -> str: except ValueError: date_str = created_at_raw[:16] - status_col = f"{symbol} {status_key.upper():<12}" + # Cancellations are always user-initiated; spell that out in the list. + status_word = "CANCELLED BY USER" if status_key == "cancelled" else status_key.upper() + status_col = f"{symbol} {status_word:<17}" date_col = f"{date_str:<14}" if date_str else " " * 14 return f"{status_col} {gid[:22]:<24} {date_col} {checkpoint_label}" diff --git a/mcp_server/tui/constants.py b/mcp_server/tui/constants.py index 1a0e59f..0410c3f 100644 --- a/mcp_server/tui/constants.py +++ b/mcp_server/tui/constants.py @@ -97,7 +97,7 @@ class PipelineStepDef: } # Terminal statuses — mirror of services.cli_service._TERMINAL_STATUSES. -TERMINAL_STATUSES: frozenset[str] = frozenset({"completed", "failed"}) +TERMINAL_STATUSES: frozenset[str] = frozenset({"completed", "failed", "cancelled"}) # Status pill text + Rich style per lifecycle status (lowercased). STATUS_PILLS: dict[str, tuple[str, str]] = { @@ -106,6 +106,7 @@ class PipelineStepDef: "running": ("⬤ RUNNING", "green"), "completed": ("✓ COMPLETED", "bold green"), "failed": ("✗ FAILED", "bold red"), + "cancelled": ("⊘ CANCELLED", "dim"), "unknown": ("? UNKNOWN", "dim"), } diff --git a/mcp_server/tui/poller.py b/mcp_server/tui/poller.py index 3de0a60..d0280b9 100644 --- a/mcp_server/tui/poller.py +++ b/mcp_server/tui/poller.py @@ -164,6 +164,11 @@ def process(self, payload: dict[str, Any] | None) -> list[Milestone]: milestones.append( Milestone("Generation completed", f"{short}... completed") ) + elif status == "cancelled": + # User-initiated stop — announce neutrally, not as a failure. + milestones.append( + Milestone("Generation cancelled", f"{short}... cancelled by user") + ) else: milestones.append(Milestone("Generation failed", f"{short}... failed")) self._last_status = status From f49d05b24db76ecb0151f8ad6a613fa6b2f6c13e Mon Sep 17 00:00:00 2001 From: Artsiom Kozak Date: Tue, 14 Jul 2026 14:39:00 +0200 Subject: [PATCH 5/5] refactor(cancellation): tighten contracts and drop needless cancel paths Three cleanups to the cancellation wiring: - raise_if_cancelled now no-ops on a None db_adapter, so callers invoke it unconditionally. db_adapter is None only in DB-less unit tests; every real run wires a generation_session_service. Removes the guard from the execute_all_phases call site and centralizes it in one place. - task_registry is now a required parameter of _run_generation_session_workflow, rerun_generation_session, and recover_interrupted_sessions, dropping the "if task_registry is not None" guards. It is always injected from app.state in production; the optionality was only test convenience, so the tests now pass a registry explicitly. - Remove the P10Y cooperative-cancellation plumbing. By the time P10Y runs the code is generated and OUTPUTS_ARCHIVED (Steel Commandment XI), task.cancel() still interrupts its poll sleeps on the owning pod, and a cross-pod cancel during P10Y lands the session in CANCELLED with no spurious failure notification (fail() rejects the terminal state before it can notify). Reverts the extra params on trigger_and_poll_p10y_metrics / _process_single_workspace, the generation_args additions, the top-level check, and the estimation parallel executor's cancellation branch. --- backend/app/api/v1/generation_sessions.py | 8 +++---- .../app/jobs/shutdown_interrupted_recovery.py | 6 ++--- backend/app/services/claude_code.py | 7 +++--- .../services/generation_workflow_runner.py | 8 +++---- backend/app/services/p10y/p10y_lib.py | 10 -------- backend/app/services/parallel_executor.py | 5 ---- backend/app/state/cancellation.py | 9 +++++++ .../multi_workspace_estimation_p10y.py | 20 +++++----------- .../test_shutdown_interrupted_recovery.py | 24 +++++++++++-------- .../test_generation_workflow_runner.py | 17 ++++++------- backend/test/state/test_cancellation.py | 6 +++++ 11 files changed, 57 insertions(+), 63 deletions(-) diff --git a/backend/app/api/v1/generation_sessions.py b/backend/app/api/v1/generation_sessions.py index 03cb041..a85159b 100644 --- a/backend/app/api/v1/generation_sessions.py +++ b/backend/app/api/v1/generation_sessions.py @@ -1084,7 +1084,7 @@ async def _run_generation_session_workflow( spec_path: str, session, generation_session_service, - task_registry: GenerationTaskRegistry | None = None, + task_registry: GenerationTaskRegistry, ) -> None: """Background task: run app generation + P10Y estimation, release session slot when done.""" agent_logger = create_agent_logger( @@ -1099,8 +1099,7 @@ async def _run_generation_session_workflow( normalized_src_dir = generation_session_params.get("src_dir", "src") try: - if task_registry is not None: - task_registry.register_current_task(generation_id) + task_registry.register_current_task(generation_id) async with session.task_slot(generation_id=generation_id): logger.info("Starting app generation + P10Y workflow for %s", generation_id) agent_logger.info("Starting app generation + P10Y workflow for %s", generation_id) @@ -1137,8 +1136,7 @@ async def _run_generation_session_workflow( # else → fail (FAILED). GenerationCancelledError is handled as a silent no-op there. await _handle_workflow_exception(e, generation_id, generation_session_service) finally: - if task_registry is not None: - task_registry.deregister_task(generation_id) + task_registry.deregister_task(generation_id) TelemetryContext.set_agent_query_totals_handler(None) diff --git a/backend/app/jobs/shutdown_interrupted_recovery.py b/backend/app/jobs/shutdown_interrupted_recovery.py index e6c5dd9..5a52185 100644 --- a/backend/app/jobs/shutdown_interrupted_recovery.py +++ b/backend/app/jobs/shutdown_interrupted_recovery.py @@ -61,15 +61,15 @@ async def _workspaces_reusable(db, workspace_ids: list, generation_id: str) -> b async def recover_interrupted_sessions( db, generation_session_service: GenerationSessionService, + task_registry: GenerationTaskRegistry, *, window_minutes: int | None = None, - task_registry: GenerationTaskRegistry | None = None, ) -> list[str]: """Auto-retry shutdown-interrupted sessions found within the recovery window. One-shot (not a loop). Returns the generation IDs it re-fired. Each candidate is - independently guarded; one failure never blocks the others. ``task_registry`` (when - provided) is threaded into each re-fired run so it stays cancellable. + independently guarded; one failure never blocks the others. ``task_registry`` is + threaded into each re-fired run so it stays cancellable. """ if not settings.AUTO_RECOVER_INTERRUPTED_SESSIONS: logger.info("shutdown_interrupted_recovery: disabled by config — skipping") diff --git a/backend/app/services/claude_code.py b/backend/app/services/claude_code.py index 884c39c..f460c63 100644 --- a/backend/app/services/claude_code.py +++ b/backend/app/services/claude_code.py @@ -1200,9 +1200,10 @@ async def execute_all_phases( consecutive_errors = 0 for phase_num in range(phase_start, phase_end + 1): # Cooperative cancellation: stop between phases if the user cancelled - # (cross-pod-safe fallback to the local task.cancel()). - if db_adapter is not None and request.generation_id: - await raise_if_cancelled(db_adapter, request.generation_id) + # (cross-pod-safe fallback to the local task.cancel()). raise_if_cancelled + # no-ops when db_adapter is None (only in DB-less unit tests; every real run + # wires a generation_session_service). + await raise_if_cancelled(db_adapter, request.generation_id) phase_info = planning_data.phases[phase_num - 1] diff --git a/backend/app/services/generation_workflow_runner.py b/backend/app/services/generation_workflow_runner.py index 0016645..23a9222 100644 --- a/backend/app/services/generation_workflow_runner.py +++ b/backend/app/services/generation_workflow_runner.py @@ -217,8 +217,8 @@ async def rerun_generation_session( generation_id: str, *, generation_session_service: GenerationSessionService, + task_registry: GenerationTaskRegistry, user_email: str | None = None, - task_registry: GenerationTaskRegistry | None = None, ) -> None: """Re-fire a generation session from its persisted parameters. @@ -252,8 +252,7 @@ async def rerun_generation_session( } try: - if task_registry is not None: - task_registry.register_current_task(generation_id) + task_registry.register_current_task(generation_id) TelemetryContext.set_user_context( user_email=user_email, generation_id=generation_id, @@ -312,6 +311,5 @@ async def rerun_generation_session( # GenerationCancelledError is handled as a silent no-op there. await _handle_workflow_exception(e, generation_id, generation_session_service) finally: - if task_registry is not None: - task_registry.deregister_task(generation_id) + task_registry.deregister_task(generation_id) TelemetryContext.set_agent_query_totals_handler(None) diff --git a/backend/app/services/p10y/p10y_lib.py b/backend/app/services/p10y/p10y_lib.py index 4cf8637..62473e2 100644 --- a/backend/app/services/p10y/p10y_lib.py +++ b/backend/app/services/p10y/p10y_lib.py @@ -8,7 +8,6 @@ from typing import Any, Dict, List, Optional from app.services.p10y.p10y_api_client import P10YInternalAPIClient -from app.state.cancellation import raise_if_cancelled # Commits whose first line starts with this prefix are excluded from P10Y / component breakdown # (e.g. user-provided initial seed: SKIP_initial_user_source, SKIP_generation_baseline). @@ -371,8 +370,6 @@ async def trigger_and_poll_p10y_metrics( organisation_id: int, workspace_name: str, logger: logging.Logger, - db_adapter=None, - generation_id: Optional[str] = None, ) -> None: """ Trigger P10Y metrics calculation and poll until processing is complete. @@ -383,8 +380,6 @@ async def trigger_and_poll_p10y_metrics( organisation_id: P10Y organisation ID workspace_name: Workspace name for logging logger: Logger instance - db_adapter: Optional state-machine DB adapter for cooperative cancellation checks - generation_id: Optional generation id for cooperative cancellation checks """ # Trigger P10Y metrics calculation logger.info(f"Triggering P10Y metrics for repository {repository_id}") @@ -401,11 +396,6 @@ async def trigger_and_poll_p10y_metrics( poll_limit = 10 while not all_commits_processed and poll_counter < poll_limit: - # Cooperative cancellation: stop polling if the user cancelled - # (cross-pod-safe fallback to the local task.cancel()). - if db_adapter is not None and generation_id: - await raise_if_cancelled(db_adapter, generation_id) - logger.info(f"Polling attempt {poll_counter + 1} / {poll_limit}") commit_stats_polled = await client.get_commit_stats( diff --git a/backend/app/services/parallel_executor.py b/backend/app/services/parallel_executor.py index 12648b8..f04070d 100644 --- a/backend/app/services/parallel_executor.py +++ b/backend/app/services/parallel_executor.py @@ -376,11 +376,6 @@ async def process_workspace(workspace: WorkspaceSettings) -> ParallelGenerationR duration_ms=duration_ms ) - except GenerationCancelledError: - # User cancellation aborts the whole estimation fan-out (gather uses - # return_exceptions=False, so this propagates out immediately). - raise - except Exception as e: # Handle other errors duration_ms = (time.time() - start_time) * 1000 diff --git a/backend/app/state/cancellation.py b/backend/app/state/cancellation.py index 64a9593..5e30e05 100644 --- a/backend/app/state/cancellation.py +++ b/backend/app/state/cancellation.py @@ -32,6 +32,12 @@ async def raise_if_cancelled( ) -> None: """Raise :class:`GenerationCancelledError` if the session is CANCELLED. + A ``None`` ``db_adapter`` is a no-op: cancellation can't be observed without a DB, + and in practice that only happens in unit tests that run phases without a + ``generation_session_service`` (every real run wires one). Centralizing the guard + here lets call sites invoke this unconditionally. ``task.cancel()`` remains the + local stop path regardless. + The first check for a generation always reads; thereafter, if called again for the same ``generation_id`` within ``min_interval_s`` seconds, it returns immediately without touching the DB. (``time.monotonic()`` has no fixed epoch, so a "never @@ -40,6 +46,9 @@ async def raise_if_cancelled( ``db_adapter`` is the async state-machine DB adapter (``service.db_adapter`` / ``orchestrator._db``). """ + if db_adapter is None: + return + now = time.monotonic() last = _last_check.get(generation_id) if last is not None and now - last < min_interval_s: diff --git a/backend/app/workflows/multi_workspace_estimation_p10y.py b/backend/app/workflows/multi_workspace_estimation_p10y.py index 904b599..f8377c0 100644 --- a/backend/app/workflows/multi_workspace_estimation_p10y.py +++ b/backend/app/workflows/multi_workspace_estimation_p10y.py @@ -72,7 +72,6 @@ is_skip_mode_enabled, ) from app.services.workflow_stats import workflow_metrics -from app.state.cancellation import raise_if_cancelled from app.state.workspace_models import load_workspace_model_overrides, resolve_workspace_model from app.services.workspace_manager import WorkspaceManager from app.state.db_adapter import COL_GENERATION_SESSIONS, StateMachineDBAdapter @@ -84,8 +83,6 @@ async def _process_single_workspace( client: P10YInternalAPIClient, settings: Settings, logger: logging.Logger, - db_adapter=None, - generation_id: Optional[str] = None, ) -> WorkspaceEstimation | None: """ Process estimation for a single workspace. @@ -144,8 +141,6 @@ async def _process_single_workspace( organisation_id=settings.P10Y_ORGANISATION_ID, workspace_name=workspace.name, logger=logger, - db_adapter=db_adapter, - generation_id=generation_id, ) # Fetch and filter commit stats @@ -516,13 +511,12 @@ async def multi_workspace_estimation_p10y_workflow( # Create workspace manager (needed for parallel executor) workspace_manager = WorkspaceManager(settings, logger) - # Cooperative cancellation: skip the whole estimation phase if the user cancelled - # (cross-pod-safe fallback to the local task.cancel()). Code is already archived by - # this point (STEEL XI), so aborting here is safe. - if db_adapter is not None and request.generation_id: - await raise_if_cancelled(db_adapter, request.generation_id) - - # Process workspaces in parallel (data collection only, no AI agents) + # Process workspaces in parallel (data collection only, no AI agents). + # No cancellation check here: by the time P10Y runs, code generation is done and + # OUTPUTS_ARCHIVED (STEEL XI) has preserved the outputs, and task.cancel() still + # interrupts this phase's sleeps on the owning pod. A cross-pod cancel during P10Y + # simply lets this bounded, read-only phase finish and lands the session in CANCELLED + # without a spurious failure notification (fail() rejects the terminal state). logger.info("Executing estimations in parallel...") parallel_results = await execute_generation_parallel( workspaces=workspaces, @@ -533,8 +527,6 @@ async def multi_workspace_estimation_p10y_workflow( "client": client, "settings": settings, "logger": logger, - "db_adapter": db_adapter, - "generation_id": request.generation_id, }, logger=logger, ) diff --git a/backend/test/jobs/test_shutdown_interrupted_recovery.py b/backend/test/jobs/test_shutdown_interrupted_recovery.py index 615991b..6f21692 100644 --- a/backend/test/jobs/test_shutdown_interrupted_recovery.py +++ b/backend/test/jobs/test_shutdown_interrupted_recovery.py @@ -11,8 +11,12 @@ from app.jobs import shutdown_interrupted_recovery as job_mod from app.jobs.shutdown_interrupted_recovery import recover_interrupted_sessions from app.schemas.generation_workflow_enums import GenerationStatus, WorkspaceStatus +from app.services.generation_task_registry import GenerationTaskRegistry from app.state.transitions import TriggeredBy +# rerun is mocked in every test here, so the registry is only forwarded, never used. +_REGISTRY = GenerationTaskRegistry() + def _seed_interrupted(db, gid, *, failed_minutes_ago=5, workspace_ids=None, code_archived=False): failed_at = datetime.now(timezone.utc) - timedelta(minutes=failed_minutes_ago) @@ -62,7 +66,7 @@ async def test_recovers_marked_session_within_window(db, gss, _patch_spawn, noti _seed_interrupted(db, "gen-1", workspace_ids=["ws-1"]) _seed_allocated_ws(db, "gen-1", ["ws-1"]) - recovered = await recover_interrupted_sessions(db, gss) + recovered = await recover_interrupted_sessions(db, gss, _REGISTRY) assert recovered == ["gen-1"] # atomic reset happened: FAILED → PENDING, marker cleared, retry_count incremented @@ -86,7 +90,7 @@ async def test_no_notification_when_nothing_recovered(db, gss, notify_mock): "failed_at": datetime.now(timezone.utc), "state_history": [], }) - recovered = await recover_interrupted_sessions(db, gss) + recovered = await recover_interrupted_sessions(db, gss, _REGISTRY) assert recovered == [] notify_mock.assert_not_called() @@ -99,7 +103,7 @@ async def test_ignores_genuine_failure_without_marker(db, gss, _patch_spawn): "failed_at": datetime.now(timezone.utc), "state_history": [], }) - recovered = await recover_interrupted_sessions(db, gss) + recovered = await recover_interrupted_sessions(db, gss, _REGISTRY) assert recovered == [] rerun.assert_not_called() # untouched — still FAILED @@ -112,7 +116,7 @@ async def test_ignores_marked_session_outside_window(db, gss, _patch_spawn): _seed_interrupted(db, "gen-old", failed_minutes_ago=120, workspace_ids=["ws-1"]) _seed_allocated_ws(db, "gen-old", ["ws-1"]) - recovered = await recover_interrupted_sessions(db, gss, window_minutes=30) + recovered = await recover_interrupted_sessions(db, gss, _REGISTRY, window_minutes=30) assert recovered == [] rerun.assert_not_called() @@ -129,7 +133,7 @@ async def test_ignores_cancelled_session(db, gss, _patch_spawn): "cancelled_at": datetime.now(timezone.utc), "state_history": [], }) - recovered = await recover_interrupted_sessions(db, gss) + recovered = await recover_interrupted_sessions(db, gss, _REGISTRY) assert recovered == [] rerun.assert_not_called() assert db.get_generation_session_data("gen-cancelled")["status"] == GenerationStatus.CANCELLED @@ -145,7 +149,7 @@ async def test_non_failed_marked_session_ignored(db, gss, _patch_spawn): "failed_at": datetime.now(timezone.utc), "state_history": [], }) - recovered = await recover_interrupted_sessions(db, gss) + recovered = await recover_interrupted_sessions(db, gss, _REGISTRY) assert recovered == [] rerun.assert_not_called() @@ -158,7 +162,7 @@ async def test_max_retries_exhausted_left_failed(db, gss, _patch_spawn): # already at the limit db.update("generation_sessions", "gen-max", {"retry_count": 3, "max_retries": 3}) - recovered = await recover_interrupted_sessions(db, gss) + recovered = await recover_interrupted_sessions(db, gss, _REGISTRY) assert recovered == [] rerun.assert_not_called() @@ -173,7 +177,7 @@ async def test_non_reusable_workspaces_skipped(db, gss, _patch_spawn): # workspace is CLEANING, not ALLOCATED → not reusable db.seed_workspace("ws-1", {"status": WorkspaceStatus.CLEANING.value, "locked_by": "gen-ws"}) - recovered = await recover_interrupted_sessions(db, gss) + recovered = await recover_interrupted_sessions(db, gss, _REGISTRY) assert recovered == [] rerun.assert_not_called() @@ -188,7 +192,7 @@ async def test_disabled_by_config(db, gss, _patch_spawn): _seed_allocated_ws(db, "gen-1", ["ws-1"]) with patch.object(job_mod.settings, "AUTO_RECOVER_INTERRUPTED_SESSIONS", False): - recovered = await recover_interrupted_sessions(db, gss) + recovered = await recover_interrupted_sessions(db, gss, _REGISTRY) assert recovered == [] rerun.assert_not_called() @@ -200,7 +204,7 @@ async def test_uses_shutdown_recovery_triggered_by(db, gss, _patch_spawn): _seed_interrupted(db, "gen-1", workspace_ids=["ws-1"]) _seed_allocated_ws(db, "gen-1", ["ws-1"]) - await recover_interrupted_sessions(db, gss) + await recover_interrupted_sessions(db, gss, _REGISTRY) history = db.get_generation_session_data("gen-1")["state_history"] assert history[-1]["triggered_by"] == TriggeredBy.SHUTDOWN_RECOVERY diff --git a/backend/test/services/test_generation_workflow_runner.py b/backend/test/services/test_generation_workflow_runner.py index 95a55ae..148104e 100644 --- a/backend/test/services/test_generation_workflow_runner.py +++ b/backend/test/services/test_generation_workflow_runner.py @@ -10,6 +10,7 @@ import pytest from app.services import generation_workflow_runner as runner +from app.services.generation_task_registry import GenerationTaskRegistry from app.services.generation_workflow_runner import rerun_generation_session from app.state.api_key_session_concurrency import OperationKind, SessionBeginOutcome @@ -57,7 +58,7 @@ def _patch_workflow_steps(): async def test_happy_path_registers_and_completes(): gss, sessions = _make_gss(parameters={"spec_path": "specs/", "outputs_dir": "docs"}) - await rerun_generation_session("gen-1", generation_session_service=gss, user_email="u@x.io") + await rerun_generation_session("gen-1", generation_session_service=gss, user_email="u@x.io", task_registry=GenerationTaskRegistry()) # Session re-registered in active_generation_sessions so the TUI shows the resumed run. sessions.try_begin_for_generation.assert_awaited_once_with( @@ -75,7 +76,7 @@ async def test_already_held_refreshes_lease(): begin_outcome=SessionBeginOutcome.ALREADY_HELD, ) - await rerun_generation_session("gen-held", generation_session_service=gss) + await rerun_generation_session("gen-held", generation_session_service=gss, task_registry=GenerationTaskRegistry()) sessions.refresh_lease.assert_awaited_once_with( generation_id="gen-held", operation=OperationKind.GENERATION @@ -92,7 +93,7 @@ async def test_refresh_lease_failure_still_completes(): ) sessions.refresh_lease = AsyncMock(side_effect=RuntimeError("firestore blip")) - await rerun_generation_session("gen-held", generation_session_service=gss) + await rerun_generation_session("gen-held", generation_session_service=gss, task_registry=GenerationTaskRegistry()) sessions.refresh_lease.assert_awaited_once() gss.complete_generation_session.assert_awaited_once() @@ -107,7 +108,7 @@ async def test_at_capacity_still_runs_workflow(): begin_outcome=SessionBeginOutcome.AT_CAPACITY, ) - await rerun_generation_session("gen-cap", generation_session_service=gss) + await rerun_generation_session("gen-cap", generation_session_service=gss, task_registry=GenerationTaskRegistry()) sessions.try_begin_for_generation.assert_awaited_once() sessions.task_slot.assert_called_once_with(generation_id="gen-cap") @@ -123,7 +124,7 @@ async def test_unresolvable_registration_still_runs_workflow(): begin_outcome=None, ) - await rerun_generation_session("gen-none", generation_session_service=gss) + await rerun_generation_session("gen-none", generation_session_service=gss, task_registry=GenerationTaskRegistry()) sessions.try_begin_for_generation.assert_awaited_once() sessions.task_slot.assert_called_once_with(generation_id="gen-none") @@ -136,7 +137,7 @@ async def test_failure_routed_not_raised(): gss.start_generation_session.side_effect = RuntimeError("boom") # failure routes through _handle_workflow_exception (fail_generation_session), never propagates - await rerun_generation_session("gen-err", generation_session_service=gss) + await rerun_generation_session("gen-err", generation_session_service=gss, task_registry=GenerationTaskRegistry()) sessions.try_begin_for_generation.assert_awaited_once() sessions.task_slot.assert_called_once_with(generation_id="gen-err") @@ -147,7 +148,7 @@ async def test_failure_routed_not_raised(): async def test_missing_parameters_fail_routed_not_raised(): gss, sessions = _make_gss(parameters={}) # no spec_path - await rerun_generation_session("gen-bad", generation_session_service=gss) + await rerun_generation_session("gen-bad", generation_session_service=gss, task_registry=GenerationTaskRegistry()) # registration + start_generation_session never reached; failure handled, not raised sessions.try_begin_for_generation.assert_not_awaited() @@ -176,6 +177,6 @@ async def test_cancellation_during_rerun_not_failed(): gss, sessions = _make_gss(parameters={"spec_path": "specs/", "outputs_dir": "docs"}) gss.start_generation_session.side_effect = GenerationCancelledError("gen-cancel") - await rerun_generation_session("gen-cancel", generation_session_service=gss) + await rerun_generation_session("gen-cancel", generation_session_service=gss, task_registry=GenerationTaskRegistry()) gss.fail_generation_session.assert_not_called() diff --git a/backend/test/state/test_cancellation.py b/backend/test/state/test_cancellation.py index 395c447..85419bd 100644 --- a/backend/test/state/test_cancellation.py +++ b/backend/test/state/test_cancellation.py @@ -50,6 +50,12 @@ async def test_noop_when_missing_doc(): await raise_if_cancelled(adapter, "est-1", min_interval_s=0.0) +@pytest.mark.asyncio +async def test_noop_when_db_adapter_is_none(): + # No DB wired (DB-less unit-test path) — must be a silent no-op, never raise. + await raise_if_cancelled(None, "est-1", min_interval_s=0.0) + + @pytest.mark.asyncio async def test_throttles_repeated_reads(): adapter = FakeAdapter(status=GenerationStatus.RUNNING.value)