feat: add node-level timeouts to prevent stuck queued states

state-manager/app/config/settings.py

@@ -13,6 +13,7 @@ class Settings(BaseModel):
     state_manager_secret: str = Field(..., description="Secret key for API authentication")
     secrets_encryption_key: str = Field(..., description="Key for encrypting secrets")
     trigger_workers: int = Field(default=1, description="Number of workers to run the trigger cron")
+    node_timeout_minutes: int = Field(default=60, gt=0, description="Timeout in minutes for nodes stuck in QUEUED status")
     trigger_retention_hours: int = Field(default=720, description="Number of hours to retain completed/failed triggers before cleanup")
 
     @classmethod
@@ -22,6 +23,7 @@ def from_env(cls) -> "Settings":
             mongo_database_name=os.getenv("MONGO_DATABASE_NAME", "exosphere-state-manager"), # type: ignore
             state_manager_secret=os.getenv("STATE_MANAGER_SECRET"), # type: ignore
             secrets_encryption_key=os.getenv("SECRETS_ENCRYPTION_KEY"), # type: ignore
+            node_timeout_minutes=int(os.getenv("NODE_TIMEOUT_MINUTES", 60)), # type: ignore
             trigger_workers=int(os.getenv("TRIGGER_WORKERS", 1)), # type: ignore
             trigger_retention_hours=int(os.getenv("TRIGGER_RETENTION_HOURS", 720)) # type: ignore
         )

state-manager/app/controller/enqueue_states.py

@@ -7,24 +7,47 @@
 from ..models.state_status_enum import StateStatusEnum
 
 from app.singletons.logs_manager import LogsManager
+from app.config.settings import get_settings
 from pymongo import ReturnDocument
 
 logger = LogsManager().get_logger()
 
 
 async def find_state(namespace_name: str, nodes: list[str]) -> State | None:
+    current_time_ms = int(time.time() * 1000)
+    settings = get_settings()
+
     data = await State.get_pymongo_collection().find_one_and_update(
         {
             "namespace_name": namespace_name,
             "status": StateStatusEnum.CREATED,
-            "node_name": {
-                "$in": nodes
-            },
-            "enqueue_after": {"$lte": int(time.time() * 1000)}
-        },
-        {
-            "$set": {"status": StateStatusEnum.QUEUED}
+            "node_name": {"$in": nodes},
+            "enqueue_after": {"$lte": current_time_ms}
         },
+        [
+            {
+                "$set": {
+                    "status": StateStatusEnum.QUEUED,
+                    "queued_at": current_time_ms,
+                    "timeout_at": {
+                        "$add": [
+                            current_time_ms,
+                            {
+                                "$multiply": [
+                                    {
+                                        "$ifNull": [
+                                            "$timeout_minutes", 
+                                            settings.node_timeout_minutes
+                                        ]
+                                    },
+                                    60000  # Convert minutes to milliseconds
+                                ]
+                            }
+                        ]
+                    }
+                }
+            }
+        ],
         return_document=ReturnDocument.AFTER
     )
     return State(**data) if data else None

state-manager/app/controller/errored_state.py

@@ -53,7 +53,8 @@ async def errored_state(namespace_name: str, state_id: PydanticObjectId, body: E
                     does_unites=state.does_unites,
                     enqueue_after= int(time.time() * 1000) + graph_template.retry_policy.compute_delay(state.retry_count + 1),
                     retry_count=state.retry_count + 1,
-                    fanout_id=state.fanout_id
+                    fanout_id=state.fanout_id,
+                    timeout_at=state.timeout_at
                 )
                 retry_state = await retry_state.insert()
                 logger.info(f"Retry state {retry_state.id} created for state {state_id}", x_exosphere_request_id=x_exosphere_request_id)

state-manager/app/controller/manual_retry_state.py

@@ -31,7 +31,8 @@ async def manual_retry_state(namespace_name: str, state_id: PydanticObjectId, bo
                 parents=state.parents,
                 does_unites=state.does_unites,
                 fanout_id=body.fanout_id, # this will ensure that multiple unwanted retries are not formed because of index in database
-                manual_retry_fanout_id=body.fanout_id # This is included in the state fingerprint to allow unique manual retries of unite nodes.
+                manual_retry_fanout_id=body.fanout_id, # This is included in the state fingerprint to allow unique manual retries of unite nodes.
+                timeout_minutes=state.timeout_minutes
             )
             retry_state = await retry_state.insert()
             logger.info(f"Retry state {retry_state.id} created for state {state_id}", x_exosphere_request_id=x_exosphere_request_id)

state-manager/app/controller/re_queue_after_signal.py

