From 5824422c424c083cd4b5b8acce7f8fc360c686b2 Mon Sep 17 00:00:00 2001 From: Shudipto Trafder Date: Sat, 7 Nov 2026 11:04:04 +0600 Subject: [PATCH 1/4] feat: Implement worker module with task management and middleware integration - Added `worker.py` to handle worker startup and shutdown, including task importing and broker configuration. - Enhanced `GraphConfig` to parse worker configuration from `10xgraph.json`. - Created unit tests for worker configuration, context management, and middleware functionality. - Developed integration tests for Redis broker functionality, including idle queue reclaim and lease renewal. - Established a testing environment with fixtures for task execution and state management. - Implemented error handling and logging for task execution and middleware processes. --- JOBS_TASKIQ_PLAN.md | 232 ++++++++++++++++++ JOBS_TASKIQ_SPRINT0.md | 72 ++++++ pyproject.toml | 8 +- tenxgraph_api/jobs/__init__.py | 39 +++ tenxgraph_api/jobs/arguments.py | 92 +++++++ tenxgraph_api/jobs/config.py | 172 +++++++++++++ tenxgraph_api/jobs/context.py | 114 +++++++++ tenxgraph_api/jobs/middleware.py | 133 ++++++++++ tenxgraph_api/jobs/monitoring.py | 70 ++++++ tenxgraph_api/jobs/redis_broker.py | 144 +++++++++++ tenxgraph_api/jobs/schedules.py | 128 ++++++++++ tenxgraph_api/jobs/state.py | 139 +++++++++++ tenxgraph_api/jobs/worker.py | 140 +++++++++++ .../src/app/core/config/graph_config.py | 14 +- tests/unit_tests/jobs/__init__.py | 0 tests/unit_tests/jobs/conftest.py | 65 +++++ tests/unit_tests/jobs/test_config.py | 101 ++++++++ tests/unit_tests/jobs/test_context.py | 176 +++++++++++++ tests/unit_tests/jobs/test_middleware.py | 217 ++++++++++++++++ tests/unit_tests/jobs/test_redis_broker.py | 175 +++++++++++++ tests/unit_tests/jobs/test_worker.py | 197 +++++++++++++++ 21 files changed, 2426 insertions(+), 2 deletions(-) create mode 100644 JOBS_TASKIQ_PLAN.md create mode 100644 JOBS_TASKIQ_SPRINT0.md create mode 100644 tenxgraph_api/jobs/__init__.py create mode 100644 tenxgraph_api/jobs/arguments.py create mode 100644 tenxgraph_api/jobs/config.py create mode 100644 tenxgraph_api/jobs/context.py create mode 100644 tenxgraph_api/jobs/middleware.py create mode 100644 tenxgraph_api/jobs/monitoring.py create mode 100644 tenxgraph_api/jobs/redis_broker.py create mode 100644 tenxgraph_api/jobs/schedules.py create mode 100644 tenxgraph_api/jobs/state.py create mode 100644 tenxgraph_api/jobs/worker.py create mode 100644 tests/unit_tests/jobs/__init__.py create mode 100644 tests/unit_tests/jobs/conftest.py create mode 100644 tests/unit_tests/jobs/test_config.py create mode 100644 tests/unit_tests/jobs/test_context.py create mode 100644 tests/unit_tests/jobs/test_middleware.py create mode 100644 tests/unit_tests/jobs/test_redis_broker.py create mode 100644 tests/unit_tests/jobs/test_worker.py diff --git a/JOBS_TASKIQ_PLAN.md b/JOBS_TASKIQ_PLAN.md new file mode 100644 index 00000000..089ef125 --- /dev/null +++ b/JOBS_TASKIQ_PLAN.md @@ -0,0 +1,232 @@ +# Jobs on TaskIQ: plan + +Status: agreed 2026-11-06. Branch `feat/jobs-taskiq` (from `main`; old code stays on `worker`). Sprint 0 done (findings in JOBS_TASKIQ_SPRINT0.md); Sprint 1 done (Redis only; other brokers later); Sprint 2 next. Replaces JOBS_PLAN.md / JOBS_DESIGN.md (own work +manager), which are kept only for history. + +## Idea + +We do not build a job system. 10xgraph owns one TaskIQ broker, built from `10xgraph.json`. The +developer imports it and writes tasks. We add what TaskIQ lacks: the broker setup (InjectQ, +monitoring, retry, OTEL already on), `worker` and `scheduler` commands, HTTP routes with our auth, +owner-scoped status, the run config in tasks, and cancel by signal. + +Target size: about 500 lines of code in `tenxgraph_api/jobs/`, plus tests. + +## What the developer writes + +### Tasks + +```python +# app/report_tasks.py +from typing import Annotated + +from injectq.integrations.taskiq import InjectTaskiq +from taskiq import TaskiqDepends +from tenxgraph.core.graph import CompiledGraph +from tenxgraph_api.jobs import broker, TaskContext + + +@broker.task(http=True) # http=True: can be started or scheduled through the API; default False +async def save_report( + report_id: int, # caller sets this + config: dict, # we fill it; optional + graph: Annotated[CompiledGraph, InjectTaskiq(CompiledGraph)], # InjectQ fills it + ctx: Annotated[TaskContext, TaskiqDepends()], # we fill it; optional +) -> dict: + await ctx.progress({"step": "loading"}) + if await ctx.cancelled(): + return {"cancelled": True} + await graph.ainvoke({"messages": [...]}, config=config) # runs as the caller + return {"ok": True} + + +@broker.task(schedule=[{"cron": "0 2 * * *"}]) # static cron, plain TaskIQ +async def nightly_cleanup() -> None: ... +``` + +The developer surface is `broker` and `TaskContext` from us; injection comes from InjectQ itself. +No broker classes, no middlewares, no URLs in code. + +`TaskContext` (optional): + +- `task_id`, `task_name` +- `config`: the run config (same value a `config` parameter gets; always has `task_id`) +- `progress(meta)`: save progress, shown by `GET /v1/tasks/{id}` +- `cancelled()`: true after `POST /v1/tasks/{id}/cancel` + +### The run config (same rule as core nodes and tools) + +Core reads the function signature and passes `state` / `config` only when the function declares +them (`invoke_node_handler._prepare_input_data` for nodes, `KwargsResolverMixin._prepare_kwargs` +with `INJECTABLE_PARAMS` for tools). A function that does not declare them is called without them. +Tasks follow the same rule: + +- A task that declares `config` gets the run config; one that does not is called without it. + `ctx.config` is the same value for tasks that take `TaskContext`. +- The config is built on the server at call time, like a graph run: `user_id`, `user` (what + `BaseAuth` returned), `authz` (the authorization result), plus `task_id`. It travels in a + message label (`tenxgraph_config`, a JSON string: labels keep only int/str/float/bool) that only + the server writes. +- Schedules store the config of the user who created them, so each scheduled run acts as them. +- Python callers pass it with `await jobs.enqueue(task, *args, config=config, **kwargs)`. + Plain `.kiq()` works too; then `config` is `{}`. + +### Reserved arguments + +A caller (HTTP body or `jobs.enqueue`) can set only the task's plain arguments. These are refused +with 422: + +- `config`, and the keys `user`, `user_id`, `authz` anywhere at the top level of `kwargs` +- any parameter that is injected (`InjectTaskiq(...)`, `TaskiqDepends()`, `TaskContext`) +- names the task does not have, and extra positional args + +This is a security check, not tidiness: TaskIQ applies caller kwargs over injected values +(`kwargs.update(message.kwargs)` in its receiver), so `graph=` from a caller would replace the +injected graph. Checked in the API and `enqueue`, and again in the worker middleware. + +Only `InjectTaskiq` injects; a plain `Inject[...]` default passes the marker object silently. +`TaskContext` is a plain class (TaskIQ reads `__init__`); progress and cancel need `async def`. + +### Config + +```json +"worker": { + "broker": {"type": "redis_stream", "url": "${REDIS_URL}", "options": {}}, + "result_backend": {"type": "redis", "url": "${REDIS_URL}", "options": {}}, + "tasks": ["app/**/*_tasks.py"], + "schedules": true, + "result_ttl": 86400, + "retries": 3, + "prometheus": {"host": "0.0.0.0", "port": 9000} +} +``` + +- `broker.type`: a name for a TaskIQ-supported broker (`redis_stream` default, `redis_list`, + `rabbitmq`, `nats`, `kafka`, `memory`) or `"module:attr"` for any other TaskIQ broker class. + `url` is passed to it; `options` are extra constructor arguments (for example `idle_timeout`). +- `result_backend.type`: `redis` (default), `memory`, or `"module:attr"`. Omitted: Redis on the + broker URL when the broker is Redis, else required. +- Each named type needs its TaskIQ package (`taskiq-redis`, `taskiq-aio-pika`, ...). A missing + package refuses to boot and prints the pip command. +- `tasks`: module names or file globs; worker, scheduler and API import them so tasks register. +- `schedules`: turns on dynamic schedules (TaskIQ `ListRedisScheduleSource`, needs Redis). +- `result_ttl`: how long status, progress and results are kept. +- `retries`: the cap for TaskIQ's retry middleware. Retries are opt-in per task + (`@broker.task(retry_on_error=True)`), so tasks that are not safe to repeat never re-run. +- Status records, cancel flags and schedule records live in the API's cache store (`jobs:` keys). + Schedule listing needs an owner index; the cache store has no list/scan, designed in Sprint 3. +- Known limit: the cancel and argument guard does not reach `--use-process-pool` workers. +- Prometheus is off unless set. OpenTelemetry follows the existing `observability` block. +- `10xgraph init` (Production) and `10xgraph config` ask for broker type and URL, result backend + type and URL, and the task file pattern. + +### Commands + +```bash +10xgraph dev # API; starts and schedules tasks, does not run them +10xgraph worker # runs tasks +10xgraph scheduler # sends scheduled tasks; run exactly ONE (TaskIQ rule) +``` + +### From Python + +```python +from tenxgraph_api import jobs + +await jobs.enqueue(save_report, 42, config=config) # run now +sid = await jobs.schedule(save_report, 42, cron="0 9 * * 1", config=config) +await jobs.reschedule(sid, cron="0 10 * * 1") +await jobs.unschedule(sid) +``` + +## HTTP (same auth and authorization as the graph routes) + +One route family for tasks and schedules. + +| Route | Does | Permission | +|---|---|---| +| `POST /v1/tasks/{name}` | body `{args, kwargs}`: run now, returns `{task_id}`. Body `{args, kwargs, schedule}`: schedule it, returns `{schedule_id}`. `schedule` is one of `{"cron": "..."}`, `{"at": ""}`, `{"every": }` | `tasks:invoke` | +| `GET /v1/tasks/{id}` | task: `{status, progress, result, error, created_at}`. Schedule: `{name, schedule, args, kwargs, created_at}` | `tasks:read` | +| `GET /v1/tasks?kind=schedule` | the caller's schedules | `tasks:read` | +| `PATCH /v1/tasks/{schedule_id}` | change the time only: body `{schedule}` (delete + add with the same id; TaskIQ has no update) | `tasks:invoke` | +| `POST /v1/tasks/{id}/cancel` | task: set the cancel flag. Schedule: delete it | `tasks:cancel` | + +- Only tasks marked `http=True` can be started or scheduled over HTTP; others answer 404. +- Webhooks use `POST /v1/tasks/{name}`, or the developer's own route calling `jobs.enqueue`. +- Status values: `queued`, `running`, `success`, `failure`, `cancelled`. Kept in a small metadata + record per task (owner, task name, created time, status), written at start and updated by our + middleware; result and progress come from the TaskIQ result backend. All of it expires after + `result_ttl`. +- Owner scope: a user sees and changes only their own tasks and schedules; anything else is 404. +- Cancel is a signal: a queued task is not run (our middleware marks it and a guard around the + task function raises `TaskCancelled` before the body; a middleware cannot skip a task and must + never raise); a running task stops when it checks `ctx.cancelled()`. +- Schedule listing uses our own owner index (TaskIQ `get_schedules()` hides future one-time + schedules). The scheduler picks up changes within `--update-interval` (we default it to 10 s). + +## How it works inside: one worker file we own + +`tenxgraph_api/jobs/worker.py`, the same shape as a hand-written TaskIQ `worker.py`: + +1. **Broker, built on import from the `worker` config.** + - `redis_stream`: taskiq-redis `RedisStreamBroker` + `RedisAsyncResultBackend` + (`result_ex_time` = `result_ttl`). It acks after the task and `XAUTOCLAIM`s messages idle + longer than `idle_timeout` (default 10 min), so a dead worker's task runs again elsewhere + (verified in `taskiq_redis/redis_broker.py`). Ours is a small subclass: upstream runs the + autoclaim only after a read that returned messages, so on an idle queue a dead worker's task + waits until new work arrives; our `listen()` also claims when nothing was read. We set + `unacknowledged_lock_timeout`. It also renews the lease of every message it holds every + `idle_timeout/3` (Lua: reset idle only while this consumer still owns it), so a long task is + never run twice while its worker lives. `ListQueueBroker` pops the message, so a task + dies with its worker; it stays available as `redis_list` but is not the default. + - `memory`: `InMemoryBroker`, for tests and local dev. +2. **Middlewares, always on:** `SimpleRetryMiddleware` (count from config); our + `MonitoringMiddleware` (task started / finished / failed events to the core publisher); our + `TaskMiddleware` (status record, cancel check, `config` from the label into kwargs when + declared). Optional from config: OpenTelemetry and Prometheus (both ship with TaskIQ). +3. **InjectQ binding:** `setup_taskiq(container, broker)` with the container built from the same + `10xgraph.json` as the API (`agent`, `checkpointer`, `store`, `injectq`, `cache_store`, + `id_generator`). +4. **`WORKER_STARTUP`:** load the graph, open checkpointer / store / cache store, ping them. + **`WORKER_SHUTDOWN`:** close them in reverse order, each in its own `try`. +5. **Task discovery:** import the modules or globs in `worker.tasks`. + +Project-specific setup (their DB, their own sweeps) stays plain TaskIQ in their code: +`@broker.on_event(TaskiqEvents.WORKER_STARTUP)`. We add no hook API of our own. + +The API process imports the same `broker` and only sends. + +## Files + +| File | Holds | Size goal | +|---|---|---| +| `jobs/__init__.py` | exports `broker`, `TaskContext`, `enqueue`, `schedule`, `reschedule`, `unschedule` | small | +| `jobs/config.py` | pydantic `WorkerConfig`, broker / result backend factory by type | ~120 | +| `jobs/worker.py` | broker, middlewares, InjectQ, startup / shutdown, discovery | ~120 | +| `jobs/middleware.py` | `TaskMiddleware`, `MonitoringMiddleware` | ~120 | +| `jobs/context.py` | `TaskContext`, `enqueue`, schedule helpers, reserved-argument check | ~100 | +| `src/app/routers/tasks/` | the HTTP routes | ~150 | +| `cli/commands/worker.py` | `worker` and `scheduler` commands | ~60 | + +## What we do NOT build (add only when users ask) + +Own run store, leases, fencing, sweeper, thread / process executors, run listing for the +playground, SQLite / Postgres backends, killing a running task. + +## Sprints + +| # | Scope | Exit | Status | +|---|---|---|---| +| 0 | Live checks on test Redis (port 16379, never 6379): `RedisStreamBroker` redelivery after `kill -9` of a worker; progress with taskiq-redis; dynamic schedule add / change time / delete with `ListRedisScheduleSource`; InjectQ `setup_taskiq` + `InjectTaskiq`; adding `config` to kwargs in `pre_execute`; how a middleware skips a cancelled queued task; constructors of each named broker and result-backend type; whether `TaskContext` works as a `TaskiqDepends` class | `JOBS_TASKIQ_SPRINT0.md` with findings; plan fixed where a check failed | Done 2026-11-06: 11 plan changes, all applied below | +| 1 | Old code not on this branch (it is on `worker`). Redis only. `redis_broker.py` (reclaim on idle queue + lease renewal), `config.py`, `worker.py`, `middleware.py`, `context.py`; extras `jobs = ["taskiq>=0.13", "taskiq-redis>=1.2.4"]`; raise the `redis` extra cap (taskiq-redis needs redis-py 8) | unit tests on `memory` broker; mypy and ruff clean | Done 2026-11-06: 1,171 lines in 10 files, 51 jobs tests (Redis incl.), full suite 2907 passed, live kill check passed | +| 2 | `10xgraph worker` and `10xgraph scheduler` commands; `10xgraph init` / `10xgraph config` questions | live run: API process enqueues, worker runs, scheduler fires, worker kill redelivers | Not started | +| 3 | HTTP routes, reserved-argument check, owner scope, permissions, route guard | route tests + live run through the API with JWT (run, schedule, change time, cancel, other user gets 404) | Not started | +| 4 | Docs page and an example project | docs build; example runs | Not started | + +Sprints run in order; 1 and the docs draft can overlap. Each sprint: small agents, one review, +fix only high / medium findings, then a live check before it is marked done. + +## Housekeeping + +- `uv.lock` is stale (regenerate after Sprint 1 changes the extras). +- Core `ddl_session` import is gone with the old code, so the rc5 pin note no longer applies. diff --git a/JOBS_TASKIQ_SPRINT0.md b/JOBS_TASKIQ_SPRINT0.md new file mode 100644 index 00000000..db9840aa --- /dev/null +++ b/JOBS_TASKIQ_SPRINT0.md @@ -0,0 +1,72 @@ +# Jobs on TaskIQ: Sprint 0 findings (2026-11-06) + +Live checks on Redis 8 (port 16379) with taskiq 0.13.0, taskiq-redis 1.2.4, redis-py 8.1.0, +injectq 0.5.0, plus source reading. Key claims re-checked in the installed source. + +## Works as planned + +| Check | Result | +|---|---| +| Stream broker + result backend | `RedisStreamBroker(url).with_result_backend(RedisAsyncResultBackend(redis_url, result_ex_time=...))`; worker default ack `when_saved` (`--ack-type`), default `--workers 2` | +| Progress | `ProgressTracker` / `TaskProgress` from `taskiq.depends.progress_tracker`; `await task.get_progress()` works on Redis. TaskIQ never sets progress itself | +| Dynamic schedules | `kicker().with_schedule_id(id).with_labels(...).schedule_by_cron / _by_interval / _by_time(source, ...)` -> `CreatedSchedule` (`schedule_id`, `unschedule()`); all three fire; delete stops it | +| InjectQ | `from injectq.integrations.taskiq import setup_taskiq, InjectTaskiq`; `setup_taskiq(container, broker)` | +| Config kwarg | middleware: `self.broker.find_task(message.task_name).original_func`, `inspect.signature`; `config` passed only when declared; tasks without it run fine; sync tasks too | +| Injected parameter list | `taskiq_dependencies.DependencyGraph(fn)` or scanning `Annotated` metadata for `Dependency`; plain args = signature minus that set | +| Middleware order | `pre_execute` in registration order, before dependency resolution and argument parsing; post hooks in reverse | +| `TaskContext` | works as `Annotated[TaskContext, TaskiqDepends()]` with `__init__(self, context: Annotated[Context, TaskiqDepends()])` | + +## Changes the plan needs + +1. **Dead-worker redelivery only happens while new messages arrive.** `RedisStreamBroker.listen()` + runs `XAUTOCLAIM` only after an `xreadgroup` that returned messages (`if not fetched: + continue`, `taskiq_redis/redis_broker.py`). Live: a killed worker's task stayed pending 50 s on + an idle queue, then ran the moment an unrelated task arrived. Fix: our `redis_stream` type is a + small subclass whose `listen()` also runs the autoclaim when nothing was fetched. Also set + `unacknowledged_lock_timeout` (default None: a worker killed while holding the autoclaim lock + blocks it forever). Worth an upstream issue. +2. **`idle_timeout` must be longer than the longest task** (milliseconds, default 600000). A + still-running task older than that is claimed by another worker and runs twice. There is no + lease renewal. Redelivery restarts the task from the beginning (at least once). +3. **Caller kwargs override injected dependencies.** The receiver does + `kwargs = await dep_ctx.resolve_kwargs(); kwargs.update(message.kwargs)` (`receiver.py` + 320-322). `kiq(1, ctx="x")` replaced `TaskContext`; `graph=` would replace the graph. The + reserved-argument check is a security check: reject injected names, `config`, `user`, + `user_id`, `authz`, unknown names and extra positional args, in the API and `enqueue`, and + again in the worker middleware. +4. **Labels keep only int, str, float, bool.** A dict becomes its Python repr + (`taskiq/labels.py::prepare_label`). The config label is `json.dumps(cfg)`, parsed in + `pre_execute`. +5. **A middleware cannot skip a task.** `pre_execute` is called outside any `try` + (`receiver.py` 189-195); a raise leaves the message unacked (redelivered forever on a stream) + and the caller waiting. Cancel design (tested, async and sync): `pre_execute` sets a + ContextVar for a cancelled id; a guard wrapped around each task function raises + `TaskCancelled` before the body runs; the caller gets an error result; our status record says + `cancelled`. Our middleware must never raise. +6. **Only `InjectTaskiq` injects.** A plain InjectQ `Inject[...]` default silently passes the + marker object. +7. **`TaskContext` is a plain class.** A pydantic model fails: TaskIQ reads `__init__`, and + pydantic's is `(**data)`. +8. **Progress and cancel need `async def`.** `ctx.progress()` / `ctx.cancelled()` are async. +9. **Schedules:** + - No update API. Re-adding the same id leaves a duplicate in the type index, and a type change + leaves the old entry. Change = delete + add (verified clean). + - `get_schedules()` hides future one-time schedules, so listing uses our own owner index. + - The scheduler reloads every `--update-interval` (default 60 s): changes take up to that long. + - Do not use `with_task_id()` on schedules (every firing would reuse the id). + - The schedule source's default serializer is pickle; use a JSON serializer (check in + Sprint 1). +10. **Constructor parameter names differ per type**, so the factory keeps a per-type map: + - `RedisStreamBroker(url, ..., idle_timeout ms, unacknowledged_lock_timeout s)` + - `ListQueueBroker(url, ...)` + - `RedisAsyncResultBackend(redis_url, result_ex_time, ...)` + - `taskiq_aio_pika.AioPikaBroker(url, ...)` + - `taskiq_nats.PullBasedJetStreamBroker(servers, ...)`; `NatsBroker` (core NATS) does not + persist; the result backend is `NATSObjectStoreResultBackend(servers, ...)` + - `taskiq_aio_kafka.AioKafkaBroker(bootstrap_servers, ...)`, no `**kwargs` + - `InMemoryBroker()`, no URL, stores its own results +11. **Dependencies:** taskiq-redis 1.2.4 needs `redis>=8,<9` and `taskiq>=0.13`. `main` still + has the `redis` extra capped at `<7` (`pyproject.toml` line 97): raise it in Sprint 1. The + API venv has injectq 0.4.2; check the `setup_taskiq` signature there too. + +Scripts and raw notes: scratchpad `s0-taskiq/redis_checks/` and `s0-taskiq/mw_checks/`. diff --git a/pyproject.toml b/pyproject.toml index 5b42733d..d1c270a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,7 +94,11 @@ snowflakekit = [ "snowflakekit>=0.1.1,<1.0", ] redis = [ - "redis>=5.0.7,<7", + "redis>=5.0.7,<9", +] +jobs = [ + "taskiq>=0.13.0", + "taskiq-redis>=1.2.4", ] jwt = [ "PyJWT>=2.15.0,<3", @@ -389,6 +393,8 @@ dev = [ "httpx==0.27.0", "requests==2.32.3", "redis>=5.0.7", + "taskiq>=0.13.0", + "taskiq-redis>=1.2.4", "PyJWT", "textxtract[pdf,docx,html,xml,md]>=0.2.0", # The ag-ui extra, so the AG-UI router tests run in CI instead of failing. diff --git a/tenxgraph_api/jobs/__init__.py b/tenxgraph_api/jobs/__init__.py new file mode 100644 index 00000000..8e2e8ff0 --- /dev/null +++ b/tenxgraph_api/jobs/__init__.py @@ -0,0 +1,39 @@ +"""Background jobs on TaskIQ: ``from tenxgraph_api.jobs import broker`` and write tasks. + +Names load on first use, so importing ``tenxgraph_api.jobs.config`` stays light and does not +build the broker. +""" + +import importlib +from typing import TYPE_CHECKING, Any + + +if TYPE_CHECKING: + from tenxgraph_api.jobs.context import ( + TaskContext, + cancel, + enqueue, + reschedule, + schedule, + unschedule, + ) + from tenxgraph_api.jobs.worker import broker + +_SOURCES = { + "broker": "worker", + "TaskContext": "context", + "enqueue": "context", + "schedule": "context", + "reschedule": "context", + "unschedule": "context", + "cancel": "context", +} + +__all__ = list(_SOURCES) + + +def __getattr__(name: str) -> Any: + if name not in _SOURCES: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module = importlib.import_module(f"{__name__}.{_SOURCES[name]}") + return getattr(module, name) diff --git a/tenxgraph_api/jobs/arguments.py b/tenxgraph_api/jobs/arguments.py new file mode 100644 index 00000000..bfb5acf4 --- /dev/null +++ b/tenxgraph_api/jobs/arguments.py @@ -0,0 +1,92 @@ +"""The reserved-argument check: a caller may set only a task's plain arguments. + +TaskIQ applies caller kwargs over injected values (``kwargs.update(message.kwargs)`` in its +receiver), so ``graph=`` from a caller would replace the injected graph. This is a security +check, used by ``jobs.enqueue`` / ``jobs.schedule``, the HTTP routes and the worker middleware. +""" + +from __future__ import annotations + +import inspect +import typing +from collections.abc import Callable +from typing import Any + +from taskiq import AsyncTaskiqDecoratedTask +from taskiq_dependencies.dependency import Dependency + + +# Run-config keys only the server writes (``config`` itself is filled from the message label). +RESERVED_NAMES = frozenset({"config", "user", "user_id", "authz"}) + +_POSITIONAL = (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) + + +class ReservedArgumentError(ValueError): + """A caller tried to set an injected or reserved argument, or one the task does not have.""" + + +def injected_names(func: Callable[..., Any]) -> set[str]: + """Parameters TaskIQ fills: ``TaskiqDepends()`` / ``InjectTaskiq(...)`` in ``Annotated`` + metadata or as the default (``Context`` and ``TaskContext`` params use one of these).""" + try: + hints = typing.get_type_hints(func, include_extras=True) + except Exception: # unresolvable forward reference: fall back to raw annotations + hints = {} + names = set() + for name, param in inspect.signature(func).parameters.items(): + annotation = hints.get(name, param.annotation) + metadata = getattr(annotation, "__metadata__", ()) + if isinstance(param.default, Dependency) or any( + isinstance(item, Dependency) for item in metadata + ): + names.add(name) + return names + + +def _max_positional(params: list[inspect.Parameter], blocked: set[str]) -> int | None: + """How many positional args can be given before one would land on a blocked parameter + (None: unlimited, through ``*args``).""" + count = 0 + for param in params: + if param.kind in _POSITIONAL: + if param.name in blocked: + return count + count += 1 + elif param.kind is inspect.Parameter.VAR_POSITIONAL: + return None + return count + + +def check_arguments( + task: AsyncTaskiqDecoratedTask[Any, Any], args: tuple[Any, ...] | list[Any], kwargs: dict +) -> None: + """Raise ``ReservedArgumentError`` unless ``args`` / ``kwargs`` set only plain arguments.""" + func = task.original_func + params = list(inspect.signature(func).parameters.values()) + by_name = {param.name: param for param in params} + blocked = injected_names(func) | RESERVED_NAMES + accepts_any = any(param.kind is inspect.Parameter.VAR_KEYWORD for param in params) + name = task.task_name + + bad = sorted(set(kwargs) & blocked) + if bad: + raise ReservedArgumentError(f"Task {name!r}: arguments {bad} are set by the server") + unknown = sorted( + key + for key in kwargs + if (key not in by_name and not accepts_any) + or (key in by_name and by_name[key].kind is inspect.Parameter.POSITIONAL_ONLY) + ) + if unknown: + raise ReservedArgumentError(f"Task {name!r} has no arguments {unknown}") + + limit = _max_positional(params, blocked) + if limit is not None and len(args) > limit: + raise ReservedArgumentError( + f"Task {name!r} takes at most {limit} positional arguments, got {len(args)}" + ) + positional = [param.name for param in params if param.kind in _POSITIONAL] + twice = sorted(set(positional[: len(args)]) & set(kwargs)) + if twice: + raise ReservedArgumentError(f"Task {name!r}: arguments {twice} are given twice") diff --git a/tenxgraph_api/jobs/config.py b/tenxgraph_api/jobs/config.py new file mode 100644 index 00000000..82bd497c --- /dev/null +++ b/tenxgraph_api/jobs/config.py @@ -0,0 +1,172 @@ +"""The ``worker`` key of 10xgraph.json and the factories that turn it into TaskIQ objects.""" + +import importlib +import os +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from tenxgraph_api.src.app.core.config.graph_config import _expand_env + + +BROKER_TYPES = ("redis_stream", "redis_list", "memory") +RESULT_BACKEND_TYPES = ("redis", "memory") +KEY = "worker.broker.url" +INSTALL_HINT = 'The jobs layer needs TaskIQ: pip install "10xgraph-api[jobs]"' + + +class _Section(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class BrokerConfig(_Section): + """``worker.broker``: ``type`` is a built-in name or ``"module:attr"``.""" + + type: str = "redis_stream" + url: str | None = None + options: dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode="after") + def _check_type(self) -> "BrokerConfig": + if self.type not in BROKER_TYPES and ":" not in self.type: + raise ValueError( + f'worker.broker.type must be one of {", ".join(BROKER_TYPES)} or "module:attr"; ' + f"got {self.type!r}" + ) + return self + + @property + def is_redis(self) -> bool: + return self.type in ("redis_stream", "redis_list") + + +class ResultBackendConfig(_Section): + """``worker.result_backend``; when omitted, Redis on the broker URL for a Redis broker.""" + + type: str = "redis" + url: str | None = None + options: dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode="after") + def _check_type(self) -> "ResultBackendConfig": + if self.type not in RESULT_BACKEND_TYPES and ":" not in self.type: + raise ValueError( + f'worker.result_backend.type must be one of {", ".join(RESULT_BACKEND_TYPES)} ' + f'or "module:attr"; got {self.type!r}' + ) + return self + + +class PrometheusConfig(_Section): + host: str = "0.0.0.0" # noqa: S104 + port: int = 9000 + + +class WorkerConfig(_Section): + broker: BrokerConfig = Field(default_factory=BrokerConfig) + result_backend: ResultBackendConfig | None = None + tasks: list[str] = Field(default_factory=list) + schedules: bool = False + result_ttl: int = Field(default=86400, gt=0) + retries: int = Field(default=3, ge=0) + scheduler_interval: int = Field(default=10, gt=0) + prometheus: PrometheusConfig | None = None + + def broker_url(self) -> str | None: + """The broker URL with ``${VAR}`` expanded; ``REDIS_URL`` for a Redis broker without one.""" + if self.broker.is_redis: + return _redis_url(self.broker.url, "worker.broker.url") + return None if self.broker.url is None else _expand(self.broker.url, "worker.broker.url") + + +def _expand(raw: str, key: str) -> str: + try: + return _expand_env(raw) or raw + except ValueError: + raise ValueError(f"{key}: {raw} uses an environment variable that is not set") from None + + +def _redis_url(raw: str | None, key: str) -> str: + """``raw`` with ``${VAR}`` expanded, else the ``REDIS_URL`` variable.""" + if raw is not None: + return _expand(raw, key) + env_url = os.environ.get("REDIS_URL", "").strip() + if not env_url: + raise ValueError(f"{key} is required: set it in 10xgraph.json or set REDIS_URL") + return env_url + + +def _require(package: str) -> None: + try: + importlib.import_module(package) + except ImportError: + raise RuntimeError(INSTALL_HINT) from None + + +def _load(path: str) -> Any: + module, attr = path.split(":", 1) + return getattr(importlib.import_module(module), attr) + + +def _make(target: Any, url: str | None, options: dict[str, Any]) -> Any: + """A class is called with the URL (when there is one) and the options; an instance is kept.""" + if not isinstance(target, type): + return target + return target(url, **options) if url else target(**options) + + +def build_broker(cfg: WorkerConfig) -> Any: + """The TaskIQ broker for ``cfg``, with its result backend attached.""" + _require("taskiq") + kind, options = cfg.broker.type, dict(cfg.broker.options) + broker: Any + if kind == "memory": + from taskiq import InMemoryBroker + + return InMemoryBroker(**options) + if kind == "redis_stream": + _require("taskiq_redis") + from tenxgraph_api.jobs.redis_broker import ReclaimingRedisStreamBroker + + options.setdefault("unacknowledged_lock_timeout", 60) + broker = ReclaimingRedisStreamBroker(_redis_url(cfg.broker.url, KEY), **options) + elif kind == "redis_list": + _require("taskiq_redis") + from taskiq_redis import ListQueueBroker + + broker = ListQueueBroker(_redis_url(cfg.broker.url, KEY), **options) + else: + broker = _make(_load(kind), cfg.broker_url(), options) + backend = build_result_backend(cfg) + return broker if backend is None else broker.with_result_backend(backend) + + +def build_result_backend(cfg: WorkerConfig) -> Any | None: + """The result backend for ``cfg``; None when the broker stores its own results.""" + section = cfg.result_backend + if section is None: + if not cfg.broker.is_redis: + return None + section = ResultBackendConfig() + if section.type == "memory": + return None + options = dict(section.options) + if section.type == "redis": + _require("taskiq_redis") + from taskiq_redis import RedisAsyncResultBackend + + url = _redis_url(section.url or cfg.broker.url, "worker.result_backend.url") + options.setdefault("result_ex_time", cfg.result_ttl) + return RedisAsyncResultBackend(url, **options) + other = section.url and _expand(section.url, "worker.result_backend.url") + return _make(_load(section.type), other, options) + + +def build_schedule_source(cfg: WorkerConfig) -> Any: + """The Redis schedule source (JSON, not the default pickle) for dynamic schedules.""" + _require("taskiq_redis") + from taskiq.serializers import JSONSerializer + from taskiq_redis import ListRedisScheduleSource + + url = _redis_url(cfg.broker.url, "worker.broker.url") + return ListRedisScheduleSource(url, serializer=JSONSerializer()) diff --git a/tenxgraph_api/jobs/context.py b/tenxgraph_api/jobs/context.py new file mode 100644 index 00000000..7f1a9c43 --- /dev/null +++ b/tenxgraph_api/jobs/context.py @@ -0,0 +1,114 @@ +"""What tasks and callers use: ``TaskContext``, ``enqueue``, ``cancel`` and the schedule helpers. + +Python callers start a task with ``await jobs.enqueue(task, *args, config=config, **kwargs)``: +the run config travels in the ``tenxgraph_config`` label (a JSON string, labels keep only +scalars) and our worker middleware passes it to tasks that declare ``config``. +""" + +from __future__ import annotations + +import json +import logging +from typing import Annotated, Any + +from taskiq import AsyncBroker, AsyncTaskiqDecoratedTask, AsyncTaskiqTask, Context, TaskiqDepends +from taskiq.depends.progress_tracker import TaskProgress, TaskState + +from tenxgraph_api.jobs import state +from tenxgraph_api.jobs.arguments import ReservedArgumentError, check_arguments +from tenxgraph_api.jobs.schedules import ( + ScheduleNotFoundError, + reschedule, + schedule, + set_schedule_source, + unschedule, +) + + +logger = logging.getLogger("tenxgraph_api.jobs") + +__all__ = [ + "ReservedArgumentError", + "ScheduleNotFoundError", + "TaskCancelled", + "TaskContext", + "cancel", + "check_arguments", + "enqueue", + "reschedule", + "schedule", + "set_schedule_source", + "unschedule", +] + + +class TaskCancelled(Exception): # - public name from the plan + """The task was cancelled before it started (the caller sees it as the error result).""" + + +class TaskContext: + """Optional task parameter: ``ctx: Annotated[TaskContext, TaskiqDepends()]``. + + A plain class on purpose: TaskIQ builds it from ``__init__`` (a pydantic model does not + work). ``progress`` and ``cancelled`` are async, so use them from ``async def`` tasks. + """ + + def __init__(self, context: Annotated[Context, TaskiqDepends()]) -> None: + self._context = context + self.task_id: str = context.message.task_id + self.task_name: str = context.message.task_name + # The same value a ``config`` parameter gets (see TaskMiddleware). + self.config: dict[str, Any] = state.parse_config(context.message.labels) + self.config["task_id"] = self.task_id + + async def progress(self, meta: Any) -> None: + """Save progress; ``GET /v1/tasks/{id}`` shows it.""" + progress: TaskProgress[Any] = TaskProgress(state=TaskState.STARTED, meta=meta) + await self._context.broker.result_backend.set_progress(self.task_id, progress) + + async def cancelled(self) -> bool: + """True after ``jobs.cancel`` / ``POST /v1/tasks/{id}/cancel`` for this task.""" + store = state.cache_store_for(self._context.broker) + return store is not None and await state.is_cancel_flag_set(store, self.task_id) + + +async def enqueue( + task: AsyncTaskiqDecoratedTask[Any, Any], + *args: Any, + config: dict[str, Any] | None = None, + **kwargs: Any, +) -> AsyncTaskiqTask[Any]: + """Check the arguments, record the task as ``queued`` and send it; return TaskIQ's handle. + + The record is written before sending, so a fast worker's ``running`` is never overwritten. + """ + check_arguments(task, args, kwargs) + config = config or {} + task_id = task.broker.id_generator() + store = state.cache_store_for(task.broker) + if store is not None: + record = state.TaskStatus( + task_id=task_id, + task_name=task.task_name, + status="queued", + owner=state.owner_of(config), + ) + await state.write_status(store, record) + kicker = task.kicker().with_task_id(task_id) + kicker = kicker.with_labels(**{state.CONFIG_LABEL: json.dumps(config)}) + return await kicker.kiq(*args, **kwargs) + + +async def cancel(task_id: str, broker: AsyncBroker | None = None) -> None: + """Set the cancel flag; a queued task is not run, a running one sees ``ctx.cancelled()``. + + ``broker`` picks the container (worker side); the API passes nothing. + """ + store = state.cache_store_for(broker) + if store is None: + raise RuntimeError("No cache store is bound: cancel needs one") + await state.set_cancel_flag(store, task_id) + record = await state.read_status(store, task_id) + if record is not None and record.status == "queued": + record.status = "cancelled" + await state.write_status(store, record) diff --git a/tenxgraph_api/jobs/middleware.py b/tenxgraph_api/jobs/middleware.py new file mode 100644 index 00000000..ceeca07e --- /dev/null +++ b/tenxgraph_api/jobs/middleware.py @@ -0,0 +1,133 @@ +"""Worker middlewares: ``TaskMiddleware`` (config, argument check, cancel, status) and +``MonitoringMiddleware`` (logs and publisher events). + +A TaskIQ middleware cannot skip a task: ``pre_execute`` runs outside any ``try`` in the +receiver, so a raise there leaves the message unacked (redelivered forever on a stream) and the +caller waiting. Ours never raise. To stop a run, ``pre_execute`` puts the error in a ContextVar +and the guard that ``wrap_tasks`` puts around every task function raises it before the body +runs; the caller gets an error result. The ContextVar does not reach a process-pool executor. +""" + +from __future__ import annotations + +import contextvars +import functools +import inspect +import logging +from collections.abc import Callable +from typing import Any + +from taskiq import AsyncBroker, TaskiqMessage, TaskiqMiddleware, TaskiqResult +from taskiq.exceptions import NoResultError + +from tenxgraph_api.jobs import state +from tenxgraph_api.jobs.arguments import ReservedArgumentError, check_arguments, injected_names +from tenxgraph_api.jobs.context import TaskCancelled +from tenxgraph_api.jobs.monitoring import MonitoringMiddleware + + +logger = logging.getLogger("tenxgraph_api.jobs") + +__all__ = ["MonitoringMiddleware", "TaskMiddleware", "wrap_tasks"] + +# The error the guard raises instead of running the task body (None: run it). +_ABORT: contextvars.ContextVar[BaseException | None] = contextvars.ContextVar( + "tenxgraph_task_abort", default=None +) +_GUARD_MARK = "__tenxgraph_guard__" + + +def _guard(func: Callable[..., Any]) -> Callable[..., Any]: + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def async_guard(*args: Any, **kwargs: Any) -> Any: + if (error := _ABORT.get()) is not None: + raise error + return await func(*args, **kwargs) + + wrapper: Callable[..., Any] = async_guard + else: + + @functools.wraps(func) + def sync_guard(*args: Any, **kwargs: Any) -> Any: + if (error := _ABORT.get()) is not None: + raise error + return func(*args, **kwargs) + + wrapper = sync_guard + setattr(wrapper, _GUARD_MARK, True) + return wrapper + + +def wrap_tasks(broker: AsyncBroker) -> None: + """Put the guard around every registered task's function (the receiver calls + ``task.original_func``). Safe to call more than once.""" + for task in broker.get_all_tasks().values(): + if not getattr(task.original_func, _GUARD_MARK, False): + task.original_func = _guard(task.original_func) + + +class TaskMiddleware(TaskiqMiddleware): + """Run config into ``config``, the reserved-argument re-check, cancel and the status record. + + A message whose kwargs break the reserved-argument rule is not run: injected names are + removed from its kwargs (so injection wins even without the guard) and the guard raises + ``ReservedArgumentError``. Failing the run is safer than dropping arguments and running + the task with different input than the caller sent. + """ + + async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage: + _ABORT.set(None) + try: + await self._prepare(message) + except Exception: + logger.exception("Task %s (%s): pre_execute failed", message.task_name, message.task_id) + return message + + async def _prepare(self, message: TaskiqMessage) -> None: + task = self.broker.find_task(message.task_name) + if task is None: + return + config = state.parse_config(message.labels) + config["task_id"] = message.task_id + # ``config`` always comes from the label: a retry re-sends the kwargs we set here. + message.kwargs.pop("config", None) + try: + check_arguments(task, message.args, message.kwargs) + except ReservedArgumentError as error: + for name in injected_names(task.original_func) & set(message.kwargs): + message.kwargs.pop(name) + _ABORT.set(error) + if "config" in inspect.signature(task.original_func).parameters: + message.kwargs["config"] = config + + store = state.cache_store_for(self.broker) + if store is None: + return + status: state.Status = "running" + if await state.is_cancel_flag_set(store, message.task_id): + _ABORT.set(TaskCancelled("cancelled before start")) + status = "cancelled" + owner = state.owner_of(config) + await state.update_status(store, message.task_id, message.task_name, status, owner) + + async def post_execute(self, message: TaskiqMessage, result: TaskiqResult[Any]) -> None: + _ABORT.set(None) # the in-memory broker can run a task in the caller's context + try: + store = state.cache_store_for(self.broker) + if store is None: + return + owner = state.owner_of(state.parse_config(message.labels)) + status = _final_status(result) + await state.update_status(store, message.task_id, message.task_name, status, owner) + except Exception: + logger.exception("Task %s (%s): status not saved", message.task_name, message.task_id) + + +def _final_status(result: TaskiqResult[Any]) -> state.Status: + if isinstance(result.error, NoResultError): + return "queued" # SimpleRetryMiddleware sent it again + if isinstance(result.error, TaskCancelled): + return "cancelled" + return "failure" if result.is_err else "success" diff --git a/tenxgraph_api/jobs/monitoring.py b/tenxgraph_api/jobs/monitoring.py new file mode 100644 index 00000000..bb8b17e7 --- /dev/null +++ b/tenxgraph_api/jobs/monitoring.py @@ -0,0 +1,70 @@ +"""``MonitoringMiddleware``: log task started / finished / failed, and send the same as core +publisher events when a ``BasePublisher`` is bound in the container. Never raises.""" + +from __future__ import annotations + +import logging +from typing import Any + +from taskiq import TaskiqMessage, TaskiqMiddleware, TaskiqResult + +from tenxgraph_api.jobs import state + + +logger = logging.getLogger("tenxgraph_api.jobs") + + +class MonitoringMiddleware(TaskiqMiddleware): + """Task lifecycle logs and events. Core has no task event kind yet, so events use + ``Event.GRAPH_EXECUTION`` with ``metadata["source"] == "task"``.""" + + async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage: + logger.info("Task %s (%s) started", message.task_name, message.task_id) + self._publish(message, "start", {}) + return message + + async def post_execute(self, message: TaskiqMessage, result: TaskiqResult[Any]) -> None: + if result.is_err: + logger.warning( + "Task %s (%s) failed after %.2fs: %s", + message.task_name, + message.task_id, + result.execution_time, + type(result.error).__name__, + ) + data = {"error": type(result.error).__name__} + self._publish(message, "error", {**data, "duration": result.execution_time}) + else: + logger.info( + "Task %s (%s) finished in %.2fs", + message.task_name, + message.task_id, + result.execution_time, + ) + self._publish(message, "end", {"duration": result.execution_time}) + + def _publish(self, message: TaskiqMessage, phase: str, data: dict[str, Any]) -> None: + try: + from tenxgraph.runtime.publisher.base_publisher import BasePublisher + from tenxgraph.runtime.publisher.events import Event, EventModel, EventType + from tenxgraph.runtime.publisher.publish import publish_event + from tenxgraph.utils.background_task_manager import BackgroundTaskManager + + container = state.container_for(self.broker) + publisher = container.try_get(BasePublisher) + manager = container.try_get(BackgroundTaskManager) + if publisher is None or manager is None: + return + event = EventModel( + event=Event.GRAPH_EXECUTION, + event_type=EventType(phase), + node_name=message.task_name, + run_id=message.task_id, + user_id=state.parse_config(message.labels).get("user_id"), + is_error=phase == "error", + data=data, + metadata={"source": "task", "task_name": message.task_name}, + ) + publish_event(event, publisher=publisher, task_manager=manager) + except Exception: + logger.debug("Task event not published", exc_info=True) diff --git a/tenxgraph_api/jobs/redis_broker.py b/tenxgraph_api/jobs/redis_broker.py new file mode 100644 index 00000000..d09328b0 --- /dev/null +++ b/tenxgraph_api/jobs/redis_broker.py @@ -0,0 +1,144 @@ +"""Redis stream broker that reclaims on an idle queue and renews leases of running tasks. + +Two changes over taskiq-redis ``RedisStreamBroker``: + +1. ``listen()`` runs ``XAUTOCLAIM`` after every read, also when the read returned nothing. + Upstream skips it on an empty read, so a dead worker's task waits until new work arrives. +2. Every message this worker holds (fetched, not yet acked) has its idle time reset every + ``idle_timeout / 3`` by one background renewer task. A running task is therefore never + claimed by another worker, however long it runs, while a dead worker's messages go idle + and are claimed after ``idle_timeout``. The renewer stops on ``shutdown()``. +""" + +import asyncio +import logging +from collections.abc import AsyncGenerator, Awaitable, Callable +from typing import Any + +from redis.asyncio import Redis +from taskiq import AckableMessage +from taskiq_redis import RedisStreamBroker + + +logger = logging.getLogger("tenxgraph_api.jobs.redis_broker") + +# Reset the idle time of one pending entry, only while this consumer still owns it. +# An acked entry, or one another worker has claimed, is left alone (returns 0). +RENEW_SCRIPT = """ +local pending = redis.call('XPENDING', KEYS[1], ARGV[1], ARGV[3], ARGV[3], 1, ARGV[2]) +if #pending == 0 then return 0 end +redis.call('XCLAIM', KEYS[1], ARGV[1], ARGV[2], 0, ARGV[3], 'JUSTID') +return 1 +""" + +Lease = tuple[Any, Any] # (stream name, message id) as returned by Redis + + +class ReclaimingRedisStreamBroker(RedisStreamBroker): + """``RedisStreamBroker`` with idle-queue reclaim and lease renewal.""" + + def __init__(self, url: str, unacknowledged_lock_timeout: float = 30, **kwargs: Any) -> None: + super().__init__(url, unacknowledged_lock_timeout=unacknowledged_lock_timeout, **kwargs) + self._held: set[Lease] = set() + self._renewer: asyncio.Task[None] | None = None + + async def shutdown(self) -> None: + """Stop renewing (held messages go idle and are reclaimed), then close the pool.""" + renewer, self._renewer = self._renewer, None + if renewer is not None: + renewer.cancel() + await asyncio.gather(renewer, return_exceptions=True) + self._held.clear() + await super().shutdown() + + async def listen(self) -> AsyncGenerator[AckableMessage, None]: + """Yield new messages, then idle pending ones claimed from dead consumers.""" + if self._renewer is None or self._renewer.done(): + self._renewer = asyncio.create_task(self._renew_forever()) + unyielded: list[Lease] = [] + try: + async with Redis(connection_pool=self.connection_pool) as redis_conn: + while True: + fetched: Any = await redis_conn.xreadgroup( + self.consumer_group_name, + self.consumer_name, + {self.queue_name: ">", **self.additional_streams}, # type: ignore[dict-item] + block=self.block, + noack=False, + count=self.count, + ) + batch: list[Any] = [(s, i, m) for s, msgs in fetched or [] for i, m in msgs] + for stream in [self.queue_name, *self.additional_streams.keys()]: + claimed = await self._claim(redis_conn, stream) + batch.extend((stream, i, m) for i, m in claimed) + # Hold the whole batch now: a message waiting here for a free slot must + # not go idle and be claimed by another worker. + unyielded = [(stream, msg_id) for stream, msg_id, _ in batch] + self._held.update(unyielded) + for stream, msg_id, msg in batch: + unyielded.pop(0) + yield AckableMessage( + data=msg[b"data"], + ack=self._ack_generator(id=msg_id, queue_name=stream), + ) + finally: + # Closed mid-batch: release what was never handed to the worker. + self._held.difference_update(unyielded) + + async def _claim(self, redis_conn: Redis, stream: str) -> list[Any]: + """``XAUTOCLAIM`` idle pending entries of ``stream`` (same lock pattern as upstream).""" + pipe = redis_conn.pipeline() + lock = pipe.lock( + f"autoclaim:{self.consumer_group_name}:{stream}", + timeout=self.unacknowledged_lock_timeout, + ) + await lock.acquire() + await pipe.xautoclaim( + name=stream, + groupname=self.consumer_group_name, + consumername=self.consumer_name, + min_idle_time=self.idle_timeout, + count=self.unacknowledged_batch_size, + ) + await lock.release() + results = await pipe.execute() + claimed: list[Any] = results[1][1] + if claimed: + logger.info("Reclaimed %d idle message(s) from stream %s", len(claimed), stream) + return claimed + + def _ack_generator( + self, + id: str, # noqa: A002 (upstream signature) + queue_name: str, + ) -> Callable[[], Awaitable[None]]: + """Stop renewing the message, then ack it (a failed ack leaves it to be reclaimed).""" + upstream_ack = super()._ack_generator(id=id, queue_name=queue_name) + + async def _ack() -> None: + self._held.discard((queue_name, id)) + await upstream_ack() + + return _ack + + async def _renew_forever(self) -> None: + """Reset the idle time of every held message every ``idle_timeout / 3``.""" + interval = self.idle_timeout / 3000 + while True: + await asyncio.sleep(interval) + try: + await self._renew_held() + except Exception: + logger.warning("Lease renewal failed; retrying", exc_info=True) + + async def _renew_held(self) -> None: + if not self._held: + return + async with Redis(connection_pool=self.connection_pool) as redis_conn: + for lease in list(self._held): + stream, msg_id = lease + args = (self.consumer_group_name, self.consumer_name, msg_id) + owned = await redis_conn.eval(RENEW_SCRIPT, 1, stream, *args) # type: ignore[misc] + if not owned and lease in self._held: + self._held.discard(lease) + logger.warning("Lost lease on message %s in %s", msg_id, stream) diff --git a/tenxgraph_api/jobs/schedules.py b/tenxgraph_api/jobs/schedules.py new file mode 100644 index 00000000..842197a0 --- /dev/null +++ b/tenxgraph_api/jobs/schedules.py @@ -0,0 +1,128 @@ +"""Dynamic schedules: add, change the time, delete. + +TaskIQ has no update and no read-by-id (``get_schedules()`` also hides future one-time +schedules), so each schedule we create is also kept as a ``ScheduleRecord`` in the cache store. +A change deletes the TaskIQ schedule and adds it again with the same id. Never +``with_task_id`` on a schedule: every firing would reuse the id. +""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from datetime import datetime, timedelta +from typing import TYPE_CHECKING, Any + +from taskiq import AsyncBroker, AsyncTaskiqDecoratedTask, ScheduledTask, ScheduleSource + +from tenxgraph_api.jobs import state +from tenxgraph_api.jobs.arguments import check_arguments + + +if TYPE_CHECKING: + from tenxgraph_api.src.app.core.cache_store import BaseCacheStore + + +class ScheduleNotFoundError(LookupError): + """No schedule with this id (deleted, expired or never created).""" + + +_source: ScheduleSource | None = None +_source_factory: Callable[[], ScheduleSource] | None = None + + +def set_schedule_source( + source: ScheduleSource | None = None, + factory: Callable[[], ScheduleSource] | None = None, +) -> None: + """Use ``source`` (or build one lazily with ``factory``); tests pass a fake here.""" + global _source, _source_factory # - one process-wide source + _source, _source_factory = source, factory + + +def get_schedule_source() -> ScheduleSource: + """The set source, else one built once from the ``worker`` config (Redis).""" + global _source # noqa: PLW0603 + if _source is None: + _source = _source_factory() if _source_factory else _default_source() + return _source + + +def _default_source() -> ScheduleSource: + from tenxgraph_api.jobs.config import build_schedule_source + from tenxgraph_api.jobs.worker import worker_config + + if worker_config is None: + raise RuntimeError('Dynamic schedules need the "worker" key in 10xgraph.json') + source: ScheduleSource = build_schedule_source(worker_config) + return source + + +def _timing(cron: str | None, at: datetime | None, every: int | timedelta | None) -> dict[str, Any]: + given = [value for value in (cron, at, every) if value is not None] + if len(given) != 1: + raise ValueError("Give exactly one of cron, at or every") + return {"cron": cron, "cron_offset": None, "time": at, "interval": every} + + +def _store(broker: AsyncBroker | None) -> BaseCacheStore: + store = state.cache_store_for(broker) + if store is None: + raise RuntimeError("No cache store is bound: schedules need one") + return store + + +async def schedule( + task: AsyncTaskiqDecoratedTask[Any, Any], + *args: Any, + cron: str | None = None, + at: datetime | None = None, + every: int | timedelta | None = None, + config: dict[str, Any] | None = None, + **kwargs: Any, +) -> str: + """Schedule ``task``: one of ``cron``, ``at`` (a datetime) or ``every`` (seconds); returns + the schedule id. Each run acts as the user in ``config``.""" + timing = _timing(cron, at, every) + check_arguments(task, args, kwargs) + config = config or {} + store = _store(task.broker) + kicker = task.kicker().with_labels(**{state.CONFIG_LABEL: json.dumps(config)}) + source = get_schedule_source() + if cron is not None: + created = await kicker.schedule_by_cron(source, cron, *args, **kwargs) + elif at is not None: + created = await kicker.schedule_by_time(source, at, *args, **kwargs) + else: + created = await kicker.schedule_by_interval(source, timing["interval"], *args, **kwargs) + record = state.ScheduleRecord(owner=state.owner_of(config), task=created.task) + await state.write_schedule(store, record) + return created.schedule_id + + +async def reschedule( + schedule_id: str, + *, + cron: str | None = None, + at: datetime | None = None, + every: int | timedelta | None = None, + broker: AsyncBroker | None = None, +) -> None: + """Change only the time of a schedule (delete + add with the same id). ``broker`` picks + the container (worker side); the API passes nothing.""" + timing = _timing(cron, at, every) + store = _store(broker) + record = await state.read_schedule(store, schedule_id) + if record is None: + raise ScheduleNotFoundError(schedule_id) + new_task = ScheduledTask.model_validate({**record.task.model_dump(), **timing}) + source = get_schedule_source() + await source.delete_schedule(schedule_id) + await source.add_schedule(new_task) + await state.write_schedule(store, record.model_copy(update={"task": new_task})) + + +async def unschedule(schedule_id: str, broker: AsyncBroker | None = None) -> None: + """Delete a schedule (no error when it is already gone).""" + await get_schedule_source().delete_schedule(schedule_id) + await state.delete_schedule(_store(broker), schedule_id) diff --git a/tenxgraph_api/jobs/state.py b/tenxgraph_api/jobs/state.py new file mode 100644 index 00000000..b43d5235 --- /dev/null +++ b/tenxgraph_api/jobs/state.py @@ -0,0 +1,139 @@ +"""Job state kept in the API's cache store: status records, cancel flags, schedule records. + +Keys live under ``jobs:`` and expire after ``RESULT_TTL`` seconds (schedule records live as long +as the schedule). The store is the ``BaseCacheStore`` bound in the InjectQ container: the one +``setup_taskiq`` attached to the broker in a worker, else the API's container. +""" + +from __future__ import annotations + +import json +import logging +import time +from typing import TYPE_CHECKING, Any, Literal + +from pydantic import BaseModel, Field +from taskiq import ScheduledTask + + +if TYPE_CHECKING: + from injectq import InjectQ + from taskiq import AsyncBroker + + from tenxgraph_api.src.app.core.cache_store import BaseCacheStore + + +logger = logging.getLogger("tenxgraph_api.jobs") + +# How long status records and cancel flags are kept, in seconds. jobs.config sets it from +# ``worker.result_ttl``. +RESULT_TTL: int = 86400 +# Message label that carries the run config as a JSON string (labels keep only scalars). +CONFIG_LABEL = "tenxgraph_config" + +Status = Literal["queued", "running", "success", "failure", "cancelled"] + + +class TaskStatus(BaseModel): + """What ``GET /v1/tasks/{id}`` shows besides the result and progress.""" + + task_id: str + task_name: str + status: Status + owner: str | None = None + created_at: float = Field(default_factory=time.time) + updated_at: float = Field(default_factory=time.time) + + +class ScheduleRecord(BaseModel): + """A schedule as we created it: TaskIQ has no read-by-id, and hides future one-time ones.""" + + owner: str | None = None + created_at: float = Field(default_factory=time.time) + task: ScheduledTask + + +def owner_of(config: dict[str, Any]) -> str | None: + """The run config's ``user_id`` as a string (owner checks compare strings).""" + user_id = config.get("user_id") + return None if user_id is None else str(user_id) + + +def parse_config(labels: dict[str, Any]) -> dict[str, Any]: + """The run config from the message label; ``{}`` when missing or not a JSON object.""" + raw = labels.get(CONFIG_LABEL) + if not raw: + return {} + try: + value = json.loads(raw) if isinstance(raw, str) else raw + except ValueError: + logger.warning("Ignoring a %s label that is not JSON", CONFIG_LABEL) + return {} + return value if isinstance(value, dict) else {} + + +def container_for(broker: AsyncBroker | None) -> InjectQ: + """The broker's container (``setup_taskiq``), else the current API container.""" + container = getattr(broker.state, "injectq_container", None) if broker else None + if container is not None: + return container + from tenxgraph_api.src.app.core.utils.container import current_container + + return current_container() + + +def cache_store_for(broker: AsyncBroker | None) -> BaseCacheStore | None: + from tenxgraph_api.src.app.core.cache_store import BaseCacheStore + + return container_for(broker).try_get(BaseCacheStore) + + +def _key(kind: str, ident: str) -> str: + return f"jobs:{kind}:{ident}" + + +async def read_status(store: BaseCacheStore, task_id: str) -> TaskStatus | None: + raw = await store.get(_key("status", task_id)) + return TaskStatus.model_validate_json(raw) if raw else None + + +async def write_status(store: BaseCacheStore, record: TaskStatus) -> None: + record.updated_at = time.time() + await store.set(_key("status", record.task_id), record.model_dump_json().encode(), RESULT_TTL) + + +async def update_status( + store: BaseCacheStore, task_id: str, task_name: str, status: Status, owner: str | None +) -> None: + """Set ``status``, keeping owner and created time of an existing record.""" + record = await read_status(store, task_id) + if record is None: + record = TaskStatus(task_id=task_id, task_name=task_name, status=status, owner=owner) + record.status = status + await write_status(store, record) + + +async def set_cancel_flag(store: BaseCacheStore, task_id: str) -> None: + await store.set(_key("cancel", task_id), b"1", RESULT_TTL) + + +async def is_cancel_flag_set(store: BaseCacheStore, task_id: str) -> bool: + return await store.get(_key("cancel", task_id)) is not None + + +async def read_schedule(store: BaseCacheStore, schedule_id: str) -> ScheduleRecord | None: + raw = await store.get(_key("schedule", schedule_id)) + return ScheduleRecord.model_validate_json(raw) if raw else None + + +async def write_schedule(store: BaseCacheStore, record: ScheduleRecord) -> None: + """Keep a cron / interval schedule until deleted, a one-time one until RESULT_TTL after.""" + ttl: float | None = None + if record.task.time is not None: + ttl = max(record.task.time.timestamp() - time.time(), 0) + RESULT_TTL + data = record.model_dump_json().encode() + await store.set(_key("schedule", record.task.schedule_id), data, ttl) + + +async def delete_schedule(store: BaseCacheStore, schedule_id: str) -> None: + await store.delete(_key("schedule", schedule_id)) diff --git a/tenxgraph_api/jobs/worker.py b/tenxgraph_api/jobs/worker.py new file mode 100644 index 00000000..beef67d5 --- /dev/null +++ b/tenxgraph_api/jobs/worker.py @@ -0,0 +1,140 @@ +"""The broker the developer imports, and the worker's startup and shutdown. + +Built on import from the ``worker`` key of 10xgraph.json (the file the API reads: ``GRAPH_PATH``, +else ``10xgraph.json`` in the cwd). Without that key the broker is in-memory. +""" + +import asyncio +import importlib +import logging +import os +import sys +from pathlib import Path + +from injectq import InjectQ +from injectq.integrations.taskiq import setup_taskiq +from taskiq import InMemoryBroker, TaskiqEvents, TaskiqState +from taskiq.middlewares import SimpleRetryMiddleware + +from tenxgraph_api.jobs import state as job_state +from tenxgraph_api.jobs.config import WorkerConfig, build_broker +from tenxgraph_api.jobs.middleware import TaskMiddleware, wrap_tasks +from tenxgraph_api.jobs.monitoring import MonitoringMiddleware +from tenxgraph_api.src.app.core.config.graph_config import GraphConfig, default_config_path + + +logger = logging.getLogger("tenxgraph_api.jobs") + +# Same value as BOOT_CLEANUP_TIMEOUT_SECONDS in src/app/main.py (not imported: it builds the app). +BOOT_CLEANUP_TIMEOUT_SECONDS = 15.0 + + +def _load_graph_config() -> GraphConfig | None: + path = os.environ.get("GRAPH_PATH") or default_config_path() + return GraphConfig(path) if Path(path).exists() else None + + +def import_tasks(cfg: WorkerConfig) -> None: + """Import the ``worker.tasks`` modules (names or file globs) so their tasks register.""" + for entry in cfg.tasks: + if entry.endswith(".py") or any(c in entry for c in "*?/"): + if str(Path.cwd()) not in sys.path: + sys.path.insert(0, str(Path.cwd())) + for file in sorted(Path.cwd().glob(entry)): + importlib.import_module( + file.relative_to(Path.cwd()).with_suffix("").as_posix().replace("/", ".") + ) + else: + importlib.import_module(entry) + + +graph_config = _load_graph_config() +worker_config = graph_config.worker if graph_config else None + +if worker_config is None: + logger.warning('No "worker" key in 10xgraph.json: using an in-memory broker (no worker).') + broker = InMemoryBroker() +else: + broker = build_broker(worker_config) + job_state.RESULT_TTL = worker_config.result_ttl + +broker.add_middlewares( + SimpleRetryMiddleware(default_retry_count=worker_config.retries if worker_config else 3), + MonitoringMiddleware(), + TaskMiddleware(), +) +if worker_config is not None and worker_config.prometheus is not None: + from taskiq.middlewares import PrometheusMiddleware + + broker.add_middlewares( + PrometheusMiddleware( + server_addr=worker_config.prometheus.host, server_port=worker_config.prometheus.port + ) + ) + + +async def _prepare_checkpointer(container: InjectQ) -> None: + """Create or upgrade the checkpointer schema like the API does at boot. Best effort.""" + checkpointer = container.try_get("checkpointer") or container.try_get("BaseCheckpointer") + asetup = getattr(checkpointer, "asetup", None) + if not callable(asetup): + return + try: + await asyncio.wait_for(asetup(), BOOT_CLEANUP_TIMEOUT_SECONDS) + except Exception as exc: + logger.error( + "Checkpointer setup failed at startup (%s: %s); it will be retried on first use", + type(exc).__name__, + exc, + ) + + +async def worker_startup(state: TaskiqState) -> None: + """Build the container and graph like the API does, then find and wrap the tasks.""" + from tenxgraph_api.src.app.core.config.settings import get_settings + from tenxgraph_api.src.app.loader import attach_all_modules, load_container + from tenxgraph_api.src.app.routers.ping.startup_checks import check_backends_at_startup + + if graph_config is None or worker_config is None: + raise RuntimeError('The worker needs a "worker" key in 10xgraph.json') + container = load_container(graph_config.injectq_path) or InjectQ.get_instance() + container.bind_instance(GraphConfig, graph_config) + state.container = container + state.graph = await attach_all_modules(graph_config, container) + await _prepare_checkpointer(container) + settings = get_settings() + if settings.STARTUP_CHECKS: + await check_backends_at_startup(container, timeout=settings.STARTUP_CHECK_TIMEOUT_SECONDS) + setup_taskiq(container, broker) + import_tasks(worker_config) + wrap_tasks(broker) + + +async def worker_shutdown(state: TaskiqState) -> None: + """Close the graph (and with it the checkpointer and store), then the cache store.""" + from tenxgraph_api.src.app.core.cache_store import BaseCacheStore + from tenxgraph_api.src.app.core.config.settings import get_settings + + container = getattr(state, "container", None) or InjectQ.get_instance() + graph = getattr(state, "graph", None) + if graph is not None: + settings = get_settings() + try: + await graph.aclose( + grace=settings.SHUTDOWN_GRACE_SECONDS, + stop_timeout=settings.SHUTDOWN_STOP_SECONDS, + timeout=settings.SHUTDOWN_TIMEOUT_SECONDS, + ) + except Exception as exc: + logger.warning("Graph close did not finish cleanly: %s", exc) + try: + cache_store = container.try_get(BaseCacheStore) + if cache_store is not None: + await cache_store.aclose() + except Exception as exc: + logger.warning("Cache store close failed: %s", exc) + + +if worker_config is not None: + broker.add_event_handler(TaskiqEvents.WORKER_STARTUP, worker_startup) + broker.add_event_handler(TaskiqEvents.WORKER_SHUTDOWN, worker_shutdown) diff --git a/tenxgraph_api/src/app/core/config/graph_config.py b/tenxgraph_api/src/app/core/config/graph_config.py index 40339ae5..a1b2d8e8 100644 --- a/tenxgraph_api/src/app/core/config/graph_config.py +++ b/tenxgraph_api/src/app/core/config/graph_config.py @@ -4,12 +4,16 @@ import os from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from dotenv import load_dotenv from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator +if TYPE_CHECKING: + from tenxgraph_api.jobs.config import WorkerConfig + + logger = logging.getLogger("tenxgraph_api") @@ -808,6 +812,14 @@ def __init__(self, path: str | None = None): if isinstance(self.data, dict): check_cache_store_settings(self.data) + @property + def worker(self) -> "WorkerConfig | None": + """``worker`` key (the jobs layer), parsed; None when absent.""" + from tenxgraph_api.jobs.config import WorkerConfig + + raw = self.data.get("worker") + return None if raw is None else WorkerConfig.model_validate(raw) + @property def graph_path(self) -> str: agent = self.data.get("agent") diff --git a/tests/unit_tests/jobs/__init__.py b/tests/unit_tests/jobs/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/jobs/conftest.py b/tests/unit_tests/jobs/conftest.py new file mode 100644 index 00000000..325655ce --- /dev/null +++ b/tests/unit_tests/jobs/conftest.py @@ -0,0 +1,65 @@ +"""An in-memory broker with our middlewares, a fresh InjectQ container and cache store.""" + +from __future__ import annotations + +from typing import Any + +import pytest +import pytest_asyncio +from injectq import InjectQ +from injectq.integrations.taskiq import setup_taskiq +from taskiq import InMemoryBroker, ScheduledTask, ScheduleSource + +from tenxgraph_api.jobs.middleware import TaskMiddleware, wrap_tasks +from tenxgraph_api.jobs.schedules import set_schedule_source +from tenxgraph_api.src.app.core.cache_store import BaseCacheStore, InMemoryCacheStore + + +class Service: + """Something InjectQ injects into tasks.""" + + name = "bound" + + +class FakeScheduleSource(ScheduleSource): + def __init__(self) -> None: + self.items: dict[str, ScheduledTask] = {} + + async def get_schedules(self) -> list[ScheduledTask]: + return list(self.items.values()) + + async def add_schedule(self, schedule: ScheduledTask) -> None: + self.items[schedule.schedule_id] = schedule + + async def delete_schedule(self, schedule_id: str) -> None: + self.items.pop(schedule_id, None) + + +class Env: + def __init__(self, store: BaseCacheStore, *middlewares: Any) -> None: + self.store = store + self.service = Service() + self.container = InjectQ() + self.container.bind_instance(BaseCacheStore, store) + self.container.bind_instance(Service, self.service) + self.broker = InMemoryBroker().with_middlewares(*(middlewares or (TaskMiddleware(),))) + setup_taskiq(self.container, self.broker) + + def wrap(self) -> None: + wrap_tasks(self.broker) + + +@pytest_asyncio.fixture +async def env(): + environment = Env(InMemoryCacheStore()) + await environment.broker.startup() + yield environment + await environment.broker.shutdown() + + +@pytest.fixture +def source(): + fake = FakeScheduleSource() + set_schedule_source(fake) + yield fake + set_schedule_source(None) diff --git a/tests/unit_tests/jobs/test_config.py b/tests/unit_tests/jobs/test_config.py new file mode 100644 index 00000000..099ff53e --- /dev/null +++ b/tests/unit_tests/jobs/test_config.py @@ -0,0 +1,101 @@ +import pytest +from pydantic import ValidationError +from taskiq import InMemoryBroker +from taskiq_redis import ListQueueBroker, ListRedisScheduleSource, RedisAsyncResultBackend + +from tenxgraph_api.jobs import config as jobs_config +from tenxgraph_api.jobs.config import ( + WorkerConfig, + build_broker, + build_result_backend, + build_schedule_source, +) +from tenxgraph_api.src.app.core.config.graph_config import GraphConfig + + +def test_defaults(): + cfg = WorkerConfig() + assert cfg.broker.type == "redis_stream" + assert cfg.result_ttl == 86400 + assert cfg.retries == 3 + assert cfg.scheduler_interval == 10 + assert cfg.schedules is False + assert cfg.prometheus is None + + +@pytest.mark.parametrize("bad", [{"broker": {"type": "nope"}}, {"retries": -1}, {"typo": 1}]) +def test_invalid_values_refused(bad): + with pytest.raises(ValidationError): + WorkerConfig.model_validate(bad) + + +def test_url_expands_env(monkeypatch): + monkeypatch.setenv("MY_REDIS", "redis://example:1/2") + cfg = WorkerConfig.model_validate({"broker": {"url": "${MY_REDIS}"}}) + assert cfg.broker_url() == "redis://example:1/2" + + +def test_unset_env_var_refused(monkeypatch): + monkeypatch.delenv("NOT_SET_VAR", raising=False) + cfg = WorkerConfig.model_validate({"broker": {"url": "${NOT_SET_VAR}"}}) + with pytest.raises(ValueError, match="environment variable that is not set"): + cfg.broker_url() + + +def test_redis_url_falls_back_to_env(monkeypatch): + monkeypatch.setenv("REDIS_URL", "redis://fallback:6379/0") + assert WorkerConfig().broker_url() == "redis://fallback:6379/0" + monkeypatch.delenv("REDIS_URL") + with pytest.raises(ValueError, match="worker.broker.url is required"): + WorkerConfig().broker_url() + + +def test_memory_broker_has_no_url_or_backend(): + cfg = WorkerConfig.model_validate({"broker": {"type": "memory"}}) + assert isinstance(build_broker(cfg), InMemoryBroker) + assert build_result_backend(cfg) is None + + +def test_redis_list_broker_gets_default_redis_backend(): + cfg = WorkerConfig.model_validate( + {"broker": {"type": "redis_list", "url": "redis://h:1/0"}, "result_ttl": 50} + ) + broker = build_broker(cfg) + assert isinstance(broker, ListQueueBroker) + assert isinstance(broker.result_backend, RedisAsyncResultBackend) + assert broker.result_backend.result_ex_time == 50 + + +def test_module_attr_broker_instance_and_class(): + cfg = WorkerConfig.model_validate( + { + "broker": {"type": "taskiq:InMemoryBroker"}, + "result_backend": {"type": "memory"}, + } + ) + assert isinstance(build_broker(cfg), InMemoryBroker) + + +def test_schedule_source_uses_json(monkeypatch): + cfg = WorkerConfig.model_validate({"broker": {"url": "redis://h:1/0"}}) + source = build_schedule_source(cfg) + assert isinstance(source, ListRedisScheduleSource) + assert type(source._serializer).__name__ == "JSONSerializer" + + +def test_missing_package_names_the_extra(monkeypatch): + def boom(name): + raise ImportError(name) + + monkeypatch.setattr(jobs_config.importlib, "import_module", boom) + with pytest.raises(RuntimeError, match=r"10xgraph-api\[jobs\]"): + build_broker(WorkerConfig()) + + +def test_graph_config_worker_property(tmp_path): + path = tmp_path / "10xgraph.json" + path.write_text('{"agent": "a:b", "worker": {"broker": {"type": "memory"}, "retries": 1}}') + cfg = GraphConfig(str(path)).worker + assert cfg is not None and cfg.retries == 1 + path.write_text('{"agent": "a:b"}') + assert GraphConfig(str(path)).worker is None diff --git a/tests/unit_tests/jobs/test_context.py b/tests/unit_tests/jobs/test_context.py new file mode 100644 index 00000000..18e63b09 --- /dev/null +++ b/tests/unit_tests/jobs/test_context.py @@ -0,0 +1,176 @@ +"""TaskContext, the reserved-argument check, enqueue, cancel and the schedule helpers.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import Annotated + +import pytest +from injectq.integrations.taskiq import InjectTaskiq +from taskiq import Context, TaskiqDepends + +from tenxgraph_api.jobs import state +from tenxgraph_api.jobs.context import ( + ReservedArgumentError, + ScheduleNotFoundError, + TaskContext, + cancel, + check_arguments, + enqueue, + reschedule, + schedule, + unschedule, +) + +from .conftest import Service + + +pytestmark = pytest.mark.asyncio +CONFIG = {"user_id": "u1", "user": {"name": "Ann"}, "authz": {"ok": True}} + + +def make_report(env): + @env.broker.task + async def report( + report_id: int, + config: dict, + graph: Annotated[Service, InjectTaskiq(Service)], + ctx: Annotated[TaskContext, TaskiqDepends()], + *, + label: str = "x", + ) -> dict: + await ctx.progress({"step": "loading"}) + return { + "id": report_id, + "config": config, + "ctx_config": ctx.config, + "graph": graph is env.service, + "cancelled": await ctx.cancelled(), + } + + return report + + +@pytest.mark.parametrize( + ("args", "kwargs", "message"), + [ + ((1,), {"graph": "EVIL"}, "set by the server"), + ((1,), {"ctx": "EVIL"}, "set by the server"), + ((1,), {"config": {}}, "set by the server"), + ((1,), {"user_id": "u2"}, "set by the server"), + ((1,), {"authz": {}}, "set by the server"), + ((1,), {"bogus": 1}, "has no arguments"), + ((1, {}), {}, "at most 1 positional"), + ((1,), {"report_id": 2}, "given twice"), + ], +) +async def test_check_arguments_rejects(env, args, kwargs, message): + task = make_report(env) + with pytest.raises(ReservedArgumentError, match=message): + check_arguments(task, args, kwargs) + + +async def test_check_arguments_accepts_plain_and_spots_all_injection_styles(env): + @env.broker.task + async def styles( + a: int, + c: Context = TaskiqDepends(), + s: Service = InjectTaskiq(Service), + **extra: int, + ) -> None: ... + + check_arguments(make_report(env), (1,), {"label": "y"}) + check_arguments(styles, (1,), {"anything": 2}) # **extra takes unknown names + for name in ("c", "s"): + with pytest.raises(ReservedArgumentError): + check_arguments(styles, (), {name: 1}) + + +async def test_enqueue_passes_config_records_owner_and_progress(env): + task = make_report(env) + handle = await enqueue(task, 7, config=CONFIG, label="z") + result = await handle.wait_result(timeout=5) + assert not result.is_err, result.error + value = result.return_value + expected = {**CONFIG, "task_id": handle.task_id} + assert value["config"] == expected + assert value["ctx_config"] == expected + assert value["graph"] is True + assert value["cancelled"] is False + progress = await env.broker.result_backend.get_progress(handle.task_id) + assert progress.meta == {"step": "loading"} + record = await state.read_status(env.store, handle.task_id) + assert (record.status, record.owner, record.task_name) == ("success", "u1", task.task_name) + + +async def test_enqueue_refuses_before_sending(env): + task = make_report(env) + with pytest.raises(ReservedArgumentError): + await enqueue(task, 1, graph="EVIL") + with pytest.raises(ReservedArgumentError): + await enqueue(task, 1, 2) + + +async def test_enqueue_without_config_and_sync_task(env): + @env.broker.task + def plain(a: int) -> int: + return a * 2 + + @env.broker.task + def sync_cfg(a: int, config: dict, ctx: Annotated[TaskContext, TaskiqDepends()]) -> list: + return [a, config, ctx.config, ctx.task_name] + + result = await (await enqueue(plain, 4, config=CONFIG)).wait_result(timeout=5) + assert result.return_value == 8 + handle = await enqueue(sync_cfg, 1) + result = await handle.wait_result(timeout=5) + only_id = {"task_id": handle.task_id} + assert result.return_value == [1, only_id, only_id, sync_cfg.task_name] + + +async def test_cancel_marks_only_a_queued_task(env): + await state.write_status( + env.store, state.TaskStatus(task_id="q", task_name="t", status="queued") + ) + await state.write_status( + env.store, state.TaskStatus(task_id="r", task_name="t", status="running") + ) + await cancel("q", env.broker) + await cancel("r", env.broker) + assert (await state.read_status(env.store, "q")).status == "cancelled" + assert (await state.read_status(env.store, "r")).status == "running" + assert await state.is_cancel_flag_set(env.store, "r") + + +async def test_schedule_reschedule_unschedule(env, source): + task = make_report(env) + schedule_id = await schedule(task, 3, cron="0 9 * * 1", config=CONFIG, label="w") + stored = source.items[schedule_id] + assert (stored.cron, stored.args, stored.kwargs) == ("0 9 * * 1", [3], {"label": "w"}) + assert stored.task_id is None # never with_task_id on a schedule + assert state.parse_config(stored.labels) == CONFIG + assert (await state.read_schedule(env.store, schedule_id)).owner == "u1" + + await reschedule(schedule_id, every=60, broker=env.broker) + changed = source.items[schedule_id] + assert (changed.cron, changed.interval, changed.args) == (None, 60, [3]) + assert (await state.read_schedule(env.store, schedule_id)).task.interval == 60 + + at = datetime.now(UTC) + timedelta(hours=1) + await reschedule(schedule_id, at=at, broker=env.broker) + assert source.items[schedule_id].time == at + + await unschedule(schedule_id, env.broker) + assert schedule_id not in source.items + assert await state.read_schedule(env.store, schedule_id) is None + with pytest.raises(ScheduleNotFoundError): + await reschedule(schedule_id, cron="* * * * *", broker=env.broker) + + +async def test_schedule_checks_timing_and_arguments(env, source): + task = make_report(env) + with pytest.raises(ValueError, match="exactly one"): + await schedule(task, 1, cron="* * * * *", every=5) + with pytest.raises(ReservedArgumentError): + await schedule(task, 1, cron="* * * * *", graph="EVIL") + assert source.items == {} diff --git a/tests/unit_tests/jobs/test_middleware.py b/tests/unit_tests/jobs/test_middleware.py new file mode 100644 index 00000000..34e9405b --- /dev/null +++ b/tests/unit_tests/jobs/test_middleware.py @@ -0,0 +1,217 @@ +"""TaskMiddleware, the cancel guard, wrap_tasks and MonitoringMiddleware on InMemoryBroker. + +String annotations on purpose (``from __future__``): the guard must keep TaskIQ's signature, +type-hint and dependency reading working through ``functools.wraps``. +""" + +from __future__ import annotations + +import logging +import uuid +from typing import Annotated + +import pytest +import pytest_asyncio +from injectq.integrations.taskiq import InjectTaskiq +from taskiq import SimpleRetryMiddleware, TaskiqDepends +from tenxgraph.runtime.publisher.base_publisher import BasePublisher +from tenxgraph.utils.background_task_manager import BackgroundTaskManager + +from tenxgraph_api.jobs import state +from tenxgraph_api.jobs.arguments import ReservedArgumentError +from tenxgraph_api.jobs.context import TaskCancelled, TaskContext, cancel +from tenxgraph_api.jobs.middleware import MonitoringMiddleware, TaskMiddleware +from tenxgraph_api.src.app.core.cache_store import InMemoryCacheStore + +from .conftest import Env, Service + + +pytestmark = pytest.mark.asyncio +RAN: list[str] = [] + + +def add_tasks(env: Env): + @env.broker.task + async def job(a: int, config: dict, graph: Annotated[Service, InjectTaskiq(Service)]) -> list: + RAN.append("job") + return [a, config, graph is env.service] + + @env.broker.task + def sync_job(a: int) -> int: + RAN.append("sync_job") + return a + + env.wrap() + return job, sync_job + + +async def kiq_with_id(task, *args, **kwargs): + task_id = uuid.uuid4().hex + return await task.kicker().with_task_id(task_id).kiq(*args, **kwargs) + + +@pytest.fixture(autouse=True) +def _reset(): + RAN.clear() + + +async def test_cancel_before_start_gives_error_result_and_status(env): + job, sync_job = add_tasks(env) + for task in (job, sync_job): + task_id = uuid.uuid4().hex + await state.write_status( + env.store, state.TaskStatus(task_id=task_id, task_name=task.task_name, status="queued") + ) + await cancel(task_id, env.broker) + handle = await task.kicker().with_task_id(task_id).kiq(1) + result = await handle.wait_result(timeout=5) + assert result.is_err + assert isinstance(result.error, TaskCancelled) + assert (await state.read_status(env.store, task_id)).status == "cancelled" + assert RAN == [] + # The next run is not affected by the cancelled one. + assert (await (await kiq_with_id(sync_job, 5)).wait_result(timeout=5)).return_value == 5 + + +async def test_worker_recheck_blocks_injected_override(env): + job, _ = add_tasks(env) + handle = await kiq_with_id(job, 1, graph="EVIL") # plain kiq, no enqueue check + result = await handle.wait_result(timeout=5) + assert isinstance(result.error, ReservedArgumentError) + assert RAN == [] + assert (await state.read_status(env.store, handle.task_id)).status == "failure" + + +async def test_config_comes_only_from_the_label(env): + job, sync_job = add_tasks(env) + kicker = job.kicker().with_labels(**{state.CONFIG_LABEL: '{"user_id": "u1"}'}) + handle = await kicker.kiq(2, config={"user_id": "EVIL"}) + result = await handle.wait_result(timeout=5) + assert result.return_value == [2, {"user_id": "u1", "task_id": handle.task_id}, True] + record = await state.read_status(env.store, handle.task_id) + assert (record.status, record.owner) == ("success", "u1") + # A task without ``config`` never gets one; a broken label means ``{}``. + kicker = sync_job.kicker().with_labels(**{state.CONFIG_LABEL: "not json"}) + assert (await (await kicker.kiq(3, config={})).wait_result(timeout=5)).return_value == 3 + + +async def test_retry_resends_config_without_tripping_the_check(): + env = Env(InMemoryCacheStore(), SimpleRetryMiddleware(default_retry_count=3), TaskMiddleware()) + await env.broker.startup() + calls: list[dict] = [] + + @env.broker.task(retry_on_error=True) + async def flaky(config: dict) -> str: + calls.append(config) + if len(calls) == 1: + raise RuntimeError("first try fails") + return "ok" + + env.wrap() + handle = await flaky.kicker().with_labels(**{state.CONFIG_LABEL: "{}"}).kiq() + result = await handle.wait_result(timeout=5) + assert result.return_value == "ok" + assert len(calls) == 2 + assert (await state.read_status(env.store, handle.task_id)).status == "success" + await env.broker.shutdown() + + +class BrokenStore(InMemoryCacheStore): + async def get(self, key: str) -> bytes | None: + raise ConnectionError("store down") + + async def set(self, key, value, ttl, mode="always") -> bool: + raise ConnectionError("store down") + + +async def test_middleware_never_raises_when_the_store_fails(caplog): + env = Env(BrokenStore()) + await env.broker.startup() + job, _ = add_tasks(env) + with caplog.at_level(logging.ERROR, logger="tenxgraph_api.jobs"): + kicker = job.kicker().with_labels(**{state.CONFIG_LABEL: '{"user_id": "u1"}'}) + result = await (await kicker.kiq(1)).wait_result(timeout=5) + assert not result.is_err, result.error + assert result.return_value[1]["user_id"] == "u1" # config still passed + assert "pre_execute failed" in caplog.text + assert "status not saved" in caplog.text + await env.broker.shutdown() + + +async def test_wrap_tasks_is_idempotent_and_keeps_injection(env): + job, _ = add_tasks(env) + wrapped = job.original_func + env.wrap() + assert job.original_func is wrapped + assert not hasattr(wrapped.__wrapped__, "__tenxgraph_guard__") # wrapped once + + @env.broker.task + async def with_ctx(ctx: Annotated[TaskContext, TaskiqDepends()]) -> str: + return ctx.task_id + + env.wrap() + handle = await kiq_with_id(with_ctx) + assert (await handle.wait_result(timeout=5)).return_value == handle.task_id + + +class FakePublisher(BasePublisher): + def __init__(self, fail: bool = False) -> None: + super().__init__({}) + self.events: list = [] + self.fail = fail + + async def publish(self, event): + if self.fail: + raise RuntimeError("sink down") + self.events.append(event) + + async def close(self): ... + + def sync_close(self): ... + + +@pytest_asyncio.fixture +async def monitored(): + env = Env(InMemoryCacheStore(), MonitoringMiddleware(), TaskMiddleware()) + manager = BackgroundTaskManager() + env.container.bind_instance(BackgroundTaskManager, manager) + await env.broker.startup() + yield env, manager + await env.broker.shutdown() + + +async def test_monitoring_logs_and_publishes(monitored, caplog): + env, manager = monitored + publisher = FakePublisher() + env.container.bind_instance(BasePublisher, publisher) + job, sync_job = add_tasks(env) + + @env.broker.task + async def boom() -> None: + raise ValueError("bad") + + env.wrap() + with caplog.at_level(logging.INFO, logger="tenxgraph_api.jobs"): + await (await kiq_with_id(sync_job, 1)).wait_result(timeout=5) + await (await kiq_with_id(boom)).wait_result(timeout=5) + await manager.wait_for_all(timeout=5) + phases = [ + (event.node_name.rsplit(":", 1)[-1], str(event.event_type)) for event in publisher.events + ] + assert phases == [ + ("sync_job", "start"), + ("sync_job", "end"), + ("boom", "start"), + ("boom", "error"), + ] + assert publisher.events[-1].data["error"] == "ValueError" + assert ( + "started" in caplog.text and "finished in" in caplog.text and "failed after" in caplog.text + ) + + +async def test_monitoring_never_raises_on_a_failing_publisher(monitored): + env, _ = monitored + env.container.bind_instance(BasePublisher, FakePublisher(fail=True)) + _, sync_job = add_tasks(env) + assert (await (await kiq_with_id(sync_job, 9)).wait_result(timeout=5)).return_value == 9 diff --git a/tests/unit_tests/jobs/test_redis_broker.py b/tests/unit_tests/jobs/test_redis_broker.py new file mode 100644 index 00000000..694cfc4a --- /dev/null +++ b/tests/unit_tests/jobs/test_redis_broker.py @@ -0,0 +1,175 @@ +"""Live tests for ReclaimingRedisStreamBroker (idle-queue reclaim and lease renewal). + +Need a real Redis: run with ``--integration``. The URL comes from ``TEST_JOBS_REDIS_URL`` +(default ``redis://localhost:16379/7``); the test is skipped when it is not reachable. +""" + +from __future__ import annotations + +import asyncio +import os +import uuid +from collections.abc import AsyncIterator, Callable + +import pytest +import pytest_asyncio +from redis.asyncio import Redis +from taskiq import AckableMessage +from taskiq.message import BrokerMessage + +from tenxgraph_api.jobs.redis_broker import ReclaimingRedisStreamBroker + + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio] + +REDIS_URL = os.environ.get("TEST_JOBS_REDIS_URL", "redis://localhost:16379/7") +GROUP = "g" + + +@pytest_asyncio.fixture +async def redis() -> AsyncIterator[Redis]: + client = Redis.from_url(REDIS_URL) + try: + await asyncio.wait_for(client.ping(), 2) + except Exception: + await client.aclose() + pytest.skip(f"Redis not reachable at {REDIS_URL}") + yield client + await client.aclose() + + +@pytest_asyncio.fixture +async def queue(redis: Redis) -> AsyncIterator[str]: + name = f"test-jobs-{uuid.uuid4().hex}" + yield name + await redis.delete(name, f"autoclaim:{GROUP}:{name}") + + +@pytest_asyncio.fixture +async def make_broker(queue: str) -> AsyncIterator[Callable[[str], ReclaimingRedisStreamBroker]]: + brokers: list[ReclaimingRedisStreamBroker] = [] + + def make(consumer: str) -> ReclaimingRedisStreamBroker: + broker = ReclaimingRedisStreamBroker( + REDIS_URL, + queue_name=queue, + consumer_group_name=GROUP, + consumer_name=consumer, + idle_timeout=1000, + xread_block=200, + ) + brokers.append(broker) + return broker + + yield make + for broker in brokers: + await broker.shutdown() + + +async def kick(broker: ReclaimingRedisStreamBroker, data: bytes) -> None: + msg = BrokerMessage(task_id=uuid.uuid4().hex, task_name="t", message=data, labels={}) + await broker.kick(msg) + + +async def collect(broker: ReclaimingRedisStreamBroker, into: list[AckableMessage]) -> None: + async for message in broker.listen(): + into.append(message) + + +async def pending_count(redis: Redis, queue: str) -> int: + info = await redis.xpending(queue, GROUP) + return int(info["pending"]) + + +async def test_idle_queue_reclaims_dead_consumer_message(redis, queue, make_broker): + live = make_broker("live") + await live.startup() + await kick(live, b"orphan") + # A dead consumer read it and never acked. + await redis.xreadgroup(GROUP, "dead", {queue: ">"}, count=1) + + # No new message arrives; the live consumer must still claim it once it is idle. + gen = live.listen() + message = await asyncio.wait_for(anext(gen), 5) + assert message.data == b"orphan" + await message.ack() + await gen.aclose() + assert await pending_count(redis, queue) == 0 + + +async def test_renewal_keeps_long_task_from_being_reclaimed(redis, queue, make_broker): + worker_a, worker_b = make_broker("a"), make_broker("b") + await worker_a.startup() + await worker_b.startup() + await kick(worker_a, b"long") + + gen_a = worker_a.listen() + message = await asyncio.wait_for(anext(gen_a), 5) + seen_by_b: list[AckableMessage] = [] + task_b = asyncio.create_task(collect(worker_b, seen_by_b)) + await asyncio.sleep(3) # three idle timeouts while "running" + + assert seen_by_b == [] + await message.ack() + task_b.cancel() + await asyncio.gather(task_b, return_exceptions=True) + await gen_a.aclose() + assert await pending_count(redis, queue) == 0 + + +async def test_ack_stops_renewal_and_late_renew_does_not_resurrect(redis, queue, make_broker): + worker = make_broker("a") + await worker.startup() + await kick(worker, b"quick") + gen = worker.listen() + message = await asyncio.wait_for(anext(gen), 5) + assert len(worker._held) == 1 + + await message.ack() + assert worker._held == set() + # A renewal racing the ack: renew the acked id directly. + stream = queue.encode() + msg_id = (await redis.xrange(queue))[0][0] + worker._held.add((stream, msg_id)) + await worker._renew_held() + assert worker._held == set() # lost lease dropped + assert await pending_count(redis, queue) == 0 + await gen.aclose() + + +async def test_shutdown_cancels_renewal(redis, queue, make_broker): + worker_a, worker_b = make_broker("a"), make_broker("b") + await worker_a.startup() + await worker_b.startup() + await kick(worker_a, b"abandoned") + gen_a = worker_a.listen() + await asyncio.wait_for(anext(gen_a), 5) + renewer = worker_a._renewer + assert renewer is not None and not renewer.done() + + await gen_a.aclose() + await worker_a.shutdown() + assert renewer.cancelled() + assert worker_a._renewer is None and worker_a._held == set() + + # No longer renewed: it goes idle and another worker claims it. + gen_b = worker_b.listen() + message = await asyncio.wait_for(anext(gen_b), 5) + assert message.data == b"abandoned" + await message.ack() + await gen_b.aclose() + + +async def test_closing_listen_mid_batch_releases_unyielded(redis, queue, make_broker): + worker = make_broker("a") + await worker.startup() + for i in range(3): + await kick(worker, f"m{i}".encode()) + gen = worker.listen() + first = await asyncio.wait_for(anext(gen), 5) + assert len(worker._held) == 3 # whole batch held while waiting for a slot + + await gen.aclose() + assert len(worker._held) == 1 # only the one handed to the worker + await first.ack() + assert worker._held == set() diff --git a/tests/unit_tests/jobs/test_worker.py b/tests/unit_tests/jobs/test_worker.py new file mode 100644 index 00000000..ccd4cb6e --- /dev/null +++ b/tests/unit_tests/jobs/test_worker.py @@ -0,0 +1,197 @@ +import importlib +import json +import logging +import sys +from types import SimpleNamespace + +import pytest +from taskiq import InMemoryBroker +from taskiq.middlewares import SimpleRetryMiddleware + +from tenxgraph_api.jobs import middleware as mw +from tenxgraph_api.jobs import monitoring +from tenxgraph_api.jobs.config import WorkerConfig + + +def load_worker(monkeypatch, tmp_path, data): + """Import a fresh ``tenxgraph_api.jobs.worker`` reading a 10xgraph.json made from ``data``.""" + path = tmp_path / "10xgraph.json" + path.write_text(json.dumps(data)) + monkeypatch.setenv("GRAPH_PATH", str(path)) + sys.modules.pop("tenxgraph_api.jobs.worker", None) + return importlib.import_module("tenxgraph_api.jobs.worker") + + +def test_no_worker_key_gives_memory_broker_and_warns(monkeypatch, tmp_path, caplog): + with caplog.at_level(logging.WARNING, logger="tenxgraph_api.jobs"): + worker = load_worker(monkeypatch, tmp_path, {"agent": "a:b"}) + assert isinstance(worker.broker, InMemoryBroker) + assert sum("in-memory broker" in r.message for r in caplog.records) == 1 + assert not worker.broker.event_handlers + + +def test_missing_config_file_gives_memory_broker(monkeypatch, tmp_path): + monkeypatch.setenv("GRAPH_PATH", str(tmp_path / "missing.json")) + sys.modules.pop("tenxgraph_api.jobs.worker", None) + worker = importlib.import_module("tenxgraph_api.jobs.worker") + assert isinstance(worker.broker, InMemoryBroker) + + +def test_middlewares_in_order(monkeypatch, tmp_path): + data = {"agent": "a:b", "worker": {"broker": {"type": "memory"}, "retries": 5}} + worker = load_worker(monkeypatch, tmp_path, data) + kinds = [type(m) for m in worker.broker.middlewares] + assert kinds == [SimpleRetryMiddleware, monitoring.MonitoringMiddleware, mw.TaskMiddleware] + assert worker.broker.middlewares[0].default_retry_count == 5 + + +def test_import_tasks_by_module_and_glob(monkeypatch, tmp_path): + (tmp_path / "pkg_tasks.py").write_text("LOADED = True\n") + (tmp_path / "other_tasks.py").write_text("LOADED = True\n") + monkeypatch.chdir(tmp_path) + monkeypatch.syspath_prepend(str(tmp_path)) + data = {"agent": "a:b", "worker": {"broker": {"type": "memory"}}} + worker = load_worker(monkeypatch, tmp_path, data) + worker.import_tasks(WorkerConfig(tasks=["pkg_tasks", "other_*.py"])) + assert "pkg_tasks" in sys.modules and "other_tasks" in sys.modules + + +@pytest.mark.asyncio +async def test_startup_and_shutdown(monkeypatch, tmp_path): + data = {"agent": "a:b", "worker": {"broker": {"type": "memory"}}} + worker = load_worker(monkeypatch, tmp_path, data) + closed = [] + + class FakeGraph: + async def aclose(self, **kwargs): + closed.append(kwargs) + + async def fake_attach(config, container): + return FakeGraph() + + from tenxgraph_api.src.app import loader + + monkeypatch.setattr(loader, "attach_all_modules", fake_attach) + monkeypatch.setattr(loader, "load_container", lambda path: None) + monkeypatch.setattr(worker, "wrap_tasks", lambda b: closed.append("wrapped")) + monkeypatch.setenv("STARTUP_CHECKS", "false") + from tenxgraph_api.src.app.core.config.settings import get_settings + + get_settings.cache_clear() + + state = SimpleNamespace() + await worker.worker_startup(state) + assert isinstance(state.graph, FakeGraph) + assert "wrapped" in closed + assert worker.broker.state.injectq_container is not None + + await worker.worker_shutdown(state) + assert any(isinstance(c, dict) and "grace" in c for c in closed) + get_settings.cache_clear() + + +@pytest.mark.asyncio +async def test_shutdown_survives_failing_graph(monkeypatch, tmp_path): + worker = load_worker(monkeypatch, tmp_path, {"agent": "a:b"}) + + class BadGraph: + async def aclose(self, **kwargs): + raise RuntimeError("boom") + + await worker.worker_shutdown(SimpleNamespace(graph=BadGraph())) + + +def _patch_loader(monkeypatch, container=None, closed=None): + from tenxgraph_api.src.app import loader + from tenxgraph_api.src.app.core.config.settings import get_settings + + class FakeGraph: + async def aclose(self, **kwargs): + if closed is not None: + closed.append(kwargs) + + async def fake_attach(config, cont): + return FakeGraph() + + monkeypatch.setattr(loader, "attach_all_modules", fake_attach) + monkeypatch.setattr(loader, "load_container", lambda path: container) + monkeypatch.setenv("STARTUP_CHECKS", "false") + get_settings.cache_clear() + + +@pytest.mark.asyncio +async def test_shutdown_uses_container_from_startup(monkeypatch, tmp_path): + from injectq import InjectQ + + from tenxgraph_api.src.app.core.cache_store import BaseCacheStore + + data = {"agent": "a:b", "worker": {"broker": {"type": "memory"}}} + worker = load_worker(monkeypatch, tmp_path, data) + mine = InjectQ() + _patch_loader(monkeypatch, container=mine) + monkeypatch.setattr(worker, "wrap_tasks", lambda b: None) + state = SimpleNamespace() + await worker.worker_startup(state) + assert state.container is mine + + closed = [] + + class FakeCache: + async def aclose(self): + closed.append("cache") + + mine.bind_instance(BaseCacheStore, FakeCache()) + await worker.worker_shutdown(state) + assert closed == ["cache"] + + +@pytest.mark.asyncio +async def test_startup_runs_checkpointer_setup(monkeypatch, tmp_path): + from injectq import InjectQ + + data = {"agent": "a:b", "worker": {"broker": {"type": "memory"}}} + worker = load_worker(monkeypatch, tmp_path, data) + calls = [] + + class FakeCheckpointer: + async def asetup(self): + calls.append("setup") + + mine = InjectQ() + mine.bind_instance("checkpointer", FakeCheckpointer()) + _patch_loader(monkeypatch, container=mine) + monkeypatch.setattr(worker, "wrap_tasks", lambda b: None) + await worker.worker_startup(SimpleNamespace()) + assert calls == ["setup"] + + +@pytest.mark.asyncio +async def test_checkpointer_setup_failure_is_logged_not_raised(monkeypatch, tmp_path, caplog): + from injectq import InjectQ + + worker = load_worker(monkeypatch, tmp_path, {"agent": "a:b"}) + + class BadCheckpointer: + async def asetup(self): + raise RuntimeError("no db") + + mine = InjectQ() + mine.bind_instance("checkpointer", BadCheckpointer()) + with caplog.at_level(logging.ERROR, logger="tenxgraph_api.jobs"): + await worker._prepare_checkpointer(mine) + assert any("Checkpointer setup failed" in r.message for r in caplog.records) + + +def test_result_ttl_set_from_config(monkeypatch, tmp_path): + from tenxgraph_api.jobs import state as job_state + + monkeypatch.setattr(job_state, "RESULT_TTL", 86400) + data = {"agent": "a:b", "worker": {"broker": {"type": "memory"}, "result_ttl": 123}} + load_worker(monkeypatch, tmp_path, data) + assert job_state.RESULT_TTL == 123 + + +@pytest.fixture(autouse=True) +def _restore_worker_module(): + yield + sys.modules.pop("tenxgraph_api.jobs.worker", None) From 8d7d8200863d9b291f06796658c90bd2096458e2 Mon Sep 17 00:00:00 2001 From: Shudipto Trafder Date: Sat, 7 Nov 2026 11:04:19 +0600 Subject: [PATCH 2/4] Refactor tests and improve configuration handling - Removed unused imports and fixtures from unit tests in `conftest.py`. - Added new unit tests for worker commands in `test_commands.py` to validate argument parsing and environment setup. - Enhanced configuration validation in `test_config.py` to support new broker types and ensure proper error handling for unsupported configurations. - Updated task scheduling tests in `test_context.py` to streamline the enqueueing process and remove deprecated functionality. - Introduced new route tests in `test_routes.py` to cover task management endpoints with authorization checks. - Cleaned up worker tests in `test_worker.py` by removing obsolete tests related to task importing and checkpointer setup. - Added comprehensive validation for worker configuration fields in `test_config_editor.py` to ensure correct schema definitions and error messages. --- config-editor-ui/src/fields.jsx | 4 +- examples/background_tasks/10xgraph.json | 9 + examples/background_tasks/README.md | 34 +++ examples/background_tasks/app/__init__.py | 0 examples/background_tasks/app/graph.py | 16 ++ examples/background_tasks/app/tasks.py | 31 ++ pyproject.toml | 13 + tenxgraph_api/cli/commands/worker.py | 122 ++++++++ tenxgraph_api/cli/config_editor/schema.py | 77 +++++ tenxgraph_api/cli/config_editor/static/app.js | 2 +- tenxgraph_api/cli/config_editor/validation.py | 33 +++ tenxgraph_api/cli/main.py | 33 +++ tenxgraph_api/jobs/__init__.py | 14 +- tenxgraph_api/jobs/arguments.py | 19 +- tenxgraph_api/jobs/config.py | 155 +++++----- tenxgraph_api/jobs/context.py | 69 +---- tenxgraph_api/jobs/scheduler.py | 19 ++ tenxgraph_api/jobs/schedules.py | 128 --------- tenxgraph_api/jobs/state.py | 35 +-- tenxgraph_api/jobs/worker.py | 45 +-- .../src/app/core/auth/authorization.py | 3 + tenxgraph_api/src/app/main.py | 15 + tenxgraph_api/src/app/routers/setup_router.py | 10 +- .../src/app/routers/tasks/__init__.py | 7 + tenxgraph_api/src/app/routers/tasks/boot.py | 49 ++++ tenxgraph_api/src/app/routers/tasks/router.py | 246 ++++++++++++++++ .../src/app/routers/tasks/schemas.py | 75 +++++ tests/unit_tests/jobs/conftest.py | 26 +- tests/unit_tests/jobs/test_commands.py | 109 +++++++ tests/unit_tests/jobs/test_config.py | 212 +++++++++++--- tests/unit_tests/jobs/test_context.py | 78 +---- tests/unit_tests/jobs/test_middleware.py | 2 +- tests/unit_tests/jobs/test_routes.py | 266 ++++++++++++++++++ tests/unit_tests/jobs/test_worker.py | 49 ---- tests/unit_tests/test_config_editor.py | 80 +++++- 35 files changed, 1549 insertions(+), 536 deletions(-) create mode 100644 examples/background_tasks/10xgraph.json create mode 100644 examples/background_tasks/README.md create mode 100644 examples/background_tasks/app/__init__.py create mode 100644 examples/background_tasks/app/graph.py create mode 100644 examples/background_tasks/app/tasks.py create mode 100644 tenxgraph_api/cli/commands/worker.py create mode 100644 tenxgraph_api/jobs/scheduler.py delete mode 100644 tenxgraph_api/jobs/schedules.py create mode 100644 tenxgraph_api/src/app/routers/tasks/__init__.py create mode 100644 tenxgraph_api/src/app/routers/tasks/boot.py create mode 100644 tenxgraph_api/src/app/routers/tasks/router.py create mode 100644 tenxgraph_api/src/app/routers/tasks/schemas.py create mode 100644 tests/unit_tests/jobs/test_commands.py create mode 100644 tests/unit_tests/jobs/test_routes.py diff --git a/config-editor-ui/src/fields.jsx b/config-editor-ui/src/fields.jsx index b1f21ff7..718cc493 100644 --- a/config-editor-ui/src/fields.jsx +++ b/config-editor-ui/src/fields.jsx @@ -136,12 +136,12 @@ function Control({ id, field: f, value, invalid }) { ); } if (f.type === "select") { - const options = f.options.map((o) => [o, o]); + const options = f.options.map((o) => (Array.isArray(o) ? o : [o, o])); if (f.unset_label !== undefined) options.unshift(["", f.unset_label]); return (