Metadata-Version: 2.5
Name: async-task-api-worker
Version: 0.1.0-rc.2
Summary: SDK to write workers for the AsyncTaskAPI platform: protocol models, runtime and conformance kit.
Project-URL: Repository, https://github.com/IA-Generative/async-api
License: MIT
License-File: LICENSE
Keywords: async-api,rabbitmq,sdk,worker
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Typing :: Typed
Requires-Python: >=3.13
Requires-Dist: aio-pika>=9.5.5
Requires-Dist: prometheus-client>=0.22.1
Requires-Dist: pydantic-settings>=2.10.1
Requires-Dist: pydantic>=2.10
Description-Content-Type: text/markdown

# async-task-api-worker

SDK to write workers for the AsyncTaskAPI platform.

This is a work in progress: the package ships the protocol models (inbound envelope decoder and the byte-exact
outbound message encoder), the error types, the settings, the runtime engine with the `Worker` definition API, the
RabbitMQ transport, the process lifecycle (graceful drain), the probes, the metrics and the JSON logs. The conformance
kit and the extras land in later releases.

## Declaring a worker

```python
from pydantic import BaseModel, Field

from async_task_api_worker import TaskContext, Worker


class Input(BaseModel):
    text: str = Field(min_length=1, max_length=100_000)


class Output(BaseModel):
    words: int


worker = Worker("count-words")


@worker.task(timeout=30.0)
async def handle(task_input: Input, ctx: TaskContext) -> Output:
    await ctx.progress(0.5)
    return Output(words=len(task_input.text.split()))
```

- `async def` handlers run on the event loop; `def` handlers (with a `SyncTaskContext`) run in a thread pool. The
  context parameter must be annotated exactly `TaskContext` or `SyncTaskContext` (an identity check: no subclass or
  `Annotated` wrapper). Hooks and probes (`on_startup`, `on_shutdown`, `readiness_probe`) must be `async def`.
- The input and output models are read from the annotations. Exactly one task per worker.
- Raise `FatalError(code, message)` to fail a task with a code the client can match, `RetryableError` for a
  transient failure. Any other exception is reported as `INTERNAL_ERROR` without its text.
- `ctx.progress(fraction, payload)` reports the **latest state** of the task, not an event stream: updates are
  throttled (at most one per `WORKER_PROGRESS_MIN_INTERVAL_S`, default 1 s), a newer update replaces a held one,
  `fraction` is in `[0, 1)` and never goes back. The terminal success or failure message is the 100%.

## Running a worker

```bash
export BROKER_URL=amqp://user:password@broker:5672/ IN_QUEUE_NAME=count-words_queue_in OUT_QUEUE_NAME=count-words_queue_out
async-task-api-worker run my_module:worker      # or: python -m async_task_api_worker run my_module:worker
```

`my_module:worker` is a `Worker`, or a zero-argument factory function returning one; modules in the current directory can
be imported. A definition error in the module exits with code 2 and one line. From code, `worker.run()` blocks until the
process is stopped and `await worker.serve()` serves from an existing event loop.

- The settings are read and validated first: an invalid configuration, or a task timeout that does not fit under the
  broker's consumer timeout (`task timeout + WORKER_RETRY_BACKOFF_MAX_S + 60 s < WORKER_CONSUMER_TIMEOUT_S`), exits with
  code 2 before any connection. Exit codes: 0 clean stop, 1 failure, 2 configuration or usage error.
- The queues are declared durable and with no arguments (the broker decides their type and policies), and consumed with
  manual acknowledgements: a task is acked only after its result is published.
- SIGTERM or SIGINT stops consuming, lets the in-flight tasks finish within `WORKER_SHUTDOWN_GRACE_S` (default 25 s),
  requeues the ones that did not, closes the connection and exits 0. A second signal skips the rest of the grace period.
  A `def` handler whose thread is still running after the drain, or an `async def` handler whose `asyncio.to_thread`
  thread is (given 2 s after the drain), is logged and the process exits at once with code 1 (only for `run`, never for
  `serve()`). The whole shutdown is bounded: the process exits at most
  `max(WORKER_SHUTDOWN_GRACE_S, 5)` + 17 s after the signal (5 s for the nacks of the cancelled tasks, 5 s to close the
  connection, and 5 s in total, `WORKER_SHUTDOWN_HOOKS_TIMEOUT_S`, for all the shutdown hooks together, each getting what is
  left). Set a pod's `terminationGracePeriodSeconds` to at least that. A signal during startup does not wait for it: the
  connection attempt and the startup hooks are cancelled and the process exits 0 (a hook that blocks the event loop
  cannot be interrupted).
- **Retries and crash-loop protection need quorum queues** (the broker's default queue type, as in the platform's
  chart). A classic queue gives no delivery counter: there, a `RetryableError` gets one retry, then concludes with its
  own code, and a warning says so.