@@ -21,6 +21,8 @@ async def re_queue_after_signal(namespace_name: str, state_id: PydanticObjectId,
 
         state.status = StateStatusEnum.CREATED
         state.enqueue_after = int(time.time() * 1000) + body.enqueue_after
+        state.timeout_at = None
+        state.queued_at = None
         await state.save()
 
         return SignalResponseModel(status=StateStatusEnum.CREATED, enqueue_after=state.enqueue_after)

state-manager/app/controller/register_nodes.py

@@ -31,7 +31,8 @@ async def register_nodes(namespace_name: str, body: RegisterNodesRequestModel, x
                         RegisteredNode.runtime_namespace: namespace_name,
                         RegisteredNode.inputs_schema: node_data.inputs_schema, # type: ignore
                         RegisteredNode.outputs_schema: node_data.outputs_schema, # type: ignore
-                        RegisteredNode.secrets: node_data.secrets # type: ignore
+                        RegisteredNode.secrets: node_data.secrets, # type: ignore
+                        RegisteredNode.timeout_minutes: node_data.timeout_minutes # type: ignore
                 }))
                 logger.info(f"Updated existing node {node_data.name} in namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
 
@@ -44,7 +45,8 @@ async def register_nodes(namespace_name: str, body: RegisterNodesRequestModel, x
                     runtime_namespace=namespace_name,
                     inputs_schema=node_data.inputs_schema,
                     outputs_schema=node_data.outputs_schema,
-                    secrets=node_data.secrets
+                    secrets=node_data.secrets,
+                    timeout_minutes=node_data.timeout_minutes
                 )
                 await new_node.insert()
                 logger.info(f"Created new node {node_data.name} in namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
@@ -54,7 +56,8 @@ async def register_nodes(namespace_name: str, body: RegisterNodesRequestModel, x
                     name=node_data.name,
                     inputs_schema=node_data.inputs_schema,
                     outputs_schema=node_data.outputs_schema,
-                    secrets=node_data.secrets
+                    secrets=node_data.secrets,
+                    timeout_minutes=node_data.timeout_minutes
                 )
             )
 

state-manager/app/controller/trigger_graph.py

@@ -7,8 +7,10 @@
 from app.models.db.store import Store
 from app.models.db.run import Run
 from app.models.db.graph_template_model import GraphTemplate
+from app.models.db.registered_node import RegisteredNode
 from app.models.node_template_model import NodeTemplate
 from app.models.dependent_string import DependentString
+from app.config.settings import get_settings
 
 import uuid
 import time
@@ -91,6 +93,16 @@ async def trigger_graph(namespace_name: str, graph_name: str, body: TriggerGraph
         if len(new_stores) > 0:
             await Store.insert_many(new_stores)
 
+        # Get node timeout setting
+        registered_node = await RegisteredNode.get_by_name_and_namespace(root.node_name, root.namespace)
+        timeout_minutes = None
+        if registered_node and registered_node.timeout_minutes:
+            timeout_minutes = registered_node.timeout_minutes
+        else:
+            # Fall back to global setting
+            settings = get_settings()
+            timeout_minutes = settings.node_timeout_minutes
+
         new_state = State(
             node_name=root.node_name,
             namespace_name=namespace_name,
@@ -101,7 +113,8 @@ async def trigger_graph(namespace_name: str, graph_name: str, body: TriggerGraph
             enqueue_after=int(time.time() * 1000) + body.start_delay,
             inputs=inputs,
             outputs={},
-            error=None
+            error=None,
+            timeout_minutes=timeout_minutes
         )
         await new_state.insert()
 

state-manager/app/main.py

@@ -38,6 +38,7 @@
 from apscheduler.schedulers.asyncio import AsyncIOScheduler
 from apscheduler.triggers.cron import CronTrigger
 from .tasks.trigger_cron import trigger_cron
+from .tasks.check_node_timeout import check_node_timeout
 
 # init tasks
 from .tasks.init_tasks import init_tasks
@@ -83,6 +84,15 @@ async def lifespan(app: FastAPI):
         max_instances=1,
         id="every_minute_task"
     )
+    scheduler.add_job(
+        check_node_timeout,
+        CronTrigger.from_crontab("* * * * *"),
+        replace_existing=True,
+        misfire_grace_time=60,
+        coalesce=True,
+        max_instances=1,
+        id="check_node_timeout_task"
+    )
     scheduler.start()
 
     # main logic of the server

state-manager/app/models/db/registered_node.py

@@ -1,6 +1,6 @@
 from .base import BaseDatabaseModel
 from pydantic import Field
-from typing import Any
+from typing import Any, Optional
 from pymongo import IndexModel
 from ..node_template_model import NodeTemplate
 
@@ -13,6 +13,7 @@ class RegisteredNode(BaseDatabaseModel):
     inputs_schema: dict[str, Any] = Field(..., description="JSON schema for node inputs")
     outputs_schema: dict[str, Any] = Field(..., description="JSON schema for node outputs") 
     secrets: list[str] = Field(default_factory=list, description="List of secrets that the node uses")
+    timeout_minutes: Optional[int] = Field(None, gt=0, description="Timeout in minutes for this node. Falls back to global setting if not provided")
 
     class Settings:
         indexes = [

state-manager/app/models/db/state.py

@@ -28,6 +28,9 @@ class State(BaseDatabaseModel):
     retry_count: int = Field(default=0, description="Number of times the state has been retried")
     fanout_id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Fanout ID of the state")
     manual_retry_fanout_id: str = Field(default="", description="Fanout ID from a manual retry request, ensuring unique retries for unite nodes.")
+    queued_at: Optional[int] = Field(default=None, description="Unix time in milliseconds when state was queued")
+    timeout_at: Optional[int] = Field(default=None, description="Unix time in milliseconds when state times out")
+    timeout_minutes: Optional[int] = Field(default=None, gt=0, description="Timeout in minutes for this specific state, taken from node registration")
 
     @before_event([Insert, Replace, Save])
     def _generate_fingerprint(self):
@@ -102,5 +105,12 @@ class Settings:
                     ("status", 1),
                 ],
                 name="run_id_status_index"
+            ),
+            IndexModel(
+                [
+                    ("status", 1),
+                    ("timeout_at", 1),
+                ],
+                name="timeout_query_index"
             )
         ]

