feat(scribe): Add STARTUP state to Lambda to prevent race condition. (#4613)

Author: eyw520

servers/fai-lambda/fai-scribe/src/handler.py

@@ -15,9 +15,7 @@
 )
 from .utils.agent import (
     SessionInterruptedError,
-    check_if_interrupted,
     run_editing_session,
-    update_session_status,
 )
 from .utils.git import setup_editing_repo
 
@@ -77,11 +75,6 @@ async def handle_editing_request(
         )
         LOGGER.info(f"Repository ready at: {repo_path}")
 
-        if await check_if_interrupted(editing_id):
-            LOGGER.warning(f"Session interrupted after git operations: {editing_id}")
-            await update_session_status(editing_id, "waiting")
-            raise SessionInterruptedError(f"Editing session {editing_id} was interrupted")
-
         try:
             session_id, pr_url = await run_editing_session(
                 repo_path=repo_path,

servers/fai-lambda/fai-scribe/src/settings.py

@@ -6,6 +6,7 @@
 
 load_dotenv()
 LOGGER = logging.getLogger()
+LOGGER.setLevel(logging.INFO)
 
 
 class Settings:

servers/fai-lambda/fai-scribe/src/utils/agent.py

@@ -96,13 +96,6 @@ async def run_editing_session(
     setup_persistent_claude_storage(repo_path)
     configure_git_auth(repo_path)
 
-    await update_session_status(editing_id, "active")
-
-    if await check_if_interrupted(editing_id):
-        LOGGER.warning(f"Session interrupted before starting: {editing_id}")
-        await update_session_status(editing_id, "waiting")
-        raise SessionInterruptedError(f"Editing session {editing_id} was interrupted")
-
     if existing_pr_url:
         full_prompt = f"""{user_prompt}
 
@@ -144,6 +137,10 @@ async def run_editing_session(
     async with ClaudeSDKClient(options=options) as client:
         await client.query(full_prompt)
 
+        if resume_session_id:
+            await update_session_status(editing_id, "active")
+            LOGGER.info(f"Transitioned resumed session {editing_id} from STARTUP to ACTIVE")
+
         async for message in client.receive_response():
             if hasattr(message, "subtype") and message.subtype == "init":
                 if hasattr(message, "data") and isinstance(message.data, dict):
@@ -152,10 +149,13 @@ async def run_editing_session(
                         session_id = init_session_id
                         LOGGER.info(f"Captured session_id from init message: {session_id}")
                         await update_session_metadata(editing_id, session_id=session_id)
+                        await update_session_status(editing_id, "active")
+                        LOGGER.info(f"Transitioned new session {editing_id} from STARTUP to ACTIVE")
 
             if await check_if_interrupted(editing_id):
                 LOGGER.warning(f"Session interrupted: {editing_id}")
-                await update_session_status(editing_id, "waiting")
+                if session_id is not None:
+                    await update_session_status(editing_id, "waiting")
                 raise SessionInterruptedError(f"Editing session {editing_id} was interrupted")
 
             if isinstance(message, AssistantMessage):

servers/fai/alembic/versions/af951c45da91_add_startup_status_to_editing_sessions.py

@@ -0,0 +1,35 @@
+"""add_startup_status_to_editing_sessions
+
+Revision ID: af951c45da91
+Revises: 2d743e49aaa1
+Create Date: 2025-10-31 11:19:09.069432
+
+"""
+
+from typing import (
+    Sequence,
+    Union,
+)
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "af951c45da91"
+down_revision: Union[str, Sequence[str], None] = "2d743e49aaa1"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # Add 'STARTUP' value to the editingsessionstatus enum
+    op.execute("ALTER TYPE editingsessionstatus ADD VALUE IF NOT EXISTS 'STARTUP'")
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # PostgreSQL does not support removing enum values easily
+    # This would require recreating the enum type, which is complex
+    # For now, we'll leave the enum value in place
+    pass

servers/fai/src/fai/models/db/editing_session_db.py

@@ -31,7 +31,7 @@ class EditingSessionDb(Base):
     id = Column(String, primary_key=True)
 
     session_id = Column(String, nullable=True)
-    status = Column(Enum(EditingSessionStatus), nullable=False, default=EditingSessionStatus.WAITING)
+    status = Column(Enum(EditingSessionStatus), nullable=False, default=EditingSessionStatus.STARTUP)
 
     # Repository information
     repository = Column(String, nullable=False)

servers/fai/src/fai/models/types/editing_session_types.py

@@ -7,6 +7,7 @@
 class EditingSessionStatus(str, Enum):
     """Status of an editing session."""
 
+    STARTUP = "startup"
     WAITING = "waiting"
     ACTIVE = "active"
     INTERRUPTED = "interrupted"

servers/fai/src/fai/routes/editing_sessions.py

@@ -55,7 +55,7 @@ async def create_editing_session(
             base_branch=request.base_branch,
             working_branch=working_branch,
             pr_url=None,
-            status=EditingSessionStatus.WAITING,
+            status=EditingSessionStatus.STARTUP,
             created_at=datetime.now(UTC),
             updated_at=datetime.now(UTC),
         )

servers/fai/src/fai/utils/slack/edit_handler.py

@@ -96,6 +96,30 @@ async def get_editing_session_status(editing_id: str) -> EditingSessionStatus |
 async def interrupt_editing_session(editing_id: str) -> bool:
     """Interrupt an active editing session. Returns True if successful."""
     try:
+        status = await get_editing_session_status(editing_id)
+
+        if status == EditingSessionStatus.STARTUP:
+            logger.info(f"Session {editing_id} is in STARTUP state, waiting for it to become ACTIVE")
+            max_wait: int = 30
+            poll_interval: float = 0.5
+            elapsed: float = 0
+
+            while elapsed < max_wait:
+                await asyncio.sleep(poll_interval)
+                elapsed += poll_interval
+
+                status = await get_editing_session_status(editing_id)
+                if status == EditingSessionStatus.ACTIVE:
+                    logger.info(f"Session {editing_id} transitioned to ACTIVE, proceeding with interruption")
+                    break
+                elif status not in [EditingSessionStatus.STARTUP, EditingSessionStatus.ACTIVE]:
+                    logger.warning(f"Session {editing_id} in unexpected state {status}, aborting interruption")
+                    return False
+
+            if status != EditingSessionStatus.ACTIVE:
+                logger.warning(f"Timeout waiting for session {editing_id} to become ACTIVE")
+                return False
+
         async with httpx.AsyncClient(timeout=10.0) as client:
             response = await client.post(f"{CONFIG.FAI_SERVER_URL}/editing-sessions/{editing_id}/interrupt")
 
@@ -176,27 +200,6 @@ async def create_editing_session(repository: str, base_branch: str = "main") ->
         return None
 
 
-async def mark_session_as_active(editing_id: str) -> bool:
-    """Mark an editing session as ACTIVE before Lambda invocation."""
-    try:
-        async with httpx.AsyncClient(timeout=10.0) as client:
-            response = await client.put(
-                f"{CONFIG.FAI_SERVER_URL}/editing-sessions/{editing_id}",
-                json={"status": "active"},
-            )
-
-            if response.status_code != 200:
-                logger.error(f"Failed to mark session as active: {response.status_code} - {response.text}")
-                return False
-
-            logger.info(f"Marked editing session {editing_id} as ACTIVE")
-            return True
-
-    except Exception as e:
-        logger.error(f"Error marking session as active: {e}", exc_info=True)
-        return False
-
-
 async def invoke_editing_lambda(
     prompt: str,
     domain: str,

servers/fai/src/fai/utils/slack/message_handler.py

@@ -1,3 +1,4 @@
+import asyncio
 from dataclasses import dataclass
 from datetime import (
     UTC,
@@ -43,7 +44,6 @@
     get_or_create_editing_session_for_thread,
     interrupt_editing_session,
     invoke_editing_lambda,
-    mark_session_as_active,
     store_editing_session_for_thread,
     wait_for_interruption,
 )
@@ -401,7 +401,50 @@ async def handle_slack_message(
                     query_id=None,
                     user_id=context.user,
                 )
-            elif session_status and session_status.value in ["active", "waiting"]:
+            elif session_status and session_status.value == "startup":
+                LOGGER.info(f"Session {existing_editing_id} is in STARTUP state, waiting for it to become ACTIVE")
+                max_wait: int = 30
+                poll_interval: float = 0.5
+                elapsed: float = 0
+
+                while elapsed < max_wait:
+                    await asyncio.sleep(poll_interval)
+                    elapsed += poll_interval
+
+                    session_status = await get_editing_session_status(existing_editing_id)
+                    if session_status and session_status.value in ["active", "waiting"]:
+                        LOGGER.info(
+                            f"Session {existing_editing_id} transitioned to {session_status.value}, "
+                            f"proceeding with interruption"
+                        )
+                        break
+                    elif session_status and session_status.value not in ["startup", "active", "waiting"]:
+                        LOGGER.warning(
+                            f"Session {existing_editing_id} in unexpected state {session_status.value}, aborting"
+                        )
+                        error_text = "❌ Session is in an unexpected state. Please try again."
+                        return SlackMessageResponse(
+                            response_text=error_text,
+                            channel=context.channel,
+                            thread_ts=context.thread_ts,
+                            bot_token=integration.slack_bot_token,
+                            query_id=None,
+                            user_id=context.user,
+                        )
+
+                if not session_status or session_status.value not in ["active", "waiting"]:
+                    LOGGER.warning(f"Timeout waiting for session {existing_editing_id} to become ACTIVE/WAITING")
+                    error_text = "⏳ Previous edit session is still starting up. Please wait a moment and try again."
+                    return SlackMessageResponse(
+                        response_text=error_text,
+                        channel=context.channel,
+                        thread_ts=context.thread_ts,
+                        bot_token=integration.slack_bot_token,
+                        query_id=None,
+                        user_id=context.user,
+                    )
+
+            if session_status and session_status.value in ["active", "waiting"]:
                 LOGGER.info(f"Interrupting {session_status.value} session {existing_editing_id} for new edit request")
 
                 interrupt_success = await interrupt_editing_session(existing_editing_id)
@@ -508,13 +551,6 @@ async def handle_slack_message(
             except Exception as e:
                 LOGGER.error(f"Failed to store editing session mapping: {e}", exc_info=True)
 
-        if existing_editing_id:
-            mark_success = await mark_session_as_active(existing_editing_id)
-            if not mark_success:
-                LOGGER.warning(
-                    f"Failed to mark session {existing_editing_id} as ACTIVE, but proceeding with Lambda invocation"
-                )
-
         result = await invoke_editing_lambda(
             prompt=context.text,
             domain=domain_to_use,