diff --git a/backend/app/api/v1/generation_sessions.py b/backend/app/api/v1/generation_sessions.py index 9b210da..a85159b 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.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, @@ -249,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 # ============================================================ @@ -775,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. @@ -815,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, ) ) @@ -850,14 +858,21 @@ 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. + 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 +881,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 = 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( @@ -1052,6 +1084,7 @@ async def _run_generation_session_workflow( spec_path: str, session, generation_session_service, + task_registry: GenerationTaskRegistry, ) -> None: """Background task: run app generation + P10Y estimation, release session slot when done.""" agent_logger = create_agent_logger( @@ -1066,6 +1099,7 @@ async def _run_generation_session_workflow( normalized_src_dir = generation_session_params.get("src_dir", "src") try: + 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 +1133,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: + task_registry.deregister_task(generation_id) TelemetryContext.set_agent_query_totals_handler(None) @@ -1296,6 +1331,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), ): """ @@ -1489,6 +1525,7 @@ async def run_generation_session( spec_path=spec_path, session=session, generation_session_service=generation_session_service, + task_registry=task_registry, )) return RunGenerationSessionResponse( @@ -1562,6 +1599,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/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..5a52185 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 ( @@ -60,13 +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, ) -> 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`` 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/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..f460c63 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,12 @@ 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()). 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] 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..e8e31c2 --- /dev/null +++ b/backend/app/services/generation_task_registry.py @@ -0,0 +1,58 @@ +""" +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. + +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__) + + +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 0f22aba..23a9222 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.generation_task_registry import GenerationTaskRegistry 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 @@ -205,6 +217,7 @@ async def rerun_generation_session( generation_id: str, *, generation_session_service: GenerationSessionService, + task_registry: GenerationTaskRegistry, user_email: str | None = None, ) -> None: """Re-fire a generation session from its persisted parameters. @@ -239,6 +252,7 @@ async def rerun_generation_session( } try: + task_registry.register_current_task(generation_id) TelemetryContext.set_user_context( user_email=user_email, generation_id=generation_id, @@ -294,6 +308,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: + 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..62473e2 100644 --- a/backend/app/services/p10y/p10y_lib.py +++ b/backend/app/services/p10y/p10y_lib.py @@ -373,7 +373,7 @@ async def trigger_and_poll_p10y_metrics( ) -> None: """ Trigger P10Y metrics calculation and poll until processing is complete. - + Args: client: P10Y API client repository_id: P10Y repository ID @@ -397,7 +397,7 @@ async def trigger_and_poll_p10y_metrics( while not all_commits_processed and poll_counter < poll_limit: 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..f04070d 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,7 @@ async def process_workspace(workspace: WorkspaceSettings) -> ParallelGenerationR warning="Commit files not found - workspace may not have completed code generation", duration_ms=duration_ms ) - + 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..5e30e05 --- /dev/null +++ b/backend/app/state/cancellation.py @@ -0,0 +1,65 @@ +""" +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. + + 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 + 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``). + """ + 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: + 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..f8377c0 100644 --- a/backend/app/workflows/multi_workspace_estimation_p10y.py +++ b/backend/app/workflows/multi_workspace_estimation_p10y.py @@ -511,7 +511,12 @@ async def multi_workspace_estimation_p10y_workflow( # Create workspace manager (needed for parallel executor) workspace_manager = WorkspaceManager(settings, logger) - # 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, diff --git a/backend/test/api/test_generation_sessions.py b/backend/test/api/test_generation_sessions.py index dd72d76..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,11 +812,12 @@ 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 ): - """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 +835,50 @@ 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"} ) - + 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_task_registry.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..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,13 +116,29 @@ 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() 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, _REGISTRY) + 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).""" @@ -129,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() @@ -142,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() @@ -157,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() @@ -172,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() @@ -184,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_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..86b7cfd --- /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.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() + + await asyncio.create_task(body()) + + +@pytest.mark.asyncio +async def test_cancel_task_cancels_running_task(): + reg = GenerationTaskRegistry() + 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(): + reg = GenerationTaskRegistry() + 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. + reg = GenerationTaskRegistry() + sentinel = object() + 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.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 + 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..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,9 +148,35 @@ 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() 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, 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 new file mode 100644 index 0000000..85419bd --- /dev/null +++ b/backend/test/state/test_cancellation.py @@ -0,0 +1,65 @@ +""" +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_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) + # 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.""" 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