- Refused broker credentials exit with code 2 right away instead of being retried.
- The attempt number of a delivery (`ctx.attempt`) counts every time the broker handed the message out. A
  `RetryableError` requeues until attempt `WORKER_DELIVERY_LIMIT`, then concludes with its own code. A delivery past
  that limit (a crash loop) is concluded with `TASK_ABANDONED_AFTER_REDELIVERIES` without running the handler.

## Probes, metrics and logs

A small HTTP server runs on `HEALTH_CHECK_HOST:HEALTH_CHECK_PORT` (default `0.0.0.0:8083`) in both modes, from the very
start of the process until the end of the shutdown.

| Path | Answer |
|---|---|
| `GET /health` | Liveness. 200 `{"status": "ok"}`, or 503 `TASK_OVERRUN` when a `def` handler is still running `WORKER_TASK_OVERRUN_GRACE_S` after its deadline. It consults no dependency. |
| `GET /ready` | Readiness. 200 with a `components` map, or 503 `DEPENDENCIES_NOT_READY` (with the same `components` in `error`) when the broker is disconnected, the worker is draining, or one of the probes registered with `@worker.readiness_probe("name")` raises or exceeds `WORKER_PROBE_TIMEOUT_S` (3 s). Probes run in parallel. A `ProbeError` message is shown as the detail; any other exception only by its class name. |
| `GET /metrics`, `GET /internal/metrics` | Prometheus text exposition. |
| `GET /__version__` | `/app/version.json` of the image, or a `dev` document. |

Errors use the platform's format, `{"status": "error", "error": {"code", "message"}}` (404 `NOT_FOUND`, 405
`METHOD_NOT_ALLOWED`, 500 `PROBE_FAILED`). `broker` and `draining` are reserved component names. `/ready` is also 503 in
job mode once the quota is reached and the consumer cancelled.

The server is bounded: at most 32 connections are served at once (the others are closed at once), reading the request,
writing the answer and closing each take at most 5 s, and concurrent `/ready` requests share a single evaluation of the
probes.

**Startup.** `/health` answers while the startup hooks run, but only as long as the event loop is free: a hook that loads a
model or a large file synchronously blocks it, `/health` stops answering and the kubelet restarts the pod in a loop. Run
heavy loading in `asyncio.to_thread`, or give the container a `startupProbe`.

Metrics have no task, client or service label (the scrape adds the pod and service): `worker_tasks_total{status}`,
`worker_task_duration_seconds`, `worker_messages_published_total{message_type}`, `event_loop_lag_seconds{component}`,
`worker_task_failures_total{code}`, `worker_task_retries_total`, `worker_poison_messages_total{reason}`,
`worker_publish_failures_total{message_type}`, `worker_ack_failures_total`, `worker_tasks_in_flight`,
`worker_shutdown_requeued_total`, `worker_progress_dropped_total{reason}` and `worker_sdk_info`, plus the process, platform
and GC collectors.

