Document HTTP statuses that route helpers raise

airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml

@@ -2384,6 +2384,12 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/CalendarTimeRangeCollectionResponse'
+        '404':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Not Found
         '422':
           description: Validation Error
           content:

airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml

@@ -6603,6 +6603,12 @@ paths:
               schema:
                 $ref: '#/components/schemas/HTTPExceptionResponse'
           description: Forbidden
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
         '404':
           content:
             application/json:

airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py

@@ -221,7 +221,7 @@ def get_task_state_store(
 @task_state_store_router.put(
     "/{key:path}",
     status_code=status.HTTP_204_NO_CONTENT,
-    responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
+    responses=create_openapi_http_exception_doc([status.HTTP_400_BAD_REQUEST, status.HTTP_404_NOT_FOUND]),
     dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.TASK_INSTANCE))],
 )
 def set_task_state_store(

airflow-core/src/airflow/api_fastapi/core_api/routes/ui/calendar.py

@@ -18,7 +18,7 @@
 
 from typing import Annotated, Literal
 
-from fastapi import Depends
+from fastapi import Depends, status
 
 from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
 from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag
@@ -29,6 +29,7 @@
     CalendarDeadlineCollectionResponse,
     CalendarTimeRangeCollectionResponse,
 )
+from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.security import requires_access_dag
 from airflow.api_fastapi.core_api.services.ui.calendar import CalendarService
 from airflow.models.dagrun import DagRun
@@ -39,6 +40,7 @@
 
 @calendar_router.get(
     "/{dag_id}",
+    responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
     dependencies=[
         Depends(
             requires_access_dag(

airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py

@@ -123,7 +123,10 @@ def _check_hitl_detail_exists(hitl_detail_model: HITLDetail | None) -> HITLDetai
 @router.patch(
     "/{task_instance_id}",
     responses=create_openapi_http_exception_doc(
-        [(status.HTTP_409_CONFLICT, "A response has already been received for this HITLDetail")]
+        [
+            (status.HTTP_404_NOT_FOUND, "HITLDetail not found"),
+            (status.HTTP_409_CONFLICT, "A response has already been received for this HITLDetail"),
+        ]
     ),
 )
 def update_hitl_detail(
@@ -153,6 +156,7 @@ def update_hitl_detail(
 @router.get(
     "/{task_instance_id}",
     status_code=status.HTTP_200_OK,
+    responses=create_openapi_http_exception_doc([(status.HTTP_404_NOT_FOUND, "HITLDetail not found")]),
 )
 async def get_hitl_detail(
     task_instance_id: UUID,

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -1257,7 +1257,11 @@ async def get_previous_successful_dagrun(
     return PrevSuccessfulDagRunResponse.model_validate(dag_run)
 
 
-@router.get("/count", status_code=status.HTTP_200_OK)
+@router.get(
+    "/count",
+    status_code=status.HTTP_200_OK,
+    responses=create_openapi_http_exception_doc([(status.HTTP_404_NOT_FOUND, "Task group not found")]),
+)
 def get_task_instance_count(
     dag_id: str,
     session: SessionDep,
@@ -1366,7 +1370,11 @@ async def get_previous_task_instance(
     )
 
 
-@router.get("/states", status_code=status.HTTP_200_OK)
+@router.get(
+    "/states",
+    status_code=status.HTTP_200_OK,
+    responses=create_openapi_http_exception_doc([(status.HTTP_404_NOT_FOUND, "Task group not found")]),
+)
 def get_task_instance_states(
     dag_id: str,
     session: SessionDep,

airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts

@@ -4297,6 +4297,7 @@ export class TaskStateStoreService {
             body: data.requestBody,
             mediaType: 'application/json',
             errors: {
+                400: 'Bad Request',
                 401: 'Unauthorized',
                 403: 'Forbidden',
                 404: 'Not Found',
@@ -5535,6 +5536,7 @@ export class CalendarService {
                 partition_date_lt: data.partitionDateLt
             },
             errors: {
+                404: 'Not Found',
                 422: 'Validation Error'
             }
         });

airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts

@@ -8592,6 +8592,10 @@ export type $OpenApiTs = {
                  * Successful Response
                  */
                 204: void;
+                /**
+                 * Bad Request
+                 */
+                400: HTTPExceptionResponse;
                 /**
                  * Unauthorized
                  */
@@ -9450,6 +9454,10 @@ export type $OpenApiTs = {
                  * Successful Response
                  */
                 200: CalendarTimeRangeCollectionResponse;
+                /**
+                 * Not Found
+                 */
+                404: HTTPExceptionResponse;
                 /**
                  * Validation Error
                  */