From d4e59381c05c6d7c70232e82cc53c058ac049159 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 28 Aug 2026 20:53:20 +0000 Subject: [PATCH] Log callback delivery timing and failures in finish_query The async callback POST was a black box: one INFO line on success with no duration, and a failure line whose only detail was the raw exception. Nothing this worker logged was persisted either -- it acks directly rather than going through wrap_up_task, so its records never reached the query's log list. Now, after each attempt completes: - success logs the elapsed time, payload size and attempt number, plus a summary line when it took more than one attempt; - failure logs the elapsed time and a bounded reason -- HTTP status with the head of the server's response body, or the httpx exception class (connect and timeout errors routinely stringify to nothing); - a single "gave up" line records that the response was never delivered; - the same numbers go on the otel span as attributes, so timing is queryable without adding per-query log storage. A failed attempt's note is also spliced into the next attempt's payload, so a receiver that eventually gets the response can see the attempts that didn't make it. Splicing rebuilds the buffer, so it's skipped above 64 MiB -- the transient second copy would cost more than the note is worth, and the note is in the query's logs either way. Memory-wise this is a net improvement: the logs list and its serialization are dropped once they're inside the payload instead of being held for the length of the POST, the payload is released before the closing db round trips, and the log flush drains a handler queue that this worker previously left to accumulate. Also stop sleeping after the final failed attempt, and make the fallback serialization put "logs" last so the retry splice has a predictable tail. --- tests/unit/test_finish_query_callback_logs.py | 246 ++++++++++++++++++ workers/finish_query/worker.py | 203 +++++++++++++-- 2 files changed, 428 insertions(+), 21 deletions(-) create mode 100644 tests/unit/test_finish_query_callback_logs.py diff --git a/tests/unit/test_finish_query_callback_logs.py b/tests/unit/test_finish_query_callback_logs.py new file mode 100644 index 0000000..a0b47e2 --- /dev/null +++ b/tests/unit/test_finish_query_callback_logs.py @@ -0,0 +1,246 @@ +"""Tests for the callback timing/failure logging in ``finish_query``. + +Covers what the delivery logs have to answer after the fact: how long the +callback POST took, and -- when it didn't land -- why, both in the worker's +logs and in the payload the receiver eventually gets. +""" + +import json +import logging + +import httpx +import orjson +import pytest + +from workers.finish_query.worker import ( + CALLBACK_ERROR_BODY_BYTES, + CALLBACK_RETRIES, + _append_log_entry, + _describe_callback_failure, + finish_query, +) + +logger = logging.getLogger(__name__) + +TASK = [ + "test", + { + "query_id": "test", + "response_id": "rid", + "workflow": json.dumps([]), + "log_level": "20", + }, +] + + +def _patch_async_query(mocker, message=None, logs=None): + """Patch the db reads for an async (callback) query.""" + mocker.patch( + "workers.finish_query.worker.get_query_state", + new_callable=mocker.AsyncMock, + return_value=["", "", "", "", "", "", "", "rid", "http://callback"], + ) + mocker.patch( + "workers.finish_query.worker.set_query_completed", + new_callable=mocker.AsyncMock, + ) + mocker.patch( + "workers.finish_query.worker.get_message", + new_callable=mocker.AsyncMock, + return_value=orjson.dumps(message if message is not None else {"message": {}}), + ) + mocker.patch( + "workers.finish_query.worker.get_logs", + new_callable=mocker.AsyncMock, + return_value=logs if logs is not None else [], + ) + + +def _http_error_response(status_code: int, body: bytes) -> httpx.Response: + """A real httpx response, so ``raise_for_status`` raises the real error.""" + return httpx.Response( + status_code, + content=body, + request=httpx.Request("POST", "http://callback"), + ) + + +@pytest.mark.asyncio +async def test_successful_callback_logs_duration_and_size(redis_mock, mocker, caplog): + """A delivered callback logs how long the POST took and how big it was.""" + _patch_async_query(mocker) + mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=_http_error_response(200, b"ok"), + ) + + with caplog.at_level(logging.INFO): + await finish_query(TASK, logger) + + sent = [r.message for r in caplog.records if "Sent response back" in r.message] + assert len(sent) == 1 + assert "http://callback" in sent[0] + # ".s" duration and a byte count, both after the send completed. + assert "s (" in sent[0] and "bytes" in sent[0] + assert f"attempt 1/{CALLBACK_RETRIES}" in sent[0] + + +@pytest.mark.asyncio +async def test_failed_callback_logs_status_and_body(redis_mock, mocker, caplog): + """A rejected callback logs the status code and the server's explanation.""" + _patch_async_query(mocker) + mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=_http_error_response(502, b"upstream exploded"), + ) + mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock) + + with caplog.at_level(logging.INFO): + await finish_query(TASK, logger) + + failures = [ + r.message for r in caplog.records if "Failed to send callback" in r.message + ] + assert len(failures) == CALLBACK_RETRIES + assert "HTTP 502" in failures[0] + assert "upstream exploded" in failures[0] + assert f"attempt 1/{CALLBACK_RETRIES}" in failures[0] + # And a single summary line once the retries are spent. + gave_up = [ + r.message for r in caplog.records if "Gave up sending callback" in r.message + ] + assert len(gave_up) == 1 + assert "http://callback" in gave_up[0] + + +@pytest.mark.asyncio +async def test_failure_is_spliced_into_the_retry_payload(redis_mock, mocker): + """The receiver sees the attempts that didn't make it. + + A callback that fails and is retried can't carry its own failure, but it + can carry the previous attempt's -- so the eventual recipient knows the + response is late and why. + """ + _patch_async_query(mocker, logs=[{"message": "earlier", "level": "INFO"}]) + payloads = [] + + async def record(*args, **kwargs): + payloads.append(kwargs["content"]) + if len(payloads) < 3: + return _http_error_response(503, b"try again later") + return _http_error_response(200, b"ok") + + mocker.patch("httpx.AsyncClient.post", side_effect=record) + mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock) + + await finish_query(TASK, logger) + + assert len(payloads) == 3 + first, second, third = (orjson.loads(p) for p in payloads) + assert [entry["message"] for entry in first["logs"]] == ["earlier"] + assert len(second["logs"]) == 2 + assert "HTTP 503" in second["logs"][1]["message"] + assert second["logs"][1]["level"] == "ERROR" + assert second["logs"][1]["timestamp"] + # Each attempt adds only its own failure, and the message survives intact. + assert len(third["logs"]) == 3 + assert third["message"] == {} + + +@pytest.mark.asyncio +async def test_oversized_payload_skips_the_inline_retry_note(redis_mock, mocker): + """A huge payload isn't rebuilt just to carry a note that's in the logs.""" + _patch_async_query(mocker) + mocker.patch("workers.finish_query.worker.RETRY_LOG_SPLICE_MAX_BYTES", 1) + payloads = [] + + async def record(*args, **kwargs): + payloads.append(kwargs["content"]) + return _http_error_response(500, b"nope") + + mocker.patch("httpx.AsyncClient.post", side_effect=record) + mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock) + + await finish_query(TASK, logger) + + assert len(payloads) == CALLBACK_RETRIES + assert len(set(payloads)) == 1 + + +@pytest.mark.asyncio +async def test_callback_logs_are_persisted_for_the_query(redis_mock, mocker): + """The delivery outcome lands in the query's logs, not just stdout. + + ``finish_query`` acks directly instead of going through ``wrap_up_task``, + so it has to flush its own logs -- otherwise a failed callback is invisible + to anyone reading the query back. + """ + _patch_async_query(mocker) + mock_save_logs = mocker.patch( + "workers.finish_query.worker.save_logs", + new_callable=mocker.AsyncMock, + ) + mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=_http_error_response(200, b"ok"), + ) + + await finish_query(TASK, logger) + + mock_save_logs.assert_awaited_once_with("rid", logger) + + +@pytest.mark.asyncio +async def test_failing_to_save_logs_does_not_fail_the_task(redis_mock, mocker): + """A log flush that blows up must not take the whole wrap-up with it.""" + _patch_async_query(mocker) + mocker.patch( + "workers.finish_query.worker.save_logs", + new_callable=mocker.AsyncMock, + side_effect=Exception("redis down"), + ) + mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=_http_error_response(200, b"ok"), + ) + + await finish_query(TASK, logger) + + +def test_describe_callback_failure_truncates_the_body(): + """A rejecting server can return anything; only the head of it is logged.""" + error = httpx.HTTPStatusError( + "boom", + request=httpx.Request("POST", "http://callback"), + response=_http_error_response(413, b"x" * (CALLBACK_ERROR_BODY_BYTES * 10)), + ) + described = _describe_callback_failure(error) + assert described.startswith("HTTP 413: ") + assert len(described) < CALLBACK_ERROR_BODY_BYTES + 50 + + +def test_describe_callback_failure_names_silent_exceptions(): + """httpx's connect/timeout errors often stringify to nothing useful.""" + assert "ConnectTimeout" in _describe_callback_failure(httpx.ConnectTimeout("")) + assert "ConnectError" in _describe_callback_failure(httpx.ConnectError("")) + + +def test_append_log_entry_handles_both_array_shapes(): + """Empty and populated logs arrays both stay valid JSON.""" + entry = {"message": "late", "level": "ERROR"} + empty = orjson.dumps({"message": {}, "logs": []}) + assert orjson.loads(_append_log_entry(empty, entry))["logs"] == [entry] + + populated = orjson.dumps({"message": {}, "logs": [{"message": "first"}]}) + appended = orjson.loads(_append_log_entry(populated, entry))["logs"] + assert [e["message"] for e in appended] == ["first", "late"] + + +def test_append_log_entry_leaves_an_unexpected_tail_alone(): + """Rather than corrupt a payload we don't recognize, send it as-is.""" + payload = orjson.dumps({"logs": [], "message": {}}) + assert _append_log_entry(payload, {"message": "late"}) == payload diff --git a/workers/finish_query/worker.py b/workers/finish_query/worker.py index 2cbed87..98fb131 100644 --- a/workers/finish_query/worker.py +++ b/workers/finish_query/worker.py @@ -7,8 +7,10 @@ import uuid import orjson +from datetime import datetime, timezone + from opentelemetry.propagate import inject -from opentelemetry.trace import Status, StatusCode +from opentelemetry.trace import Status, StatusCode, get_current_span from shepherd_utils.broker import mark_task_as_complete from shepherd_utils.db import ( @@ -16,6 +18,7 @@ get_logs, get_message, get_query_state, + save_logs, set_query_completed, ) from shepherd_utils.shared import get_tasks @@ -30,6 +33,160 @@ tracer = setup_tracer(STREAM) LOGGER = get_worker_logger(STREAM) CALLBACK_RETRIES = 3 +CALLBACK_TIMEOUT = 120 +# How much of a rejecting server's response body goes into the failure log. The +# body is already in memory (we don't stream the response), but it can be an +# arbitrarily large HTML error page, and this string is copied into the query's +# logs -- and possibly into the retry payload -- so keep only the head of it. +CALLBACK_ERROR_BODY_BYTES = 500 +# Ceiling on the payload size for which we splice a failed attempt's note into +# the *next* attempt's body. Splicing rebuilds the whole buffer, so for a large +# response the transient second copy costs far more than the note is worth. The +# note is in the query's own logs either way, so oversized payloads just skip +# the inline copy. +RETRY_LOG_SPLICE_MAX_BYTES = 64 * 1024 * 1024 + + +def _log_entry(message: str, level: str = "ERROR") -> dict: + """Build a TRAPI LogEntry, matching ReasonerLogEntryFormatter's shape. + + Used for entries we splice straight into an outgoing payload, which never + pass through the logging handler that would otherwise format them. + """ + return { + "message": message, + "timestamp": datetime.now(timezone.utc).isoformat(), + "level": level, + } + + +def _describe_callback_failure(e: Exception) -> str: + """One bounded line explaining why a callback POST failed. + + The reason is the whole point of logging the failure -- "callback failed" + alone doesn't say whether the receiver is down, slow, or rejecting the + payload -- so pull out the status code and the head of the response body + for an HTTP error, and the exception type otherwise (httpx reports connect + failures, TLS errors and timeouts as distinct classes, and several of them + stringify to an empty message). + """ + if isinstance(e, httpx.HTTPStatusError): + detail = "" + try: + body = e.response.content[:CALLBACK_ERROR_BODY_BYTES] + if body: + detail = f": {body.decode('utf-8', 'replace')}" + except Exception: + # Body not readable (streamed/closed response) -- the status code + # is still worth reporting on its own. + pass + return f"HTTP {e.response.status_code}{detail}" + if isinstance(e, httpx.TimeoutException): + return f"{type(e).__name__} (no response within {CALLBACK_TIMEOUT}s)" + return f"{type(e).__name__}: {e}" + + +def _append_log_entry(payload: bytes, entry: dict) -> bytes: + """Return ``payload`` with ``entry`` appended to its trailing logs array. + + Only sound for a payload this worker built, which always ends with the logs + array followed by the closing brace. Rebuilding costs a transient second + copy of the payload, so callers guard on size; the rebind releases the old + buffer immediately. If the payload doesn't have the expected tail, hand it + back untouched rather than risk shipping malformed JSON. + """ + entry_bytes = orjson.dumps(entry) + if payload.endswith(b"[]}"): + return payload[:-3] + b"[" + entry_bytes + b"]}" + if payload.endswith(b"]}"): + return payload[:-2] + b"," + entry_bytes + b"]}" + return payload + + +async def send_callback( + callback_url: str, + message_bytes: bytes, + logger: logging.Logger, +) -> bool: + """POST the finished response to the caller's callback URL. + + Every attempt is timed and logged *after* the send completes -- how long a + callback takes is a property of the receiver we otherwise have no record + of, and a failure is only actionable with the reason attached. Failures are + also spliced into the next attempt's payload (size permitting), so a + receiver that eventually gets the response can see the attempts that didn't + make it. + + Returns True if the response was delivered. + """ + headers = {"Content-Type": "application/json"} + # Propagate the otel trace context through the callback. + # Matches the inject() carrier pattern used by the + # lookup workers; the active span comes from process_task's + # start_as_current_span. + inject(headers) + span = get_current_span() + started = time.time() + payload_size = len(message_bytes) + delivered = False + attempts = 0 + for attempt in range(1, CALLBACK_RETRIES + 1): + attempts = attempt + attempt_start = time.time() + try: + async with httpx.AsyncClient(timeout=CALLBACK_TIMEOUT) as client: + response = await client.post( + callback_url, + content=message_bytes, + headers=headers, + ) + response.raise_for_status() + elapsed = time.time() - attempt_start + logger.info( + f"Sent response back to {callback_url} in {elapsed:.3f}s " + f"({len(message_bytes)} bytes, " + f"attempt {attempt}/{CALLBACK_RETRIES})" + ) + delivered = True + break + except Exception as e: + elapsed = time.time() - attempt_start + failure = ( + f"Failed to send callback to {callback_url} after {elapsed:.3f}s " + f"(attempt {attempt}/{CALLBACK_RETRIES}, " + f"{len(message_bytes)} bytes): {_describe_callback_failure(e)}" + ) + logger.error(failure) + span.add_event( + "callback_attempt_failed", + {"attempt": attempt, "duration_ms": int(elapsed * 1000)}, + ) + if attempt < CALLBACK_RETRIES: + if len(message_bytes) <= RETRY_LOG_SPLICE_MAX_BYTES: + message_bytes = _append_log_entry( + message_bytes, _log_entry(failure) + ) + await asyncio.sleep(1 * (2 ** (attempt - 1))) + + total = time.time() - started + if not delivered: + logger.error( + f"Gave up sending callback to {callback_url} after " + f"{CALLBACK_RETRIES} attempts and {total:.3f}s. The response was " + "not delivered." + ) + elif attempts > 1: + logger.info( + f"Callback to {callback_url} succeeded on attempt {attempts} " + f"after {total:.3f}s total." + ) + # Attributes rather than another log line: same numbers, no per-query log + # storage, and they're queryable alongside the rest of the trace. + span.set_attribute("callback.duration_ms", int(total * 1000)) + span.set_attribute("callback.attempts", attempts) + span.set_attribute("callback.payload_bytes", payload_size) + span.set_attribute("callback.delivered", delivered) + return delivered async def finish_query(task, logger: logging.Logger): @@ -63,29 +220,22 @@ async def finish_query(task, logger: logging.Logger): ) else: message = orjson.loads(message_bytes) + # Re-insert rather than assign in place so "logs" is last in + # the serialized payload -- send_callback appends retry notes + # by rewriting the payload's tail. + message.pop("logs", None) message["logs"] = logs message_bytes = orjson.dumps(message) del message - headers = {"Content-Type": "application/json"} - # Propagate the otel trace context through the callback. - # Matches the inject() carrier pattern used by the - # lookup workers; the active span comes from process_task's - # start_as_current_span. - inject(headers) - for attempt in range(CALLBACK_RETRIES): - try: - async with httpx.AsyncClient(timeout=120) as client: - response = await client.post( - callback_url, - content=message_bytes, - headers=headers, - ) - response.raise_for_status() - logger.info(f"Sent response back to {callback_url}") - break - except Exception as e: - logger.error(f"Failed to send callback to {callback_url}: {e}") - await asyncio.sleep(1 * (2**attempt)) + # The logs list and its serialization are a full second copy of + # every log line the query produced; they're inside the payload + # now, so drop them before the send rather than holding them for + # its duration. + del logs, logs_bytes + + await send_callback(callback_url, message_bytes, logger) + # Release the payload before the remaining db round trips. + del message_bytes await set_query_completed(query_id, status, logger) @@ -98,6 +248,17 @@ async def finish_query(task, logger: logging.Logger): logger.info(f"Finished task {task[0]} in {time.time() - start}") + # This worker acks directly instead of going through wrap_up_task, so + # nothing else flushes what it logged. Persist here so the callback + # outcome -- how long delivery took, or why it failed -- survives in the + # query's logs (GET /response/{query_id}) instead of only in the pod's + # stdout. Draining also clears the handler's queue, which for this + # process-wide logger would otherwise just accumulate. + try: + await save_logs(response_id, logger) + except Exception as e: + logger.error(f"Failed to save logs for {response_id}: {e}") + async def process_task(task, parent_ctx, logger: logging.Logger, limiter): """Process a given task and ACK in redis."""