feat: add manual retry state apis (#405)

* feat: add manual retry state functionality

- Introduced a new API endpoint for manual state retries, allowing users to trigger retries for specific states.
- Implemented request and response models for manual retry operations.
- Added error handling for duplicate retry states and invalid API keys.
- Enhanced logging for better traceability of retry operations.

This feature improves the system's ability to manage state retries effectively.

* fix: correct import path for manual retry state and update route definition

- Fixed the import path for the manual retry state controller in routes.py.
- Updated the route definition for manual retry to include a leading slash for consistency.
- Added a new controller file for manual retry state functionality, implementing the logic for handling manual retries.
- Introduced unit tests for the manual retry request and response models, ensuring validation and functionality.
- Enhanced unit tests for the manual retry state route, covering various scenarios including valid and invalid API keys.

These changes improve the structure and reliability of the manual retry feature.

* fix: refine manual retry state error handling and query logic

- Updated the query for fetching the state to include the namespace name, ensuring accurate state retrieval.
- Changed the HTTP status code for duplicate retry state errors from 400 to 409 to better reflect the conflict nature of the error.
- Simplified exception handling by removing the unnecessary re-raise of the caught exception.

These changes enhance the reliability and clarity of the manual retry state functionality.

* feat: add unit tests for manual retry state functionality

- Introduced a new test file for the manual retry state, covering various scenarios including successful state creation, error handling for not found states, and duplicate key errors.
- Enhanced tests to verify logging, database error handling, and preservation of original state fields during retries.
- Updated the README to include the new test file and detailed coverage of the manual retry state functionality.

These changes improve the test coverage and reliability of the manual retry state feature.

Author: NiveditJain

state-manager/.coverage

@@ -0,0 +1,49 @@
+from pymongo.errors import DuplicateKeyError
+from app.models.manual_retry import ManualRetryRequestModel, ManualRetryResponseModel
+from beanie import PydanticObjectId
+from app.singletons.logs_manager import LogsManager
+from app.models.state_status_enum import StateStatusEnum
+from fastapi import HTTPException, status
+from app.models.db.state import State
+
+
+logger = LogsManager().get_logger()
+
+async def manual_retry_state(namespace_name: str, state_id: PydanticObjectId, body: ManualRetryRequestModel, x_exosphere_request_id: str):
+    try:
+        logger.info(f"Manual retry state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+
+        state = await State.find_one(State.id == state_id, State.namespace_name == namespace_name)
+        if not state:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="State not found")
+        
+        try:
+            retry_state = State(
+                node_name=state.node_name,
+                namespace_name=state.namespace_name,
+                identifier=state.identifier,
+                graph_name=state.graph_name,
+                run_id=state.run_id,
+                status=StateStatusEnum.CREATED,
+                inputs=state.inputs,
+                outputs={},
+                error=None,
+                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
+            )
+            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.status = StateStatusEnum.RETRY_CREATED
+            await state.save()
+
+            return ManualRetryResponseModel(id=str(retry_state.id), status=retry_state.status)
+        except DuplicateKeyError:
+            logger.info(f"Duplicate retry state detected for state {state_id}. A retry state with the same unique key already exists.", x_exosphere_request_id=x_exosphere_request_id)
+            raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="Duplicate retry state detected")
+
+
+    except Exception as _:
+        logger.error(f"Error manual retry state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        raise

state-manager/app/controller/manual_retry_state.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel, Field
+from .state_status_enum import StateStatusEnum
+
+
+class ManualRetryRequestModel(BaseModel):
+    fanout_id: str = Field(..., description="Fanout ID of the state")
+
+
+class ManualRetryResponseModel(BaseModel):
+    id: str = Field(..., description="ID of the state")
+    status: StateStatusEnum = Field(..., description="Status of the state")

state-manager/app/models/manual_retry.py

@@ -50,6 +50,10 @@
 from .models.signal_models import ReEnqueueAfterRequestModel
 from .controller.re_queue_after_signal import re_queue_after_signal
 
+# manual_retry
+from .models.manual_retry import ManualRetryRequestModel, ManualRetryResponseModel
+from .controller.manual_retry_state import manual_retry_state
+
 
 logger = LogsManager().get_logger()
 
@@ -176,6 +180,24 @@ async def re_enqueue_after_state_route(namespace_name: str, state_id: str, body:
 
     return await re_queue_after_signal(namespace_name, PydanticObjectId(state_id), body, x_exosphere_request_id)
 
+@router.post(
+    "/state/{state_id}/manual-retry",
+    response_model=ManualRetryResponseModel,
+    status_code=status.HTTP_200_OK,
+    response_description="State manual retry successfully",
+    tags=["state"]
+)
+async def manual_retry_state_route(namespace_name: str, state_id: str, body: ManualRetryRequestModel, request: Request, api_key: str = Depends(check_api_key)):
+    x_exosphere_request_id = getattr(request.state, "x_exosphere_request_id", str(uuid4()))
+    
+    if api_key:
+        logger.info(f"API key is valid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+    else:
+        logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key")
+    
+    return await manual_retry_state(namespace_name, PydanticObjectId(state_id), body, x_exosphere_request_id)
+
 
 @router.put(
     "/graph/{graph_name}",

state-manager/app/routes.py

@@ -14,6 +14,7 @@ tests/
 │       ├── test_errored_state.py
 │       ├── test_get_graph_template.py
 │       ├── test_get_secrets.py
+│       ├── test_manual_retry_state.py
 │       ├── test_register_nodes.py
 │       └── test_upsert_graph_template.py
 └── README.md
@@ -80,7 +81,21 @@ The unit tests cover all controller functions in the state-manager:
 - ✅ Complex schema handling
 - ✅ Database error handling
 
-### 8. `upsert_graph_template.py`
+### 8. `manual_retry_state.py`
+- ✅ Successful manual retry state creation
+- ✅ State not found scenarios
+- ✅ Duplicate retry state detection (DuplicateKeyError)
+- ✅ Different fanout_id handling
+- ✅ Complex inputs and multiple parents preservation
+- ✅ Database errors during state lookup
+- ✅ Database errors during state save
+- ✅ Database errors during retry state insert
+- ✅ Empty inputs and parents handling
+- ✅ Namespace mismatch scenarios
+- ✅ Field preservation and reset logic
+- ✅ Logging verification
+
+### 9. `upsert_graph_template.py`
 - ✅ Existing template updates
 - ✅ New template creation
 - ✅ Empty nodes handling