Allow access to task arguments while pending

Author: mxsrc

src/api/organization/project/branch/__init__.py

@@ -96,9 +96,8 @@
 from .auth import api as auth_api
 from .control_tasks import (
     _CONTROL_TO_POWER_STATE,
-    _CONTROL_TRANSITION_INITIAL,
     dispatch_control,
-    perform_control,
+    get_control_in_progress_status,
 )
 from .resize_tasks import dispatch_resize
 from .tasks import task_api
@@ -234,29 +233,34 @@ def _should_update_branch_status(
 def _adjust_derived_status_for_stuck_creation(
     branch: Branch, current: BranchServiceStatus, derived: BranchServiceStatus
 ) -> BranchServiceStatus:
-    if derived != BranchServiceStatus.STOPPED:
+    if derived not in {BranchServiceStatus.STOPPED, BranchServiceStatus.ERROR}:
+        return derived
+    if current not in {BranchServiceStatus.CREATING, BranchServiceStatus.STARTING}:
         return derived
 
     status_timestamp = branch.status_updated_at or branch.created_datetime
     elapsed = datetime.now(UTC) - status_timestamp
 
-    if current == BranchServiceStatus.CREATING and elapsed >= _CREATING_STATUS_ERROR_GRACE_PERIOD:
-        logger.warning(
-            "Branch %s still CREATING after %s with STOPPED services; marking ERROR",
-            branch.id,
-            elapsed,
-        )
-        return BranchServiceStatus.ERROR
+    if current == BranchServiceStatus.CREATING:
+        if derived == BranchServiceStatus.STOPPED and elapsed >= _CREATING_STATUS_ERROR_GRACE_PERIOD:
+            logger.warning(
+                "Branch %s still CREATING after %s with STOPPED services; marking ERROR",
+                branch.id,
+                elapsed,
+            )
+            return BranchServiceStatus.ERROR
+        return derived
 
-    if current == BranchServiceStatus.STARTING and elapsed >= _STARTING_STATUS_ERROR_GRACE_PERIOD:
+    # current == STARTING
+    if elapsed >= _STARTING_STATUS_ERROR_GRACE_PERIOD:
         logger.warning(
-            "Branch %s still STARTING after %s with STOPPED services; marking ERROR",
+            "Branch %s still STARTING after %s; marking ERROR",
             branch.id,
             elapsed,
         )
         return BranchServiceStatus.ERROR
-
-    return derived
+    # Within grace period: suppress both STOPPED and ERROR — stay STARTING
+    return current
 
 
 async def refresh_branch_status(branch_id: Identifier) -> BranchServiceStatus:
@@ -280,10 +284,9 @@ async def refresh_branch_status(branch_id: Identifier) -> BranchServiceStatus:
 def _active_task_status(branch: Branch) -> BranchServiceStatus | None:
     """Return the in-progress status for the running task, or None if idle."""
     if branch.control_task_id is not None:
-        result = perform_control.AsyncResult(str(branch.control_task_id))
-        action = (result.kwargs or {}).get("action")
-        if action in _CONTROL_TRANSITION_INITIAL:
-            return _CONTROL_TRANSITION_INITIAL[action]
+        status = get_control_in_progress_status(branch.control_task_id)
+        if status is not None:
+            return status
     if branch.resize_task_id is not None:
         return BranchServiceStatus.RESIZING
     return None

src/api/organization/project/branch/control_tasks.py

@@ -6,6 +6,7 @@
 
 import asyncio
 import logging
+from typing import Any
 from uuid import UUID
 
 from asgiref.sync import async_to_sync
@@ -51,6 +52,24 @@
 _TIMEOUT_SEC = 600  # 10 minutes
 
 
+class _TaskRequest:
+    """Minimal Celery request stub used to pre-populate the result backend.
+
+    Celery only stores task kwargs in the backend when ``store_result`` is
+    called with a request object (requires ``result_extended = True``).  By
+    constructing this stub at dispatch time we make ``AsyncResult.kwargs``
+    readable immediately — before the worker has even picked up the message.
+    """
+
+    __slots__ = ("id", "task", "args", "kwargs")
+
+    def __init__(self, task_id: str, kwargs: dict[str, Any]) -> None:
+        self.id = task_id
+        self.task: str = perform_control.name
+        self.args: tuple[()] = ()
+        self.kwargs = kwargs
+
+
 async def _async_perform_control(branch_id: str, action: str) -> dict:
     ulid = ULID.from_str(branch_id)
     namespace, name = get_autoscaler_vm_identity(ulid)
@@ -103,7 +122,24 @@ def perform_control(branch_id: str, action: str) -> dict:
     return async_to_sync(_async_perform_control_with_error_handling)(branch_id, action)
 
 
+def get_control_in_progress_status(task_id: UUID) -> BranchServiceStatus | None:
+    """Return the in-progress status for a running control task, or None."""
+    result = perform_control.AsyncResult(str(task_id))
+    action = (result.kwargs or {}).get("action")
+    if action in _CONTROL_TRANSITION_INITIAL:
+        return _CONTROL_TRANSITION_INITIAL[action]
+    return None
+
+
 def dispatch_control(branch_id: str, action: str) -> UUID:
-    """Dispatch perform_control asynchronously; return the Celery task UUID."""
-    result = perform_control.apply_async(kwargs={"branch_id": branch_id, "action": action})
-    return UUID(result.id)
+    """Dispatch perform_control asynchronously; return the Celery task UUID.
+
+    The task kwargs are pre-stored in the result backend immediately so that
+    ``AsyncResult.kwargs`` (and therefore ``get_control_in_progress_status``)
+    works even before the worker picks up the message.
+    """
+    kwargs = {"branch_id": branch_id, "action": action}
+    result = perform_control.apply_async(kwargs=kwargs)
+    task_id = str(result.id)
+    app.backend.store_result(task_id, None, "PENDING", request=_TaskRequest(task_id, kwargs))
+    return UUID(task_id)