From 226356ddb74102777aa196c2978363aacac4705d Mon Sep 17 00:00:00 2001 From: Maxwell Wang Date: Tue, 1 Sep 2026 21:34:31 +0000 Subject: [PATCH 1/4] Report ARAX failures that arrive inside a 200, and stop faking query status Propagating the HTTP status code was not enough: ARAX answers HTTP 200 for most of its own failures and reports them in the TRAPI status field, so a service having internal issues still looked completely healthy to a check of the HTTP code alone, and the query completed OK. The ARAX worker now validates the body of a 2xx as well as its code. A body that isn't a TRAPI response (an error envelope from ARAX or from something in front of it) and a TRAPI status that names an error or a failure both raise ARAXServiceError, which puts ARAX's status on the span as arax.trapi_status, in the query's logs, and marks the query ERROR. Statuses that describe a non-failure outcome are left alone, so a query ARAX answered fine is not failed on an unrecognized status. When ARAX did send a usable response its own body is what the caller gets back -- its status, description and logs say more about the failure than anything synthesized. Two places then reported that failure as success anyway: - /asyncquery_status returned a hardcoded {"status": "Queued"} for every query it was ever asked about (a stub with a TODO), so a failed query and a healthy one were indistinguishable. It now reports Running, Completed or Failed from the query's own row, with the query's logs -- where the upstream status code is recorded -- instead of an empty list. - The sync path returned the stored response with no hint that the query had finished with anything but OK; a response an operation never got to write looks exactly like one that legitimately found nothing. The stored status is now stamped onto the TRAPI response, leaving an error the ARA reported itself alone, since that one is more specific. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BmsMzWR6tNbDM2uSxh47n8 --- shepherd_server/base_routes.py | 76 +++++++++++++++++-- tests/unit/test_arax.py | 71 ++++++++++++++++++ tests/unit/test_query_status.py | 125 ++++++++++++++++++++++++++++++++ workers/arax/worker.py | 60 ++++++++++++++- 4 files changed, 322 insertions(+), 10 deletions(-) create mode 100644 tests/unit/test_query_status.py diff --git a/shepherd_server/base_routes.py b/shepherd_server/base_routes.py index a424038..5c3d250 100644 --- a/shepherd_server/base_routes.py +++ b/shepherd_server/base_routes.py @@ -186,6 +186,26 @@ async def run_query( return query_id, response_id, logger +def apply_query_status(response: dict, status: Optional[str]) -> None: + """Stamp a non-OK query status onto the TRAPI response, in place. + + ``status``/``description`` are TRAPI Response fields, so a caller reading + the body it already parses can tell a failed query from an empty one. An + error the ARA reported itself is left alone -- it is more specific than the + query-level status -- but a body that claims nothing, or claims success for + a query that failed, is corrected here. + """ + if not status or status == "OK": + return + current = response.get("status") + if isinstance(current, str) and ( + "error" in current.lower() or "fail" in current.lower() + ): + return + response["status"] = "Error" + response.setdefault("description", f"Query finished with status {status}.") + + async def run_sync_query( target: ARATargetEnum, query: dict = Body(..., examples=[default_input_query]), @@ -224,6 +244,11 @@ async def run_sync_query( ) logs = await get_logs(response_id, logger) response["logs"] = logs + # The stored status is the one thing that knows the query + # failed -- a response an operation never got to write looks + # exactly like one that legitimately found nothing. Report it + # rather than handing back a body that only says "here you go". + apply_query_status(response, query_state[10]) return ORJSONResponse(content=response) else: # Debug, not warning: this fires every 0.5s while a query is still @@ -455,17 +480,54 @@ async def callback( return Response("Callback received.", 200) +# shepherd_brain.state/status (see shepherd_db/init_db.sql) mapped onto the +# TRAPI AsyncQueryStatusResponse vocabulary. A query is inserted QUEUED/OK and +# only leaves that state when it finishes, is abandoned, or times out. +TERMINAL_QUERY_STATES = {"COMPLETED", "ABANDONED"} +OK_QUERY_STATUS = "OK" + + @base_router.get("/asyncquery_status/{qid}", status_code=200) async def query_status( qid: str, -) -> dict: +): """Handle query status requests.""" - # TODO: get query status from db - return { - "status": "Queued", - "description": "Query is currently waiting to be run.", - "logs": [], - } + logger = logging.getLogger("shepherd.query_status") + logger.setLevel(logging.INFO) + attach_query_handler(logger) + query_state = await get_query_state(qid, logger) + if query_state is None: + return JSONResponse(content={"error": "Not found"}, status_code=404) + + response_id = query_state[7] + state = query_state[9] + status = query_state[10] + description = query_state[11] + logs = await get_logs(response_id, logger) if response_id else [] + + if state not in TERMINAL_QUERY_STATES: + # Shepherd doesn't track a separate running state: a query is in the + # pipeline from the moment it is accepted until it finishes. + trapi_status = "Running" + default_description = "Query is currently running." + elif status == OK_QUERY_STATUS: + trapi_status = "Completed" + default_description = "Query has finished." + else: + # The query reached the end of the line with something other than OK + # (ERROR from a failed operation, TIMEOUT, ABANDONED). Previously this + # endpoint answered "Queued" for every query it was ever asked about, + # so a failed query and a healthy one looked exactly alike here. + trapi_status = "Failed" + default_description = f"Query finished with status {status}." + + return ORJSONResponse( + content={ + "status": trapi_status, + "description": description or default_description, + "logs": logs, + } + ) @base_router.get("/response/{query_id}", status_code=200) diff --git a/tests/unit/test_arax.py b/tests/unit/test_arax.py index e9dbd41..60b851e 100644 --- a/tests/unit/test_arax.py +++ b/tests/unit/test_arax.py @@ -216,3 +216,74 @@ async def test_pathfinder_query_is_routed_without_calling_arax(mocker): assert not post.called assert not save.called assert json.loads(task[1]["workflow"]) == [{"id": "arax.pathfinder"}] + + +# --- TRAPI-level failures reported inside a 200 ----------------------------- +# +# ARAX answers HTTP 200 for most of its own failures and puts the failure in +# the TRAPI status field, so checking the HTTP code alone reports every one of +# them as a healthy query. + + +def _trapi(status=None, description=None): + body = {"message": {"query_graph": {}, "knowledge_graph": {}, "results": []}} + if status is not None: + body["status"] = status + if description is not None: + body["description"] = description + return body + + +@pytest.mark.parametrize("status", ["Error", "ERROR", "InternalError", "Failed"]) +@pytest.mark.asyncio +async def test_trapi_error_status_in_a_200_fails_the_query(mocker, status): + save = _patch_db(mocker) + span = _patch_span(mocker) + _patch_post( + mocker, + _http_response(200, json_body=_trapi(status, "internal issues upstream")), + ) + task = _task() + + with pytest.raises(ARAXServiceError) as excinfo: + await arax(task, logger) + + assert excinfo.value.status_code == 200 + assert status in str(excinfo.value) + assert "internal issues upstream" in str(excinfo.value) + span.set_attribute.assert_any_call("arax.trapi_status", status) + # ARAX's own body is what the caller gets -- its status and description say + # more than anything we could synthesize. + assert save.await_args.args[1]["status"] == status + assert save.await_args.args[1]["description"] == "internal issues upstream" + + +@pytest.mark.parametrize("status", ["Success", "OK", "QueryNotTraversable", None]) +@pytest.mark.asyncio +async def test_non_error_trapi_status_still_succeeds(mocker, status): + """Only statuses naming an error fail the query; the rest are outcomes.""" + save = _patch_db(mocker) + _patch_span(mocker) + _patch_post(mocker, _http_response(200, json_body=_trapi(status))) + task = _task() + + await arax(task, logger) + + assert save.await_args.args[1].get("status") == status + assert json.loads(task[1]["workflow"]) == [{"id": "arax"}] + + +@pytest.mark.asyncio +async def test_json_body_that_is_not_trapi_fails_the_query(mocker): + save = _patch_db(mocker) + _patch_span(mocker) + _patch_post(mocker, _http_response(200, json_body={"detail": "Internal Error"})) + + with pytest.raises(ARAXServiceError) as excinfo: + await arax(_task(), logger) + + assert excinfo.value.status_code == 200 + assert "not a TRAPI response" in str(excinfo.value) + # Nothing usable came back, so the caller gets one we build. + assert save.await_args.args[1]["status"] == "Error" + assert "[HTTP 200]" in save.await_args.args[1]["description"] diff --git a/tests/unit/test_query_status.py b/tests/unit/test_query_status.py new file mode 100644 index 0000000..a68ff7b --- /dev/null +++ b/tests/unit/test_query_status.py @@ -0,0 +1,125 @@ +"""Tests for how a query's outcome is reported back to the caller. + +``/asyncquery_status/{qid}`` used to answer ``{"status": "Queued"}`` for every +query it was ever asked about (a hardcoded stub with a TODO), so a query that +failed was indistinguishable from a healthy one. The sync path had the milder +version of the same problem: it returned the stored response with no hint that +the query had finished with anything other than OK. +""" + +import json +import logging + +import pytest + +from shepherd_server.base_routes import apply_query_status, query_status + +logger = logging.getLogger(__name__) + + +def _row(state="QUEUED", status="OK", description=None): + """A shepherd_brain row, in the column order get_query_state returns.""" + return ( + "qid", + "start", + "stop", + "submitter", + "ip", + "domain", + "hostname", + "response_id", + None, + state, + status, + description, + ) + + +def _patch_state(mocker, row, logs=None): + mocker.patch( + "shepherd_server.base_routes.get_query_state", + new_callable=mocker.AsyncMock, + return_value=row, + ) + mocker.patch( + "shepherd_server.base_routes.get_logs", + new_callable=mocker.AsyncMock, + return_value=logs if logs is not None else [], + ) + + +def _body(response): + return json.loads(bytes(response.body)) + + +@pytest.mark.asyncio +async def test_status_unknown_query_is_not_found(mocker): + _patch_state(mocker, None) + response = await query_status("qid") + assert response.status_code == 404 + + +@pytest.mark.asyncio +async def test_status_in_flight_query_is_running(mocker): + _patch_state(mocker, _row(state="QUEUED", status="OK")) + assert _body(await query_status("qid"))["status"] == "Running" + + +@pytest.mark.asyncio +async def test_status_finished_query_is_completed(mocker): + _patch_state(mocker, _row(state="COMPLETED", status="OK")) + assert _body(await query_status("qid"))["status"] == "Completed" + + +@pytest.mark.parametrize( + "state,status", + [ + ("COMPLETED", "ERROR"), + ("COMPLETED", "TIMEOUT"), + ("ABANDONED", "Abandoned: no completion within budget"), + ], +) +@pytest.mark.asyncio +async def test_status_failed_query_is_reported_as_failed(mocker, state, status): + _patch_state(mocker, _row(state=state, status=status)) + body = _body(await query_status("qid")) + assert body["status"] == "Failed" + assert status in body["description"] + + +@pytest.mark.asyncio +async def test_status_carries_the_query_logs(mocker): + """The logs are where the upstream status code is recorded.""" + logs = [{"level": "ERROR", "message": "ARAX service returned HTTP 500"}] + _patch_state(mocker, _row(state="COMPLETED", status="ERROR"), logs=logs) + assert _body(await query_status("qid"))["logs"] == logs + + +# --- apply_query_status ---------------------------------------------------- + + +def test_ok_query_is_left_untouched(): + response = {"message": {}} + apply_query_status(response, "OK") + assert response == {"message": {}} + + +def test_failed_query_is_marked_on_the_response(): + response = {"message": {}} + apply_query_status(response, "ERROR") + assert response["status"] == "Error" + assert "ERROR" in response["description"] + + +def test_ara_reported_error_is_not_overwritten(): + """ARAX's own status is more specific than the query-level one.""" + response = {"message": {}, "status": "InternalError", "description": "upstream"} + apply_query_status(response, "ERROR") + assert response["status"] == "InternalError" + assert response["description"] == "upstream" + + +def test_success_claimed_for_a_failed_query_is_corrected(): + response = {"message": {}, "status": "Success"} + apply_query_status(response, "TIMEOUT") + assert response["status"] == "Error" diff --git a/workers/arax/worker.py b/workers/arax/worker.py index e501116..98fcbf3 100644 --- a/workers/arax/worker.py +++ b/workers/arax/worker.py @@ -44,11 +44,35 @@ class ARAXServiceError(Exception): non-TRAPI ``{"status": "error"}`` blob that the workflow then reported as a successful response, which left nothing downstream with a status code to report (see the same fix in ``arax_pathfinder``). + + ``trapi_response`` holds ARAX's own response when it sent a usable one -- + an HTTP 200 that reports the failure in the TRAPI ``status`` field. That + body, not one we synthesize, is what the caller should get back. """ - def __init__(self, message: str, status_code: int): + def __init__(self, message: str, status_code: int, trapi_response: dict = None): super().__init__(message) self.status_code = status_code + self.trapi_response = trapi_response + + +def trapi_error_status(result: dict) -> str: + """ARAX's TRAPI ``status``, when it reports an error. Empty string if not. + + ARAX answers HTTP 200 for most of its own failures and reports them in the + TRAPI ``status`` field, so a service having internal issues still looks + perfectly healthy to anything that only checks the HTTP code. Only statuses + that name an error or a failure count: the rest of the TRAPI status + vocabulary describes outcomes that aren't service failures, and treating an + unrecognized status as one would fail queries ARAX answered fine. + """ + status = result.get("status") + if not isinstance(status, str): + return "" + lowered = status.lower() + if "error" in lowered or "fail" in lowered: + return status + return "" def body_head(response: httpx.Response) -> str: @@ -139,7 +163,32 @@ async def call_arax(message: dict, logger: logging.Logger) -> dict: status_code, ) from e - return add_shepherd_arax_to_edge_sources(result) + if not isinstance(result, dict) or not isinstance(result.get("message"), dict): + # Parseable JSON, but not a TRAPI response -- an error envelope from + # ARAX or from something in front of it ({"detail": "..."}), which + # nothing downstream can merge, score or hand back. + raise ARAXServiceError( + f"ARAX service at {settings.arax_url} returned HTTP {status_code} " + "with a body that is not a TRAPI response: " + f"{str(result)[:ERROR_BODY_BYTES]}", + status_code, + ) + + result = add_shepherd_arax_to_edge_sources(result) + + error_status = trapi_error_status(result) + if error_status: + span.set_attribute("arax.trapi_status", error_status) + description = result.get("description") or "no description given" + raise ARAXServiceError( + f"ARAX service at {settings.arax_url} returned HTTP {status_code} " + f'reporting TRAPI status "{error_status}": ' + f"{str(description)[:ERROR_BODY_BYTES]}", + status_code, + trapi_response=result, + ) + + return result def error_response(message: dict, error: ARAXServiceError) -> dict: @@ -181,7 +230,12 @@ async def arax(task, logger: logging.Logger): # status. Without this the response id still holds the echo of the # incoming query, so the caller gets a query that looks like it # simply found nothing. - await save_message(response_id, error_response(message, e), logger) + # ARAX's own body wins when it sent one: its status, description + # and logs say more about the failure than anything we can build. + saved = e.trapi_response + if saved is None: + saved = error_response(message, e) + await save_message(response_id, saved, logger) raise await save_message(response_id, result, logger) task[1]["workflow"] = json.dumps([{"id": "arax"}]) From 70ea268a2a2019cc49c9fe1c90eccf358ab48aae Mon Sep 17 00:00:00 2001 From: Maxwell Wang Date: Wed, 2 Sep 2026 02:14:40 +0000 Subject: [PATCH 2/4] Revert the TRAPI-status handling in the ARAX worker Treating a TRAPI status of "Error" inside an HTTP 200 as a failed query was speculation about how ARAX reports its internal issues, and it isn't what was hiding them: the sync path returned the stored response without ever mentioning that the query had finished with anything but OK. That fix, and the /asyncquery_status one alongside it, stay. This restores the ARAX worker to propagating ARAX's HTTP status code, and drops the body inspection, the TRAPI-status check and the response the exception carried for it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BmsMzWR6tNbDM2uSxh47n8 --- tests/unit/test_arax.py | 71 ----------------------------------------- workers/arax/worker.py | 60 ++-------------------------------- 2 files changed, 3 insertions(+), 128 deletions(-) diff --git a/tests/unit/test_arax.py b/tests/unit/test_arax.py index 60b851e..e9dbd41 100644 --- a/tests/unit/test_arax.py +++ b/tests/unit/test_arax.py @@ -216,74 +216,3 @@ async def test_pathfinder_query_is_routed_without_calling_arax(mocker): assert not post.called assert not save.called assert json.loads(task[1]["workflow"]) == [{"id": "arax.pathfinder"}] - - -# --- TRAPI-level failures reported inside a 200 ----------------------------- -# -# ARAX answers HTTP 200 for most of its own failures and puts the failure in -# the TRAPI status field, so checking the HTTP code alone reports every one of -# them as a healthy query. - - -def _trapi(status=None, description=None): - body = {"message": {"query_graph": {}, "knowledge_graph": {}, "results": []}} - if status is not None: - body["status"] = status - if description is not None: - body["description"] = description - return body - - -@pytest.mark.parametrize("status", ["Error", "ERROR", "InternalError", "Failed"]) -@pytest.mark.asyncio -async def test_trapi_error_status_in_a_200_fails_the_query(mocker, status): - save = _patch_db(mocker) - span = _patch_span(mocker) - _patch_post( - mocker, - _http_response(200, json_body=_trapi(status, "internal issues upstream")), - ) - task = _task() - - with pytest.raises(ARAXServiceError) as excinfo: - await arax(task, logger) - - assert excinfo.value.status_code == 200 - assert status in str(excinfo.value) - assert "internal issues upstream" in str(excinfo.value) - span.set_attribute.assert_any_call("arax.trapi_status", status) - # ARAX's own body is what the caller gets -- its status and description say - # more than anything we could synthesize. - assert save.await_args.args[1]["status"] == status - assert save.await_args.args[1]["description"] == "internal issues upstream" - - -@pytest.mark.parametrize("status", ["Success", "OK", "QueryNotTraversable", None]) -@pytest.mark.asyncio -async def test_non_error_trapi_status_still_succeeds(mocker, status): - """Only statuses naming an error fail the query; the rest are outcomes.""" - save = _patch_db(mocker) - _patch_span(mocker) - _patch_post(mocker, _http_response(200, json_body=_trapi(status))) - task = _task() - - await arax(task, logger) - - assert save.await_args.args[1].get("status") == status - assert json.loads(task[1]["workflow"]) == [{"id": "arax"}] - - -@pytest.mark.asyncio -async def test_json_body_that_is_not_trapi_fails_the_query(mocker): - save = _patch_db(mocker) - _patch_span(mocker) - _patch_post(mocker, _http_response(200, json_body={"detail": "Internal Error"})) - - with pytest.raises(ARAXServiceError) as excinfo: - await arax(_task(), logger) - - assert excinfo.value.status_code == 200 - assert "not a TRAPI response" in str(excinfo.value) - # Nothing usable came back, so the caller gets one we build. - assert save.await_args.args[1]["status"] == "Error" - assert "[HTTP 200]" in save.await_args.args[1]["description"] diff --git a/workers/arax/worker.py b/workers/arax/worker.py index 98fcbf3..e501116 100644 --- a/workers/arax/worker.py +++ b/workers/arax/worker.py @@ -44,35 +44,11 @@ class ARAXServiceError(Exception): non-TRAPI ``{"status": "error"}`` blob that the workflow then reported as a successful response, which left nothing downstream with a status code to report (see the same fix in ``arax_pathfinder``). - - ``trapi_response`` holds ARAX's own response when it sent a usable one -- - an HTTP 200 that reports the failure in the TRAPI ``status`` field. That - body, not one we synthesize, is what the caller should get back. """ - def __init__(self, message: str, status_code: int, trapi_response: dict = None): + def __init__(self, message: str, status_code: int): super().__init__(message) self.status_code = status_code - self.trapi_response = trapi_response - - -def trapi_error_status(result: dict) -> str: - """ARAX's TRAPI ``status``, when it reports an error. Empty string if not. - - ARAX answers HTTP 200 for most of its own failures and reports them in the - TRAPI ``status`` field, so a service having internal issues still looks - perfectly healthy to anything that only checks the HTTP code. Only statuses - that name an error or a failure count: the rest of the TRAPI status - vocabulary describes outcomes that aren't service failures, and treating an - unrecognized status as one would fail queries ARAX answered fine. - """ - status = result.get("status") - if not isinstance(status, str): - return "" - lowered = status.lower() - if "error" in lowered or "fail" in lowered: - return status - return "" def body_head(response: httpx.Response) -> str: @@ -163,32 +139,7 @@ async def call_arax(message: dict, logger: logging.Logger) -> dict: status_code, ) from e - if not isinstance(result, dict) or not isinstance(result.get("message"), dict): - # Parseable JSON, but not a TRAPI response -- an error envelope from - # ARAX or from something in front of it ({"detail": "..."}), which - # nothing downstream can merge, score or hand back. - raise ARAXServiceError( - f"ARAX service at {settings.arax_url} returned HTTP {status_code} " - "with a body that is not a TRAPI response: " - f"{str(result)[:ERROR_BODY_BYTES]}", - status_code, - ) - - result = add_shepherd_arax_to_edge_sources(result) - - error_status = trapi_error_status(result) - if error_status: - span.set_attribute("arax.trapi_status", error_status) - description = result.get("description") or "no description given" - raise ARAXServiceError( - f"ARAX service at {settings.arax_url} returned HTTP {status_code} " - f'reporting TRAPI status "{error_status}": ' - f"{str(description)[:ERROR_BODY_BYTES]}", - status_code, - trapi_response=result, - ) - - return result + return add_shepherd_arax_to_edge_sources(result) def error_response(message: dict, error: ARAXServiceError) -> dict: @@ -230,12 +181,7 @@ async def arax(task, logger: logging.Logger): # status. Without this the response id still holds the echo of the # incoming query, so the caller gets a query that looks like it # simply found nothing. - # ARAX's own body wins when it sent one: its status, description - # and logs say more about the failure than anything we can build. - saved = e.trapi_response - if saved is None: - saved = error_response(message, e) - await save_message(response_id, saved, logger) + await save_message(response_id, error_response(message, e), logger) raise await save_message(response_id, result, logger) task[1]["workflow"] = json.dumps([{"id": "arax"}]) From 0bdd067cbf25f87ef268feedd3a32cc1b2ee16c8 Mon Sep 17 00:00:00 2001 From: Maxwell Wang Date: Wed, 2 Sep 2026 02:18:47 +0000 Subject: [PATCH 3/4] Answer /query with an error code when the query failed The body has carried a TRAPI error status since apply_query_status went in, but every answer /query gave was still HTTP 200, so a caller that checks the status code rather than parsing the payload for a status field saw a failed query as a successful one. All four ways a query can fail after intake now answer 500: an operation that errored, a query that ran out of budget or was abandoned, a response that isn't in the datastore (which had no status code at all on its ORJSONResponse, so it defaulted to 200), and the caller's own timeout elapsing with the query still in flight. TRAPI 1.5 documents 200, 400, 429, 500 and 501 for this operation, so they all report the spec's InternalServerError rather than a more precise code the schema doesn't allow -- 504 for the two timeout cases. Which kind of failure it was stays in the response's own status and description. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BmsMzWR6tNbDM2uSxh47n8 --- shepherd_server/base_routes.py | 48 +++++++++++++------ tests/unit/test_query_status.py | 83 ++++++++++++++++++++++++++++++++- 2 files changed, 117 insertions(+), 14 deletions(-) diff --git a/shepherd_server/base_routes.py b/shepherd_server/base_routes.py index 5c3d250..80e59f8 100644 --- a/shepherd_server/base_routes.py +++ b/shepherd_server/base_routes.py @@ -43,6 +43,19 @@ base_router = APIRouter() +# shepherd_brain.state/status (see shepherd_db/init_db.sql). A query is +# inserted QUEUED/OK and only leaves that state when it finishes, is +# abandoned, or times out. +TERMINAL_QUERY_STATES = {"COMPLETED", "ABANDONED"} +OK_QUERY_STATUS = "OK" +# What /query answers with once a query it accepted has failed. TRAPI 1.5 +# documents 200, 400, 429, 500 and 501 for this operation, so every +# post-intake failure -- an operation that errored, a query that ran out of +# budget, a response that isn't there -- reports the spec's InternalServerError +# rather than a more precise code (504) the schema doesn't allow. Which of +# those it was is in the response's own status/description. +QUERY_ERROR_CODE = 500 + class QueryIntakeError(Exception): """Raised when a query can't be accepted because its initial state could @@ -195,7 +208,7 @@ def apply_query_status(response: dict, status: Optional[str]) -> None: query-level status -- but a body that claims nothing, or claims success for a query that failed, is corrected here. """ - if not status or status == "OK": + if not status or status == OK_QUERY_STATUS: return current = response.get("status") if isinstance(current, str) and ( @@ -218,7 +231,7 @@ async def run_sync_query( except QueryIntakeError as e: return ORJSONResponse( content={"status": "ERROR", "description": str(e)}, - status_code=500, + status_code=QUERY_ERROR_CODE, ) start = time.time() now = start @@ -240,7 +253,8 @@ async def run_sync_query( content={ "status": "ERROR", "description": "Unable to get response", - } + }, + status_code=QUERY_ERROR_CODE, ) logs = await get_logs(response_id, logger) response["logs"] = logs @@ -248,8 +262,20 @@ async def run_sync_query( # failed -- a response an operation never got to write looks # exactly like one that legitimately found nothing. Report it # rather than handing back a body that only says "here you go". - apply_query_status(response, query_state[10]) - return ORJSONResponse(content=response) + status = query_state[10] + apply_query_status(response, status) + # The body has said "status": "Error" since apply_query_status + # went in, but the HTTP code said 200 -- so a caller that + # checks the code (rather than parsing the payload for a status + # field) saw every failed query as a successful one. + return ORJSONResponse( + content=response, + status_code=( + 200 + if not status or status == OK_QUERY_STATUS + else QUERY_ERROR_CODE + ), + ) else: # Debug, not warning: this fires every 0.5s while a query is still # in flight (the row just isn't COMPLETED yet) and would otherwise @@ -258,7 +284,10 @@ async def run_sync_query( await asyncio.sleep(0.5) logger.error("Query timed out") - return ORJSONResponse(content={"status": "TIMEOUT", "description": "Query timeout"}) + return ORJSONResponse( + content={"status": "TIMEOUT", "description": "Query timeout"}, + status_code=QUERY_ERROR_CODE, + ) async def run_async_query( @@ -480,13 +509,6 @@ async def callback( return Response("Callback received.", 200) -# shepherd_brain.state/status (see shepherd_db/init_db.sql) mapped onto the -# TRAPI AsyncQueryStatusResponse vocabulary. A query is inserted QUEUED/OK and -# only leaves that state when it finishes, is abandoned, or times out. -TERMINAL_QUERY_STATES = {"COMPLETED", "ABANDONED"} -OK_QUERY_STATUS = "OK" - - @base_router.get("/asyncquery_status/{qid}", status_code=200) async def query_status( qid: str, diff --git a/tests/unit/test_query_status.py b/tests/unit/test_query_status.py index a68ff7b..37d9747 100644 --- a/tests/unit/test_query_status.py +++ b/tests/unit/test_query_status.py @@ -12,7 +12,13 @@ import pytest -from shepherd_server.base_routes import apply_query_status, query_status +from shepherd_server.base_routes import ( + ARATargetEnum, + QueryIntakeError, + apply_query_status, + query_status, + run_sync_query, +) logger = logging.getLogger(__name__) @@ -123,3 +129,78 @@ def test_success_claimed_for_a_failed_query_is_corrected(): response = {"message": {}, "status": "Success"} apply_query_status(response, "TIMEOUT") assert response["status"] == "Error" + + +# --- /query ---------------------------------------------------------------- +# +# The body has carried a TRAPI error status since apply_query_status went in, +# but the HTTP code stayed 200, so a caller checking the code rather than +# parsing the payload saw every failed query as a successful one. + + +def _patch_sync_query(mocker, row, response=None): + mocker.patch( + "shepherd_server.base_routes.run_query", + new_callable=mocker.AsyncMock, + return_value=("qid", "response_id", logger), + ) + mocker.patch( + "shepherd_server.base_routes.get_message", + new_callable=mocker.AsyncMock, + return_value=response, + ) + _patch_state(mocker, row) + + +@pytest.mark.asyncio +async def test_query_returns_200_for_a_healthy_query(mocker): + _patch_sync_query( + mocker, _row(state="COMPLETED", status="OK"), response={"message": {}} + ) + response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) + assert response.status_code == 200 + assert "status" not in _body(response) + + +@pytest.mark.parametrize("status", ["ERROR", "TIMEOUT", "Abandoned: no completion"]) +@pytest.mark.asyncio +async def test_query_returns_an_error_code_for_a_failed_query(mocker, status): + _patch_sync_query( + mocker, _row(state="COMPLETED", status=status), response={"message": {}} + ) + response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) + assert response.status_code == 500 + # The body still says which kind of failure it was. + assert _body(response)["status"] == "Error" + assert status in _body(response)["description"] + + +@pytest.mark.asyncio +async def test_query_returns_an_error_code_when_the_response_is_missing(mocker): + _patch_sync_query(mocker, _row(state="COMPLETED", status="OK"), response=None) + response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) + assert response.status_code == 500 + assert _body(response)["description"] == "Unable to get response" + + +@pytest.mark.asyncio +async def test_query_returns_an_error_code_when_it_times_out(mocker): + """The caller's own timeout elapsed with the query still in flight.""" + _patch_sync_query(mocker, _row(state="QUEUED", status="OK")) + response = await run_sync_query( + ARATargetEnum.ARAX, {"message": {}, "parameters": {"timeout": 0}} + ) + assert response.status_code == 500 + assert _body(response)["status"] == "TIMEOUT" + + +@pytest.mark.asyncio +async def test_query_returns_an_error_code_when_intake_fails(mocker): + mocker.patch( + "shepherd_server.base_routes.run_query", + new_callable=mocker.AsyncMock, + side_effect=QueryIntakeError("datastore unavailable"), + ) + response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) + assert response.status_code == 500 + assert "datastore unavailable" in _body(response)["description"] From 07f18b3b737413f4d47599863fd619d38b8ec95c Mon Sep 17 00:00:00 2001 From: Maxwell Wang Date: Wed, 2 Sep 2026 02:23:34 +0000 Subject: [PATCH 4/4] Answer /query with the code that describes the failure Reporting every post-intake failure as 500 kept /query inside the set of codes TRAPI 1.5 documents, but it told the caller less than it knows. Each way a query can fail now answers with the code that actually describes it: - an operation that errored stays 500, a genuine internal error; - a query that ran out of its budget (TIMEOUT), one reaped without ever completing (Abandoned: ...), and the caller's own timeout elapsing with the query still in flight are 504 -- Shepherd is fine, the work behind it didn't finish in time; - a query that was never accepted because the datastore was unavailable is 503, which says the retry its description asks for is worth making. The mapping from a stored query status to its code lives in query_status_code, so the two places that answer from a finished query's row agree. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01BmsMzWR6tNbDM2uSxh47n8 --- shepherd_server/base_routes.py | 48 ++++++++++++++++++++++----------- tests/unit/test_query_status.py | 39 ++++++++++++++++++++++----- 2 files changed, 66 insertions(+), 21 deletions(-) diff --git a/shepherd_server/base_routes.py b/shepherd_server/base_routes.py index 80e59f8..e93a015 100644 --- a/shepherd_server/base_routes.py +++ b/shepherd_server/base_routes.py @@ -35,7 +35,11 @@ setup_logging, ) from shepherd_utils.otel import setup_tracer -from shepherd_utils.task_deadline import deadline_field, query_deadline +from shepherd_utils.task_deadline import ( + TIMEOUT_STATUS, + deadline_field, + query_deadline, +) setup_logging() @@ -48,13 +52,17 @@ # abandoned, or times out. TERMINAL_QUERY_STATES = {"COMPLETED", "ABANDONED"} OK_QUERY_STATUS = "OK" -# What /query answers with once a query it accepted has failed. TRAPI 1.5 -# documents 200, 400, 429, 500 and 501 for this operation, so every -# post-intake failure -- an operation that errored, a query that ran out of -# budget, a response that isn't there -- reports the spec's InternalServerError -# rather than a more precise code (504) the schema doesn't allow. Which of -# those it was is in the response's own status/description. +# The status prefix the janitor writes when a query never completed within its +# budget (see the ABANDONED update in shepherd_utils.db). +ABANDONED_STATUS_PREFIX = "abandoned" +# What /query answers with for each way a query can fail. TRAPI 1.5 only +# documents 200/400/429/500/501 for this operation, but a caller is better +# served by the code that actually describes what happened: a query that ran +# out of time is not an internal error, and one that was never accepted because +# the datastore was unavailable is worth retrying. QUERY_ERROR_CODE = 500 +QUERY_TIMEOUT_CODE = 504 +QUERY_UNAVAILABLE_CODE = 503 class QueryIntakeError(Exception): @@ -199,6 +207,21 @@ async def run_query( return query_id, response_id, logger +def query_status_code(status: Optional[str]) -> int: + """The HTTP code describing how a query ended, from its stored status. + + A query that ran out of its budget (``TIMEOUT``) or was reaped without ever + completing (``Abandoned: ...``) is a gateway timeout: Shepherd is fine, the + work behind it didn't finish in time. Anything else non-OK is an operation + that failed, which is a genuine internal error. + """ + if not status or status == OK_QUERY_STATUS: + return 200 + if status == TIMEOUT_STATUS or status.lower().startswith(ABANDONED_STATUS_PREFIX): + return QUERY_TIMEOUT_CODE + return QUERY_ERROR_CODE + + def apply_query_status(response: dict, status: Optional[str]) -> None: """Stamp a non-OK query status onto the TRAPI response, in place. @@ -231,7 +254,7 @@ async def run_sync_query( except QueryIntakeError as e: return ORJSONResponse( content={"status": "ERROR", "description": str(e)}, - status_code=QUERY_ERROR_CODE, + status_code=QUERY_UNAVAILABLE_CODE, ) start = time.time() now = start @@ -269,12 +292,7 @@ async def run_sync_query( # checks the code (rather than parsing the payload for a status # field) saw every failed query as a successful one. return ORJSONResponse( - content=response, - status_code=( - 200 - if not status or status == OK_QUERY_STATUS - else QUERY_ERROR_CODE - ), + content=response, status_code=query_status_code(status) ) else: # Debug, not warning: this fires every 0.5s while a query is still @@ -286,7 +304,7 @@ async def run_sync_query( logger.error("Query timed out") return ORJSONResponse( content={"status": "TIMEOUT", "description": "Query timeout"}, - status_code=QUERY_ERROR_CODE, + status_code=QUERY_TIMEOUT_CODE, ) diff --git a/tests/unit/test_query_status.py b/tests/unit/test_query_status.py index 37d9747..5d53bec 100644 --- a/tests/unit/test_query_status.py +++ b/tests/unit/test_query_status.py @@ -17,6 +17,7 @@ QueryIntakeError, apply_query_status, query_status, + query_status_code, run_sync_query, ) @@ -162,14 +163,24 @@ async def test_query_returns_200_for_a_healthy_query(mocker): assert "status" not in _body(response) -@pytest.mark.parametrize("status", ["ERROR", "TIMEOUT", "Abandoned: no completion"]) +@pytest.mark.parametrize( + "status,code", + [ + # An operation failed: a genuine internal error. + ("ERROR", 500), + # Out of budget, or reaped without ever completing: the work behind + # Shepherd didn't finish in time. + ("TIMEOUT", 504), + ("Abandoned: no completion within budget", 504), + ], +) @pytest.mark.asyncio -async def test_query_returns_an_error_code_for_a_failed_query(mocker, status): +async def test_query_returns_the_code_for_how_it_failed(mocker, status, code): _patch_sync_query( mocker, _row(state="COMPLETED", status=status), response={"message": {}} ) response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) - assert response.status_code == 500 + assert response.status_code == code # The body still says which kind of failure it was. assert _body(response)["status"] == "Error" assert status in _body(response)["description"] @@ -190,17 +201,33 @@ async def test_query_returns_an_error_code_when_it_times_out(mocker): response = await run_sync_query( ARATargetEnum.ARAX, {"message": {}, "parameters": {"timeout": 0}} ) - assert response.status_code == 500 + assert response.status_code == 504 assert _body(response)["status"] == "TIMEOUT" @pytest.mark.asyncio -async def test_query_returns_an_error_code_when_intake_fails(mocker): +async def test_query_returns_unavailable_when_intake_fails(mocker): + """The query was never accepted, so the caller can retry it as-is.""" mocker.patch( "shepherd_server.base_routes.run_query", new_callable=mocker.AsyncMock, side_effect=QueryIntakeError("datastore unavailable"), ) response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) - assert response.status_code == 500 + assert response.status_code == 503 assert "datastore unavailable" in _body(response)["description"] + + +@pytest.mark.parametrize( + "status,code", + [ + (None, 200), + ("OK", 200), + ("ERROR", 500), + ("TIMEOUT", 504), + ("Abandoned: no completion within budget", 504), + ("something nobody writes today", 500), + ], +) +def test_query_status_code_mapping(status, code): + assert query_status_code(status) == code