feat(generation): support user cancellation of a running session

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:

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

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)

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

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:

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.

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,