state-manager/app/models/register_nodes_request.py

@@ -1,12 +1,13 @@
 from pydantic import BaseModel, Field
-from typing import Any, List
+from typing import Any, List, Optional
 
 
 class NodeRegistrationModel(BaseModel):
     name: str = Field(..., description="Unique name of the node")
     inputs_schema: dict[str, Any] = Field(..., description="JSON schema for node inputs")
     outputs_schema: dict[str, Any] = Field(..., description="JSON schema for node outputs")
     secrets: List[str] = Field(..., description="List of secrets that the node uses")
+    timeout_minutes: Optional[int] = Field(None, gt=0, description="Timeout in minutes for this node. Falls back to global setting if not provided")
 
 
 class RegisterNodesRequestModel(BaseModel):

state-manager/app/models/register_nodes_response.py

@@ -1,12 +1,13 @@
 from pydantic import BaseModel, Field
-from typing import Any, List
+from typing import Any, List, Optional
 
 
 class RegisteredNodeModel(BaseModel):
     name: str = Field(..., description="Name of the registered node")
     inputs_schema: dict[str, Any] = Field(..., description="Inputs for the registered node")
     outputs_schema: dict[str, Any] = Field(..., description="Outputs for the registered node")
     secrets: List[str] = Field(..., description="List of secrets that the node uses")
+    timeout_minutes: Optional[int] = Field(None, gt=0, description="Timeout in minutes for this node. Falls back to global setting if not provided")
 
 
 class RegisterNodesResponseModel(BaseModel):

state-manager/app/models/state_status_enum.py

@@ -11,6 +11,7 @@ class StateStatusEnum(str, Enum):
     # Errored
     ERRORED = 'ERRORED'
     NEXT_CREATED_ERROR = 'NEXT_CREATED_ERROR'
+    TIMEDOUT = 'TIMEDOUT'
 
     # Success
     SUCCESS = 'SUCCESS'

state-manager/app/tasks/check_node_timeout.py

@@ -0,0 +1,33 @@
+import time
+from app.models.db.state import State
+from app.models.state_status_enum import StateStatusEnum
+from app.singletons.logs_manager import LogsManager
+
+logger = LogsManager().get_logger()
+
+
+async def check_node_timeout():
+    try:
+        current_time_ms = int(time.time() * 1000)
+
+        logger.info(f"Checking for timed out nodes at {current_time_ms}")
+
+        # Use database query to find and update timed out states in one operation
+        result = await State.get_pymongo_collection().update_many(
+            {
+                "status": StateStatusEnum.QUEUED,
+                "timeout_at": {"$ne": None, "$lte": current_time_ms}
+            },
+            {
+                "$set": {
+                    "status": StateStatusEnum.TIMEDOUT,
+                    "error": "Node execution timed out"
+                }
+            }
+        )
+
+        if result.modified_count > 0:
+            logger.info(f"Marked {result.modified_count} states as TIMEDOUT")
+
+    except Exception:
+        logger.error("Error checking node timeout", exc_info=True)

state-manager/app/tasks/create_next_states.py

@@ -10,6 +10,7 @@
 from app.models.db.store import Store
 from app.models.dependent_string import DependentString
 from app.models.node_template_model import UnitesStrategyEnum
+from app.config.settings import get_settings
 from json_schema_to_pydantic import create_model
 from pydantic import BaseModel
 from typing import Type
@@ -162,6 +163,16 @@ async def generate_next_state(next_state_input_model: Type[BaseModel], next_stat
                 current_state.identifier: current_state.id
             }
 
+            # Get timeout for this node
+            registered_node = await get_registered_node(next_state_node_template)
+            timeout_minutes = None
+            if registered_node.timeout_minutes:
+                timeout_minutes = registered_node.timeout_minutes
+            else:
+                # Fall back to global setting
+                settings = get_settings()
+                timeout_minutes = settings.node_timeout_minutes
+
             return State(
                 node_name=next_state_node_template.node_name,
                 identifier=next_state_node_template.identifier,
@@ -173,7 +184,8 @@ async def generate_next_state(next_state_input_model: Type[BaseModel], next_stat
                 outputs={},
                 does_unites=next_state_node_template.unites is not None,
                 run_id=current_state.run_id,
-                error=None
+                error=None,
+                timeout_minutes=timeout_minutes
             )
 
         current_states = await State.find(