Logs are one JSON object per line on stdout (`ts`, `level`, `logger`, `msg`, and `task_id`, `service`, `attempt` when
bound; `ctx.logger` carries them, and so do the records of any library logging inside a handler, on the SDK's handler).
`LOG_LEVEL` sets the level. If the application configured the root logger before running the worker, its handlers are
kept and the SDK adds none (no line is printed twice): the format is then the application's, and so is the level of
the `aio_pika.tools` logger, whose two ERROR lines on a dropped connection the SDK otherwise logs as warnings (the
connection then reconnects by itself). To still get the task fields on your own handler, chain a record factory:

```python
import logging
from async_task_api_worker.observability.logging import TaskContextFilter

task_fields = TaskContextFilter()
previous_factory = logging.getLogRecordFactory()


def factory(*args, **kwargs):
    record = previous_factory(*args, **kwargs)
    task_fields.filter(record)
    return record


logging.setLogRecordFactory(factory)
```

Every delivery ends with exactly one INFO decision line, for instance `Delivery settled: outcome=failure action=conclude
code=NOT_ALLOWED attempt=1/5`, with the fields `task_id`, `attempt`, `delivery_limit`, `outcome` (`success`, `failure`,
`poison`, `retry`, `publish_failed`, `shutdown_requeued`, `cancelled`, `engine_error`), `action` (`ack`, `conclude`,
`nack-requeue`) and `code`: `grep <task_id>`
tells the whole story of a task. Payloads and secrets are never logged.

## Deployment or Job

The same worker code and image run as a long-lived consumer (a Kubernetes Deployment) or as a one-shot process (a
Job, for instance one started per message by KEDA's `ScaledJob`). The deployment picks the mode with `WORKER_MODE`; the
`Worker` definition does not change.

| Variable | Default | Meaning |
|---|---|---|
| `WORKER_MODE` | `service` | `service` consumes until SIGTERM; `job` is one-shot. Any other value exits with code 2. |
| `WORKER_JOB_MAX_TASKS` | `1` | Job mode: stop consuming at this many deliveries, finish them, and exit. |
| `WORKER_JOB_IDLE_TIMEOUT_S` | `10` | Job mode: exit when no delivery is in hand and none arrives for this long (another pod probably took the message). |
| `WORKER_JOB_CONNECT_TIMEOUT_S` | `30` | Job mode: give up and exit 1 when the broker stays unreachable for this long. Service mode retries forever. |

In job mode the prefetch is forced to 1, whatever `WORKER_CONCURRENCY` says (it then only sizes the handler semaphore
and the thread pool). The quota counts deliveries **on arrival**: any delivery counts, poison ones and redeliveries past
the limit included, and consuming stops as soon as the last one is in hand. The exit code reports the **pod**, not the
task: a task that succeeds, fails (`FatalError`, timeout, invalid input), or is requeued (`RetryableError`) all exit 0,
and so does an idle exit, because a failed task is a business outcome already published to the client. Exit 1 (an
infrastructure failure such as a broker unreachable for `WORKER_JOB_CONNECT_TIMEOUT_S`, or a `def` handler stuck after
the drain) and 2 (invalid configuration) stay non-zero, so the Job's `backoffLimit` applies to them. SIGTERM and
`activeDeadlineSeconds` drain exactly as in service mode.

## Testing your worker

`async_task_api_worker.testing` runs your task in memory through the real engine (decoding, time limit, retries, error
mapping, acknowledgements), with no broker and no Docker:

```python
from pydantic import BaseModel

from async_task_api_worker import FatalError, TaskContext, Worker
from async_task_api_worker.testing import Outcome, ProgressRecorder, WorkerHarness, make_task_context


class Input(BaseModel):
    text: str


class Output(BaseModel):
    words: int


worker = Worker("count-words")


@worker.task(timeout=30.0)
async def handle(task_input: Input, ctx: TaskContext) -> Output:
    if not task_input.text.strip():
        raise FatalError("EMPTY_TEXT", "There is nothing to count.")
    await ctx.progress(0.5)
    return Output(words=len(task_input.text.split()))


async def test_counts_the_words() -> None:
    async with WorkerHarness(worker) as harness:
        run = await harness.submit({"text": "hello big world"})

    run.assert_succeeded()
    assert run.response == {"words": 3}
    assert run.message_types == ("started", "progress", "success")


async def test_an_empty_text_fails_with_its_code() -> None:
    async with WorkerHarness(worker) as harness:
        run = await harness.submit({"text": "  "})

    run.assert_failed(code="EMPTY_TEXT")
    assert run.outcome is Outcome.ACKED


async def test_the_handler_alone_with_a_context_you_build() -> None:
    recorder = ProgressRecorder()
    ctx = make_task_context(recorder=recorder)

    assert await handle(Input(text="a b"), ctx) == Output(words=2)
    assert [update.fraction for update in recorder.updates] == [0.5]
```

- `harness.submit(body, client_id=..., input_refs=..., delivery_count=...)` builds the envelope the API publishes;
  `submit_raw(bytes)` sends bytes as they are (a poison message). The result is a `TaskRun`: `messages` (decoded) and
  `raw` (the published bytes), `message_types` (all three are tuples), `outcome` (`ACKED`, `NACKED_REQUEUE`), `response`,
  `error`, and the two assertions. `task_id` is the envelope's, or `None` for bytes that do not carry one.
- `delivery_count` is the broker's `x-delivery-count`: pass `0`, then `delivery_count=n` to see a `RetryableError`
  requeue until the last attempt (`worker_delivery_limit`), which concludes with the error's own code. Retry backoffs are
  a millisecond; every other setting is the production default (progress throttling included).
- `settings_overrides={"worker_task_timeout_s": 0.1, ...}` takes `WorkerSettings` field names; no environment variable is
  read, and the time budget is checked as at startup. `harness.fail_next_publish("success")` makes the next publish (of
  that message type, or of any) fail once, to test the requeue. Use `async with` (or `await harness.aclose()`).
- `make_task_context` and `make_sync_task_context` build the context of an `async def` and of a `def` handler, to call
  the handler directly; `ProgressRecorder` collects what it reports.

## Development

From this directory:

```bash
uv sync --locked --all-extras
uv run ruff check . && uv run ruff format --check .
uv run ty check src
uv run pytest -m "not integration" -q --cov=async_task_api_worker --cov-branch --cov-fail-under=95
```

The integration tests need Docker (they start RabbitMQ with the platform's policy through testcontainers):

```bash
uv run pytest -m integration -q
```

Add dependencies with `uv add` from this directory, never by editing `pyproject.toml`.
