diff --git a/backend/agent_kv/__init__.py b/backend/agent_kv/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/agent_kv/apps.py b/backend/agent_kv/apps.py new file mode 100644 index 0000000000..0b70b303e1 --- /dev/null +++ b/backend/agent_kv/apps.py @@ -0,0 +1,6 @@ +from django.apps import AppConfig + + +class AgentKvConfig(AppConfig): + default_auto_field = "django.db.models.BigAutoField" + name = "agent_kv" diff --git a/backend/agent_kv/constants.py b/backend/agent_kv/constants.py new file mode 100644 index 0000000000..04cf856e04 --- /dev/null +++ b/backend/agent_kv/constants.py @@ -0,0 +1,70 @@ +#: The `kv` extractor's name. Kept as a constant, but NOT routable on this +#: deployment -- see `EXTRACTOR_ROUTES` below, which carries `table` only. +#: +#: An earlier version of this comment said "v1 accepts exactly one extractor +#: (`kv`), so the job row does not carry which one it ran". Both halves are now +#: wrong, on line 1 of the file that defines the routing table: `kv` is the one +#: extractor this deployment does NOT accept, and `AgentKVJob.extractor` has +#: carried the name since migration 0002. +V1_EXTRACTOR_NAME = "kv" + +STAGE_NAMES = [ + "document_processing", + "extraction", + "qa", + "challenge", + "normalize", + "constraints", + "codegen", + "code_execution", +] +EXECUTOR_NAME = "agentic_kv" +OPERATION_KV_EXTRACT = "kv_extract" +EXECUTION_SOURCE = "agent_kv_api" + +#: The table extractor, served by the cloud `agentic_table` plugin's blind-API +#: operation. Same executor as the IDE table path (and therefore the same, +#: already-wired `celery_executor_agentic_table` queue) with a second operation +#: -- a new executor name would derive a new queue needing wiring at five +#: sites, and an unwired queue accepts work and drains nothing, silently. +TABLE_EXTRACTOR_NAME = "table" + +TABLE_EXECUTOR_NAME = "agentic_table" +OPERATION_TABLE_EXTRACT_API = "table_extract_api" + +#: Which executor and operation each extractor dispatches to. `dispatch_job` +#: reads this rather than hardcoding one pair, so adding an extractor is one +#: entry here plus its options serializer and stage list. +#: +#: **`kv` is deliberately absent.** This deployment ships the `agentic_table` +#: plugin and not `agentic_kv`, so nothing drains `celery_executor_agentic_kv`. +#: `SUPPORTED_EXTRACTORS` is derived from this table's keys +#: (`execution_serializers.py`), so the omission turns a `kv` submit into a 400 +#: at the serializer. Restoring the entry below is all it takes to re-enable +#: the extractor once the plugin ships -- and restoring it WITHOUT the plugin +#: is the failure this guards: the submit would be accepted with a 202, +#: dispatched to a queue with no consumer, and sit in DISPATCHED forever with +#: no error at the producer and nothing in any log to find. +#: +#: `V1_EXTRACTOR_NAME`, `STAGE_NAMES` and the KV options serializer stay in the +#: tree, dormant and still under test, so re-enabling is one line rather than a +#: content merge against the branch that carries the engine. +EXTRACTOR_ROUTES = { + TABLE_EXTRACTOR_NAME: (TABLE_EXECUTOR_NAME, OPERATION_TABLE_EXTRACT_API), +} + +#: The table engine reports one coarse stage: it has no node-level progress +#: hooks (the IDE path gets `stream_log` only), so inventing finer stages here +#: would describe progress the executor cannot actually report. +TABLE_STAGE_NAMES = ["table_extraction"] + +#: Stage names ARE wire format -- they are returned to clients -- and they are +#: extractor-specific (`qa`/`challenge`/`codegen` mean nothing to the table +#: extractor). `_status_document` filters a job's recorded stages through the +#: list for the extractor that ran: `StageReportView` persists whatever name +#: the executor sends, so without a per-extractor list a table job's stages +#: would be stored and then silently filtered out of every status response. +STAGE_NAMES_BY_EXTRACTOR = { + V1_EXTRACTOR_NAME: STAGE_NAMES, + TABLE_EXTRACTOR_NAME: TABLE_STAGE_NAMES, +} diff --git a/backend/agent_kv/dispatch.py b/backend/agent_kv/dispatch.py new file mode 100644 index 0000000000..7f6cb4276f --- /dev/null +++ b/backend/agent_kv/dispatch.py @@ -0,0 +1,244 @@ +"""Executor dispatch glue (spec §5.3). One dispatch per job; UUID task_id.""" + +import logging +import uuid + +from celery import signature +from django.conf import settings +from django.utils import timezone + +from agent_kv.constants import EXECUTION_SOURCE, EXTRACTOR_ROUTES +from agent_kv.models import AgentKVJob, JobStatus +from unstract.sdk1.execution.context import ExecutionContext + +logger = logging.getLogger(__name__) + +CALLBACK_QUEUE = "agent_kv_callback" + + +class DispatchError(Exception): + """Enqueue failed; the caller terminalizes the job (spec §5.3).""" + + +def dispatch_cancelled_webhook(job) -> None: + """Queue the terminal webhook for a job the API just cancelled. + + Cancellation never reaches finalize -- ``JobCancelView`` and DELETE + terminalize the row themselves -- so without this the caller who supplied + ``webhook_url`` is never told, and a late executor callback cannot tell them + either (it loses the terminal guard, and the callback declines a non-fresh + finalize rather than double-notifying). Docs §8 promises delivery on + terminal states; this is the cancel half of that promise. + + Call ONLY when the guarded cancel actually won. That is what makes the two + paths mutually exclusive: a cancel that lost means a finalize won and will + send, and a cancel that won means no fresh finalize can. + + Best-effort by design. The job is already cancelled and the caller already + has their 200; failing their request because a notification could not be + QUEUED would be the wrong trade, so this logs and returns rather than + raising. Delivery itself is the worker's problem. + """ + if not job.webhook_url: + return + try: + from pg_queue.producer import enqueue_task + + enqueue_task( + task_name="agent_kv_cancelled", + queue=CALLBACK_QUEUE, + kwargs={ + "callback_kwargs": { + "job_id": str(job.id), + "webhook_url": job.webhook_url, + } + }, + org_id=str(job.organization_id), + ) + except Exception: + logger.exception( + "agent-kv: could not queue the cancellation webhook for job %s; " + "the job IS cancelled, only the notification was lost", + job.id, + ) + + +def _dispatcher(): + # No `celery_app`: UN-4046 removed that parameter when the routing + # dispatcher's Celery branch went with the pg_queue_enabled flag. Passing it + # raises TypeError, so every submit failed to dispatch -- and because it + # fails at the call rather than at import, nothing catches it until a real + # request is made. + from pg_queue.executor_rpc import get_executor_dispatcher + + return get_executor_dispatcher() + + +def _platform_api_key(job) -> str: + # Lazy import: avoids Django app registry init order (mirrors + # PromptStudioHelper._get_platform_api_key). + from platform_settings_v2.platform_auth_service import ( + PlatformAuthenticationService, + ) + + # ``get_active_platform_key`` takes the org's public *slug* + # (``Organization.organization_id``, e.g. ``org_abc123``) and resolves it + # via ``get_organization_by_org_id`` -- NOT the row's UUID primary key that + # ``job.organization_id`` holds. Passing the PK here silently resolves to + # no organization and every dispatch fails with ``ActiveKeyNotFound`` + # (caught live in the Task 13b integration run). + org_slug = job.organization.organization_id + platform_key = PlatformAuthenticationService.get_active_platform_key(org_slug) + if not platform_key: + raise DispatchError(f"No active platform key for org {org_slug}") + return str(platform_key.key) + + +def dispatch_job( + job, *, extractor: str, schema: dict, options: dict, adapters: dict | None = None +) -> None: + executor_name, operation = EXTRACTOR_ROUTES[extractor] + org_id = str(job.organization_id) + # Everything that can fail — platform-key lookup, context construction, + # and the enqueue call itself — lives inside this try so no internal + # failure (e.g. a transient DB error resolving the platform key) can + # escape as a raw, uncaught exception. Only the post-success bookkeeping + # below runs outside it. + try: + job.task_id = uuid.uuid4() + context = ExecutionContext( + executor_name=executor_name, + operation=operation, + run_id=str(job.id), + execution_source=EXECUTION_SOURCE, + organization_id=org_id, + executor_params={ + "job_id": str(job.id), + "input_ref": job.input_ref, + "schema": schema, + "options": options, + # Platform adapter instance ids, by role, already validated at + # submit against THIS job's organization and against the + # expected `AdapterTypes` (see + # `execution_serializers._validated_adapter_shape` for the + # shape half and `execution_views._resolved_adapters` for the + # tenancy half). + # + # Defence in depth, NOT the only check: the platform service + # re-scopes every lookup as + # `WHERE id=%s and organization_id=%s` + # (platform-service/.../helper/adapter_instance.py:28-30), + # where the org comes from the bearer platform key that + # `_platform_api_key(job)` below mints from THIS job's org. So + # org A cannot spend org B's credential even with the submit + # gate removed. What the submit gate buys is a clean 400 naming + # the role instead of a mid-run `SdkError`, plus the two checks + # the platform service does NOT make: `is_usable` (exhausted + # trial) and `is_available` (deprecated). + # + # Empty for an env-configured extractor (`kv`), which is why + # this is a dict rather than three params: the two credential + # models coexist, one per extractor. + "adapters": adapters or {}, + "platform_api_key": _platform_api_key(job), + # The CAP the engine must enforce (spec §6.1/§6.6), not the + # measured count -- job.pages_total is None for Excel (no + # pre-OCR page concept), which would otherwise leave the + # engine with nothing to check the post-OCR virtual-page cap + # against. The measured count still rides along separately. + "max_pages": settings.AGENT_KV_MAX_PAGES, + "pages_total": job.pages_total, + }, + ) + # Last check before spending money. A cancel can land between the + # submit's `job.save()` and this enqueue: the cancel sees a PENDING, + # never-dispatched row, so it terminalizes it AND releases its + # concurrency slot -- correctly, because nothing had been dispatched + # yet. Enqueueing anyway would then run paid work for a job the caller + # already cancelled, with its slot already handed to someone else. + # + # Re-read rather than trusting the in-memory row, which predates the + # cancel by construction. + if AgentKVJob.objects.filter( + id=job.id, status__in=list(AgentKVJob.TERMINAL) + ).exists(): + logger.info( + "agent-kv: job %s was terminalized before dispatch; not enqueueing", + job.id, + ) + return + + cb_kwargs = {"callback_kwargs": {"job_id": str(job.id), "org_id": org_id}} + _dispatcher().dispatch_with_callback( + context, + on_success=signature( + "agent_kv_complete", kwargs=cb_kwargs, queue=CALLBACK_QUEUE + ), + on_error=signature("agent_kv_error", kwargs=cb_kwargs, queue=CALLBACK_QUEUE), + task_id=str(job.task_id), + ) + except DispatchError: + raise + except Exception as e: + raise DispatchError(str(e)) from e + job.status = JobStatus.DISPATCHED + job.dispatched_at = timezone.now() + # Guarded queryset UPDATE, not job.save(): a plain save would blindly + # overwrite whatever status this job already raced to. Concretely: the + # executor can fail (or the job be cancelled) essentially instantly + # after enqueue, and its finalize callback can land -- marking the row + # FAILED/CANCELLED -- before this post-enqueue bookkeeping runs. An + # unconditional save() here would rewrite that terminal status back to + # DISPATCHED, un-terminalizing the job forever (nothing else ever + # revisits a DISPATCHED row). Only a still-PENDING row is advanced; a + # row this UPDATE doesn't match is left exactly as the winning writer + # left it. `modified_at` is stamped automatically by + # BaseModelQuerySet.update() (utils/models/base_model.py). + # Everything below is POST-ENQUEUE bookkeeping. The task is already on the + # queue, so a failure here must never be reported as a failed dispatch: + # `SubmitView` turns a DispatchError into a FAILED job, and the executor + # would then run, callback, and find a terminal row it cannot write to -- + # the caller told nothing was billed for work that did run. Wrapped rather + # than left to propagate, which is what the single pre-review UPDATE did. + # + # Losing the bookkeeping entirely is recoverable: sweep phase 1 reaps a + # still-PENDING row with no `dispatched_at`, and phase 2's + # `dispatched_at IS NULL` arm covers the non-PENDING case. + try: + _record_dispatch(job) + except Exception: + logger.exception( + "agent-kv: dispatch bookkeeping failed for job %s after enqueue " + "(task is queued; the sweep will reconcile)", + job.id, + ) + + +def _record_dispatch(job) -> None: + """Persist task_id/status/dispatched_at against whatever the row raced to.""" + advanced = AgentKVJob.objects.filter(id=job.id, status=JobStatus.PENDING).update( + task_id=job.task_id, + status=job.status, + dispatched_at=job.dispatched_at, + ) + if not advanced: + # The row moved off PENDING between the enqueue above and this write. + # The benign case is a terminal status (the guard's whole purpose) -- + # but there is a non-terminal one: StageReportView promotes + # PENDING -> RUNNING on the executor's FIRST stage report, which can + # easily land before this bookkeeping. The guard above then matches 0 + # rows and `dispatched_at` stays NULL -- and a non-terminal row with a + # NULL `dispatched_at` is invisible to BOTH sweep phases: phase 1 + # requires `status=PENDING`, phase 2 filters `dispatched_at__lt=cutoff` + # and SQL `NULL < x` is never true. The job reports `running` forever + # and `GET result` 409s for the life of the row, with nothing able to + # recover it. + # + # So stamp the dispatch bookkeeping for any still-non-terminal row, + # WITHOUT touching `status`: the row genuinely was dispatched, and + # moving RUNNING back to DISPATCHED would lose the executor's own + # progress. `dispatched_at__isnull=True` keeps this idempotent and + # stops a retry overwriting the original dispatch time. + AgentKVJob.objects.filter(id=job.id, dispatched_at__isnull=True).exclude( + status__in=list(AgentKVJob.TERMINAL) + ).update(task_id=job.task_id, dispatched_at=job.dispatched_at) diff --git a/backend/agent_kv/exceptions.py b/backend/agent_kv/exceptions.py new file mode 100644 index 0000000000..6c84a9dd19 --- /dev/null +++ b/backend/agent_kv/exceptions.py @@ -0,0 +1,33 @@ +from rest_framework.exceptions import APIException + + +class EngineUnavailable(APIException): + status_code = 501 + default_detail = "agent-kv engine not available on this deployment" + + +class RateLimited(APIException): + status_code = 429 + default_detail = "Too many requests" + + +class JobNotFound(APIException): + status_code = 404 + default_detail = "Job not found" + + +class SubscriptionGateUnavailable(APIException): + """The engine plugin is installed but exposes no subscription gate. + + Deliberately NOT a 402: the org's subscription was never evaluated, so + claiming it was denied would be a lie. This is a deployment fault — a build + that can run billable work but cannot check entitlement — and it fails + CLOSED, because the alternative is admitting unmetered paid work on a route + whose URL carries no org segment for the middleware to fall back on. + """ + + status_code = 503 + default_detail = ( + "agent-kv cannot verify subscription entitlement on this deployment; " + "the engine plugin exposes no subscription gate" + ) diff --git a/backend/agent_kv/execution_serializers.py b/backend/agent_kv/execution_serializers.py new file mode 100644 index 0000000000..01064c3c65 --- /dev/null +++ b/backend/agent_kv/execution_serializers.py @@ -0,0 +1,483 @@ +"""Submit-time validation: every §6.1 cap lives here, before any paid work.""" + +import json + +import pdfplumber +from django.conf import settings +from rest_framework import serializers + +from agent_kv.constants import EXTRACTOR_ROUTES, TABLE_EXTRACTOR_NAME, V1_EXTRACTOR_NAME +from unstract.agent_kv_schema import SchemaError, compile_schema +from unstract.sdk1.constants import AdapterTypes + +# Images are deliberately ABSENT, and that is a cross-repo contract, not an +# oversight. The cloud engine's `_build_agent_graph` (agentic_kv +# kv_extractor.py) treats only `.pdf/.xlsx/.xls` as a document; anything else +# takes a branch that skips `document_processor` entirely, and +# `ImageLoader.load_pages` -- the only thing that would populate pages for an +# image -- has no call site anywhere in the plugin (the engine's own comment +# there says images are "out of P2 scope"). +# +# Accepting them here anyway failed OPEN: an image dispatched normally with +# `pages_total=1`, every key came back not-found, and the job returned +# `success: true` with a page billed. Refusing at submit is the honest +# behaviour until the engine side is wired; re-add them in the same change that +# gives `load_pages` a call site, not before. +ALLOWED_EXTENSIONS = {".pdf", ".xlsx", ".xls"} +PDF_LIKE = {".pdf"} +EXTRACTION_MODES = ("whole-doc", "per-page") + + +# Derived, not repeated: the result and status documents key their payloads by +# these names, so a literal here that no route/stage list knows about would let +# the serializer accept a name the responses file under something else. +SUPPORTED_EXTRACTORS = tuple(EXTRACTOR_ROUTES) + + +class KVOptionsSerializer(serializers.Serializer): + """The `kv` extractor's own knobs (spec §7.1). + + These used to be top-level submit fields. They are extractor-scoped now + because they are meaningless to any other extractor -- `qa` and `challenge` + describe the KV agent pipeline, not "the request". + """ + + qa = serializers.BooleanField(required=False, default=True) + challenge = serializers.BooleanField(required=False, default=True) + extraction_mode = serializers.ChoiceField( + required=False, choices=EXTRACTION_MODES, default="whole-doc" + ) + structured_output = serializers.BooleanField(required=False, default=False) + calculations = serializers.CharField(required=False, allow_blank=True, default="") + document_class = serializers.CharField( + required=False, allow_blank=True, default="", max_length=256 + ) + key_notes = serializers.CharField( + required=False, allow_blank=True, default="", max_length=10_000 + ) + + def validate_calculations(self, v): + if v and not settings.AGENT_KV_CALCULATIONS_ENABLED: + raise serializers.ValidationError( + "calculations is not available on this deployment yet" + ) + if len(v.encode("utf-8")) > settings.AGENT_KV_MAX_CALCULATIONS_BYTES: + raise serializers.ValidationError( + f"calculations exceeds {settings.AGENT_KV_MAX_CALCULATIONS_BYTES} bytes" + ) + return v + + def validate_structured_output(self, v): + if v and not settings.AGENT_KV_STRUCTURED_OUTPUT_ENABLED: + raise serializers.ValidationError( + "structured_output is not available on this deployment yet" + ) + return v + + def validate(self, data): + # DRF silently DROPS unknown fields. For per-extractor options that is + # the wrong default: an option aimed at the wrong extractor (or a typo) + # would be discarded and the job would run with a silently different + # configuration than the caller asked for. Reject instead. + unknown = set(self.initial_data) - set(self.fields) + if unknown: + raise serializers.ValidationError( + f"unknown options for extractor 'kv': {sorted(unknown)}" + ) + return data + + +class TableOptionsSerializer(serializers.Serializer): + """The `table` extractor's own knobs (spec §7.1). + + Deliberately a subset of what the IDE path accepts: `output_path` and the + IDE callback hints (`prompt_key`, `doc_name`) are meaningless on the blind + API, and `enable_highlight` has no consumer there. + + **`enable_header_mapping` changes the result shape.** With it off, the + engine returns `output.tables` as a flat list of row dicts. With it on, the + engine wraps them as `{"header_mapping": ..., "rows": [...]}` + (`runner.py:1712`; the executor's own enrichment path unwraps the same + shape at `executor.py:577`). The API returns whichever the caller asked + for, unchanged -- so a client that sets this flag must read `.rows`. + """ + + instructions = serializers.CharField( + required=False, allow_blank=True, default="", max_length=10_000 + ) + json_structure = serializers.CharField( + required=False, allow_blank=True, default="", max_length=100_000 + ) + enable_header_mapping = serializers.BooleanField(required=False, default=False) + correct_number_separators = serializers.BooleanField(required=False, default=False) + number_format = serializers.ChoiceField( + required=False, choices=("US", "EU"), default="US" + ) + + def validate(self, data): + # Same reason KVOptionsSerializer rejects unknowns: DRF drops them + # silently, so an option aimed at the wrong extractor would be + # discarded and the job would run with a configuration the caller + # never asked for. + unknown = set(self.initial_data) - set(self.fields) + if unknown: + raise serializers.ValidationError( + f"unknown options for extractor '{TABLE_EXTRACTOR_NAME}': " + f"{sorted(unknown)}" + ) + return data + + +class TableKeysSerializer(serializers.Serializer): + """The `table` extractor's `keys`. + + The wire format gives every extractor a `keys` member; for the table + extractor the thing being asked for is a table, named by `target_table` + (the engine's one required extraction parameter). + """ + + target_table = serializers.CharField(max_length=256) + + def validate(self, data): + unknown = set(self.initial_data) - set(self.fields) + if unknown: + raise serializers.ValidationError( + f"unknown keys for extractor '{TABLE_EXTRACTOR_NAME}': {sorted(unknown)}" + ) + return data + + +# Each extractor's own validators. Keyed by the same names as EXTRACTOR_ROUTES; +# `test_every_supported_extractor_has_a_route_stage_list_and_options_serializer` +# keeps the three tables in step. +_OPTIONS_SERIALIZERS = { + V1_EXTRACTOR_NAME: KVOptionsSerializer, + TABLE_EXTRACTOR_NAME: TableOptionsSerializer, +} + + +class TableAdaptersSerializer(serializers.Serializer): + """The platform adapters the `table` extractor runs on. + + **Why the caller names adapters rather than the operator configuring env + vars.** The engine needs two LLMs and an OCR source. The IDE table path + resolves all three from platform adapter instances the user configured + (`agentic_table/executor.py` -> `LLM(adapter_instance_id=...)`, + `X2Text(adapter_instance_id=...)`), and this API now does the same. That + reverses spec D6, which specified system-level env configuration and "no + end-user model control" -- a deliberate reversal, recorded here and in + `docs/agent-kv-api.md`: the consumers are existing customers with + accounts, so they already own adapters, and letting them choose puts model + selection and LLM spend on the account that benefits from it. + + Field names follow the platform convention for naming adapters by ROLE, + not by id: `prompt_profile_manager_v2.ProfileManager` declares `llm`, + `x2text`, `embedding_model`, `vector_store`. Hence `llm` / `lite_llm` / + `x2text` rather than `llm_adapter_id` and friends. + + `lite_llm` is separate because the engine uses a cheaper model for + per-page presence detection and the advanced one for structure and + extraction; collapsing them would silently multiply the cost of the + highest-volume stage. + """ + + llm = serializers.UUIDField() + lite_llm = serializers.UUIDField() + x2text = serializers.UUIDField() + + def validate(self, data): + unknown = set(self.initial_data) - set(self.fields) + if unknown: + raise serializers.ValidationError( + f"unknown adapters for extractor '{TABLE_EXTRACTOR_NAME}': " + f"{sorted(unknown)}" + ) + return data + + +#: Each extractor's own `adapters` validator, and which `AdapterTypes` each +#: role must be. An extractor absent from this table takes no adapters -- `kv` +#: is env-configured (`extraction_seams.ExtractionConfig`) and deliberately +#: stays that way for now, so the two credential models coexist: one extractor +#: per model, never both for one extractor. +_ADAPTERS_SERIALIZERS = { + TABLE_EXTRACTOR_NAME: ( + TableAdaptersSerializer, + { + "llm": AdapterTypes.LLM, + "lite_llm": AdapterTypes.LLM, + "x2text": AdapterTypes.X2TEXT, + }, + ), +} + + +#: Each extractor's own `keys` validator. An extractor absent from this table +#: falls through to the schema compiler, which is `kv`'s contract: its `keys` IS +#: a compiled schema rather than a fixed set of fields. +_KEYS_SERIALIZERS = { + TABLE_EXTRACTOR_NAME: TableKeysSerializer, +} + + +def _validated_adapter_shape(name: str, supplied) -> dict: + """Validate an extractor's `adapters` block SHAPE only. No database. + + Returns `{role: ""}`, or `{}` for an extractor that takes none. + + Deliberately split from the tenancy check. This runs in the serializer, so + it must stay free of the ORM -- `test_submit_serializer.py` is a unit suite + with no database, and an adapter lookup here would make every submit test + require one. What it does check is everything that needs no database: + presence, that each id parses as a UUID, and that no unknown role was sent. + + **Ownership and type are checked in the VIEW** + (`execution_views._resolved_adapters`), because they need the Bearer key's + organization -- the same reason `_subscription_denial` lives there. Both + still run before anything is staged or billed. + """ + entry = _ADAPTERS_SERIALIZERS.get(name) + if entry is None: + if supplied: + raise serializers.ValidationError( + {"adapters": f"extractor '{name}' takes no adapters"} + ) + return {} + + adapters_cls, _expected_types = entry + ser = adapters_cls(data=supplied if isinstance(supplied, dict) else {}) + ser.is_valid(raise_exception=True) + return {role: str(value) for role, value in ser.validated_data.items()} + + +class ExtractorSerializer(serializers.Serializer): + """One entry of the submit's `extractors` array (spec §7.0/§7.1).""" + + name = serializers.CharField() + keys = serializers.JSONField() + #: Sibling of `keys`/`options`, not nested inside them. These are the + #: RESOURCES the extractor runs on, which `ProfileManager` likewise holds + #: as fields on the owning object rather than inside a knobs blob -- and + #: keeping them out of `options` leaves that member optional. + adapters = serializers.DictField(required=False, default=dict) + options = serializers.DictField(required=False, default=dict) + + def validate_name(self, v): + if v not in SUPPORTED_EXTRACTORS: + raise serializers.ValidationError( + f"unknown extractor '{v}'; supported: {list(SUPPORTED_EXTRACTORS)}" + ) + return v + + def _validated_keys(self, name: str, spec): + """Validate `keys` for the extractor named by the VALIDATED `name`. + + Not a `validate_keys` field validator, deliberately. That ran before + object-level validation and branched on raw ``self.initial_data["name"]`` + -- while ``validate_name`` saw the value DRF had already trimmed + (``CharField.trim_whitespace`` defaults True). So the extractor's + identity was decided twice, under two different values: + ``{"name": " table ", ...}`` passed the name check as ``table`` and then + took the **kv** branch here, never running ``TableKeysSerializer``. A + kv-shaped payload like ``{"target_table": {"description": "x"}}`` + compiles cleanly as a KV schema, so the submit returned **202** and + dispatched to ``agentic_table`` with a `target_table` that is a dict + rather than the string the binding requires -- the job staged, billed + and then failed at the executor. + + Branching on `data["name"]` in ``validate()`` means the name is trimmed + and already checked against ``SUPPORTED_EXTRACTORS`` exactly once. + """ + keys_cls = _KEYS_SERIALIZERS.get(name) + if keys_cls is not None: + keys = keys_cls(data=spec if isinstance(spec, dict) else {}) + keys.is_valid(raise_exception=True) + return keys.validated_data + + # kv (and any future extractor with no dedicated keys serializer): size + # is capped on the SERIALIZED form -- the cap exists to bound parse and + # compile cost, and `keys` arrives here already parsed out of the + # `extractors` JSON. ensure_ascii=False so this measures the SAME bytes + # the outer `extractors` cap measured. + serialized = json.dumps(spec, ensure_ascii=False).encode("utf-8") + if len(serialized) > settings.AGENT_KV_MAX_SCHEMA_BYTES: + raise serializers.ValidationError({"keys": "keys schema too large"}) + # Called for what it REFUSES, not for what it returns. The + # `CompiledSchema` used to be stashed on the serializer and collected + # into a `{name: compiled}` dict by `SubmitSerializer`, which no + # non-test code ever read: `dispatch_job` sends `schema=entry["keys"]`, + # the raw dict, and the engine re-compiles it on its own (see the + # docstring on `compile.py` -- the caps are a submit-time gate, not an + # invariant the engine re-checks). Plumbing the compiled form through + # instead is not an option the queue allows: it would have to survive + # JSON round-tripping to the executor, which is exactly why the engine + # recompiles. So the dead attribute is gone rather than left to read as + # plumbing that exists. What this call is for is the 400 below. + try: + compile_schema(spec) + except SchemaError as e: + raise serializers.ValidationError({"keys": str(e)}) + return spec + + def validate(self, data): + # Same reason the options block rejects unknowns, one level up: DRF drops + # unrecognised keys, so `"option"` or `"Options"` for `"options"` would + # be discarded whole, `options` would default to {}, and the job would + # run with qa=True/challenge=True -- roughly double the LLM spend the + # caller asked for, with a 202 and no indication anything was ignored. + unknown = set(self.initial_data) - set(self.fields) + if unknown: + raise serializers.ValidationError( + f"unknown keys on extractor entry: {sorted(unknown)}" + ) + name = data["name"] + data["keys"] = self._validated_keys(name, data.get("keys")) + opts_cls = _OPTIONS_SERIALIZERS[name] + opts = opts_cls(data=data.get("options") or {}) + opts.is_valid(raise_exception=True) + data["options"] = opts.validated_data + return data + + +class SubmitSerializer(serializers.Serializer): + """Submit-time validation: every §6.1 cap lives here, before any paid work. + + The wire format is extractor-scoped (§7.0): per-extractor schema and knobs + live inside `extractors`, and only fields describing the REQUEST stay top + level. There is deliberately no alias for the old flat shape. + """ + + file = serializers.FileField() + extractors = serializers.CharField() # JSON array string, or a file part + # Job-level: the page range drives the shared OCR pass and the §6.1 page + # cap, so it cannot differ between extractors reading the same document. + page_start = serializers.IntegerField(required=False, default=1, min_value=1) + page_end = serializers.IntegerField( + required=False, default=None, allow_null=True, min_value=1 + ) + timeout = serializers.IntegerField(required=False, default=0, min_value=0) + tags = serializers.ListField( + child=serializers.CharField(max_length=64), + required=False, + default=list, + max_length=20, + ) + custom_data = serializers.JSONField(required=False, default=None, allow_null=True) + webhook_url = serializers.URLField( + required=False, allow_blank=True, default="", max_length=1024 + ) + + #: Measured page count of the uploaded document. None for Excel, which has + #: no pre-OCR page concept. This is what metering and the status document + #: report, and it is NOT what the page cap is compared against. + pages_total = None + #: How many pages the request actually asks to process, after `page_start` / + #: `page_end` are applied. This is what the cap bounds. None for Excel. + pages_selected = None + + def validate_file(self, f): + name = (f.name or "").lower() + ext = "." + name.rsplit(".", 1)[-1] if "." in name else "" + if ext not in ALLOWED_EXTENSIONS: + raise serializers.ValidationError( + f"Unsupported file type '{ext}'; allowed: {sorted(ALLOWED_EXTENSIONS)}" + ) + max_bytes = settings.AGENT_KV_MAX_FILE_SIZE_MB * 1024 * 1024 + if f.size > max_bytes: + raise serializers.ValidationError( + f"File exceeds {settings.AGENT_KV_MAX_FILE_SIZE_MB}MB limit" + ) + return f + + def validate_timeout(self, v): + if v > settings.AGENT_KV_MAX_TIMEOUT_SECONDS: + raise serializers.ValidationError( + f"timeout must be 0..{settings.AGENT_KV_MAX_TIMEOUT_SECONDS}" + ) + return v + + def validate_extractors(self, raw): + # §7.1: `extractors` may arrive as an inline JSON string OR a file part; + # SubmitView reads a file-typed part into a string before constructing + # the serializer, exactly as it did for the old `keys` field. + if len(raw.encode("utf-8")) > settings.AGENT_KV_MAX_SCHEMA_BYTES: + raise serializers.ValidationError("extractors payload too large") + try: + entries = json.loads(raw) + except (ValueError, TypeError) as e: + raise serializers.ValidationError(f"extractors is not valid JSON: {e}") + if not isinstance(entries, list) or not entries: + raise serializers.ValidationError("extractors must be a non-empty JSON array") + if len(entries) > 1: + # The FORMAT is being fixed before launch; the fan-out execution is + # not built (one executor exists, and page images are not shared). + # Refusing loudly beats accepting a request we would silently run + # single-extractor. + raise serializers.ValidationError( + "multiple extractors are not supported yet; pass exactly one" + ) + + validated = [] + for i, entry in enumerate(entries): + if not isinstance(entry, dict): + raise serializers.ValidationError(f"extractors[{i}] must be an object") + ser = ExtractorSerializer(data=entry) + if not ser.is_valid(): + raise serializers.ValidationError({f"extractors[{i}]": ser.errors}) + data = ser.validated_data + try: + data["adapters"] = _validated_adapter_shape( + data["name"], data.get("adapters") + ) + except serializers.ValidationError as e: + raise serializers.ValidationError( + {f"extractors[{i}]": e.detail} + ) from None + validated.append(data) + return validated + + def validate(self, data): + start, end = data.get("page_start", 1), data.get("page_end") + if end is not None and end < start: + raise serializers.ValidationError( + {"page_end": "page_end must be >= page_start"} + ) + f = data["file"] + ext = "." + f.name.lower().rsplit(".", 1)[-1] + if ext in PDF_LIKE: + try: + with pdfplumber.open(f) as pdf: + self.pages_total = len(pdf.pages) + except Exception: + raise serializers.ValidationError({"file": "Unreadable PDF"}) + finally: + f.seek(0) + if start > self.pages_total: + raise serializers.ValidationError( + { + "page_start": f"page_start {start} is past the end of a " + f"{self.pages_total}-page document" + } + ) + # The cap bounds the work the job will DO, not the size of the file + # it was handed. A caller asking for pages 1-5 of a 400-page PDF is + # requesting five pages of OCR and extraction; refusing that against + # a 100-page cap rejected a request that was inside the documented + # limit. `pages_total` stays the measured document count -- it is + # what metering and the status document report -- and only the cap + # comparison moves to the selected range. + last_page = self.pages_total if end is None else min(end, self.pages_total) + self.pages_selected = last_page - start + 1 + if self.pages_selected > settings.AGENT_KV_MAX_PAGES: + raise serializers.ValidationError( + { + "file": f"Requested {self.pages_selected} pages " + f"(page_start={start}, page_end=" + f"{end if end is not None else self.pages_total}) of a " + f"{self.pages_total}-page document; max is " + f"{settings.AGENT_KV_MAX_PAGES} (§6.1)" + } + ) + # Excel: no page concept pre-OCR (spec §6.1); pages_total stays None, + # size cap already enforced; the engine enforces the post-OCR cap. + return data diff --git a/backend/agent_kv/execution_urls.py b/backend/agent_kv/execution_urls.py new file mode 100644 index 0000000000..eff4807722 --- /dev/null +++ b/backend/agent_kv/execution_urls.py @@ -0,0 +1,29 @@ +"""Public Agent-KV execution routes. + +``ValidateView`` is **deliberately not routed here.** It compiles a `kv` keys +schema through ``unstract.agent_kv_schema.compile_schema`` and does nothing +else; the `table` extractor's ``keys`` is ``{"target_table": ...}``, validated +by ``TableKeysSerializer`` at submit. Publishing an endpoint that validates +schemas for an extractor this deployment refuses with a 400 would be an +incoherent public contract. + +The view class and the schema package both stay in the tree. Deleting them +would force a content merge in the files the branch carrying the KV engine +rewrites most; leaving them makes restoring this endpoint one line. +""" + +from django.urls import path + +from agent_kv.execution_views import ( + JobCancelView, + JobResultView, + JobStatusView, + SubmitView, +) + +urlpatterns = [ + path("", SubmitView.as_view(), name="agent_kv_submit"), + path("", JobStatusView.as_view(), name="agent_kv_status"), + path("/result", JobResultView.as_view(), name="agent_kv_result"), + path("/cancel", JobCancelView.as_view(), name="agent_kv_cancel"), +] diff --git a/backend/agent_kv/execution_views.py b/backend/agent_kv/execution_views.py new file mode 100644 index 0000000000..93d1699167 --- /dev/null +++ b/backend/agent_kv/execution_views.py @@ -0,0 +1,785 @@ +import logging +import time +from datetime import timedelta + +from django.conf import settings +from django.utils import timezone +from plugins import get_plugin +from rest_framework.exceptions import ValidationError +from rest_framework.response import Response +from rest_framework.views import APIView + +from agent_kv.constants import STAGE_NAMES_BY_EXTRACTOR +from agent_kv.dispatch import DispatchError, dispatch_cancelled_webhook, dispatch_job +from agent_kv.exceptions import ( + EngineUnavailable, + JobNotFound, + RateLimited, + SubscriptionGateUnavailable, +) +from agent_kv.execution_serializers import _ADAPTERS_SERIALIZERS, SubmitSerializer +from agent_kv.execution_views_result import result_payload +from agent_kv.key_validator import AgentKVKeyValidator +from agent_kv.models import AgentKVJob, JobStatus +from agent_kv.rate_limiter import AgentKVConcurrencyLimiter, check_key_rate +from agent_kv.storage import delete_job_files, stage_input +from unstract.agent_kv_schema.compile import SchemaError, compile_schema + +logger = logging.getLogger(__name__) + + +def _get_job(agent_kv_key, job_id): + """Org-scoped lookup used by every job-scoped endpoint (spec §5.4). + + Unknown job_id and a job that belongs to a different org must be + indistinguishable to the caller, so both funnel through the same + ``DoesNotExist`` -> ``JobNotFound`` (404) path. + """ + try: + return AgentKVJob.objects.get( + id=job_id, organization_id=agent_kv_key.organization_id + ) + except AgentKVJob.DoesNotExist: + raise JobNotFound() + + +def _status_document(job) -> dict: + """Build the status document per spec §7.2.""" + stages_json = job.stages or {} + # `.get`, not a subscript. No row can hold an out-of-dict value today (the + # sole creation site is serializer-validated and the migration backfills + # "kv"), but a RETIRED extractor name with surviving rows would 500 every + # `GET status` for those jobs while `GET result` kept working -- the result + # payload keys by `job.extractor` without consulting this table at all. An + # empty stage list degrades to "no stages reported", which is honest for an + # extractor this build no longer knows how to describe. + stage_names = STAGE_NAMES_BY_EXTRACTOR.get(job.extractor) + if stage_names is None: + logger.warning( + "agent-kv job %s ran extractor %r, which has no stage list in " + "STAGE_NAMES_BY_EXTRACTOR; reporting an empty stage array", + job.id, + job.extractor, + ) + stage_names = [] + doc = { + "job_id": str(job.id), + # The JOB's state, not an extractor's: a job is not complete until every + # extractor is, so this stays top level (spec §7.2). + "status": job.status.lower(), + # Stage names are extractor-specific (`qa`/`challenge`/`codegen` mean + # nothing to the table extractor) and they ARE returned to clients, so + # they are wire format and namespaced with everything else. + "extractors": { + job.extractor: { + "stage": job.stage, + "stages": [ + {"name": name, **stages_json[name]} + for name in stage_names + if name in stages_json + ], + } + }, + # Which adapters this run spent, by role. Reported because the caller + # chose them and is paying for them: a run naming the wrong `llm` id + # produces a correct-looking result at a different price, and without + # this the only record is in `usage_v2`, which the caller cannot read. + # `{}` for an env-configured extractor. + "adapters": job.adapters or {}, + "created_at": job.created_at.isoformat() if job.created_at else None, + "started_at": job.dispatched_at.isoformat() if job.dispatched_at else None, + "completed_at": job.completed_at.isoformat() if job.completed_at else None, + "pages_total": job.pages_total, + } + if job.status == JobStatus.FAILED: + doc["error"] = job.error + return doc + + +def _never_dispatched(job) -> bool: + """True when this job never reached an executor, so no callback will come. + + **Re-read from the database, not from the in-memory row.** The row the + cancel path holds was fetched BEFORE `mark_terminal`, and a dispatch can + complete inside that gap: + + cancel reads the row -> PENDING, dispatched_at=NULL + dispatch_job -> passes its own terminal re-read, enqueues, + writes dispatched_at and DISPATCHED + cancel's mark_terminal-> wins (DISPATCHED is not terminal) + _never_dispatched -> True on the STALE row + + which releases the slot while the executor is running and billing. That is + a smaller version of the original over-release bug this narrowing was + written to fix, reintroduced by trusting the pre-cancel snapshot. + + Only `dispatched_at` is consulted, and `status` deliberately is not: by the + time this runs the cancel has already rewritten `status` to CANCELLED, so + the old `status == PENDING` test could only ever be answered from the stale + object. `dispatched_at` is the durable marker of "an executor may exist" + and the cancel never touches it. + + **Residual window, stated rather than papered over.** `dispatch_job` + enqueues and only then writes `dispatched_at` (`_record_dispatch`). A cancel + landing between those two still sees NULL and still releases early. The + consequence is bounded and self-correcting: one extra concurrent slot for + the remainder of that one run, after which the executor's finalize releases + again (`release` is an idempotent zrem). Closing it completely means + stamping a dispatch-intent marker before the enqueue, which changes what + `dispatched_at` means to both sweep phases -- a larger change than the + defect warrants. + """ + return not AgentKVJob.objects.filter(id=job.id, dispatched_at__isnull=False).exists() + + +def _fail_job_response(job, org_id: str, message: str, *, job_saved: bool) -> Response: + """Shared cleanup for any post-acquire submit failure (spec §5.3/§5.4). + + Always releases the concurrency slot acquired earlier in the request. + Only calls ``mark_terminal`` when a job row may actually exist — calling + it against a job_id with no row is harmless (the guarded UPDATE just + matches zero rows), but ``job_saved`` keeps the write guard's intent + (only terminalize rows that exist) obvious at the call site. + """ + if job_saved: + AgentKVJob.mark_terminal( + job.id, job.organization_id, JobStatus.FAILED, error=message + ) + AgentKVConcurrencyLimiter.release(org_id, str(job.id)) + return Response( + {"job_id": str(job.id), "status": JobStatus.FAILED.lower(), "error": message}, + status=500, + ) + + +def _subscription_denial(plugin, agent_kv_key, request) -> Response | None: + """Subscription admission (§6.6), or ``None`` when the org may proceed. + + A submit dispatches paid work, so it is gated exactly as an API deployment + execute is -- same policy, same 402 bodies -- via the cloud plugin's gate, + which calls the very ``SubscriptionHelper`` that cloud's + ``SubscriptionMiddleware`` calls. + + Why here and not in that middleware: it resolves the org from the URL + (``/deployment/api/{org_name}/...``). Agent-KV's URL carries no org segment + -- the org is inside the Bearer key -- so the middleware's + ``get_organization_id`` returns None for these requests, finds no + subscription row, and admits every one of them. The view is the first point + where the org is actually known. + + ``organization.organization_id`` is the org SLUG -- the CharField + ``Subscription.organization_id`` is keyed on, and what the deployment URL + supplies as ``org_name``. NOT ``agent_kv_key.organization_id``, which is the + Organization FK primary key: that matches no row, and the shared policy + reads "no row" as "nothing to enforce", so the gate would silently admit + everything while looking correctly wired. + + **Fails closed when the plugin exposes no gate.** This used to return None + -- admit -- to tolerate a cloud build predating the gate. But the admitted + request dispatches billable LLM and OCR work, and this route's URL carries + no org segment, so ``SubscriptionMiddleware`` cannot catch it downstream + either: a mixed deploy (this backend against a pre-gate cloud image) would + run unmetered paid work with nothing anywhere enforcing entitlement. + + Refusing instead is the conservative read of a deployment that can spend + money but cannot check whether it may. It surfaces as a 503 naming the + cause, not a 402 -- the subscription was never evaluated, and reporting it + as denied would send an operator to the billing system for what is an + image-pairing problem. + + Raises: + SubscriptionGateUnavailable: the engine plugin is installed but exposes + no ``service_class``. + """ + gate_factory = plugin.get("service_class") + if not gate_factory: + logger.error( + "agent-kv: engine plugin exposes no subscription gate; refusing the " + "submit rather than dispatching billable work unmetered" + ) + raise SubscriptionGateUnavailable() + return gate_factory().check(agent_kv_key.organization.organization_id, request) + + +def _dispatch_or_fail(job, org_id: str, entry: dict, options: dict) -> Response | None: + """Dispatch the job, or return the failure response. ``None`` means sent. + + Both failure paths terminalize the job and release the concurrency slot. + The bare ``except`` is belt-and-braces: ``dispatch_job`` wraps its own + internal failures as ``DispatchError``, but nothing here may rely on that + alone -- any other exception must still terminalize and release rather than + escape as an unhandled 500 with the slot still held. + """ + try: + dispatch_job( + job, + extractor=entry["name"], + schema=entry["keys"], + options=options, + adapters=entry.get("adapters") or {}, + ) + except DispatchError: + logger.exception("agent-kv dispatch failed for job %s", job.id) + except Exception: + logger.exception("agent-kv dispatch raised unexpectedly for job %s", job.id) + else: + return None + return _fail_job_response( + job, org_id, "Job could not be dispatched; nothing was billed.", job_saved=True + ) + + +def _sync_wait_response(job, wait: float) -> Response | None: + """Poll until the job is terminal or ``wait`` elapses (§7.1 sync mode). + + Returns the full result payload if it terminalized in time, else ``None`` + so the caller falls through to the normal 202 handshake. + """ + if not wait: + return None + deadline = time.monotonic() + wait + while time.monotonic() < deadline: + job.refresh_from_db() + if job.status in AgentKVJob.TERMINAL: + from agent_kv.execution_views_result import result_payload + + return Response(result_payload(job), status=200) + time.sleep(1) + return None + + +def _request_data_with_extractors_inlined(request): + """A mutable copy of the request data with a file-part `extractors` read in. + + Split out of ``SubmitView.post`` for sonar S3776 (the method was at + cognitive complexity 20 against a limit of 15); the behaviour is unchanged. + + Bounded read, which is the point. The size cap lives in + ``validate_extractors``, i.e. AFTER this point -- and + ``DATA_UPLOAD_MAX_MEMORY_SIZE`` excludes file-typed parts, so a 500 MB + `extractors` part was materialised in full (plus up to 4x that again for + the ``str``) before being rejected at 256 KiB. One such request per worker + process OOMs the pod, which makes it a cheap denial of service. + """ + data = request.data.copy() + part = data.get("extractors") + if not hasattr(part, "read"): # not a file part (§7.1); nothing to inline + return data + # Read one byte past the cap: enough to know it is over without ever + # holding more than the cap plus one. + limit = settings.AGENT_KV_MAX_SCHEMA_BYTES + raw = part.read(limit + 1) + if len(raw) > limit: + raise ValidationError({"extractors": f"extractors payload exceeds {limit} bytes"}) + try: + # `errors="strict"`, not "replace". A latin-1 key name used to decode + # to U+FFFD and then compile cleanly, so a malformed payload became a + # job that ran against a schema the caller never wrote. + data["extractors"] = raw.decode("utf-8") + except UnicodeDecodeError: + raise ValidationError({"extractors": "extractors must be valid UTF-8"}) from None + return data + + +def _discard_orphaned_input(job) -> None: + """Delete the staged object for a submit that failed before the row landed. + + Split out of ``SubmitView.post`` for sonar S3776; behaviour unchanged. + + If ``stage_input`` succeeded and ``job.save()`` then raised, the upload + exists with no row to carry its ref -- and ``run_ttl_cleanup`` selects + candidates from ``AgentKVJob`` rows, so an object with no row is + structurally unreachable by every cleanup path there is. It would sit in + the bucket forever, holding customer data nobody can find or delete. + """ + if not job.input_ref: + return + # `delete_job_files` reports which refs are confirmed gone and swallows the + # rest -- normally the ref survives as the retry handle. Here there is no + # row to hold it, so a delete that failed means the object is orphaned with + # nothing anywhere pointing at it. Log the ref itself: it is the only way + # anyone can clean it up by hand. + try: + cleared = delete_job_files(job) + except Exception: + cleared = [] + logger.exception( + "agent-kv: staged-input cleanup raised for failed submit %s (ref=%s)", + job.id, + job.input_ref, + ) + if "input_ref" not in cleared: + logger.error( + "agent-kv: the staged input for failed submit %s could NOT be " + "removed and no job row exists to carry its ref -- object %s is " + "orphaned and unreachable by TTL cleanup; remove it manually", + job.id, + job.input_ref, + ) + + +def _lookup_adapter(adapter_id, organization_id): + """One adapter belonging to `organization_id`, or ``None``. + + **`_base_manager`, NOT `.objects`.** `AdapterInstance.objects` is an + `AdapterInstanceModelManager`, which inherits + `DefaultOrganizationManagerMixin.get_queryset()` -- and that filters every + query by `UserContext.get_organization()`, a REQUEST-LOCAL thread-local the + tenant middleware sets. + + This route is deliberately whitelisted past that middleware (the org comes + from the Bearer key, not the URL), so the thread-local is unset here and the + ambient filter becomes `organization=None`, matching nothing. Through + `.objects` this gate refused EVERY submit -- including valid ones -- with + "no such adapter in this organization". + + That is NOT the same situation as `_get_job`, and reasoning by analogy to it + is how the defect got written: `AgentKVJob` uses `BaseModelManager`, which + does not auto-filter, so its explicit `organization_id=` is the only filter. + `AdapterInstance` sets a different default manager. + + `_base_manager` is Django's unfiltered manager and exists for exactly this + -- internal lookups that must not inherit a custom default manager's + filtering. It therefore makes the explicit `organization_id` below the ONLY + org filter, which is why that argument is required rather than optional. + + Not fixed by setting the thread-local instead: that would silently re-scope + every other org-managed query in the request, and `UserContext` keys on the + org SLUG while the key carries the FK pk -- the exact confusion + `dispatch._platform_api_key` documents shipping once already. + + A separate function so the view tests have an honest seam to patch: + `_base_manager` is a read-only property and cannot be patched, and mocking + `.objects` is what hid this defect in the first place. The real manager is + exercised by `tests/test_adapter_scoping.py`. + """ + from adapter_processor_v2.models import AdapterInstance + + return AdapterInstance._base_manager.filter( + id=adapter_id, organization_id=organization_id + ).first() + + +def _resolved_adapters(entry: dict, agent_kv_key) -> dict: + """Confirm every adapter the caller named is THIS org's, and of the right type. + + Returns `{role: ""}` unchanged on success -- the ids are already + UUID-shaped and role-checked by the serializer + (`_validated_adapter_shape`); what is added here is the half that needs a + database and the Bearer key. + + **Defence in depth, not the only tenancy check** -- an earlier version of + this docstring claimed it was, and was wrong. The platform service re-scopes + every adapter lookup as `WHERE id=%s and organization_id=%s` + (platform-service/.../helper/adapter_instance.py:28-30), against the org of + the bearer platform key that `dispatch._platform_api_key(job)` mints from + THIS job's org. Org A therefore cannot spend org B's credential even if this + gate were deleted outright. + + What this gate does buy, and why it is still worth having: + + - A clean 400 naming the offending role at submit, instead of a mid-run + `SdkError` on a job that already took a concurrency slot and billed for + staging. + - The two refusals the platform service does NOT make: `is_usable` (a + frictionless org's operator-funded trial is exhausted -- it hands back the + credentials anyway) and `is_available` (deprecated, removed from the SDK + registry). + + It lives in the view rather than the serializer for the same reason + `_subscription_denial` does: it needs the key. The platform's own helpers + do not fit -- `AdapterProcessor.get_adapters_by_type` filters + `.for_user(user)` and this path has no user (a key resolves to an + ORGANIZATION), and `get_adapter_by_name_and_type` is not scoped at all. + + TYPE is checked too: an `X2TEXT` id accepted into the `llm` slot resolves + fine and then fails deep in the engine as a provider error, which reads + like a broken model rather than two swapped UUIDs. + + What this deliberately does NOT honour is per-user adapter visibility + (`AdapterInstance` is `HasMembersMixin`, with `ResourceMembership` VIEWER + rows). An org-scoped key has no user to evaluate it against, so a key may + use ANY adapter in its own organization. Deliberate: a key is an + organization-level credential, and whoever can mint one can already read + the org's data. + + Raises: + ValidationError: 400 at submit -- before staging, before a slot is + taken, before anything is billed. + """ + requested = entry.get("adapters") or {} + if not requested: + return {} + + _, expected_types = _ADAPTERS_SERIALIZERS[entry["name"]] + for role, adapter_id in requested.items(): + expected = expected_types[role] + adapter = _lookup_adapter(adapter_id, agent_kv_key.organization_id) + if adapter is None: + # One message for "no such adapter" and "not yours", on purpose: + # telling them apart reveals which UUIDs exist. + raise ValidationError( + {"adapters": f"{role}: no such adapter in this organization"} + ) + if adapter.adapter_type != expected.value: + raise ValidationError( + { + "adapters": ( + f"{role}: adapter is of type '{adapter.adapter_type}', " + f"expected '{expected.value}'" + ) + } + ) + # Trial exhaustion. `is_usable` is flipped False when a frictionlessly + # onboarded org burns through its operator-funded sample allowance, and + # the platform service does NOT check it -- it returns the credentials + # regardless. So without this the IDE, workflows and Prompt Studio all + # refuse (tool_instance_helper.py, prompt_studio_helper.py) while an + # Agent-KV submit naming the SAME adapter id gets a 202 and runs on + # credentials the operator is paying for and has already cut off. + # + # Same message those surfaces use, so an exhausted trial reads + # identically whichever door the caller knocks on. + if not adapter.is_usable: + logger.error( + "agent-kv: free usage for sample adapter %s is exhausted", adapter.id + ) + raise ValidationError( + { + "adapters": ( + f"{role}: free usage for the configured trial adapter is " + "exhausted. Connect your own service account to continue: " + "https://docs.unstract.com/unstract_platform/setup_accounts/" + "whats_needed" + ) + } + ) + # Deprecated: removed from the SDK registry, so resolving it raises + # `InValidAdapterId` (adapter_processor.py:120) deep inside the engine + # and surfaces to the caller as a mid-run extraction failure on a job + # that was already billed for staging. A 400 at submit names the cause. + if not adapter.is_available: + raise ValidationError( + { + "adapters": ( + f"{role}: adapter is deprecated and no longer available; " + "pick a current one" + ) + } + ) + return requested + + +class SubmitView(APIView): + authentication_classes: list = [] + permission_classes: list = [] + + @AgentKVKeyValidator.validate_api_key + def post(self, request, *args, agent_kv_key=None, **kwargs): + plugin = get_plugin("agent_kv") + if not plugin: + raise EngineUnavailable() + if not check_key_rate(str(agent_kv_key.id)): + raise RateLimited() + + denied = _subscription_denial(plugin, agent_kv_key, request) + if denied is not None: + return denied + + data = _request_data_with_extractors_inlined(request) + serializer = SubmitSerializer(data=data) + serializer.is_valid(raise_exception=True) + v = serializer.validated_data + org_id = str(agent_kv_key.organization_id) + + # Tenancy + type + usability on the named adapters, HERE and not later: + # before the job row, before the concurrency slot, before the staging + # write. A caller who names another org's adapter, swaps two UUIDs, or + # points at an exhausted trial gets a 400 and is charged nothing -- the + # same "every cap before paid work" rule the serializer's caps follow. + # + # Assigned back rather than discarded: `_dispatch_or_fail` reads + # `entry["adapters"]` to build `executor_params`, so anything this gate + # ever normalizes (it returns the dict unchanged today) must land where + # dispatch will see it. Discarding the return made the two paths agree + # only by the accident of being the same dict object. + v["extractors"][0]["adapters"] = _resolved_adapters( + v["extractors"][0], agent_kv_key + ) + + job = AgentKVJob( + api_key=agent_kv_key, + organization_id=agent_kv_key.organization_id, + extractor=v["extractors"][0]["name"], + # The gate above validated these against this org, their expected + # types, and usability -- so what lands here is exactly what the + # run will spend. Ids only, never `adapter_metadata`: this column + # is returned to the caller. + adapters=v["extractors"][0].get("adapters") or {}, + pages_total=serializer.pages_total, + tags=v["tags"], + custom_data=v["custom_data"], + webhook_url=v["webhook_url"], + expires_at=timezone.now() + timedelta(days=settings.AGENT_KV_RESULT_TTL_DAYS), + ) + if not AgentKVConcurrencyLimiter.check_and_acquire(org_id, str(job.id)): + raise RateLimited("Concurrent job limit reached") + + job_saved = False + try: + job.input_ref = stage_input(org_id, str(job.id), v["file"]) + job.save() + job_saved = True + except Exception: + logger.exception("agent-kv staging/save failed for job %s", job.id) + _discard_orphaned_input(job) + return _fail_job_response( + job, + org_id, + "Job could not be accepted; nothing was billed.", + job_saved=job_saved, + ) + + # Unpack the single extractor entry into the FROZEN OSS<->cloud + # executor_params contract (`schema` + `options`). The wire format + # changed at the API edge only: the engine still receives exactly the + # option names it always did, so no cloud change is required. + entry = v["extractors"][0] + options = dict(entry["options"]) + # Job-level, but the engine reads them from `options` (§7.1: the page + # range drives the shared OCR pass, so it cannot be per-extractor). + options["page_start"] = v["page_start"] + options["page_end"] = v["page_end"] + + failed = _dispatch_or_fail(job, org_id, entry, options) + if failed is not None: + return failed + + inline = _sync_wait_response(job, v["timeout"]) + if inline is not None: + return inline + + return Response( + { + "job_id": str(job.id), + # Lowercased for cross-endpoint consistency (spec §7.2) -- the + # 202 body used to leak the raw uppercase status. + "status": job.status.lower(), + "status_url": f"/{settings.AGENT_KV_PATH_PREFIX}/{job.id}", + "created_at": job.created_at.isoformat(), + }, + status=202, + ) + + +class ValidateView(APIView): + """Dry-run schema compile. **NOT ROUTED ON THIS DEPLOYMENT.** + + ``execution_urls.py`` deliberately omits this view, and that file carries + the full rationale. The short version: it compiles a `kv` keys schema and + does nothing else, and `kv` is the one extractor this build refuses + (``EXTRACTOR_ROUTES`` carries `table` only), so publishing it would + advertise validation for an extractor every submit 400s. The `table` + extractor's ``keys`` is ``{"target_table": ...}``, validated by + ``TableKeysSerializer`` at submit instead. + + Kept in the tree, not deleted: restoring the endpoint is one line in + ``execution_urls.py``, and deleting it would force a content merge in the + file the branch carrying the KV engine rewrites most. + + So read the code below as dormant. Nothing reaches it but the unit tests + that keep it honest -- in particular, it is NOT a live surface for the + rate limiter or the key validator it still decorates, and a reader looking + for the deployment's public endpoints should look at ``execution_urls.py``, + which is the only place that decides. + """ + + authentication_classes: list = [] + permission_classes: list = [] + + @AgentKVKeyValidator.validate_api_key + def post(self, request, *args, agent_kv_key=None, **kwargs): + if not check_key_rate(str(agent_kv_key.id)): + raise RateLimited() + spec = request.data.get("keys") + if spec is None: + return Response({"detail": "body must include 'keys'"}, status=400) + try: + compiled = compile_schema(spec) + except SchemaError as e: + return Response({"valid": False, "error": str(e)}, status=200) + return Response( + { + "valid": True, + "leaves": len(compiled.key_specs), + "arrays": len(compiled.array_specs), + "constraints": len(compiled.constraints), + }, + status=200, + ) + + +class JobStatusView(APIView): + """GET status document; DELETE purges the job's staged/result files. + + DELETE is merged onto this class (rather than a standalone + ``JobDeleteView``) because both share the ```` URL — Django + matches a URL pattern once per request regardless of HTTP method, so two + separate ``path()`` entries for the same literal path can't coexist. + ``JobDeleteView`` below is kept as a name alias for this same class. + """ + + authentication_classes: list = [] + permission_classes: list = [] + + @AgentKVKeyValidator.validate_api_key + def get(self, request, *args, job_id=None, agent_kv_key=None, **kwargs): + job = _get_job(agent_kv_key, job_id) + return Response(_status_document(job), status=200) + + @AgentKVKeyValidator.validate_api_key + def delete(self, request, *args, job_id=None, agent_kv_key=None, **kwargs): + job = _get_job(agent_kv_key, job_id) + if job.status not in AgentKVJob.TERMINAL: + # Cancel BEFORE deleting files: a still-running job would + # otherwise keep running after its files are gone, and its + # eventual finalize call would write a fresh result_ref onto a + # job the caller already asked to delete -- resurrecting a + # result they explicitly discarded. Terminalizing first closes + # that window; a finalize call that still lands late loses the + # terminal-state guard and, on the success path, cleans up its + # own now-orphaned write (FinalizeView, storage.delete_result_file). + won = AgentKVJob.mark_terminal( + job.id, job.organization_id, JobStatus.CANCELLED + ) + if won: + # Release the concurrency slot ONLY for a job that was never + # dispatched. + # + # The slot is taken at submit and released by + # `_fail_job_response`, the finalize callback and the sweep. A + # job cancelled BEFORE dispatch gets no finalize callback, and + # the sweep's phase-1 only targets PENDING, never CANCELLED -- + # so without this release its slot would leak until the 6h TTL. + # That is what this release is for, and all it is for. + # + # It used to fire for ANY non-terminal job, including one + # mid-run. Nothing revokes a running executor -- `job.task_id` + # is written at dispatch and never read again -- so the engine + # kept running, kept calling LLMs and kept billing while its + # slot was handed to the next submit. Submit-then-cancel in a + # loop therefore ran arbitrarily many concurrent extractions + # against a ceiling of `AGENT_KV_CONCURRENT_LIMIT`, all paid for. + # + # Narrowing loses nothing: a mid-run cancel's slot is released + # by `FinalizeView`'s `finally` when the executor's callback + # lands, and `release()` is idempotent (zrem). + # + # Guarded on `won` so a lost race (a concurrent cancel or + # finalize terminalized it first) does not release a slot that + # the winner is still accounting for. + # + # `_never_dispatched` re-reads `dispatched_at` from the DB: the + # `job` here predates `mark_terminal`, and a dispatch can + # complete in that gap (see that function). + if _never_dispatched(job): + AgentKVConcurrencyLimiter.release( + str(job.organization_id), str(job.id) + ) + # This request owns the terminal notification: no fresh finalize + # can follow a cancel that won, so nothing else will send it. + dispatch_cancelled_webhook(job) + else: + # Lost the race: a finalize or cancel terminalized this job + # between our read above and the guarded UPDATE. `job` is now + # STALE, and `result_ref` is the field that matters -- a winning + # finalize has just written one. + # + # Without this refresh the cleanup below runs against the stale + # copy, where `result_ref` is still "". `delete_job_files` + # reports an already-empty ref as "cleared" (nothing to delete), + # so the save then writes "" OVER the winner's real ref. The job + # stays COMPLETED, its result 404s (`JobResultView` treats + # COMPLETED-without-a-ref as swept), and the object is orphaned + # in the bucket with nothing left pointing at it -- TTL cleanup + # selects on `result_ref > ""`, so a blanked row never comes + # back. Deleting is still the caller's intent; it just has to + # act on the refs that actually exist now. + job.refresh_from_db() + # Blank only the refs whose files are confirmed gone, so a ref whose + # delete failed survives as the handle TTL cleanup retries from. 204 + # either way: the job IS terminal and the caller's intent is recorded, + # and a sync 5xx here would only invite a retry of a DELETE that already + # did everything it could. + cleared = delete_job_files(job) + if cleared: + for field in cleared: + setattr(job, field, "") + job.save(update_fields=cleared) + return Response(status=204) + + +# Alias kept for a descriptive import name; there is no separate URL route +# (see the JobStatusView docstring above) — DELETE rides JobStatusView's URL. +JobDeleteView = JobStatusView + + +class JobResultView(APIView): + authentication_classes: list = [] + permission_classes: list = [] + + @AgentKVKeyValidator.validate_api_key + def get(self, request, *args, job_id=None, agent_kv_key=None, **kwargs): + job = _get_job(agent_kv_key, job_id) + if job.status not in AgentKVJob.TERMINAL: + return Response({"status": job.status.lower()}, status=409) + # A job's row outlives its result by design (audit trail after TTL + # cleanup blanks the refs) -- expired or a COMPLETED job whose + # result was already swept both mean "nothing left to serve". + if (job.expires_at and job.expires_at < timezone.now()) or ( + job.status == JobStatus.COMPLETED and not job.result_ref + ): + raise JobNotFound() + return Response(result_payload(job), status=200) + + +class JobCancelView(APIView): + authentication_classes: list = [] + permission_classes: list = [] + + @AgentKVKeyValidator.validate_api_key + def post(self, request, *args, job_id=None, agent_kv_key=None, **kwargs): + job = _get_job(agent_kv_key, job_id) + won = AgentKVJob.mark_terminal(job.id, job.organization_id, JobStatus.CANCELLED) + if won: + # Release the concurrency slot ONLY for a job that was never + # dispatched. + # + # The slot is taken at submit and released by + # `_fail_job_response`, the finalize callback and the sweep. A + # job cancelled BEFORE dispatch gets no finalize callback, and + # the sweep's phase-1 only targets PENDING, never CANCELLED -- + # so without this release its slot would leak until the 6h TTL. + # That is what this release is for, and all it is for. + # + # It used to fire for ANY non-terminal job, including one + # mid-run. Nothing revokes a running executor -- `job.task_id` + # is written at dispatch and never read again -- so the engine + # kept running, kept calling LLMs and kept billing while its + # slot was handed to the next submit. Submit-then-cancel in a + # loop therefore ran arbitrarily many concurrent extractions + # against a ceiling of `AGENT_KV_CONCURRENT_LIMIT`, all paid for. + # + # Narrowing loses nothing: a mid-run cancel's slot is released + # by `FinalizeView`'s `finally` when the executor's callback + # lands, and `release()` is idempotent (zrem). + if _never_dispatched(job): + AgentKVConcurrencyLimiter.release(str(job.organization_id), str(job.id)) + # Docs §8 promises a terminal notification, and cancellation never + # reaches finalize -- so this path has to send it. Guarded on `won` + # so a cancel that LOST leaves the notification to the finalize that + # beat it, and the caller is told exactly once. + dispatch_cancelled_webhook(job) + return Response({"status": "cancelled"}, status=200) + # Lowercased for cross-endpoint consistency (spec §7.2). + return Response({"status": job.status.lower()}, status=409) diff --git a/backend/agent_kv/execution_views_result.py b/backend/agent_kv/execution_views_result.py new file mode 100644 index 0000000000..5ec465a604 --- /dev/null +++ b/backend/agent_kv/execution_views_result.py @@ -0,0 +1,63 @@ +"""Result/outcome payload for a TERMINAL agent-kv job (spec §7.3). + +Single source of truth for "what does the result endpoint hand back for a +job that's done", shared by ``JobResultView`` (``execution_views.py``) and +``SubmitView``'s synchronous-wait branch (Task 8, lazily imported to avoid a +circular import at module load time) -- both call this only once they've +already established the job is terminal, so it is total over +``AgentKVJob.TERMINAL`` and never raises for a terminal job (spec §7.3: +"Failed jobs: {success: false, error, timing} with a user-safe error"). + +Everything that depends on the caller's own status code (409 for a +non-terminal job at ``JobResultView``, vs. a plain 200 for any terminal +status at ``SubmitView``'s wait loop) -- and the expired/blank-``result_ref`` +404 check -- lives in the caller instead; this function's only job is +building the right body once "terminal" is already a given. +""" + +from agent_kv.models import JobStatus +from agent_kv.storage import read_result + + +def result_payload(job) -> dict: + """Return the result/outcome payload for a terminal ``job``. + + - COMPLETED: ``{success, status, extractors, usage_summary}`` -- the stored + engine result namespaced per extractor (spec §7.3), not the bare blob. + - FAILED: ``{"success": False, "status": "failed", "error": }``. + - CANCELLED: ``{"success": False, "status": "cancelled"}``. + + Callers are responsible for confirming ``job.status`` is terminal (and, + for COMPLETED, that ``result_ref`` is non-blank and unexpired) before + calling this -- it never raises for a job in any of the three states + above. + """ + if job.status == JobStatus.COMPLETED: + # Wrapped at READ time, not stored wrapped: the stored blob is the + # engine's own result object, and namespacing it here keeps the cloud + # engine and every already-written result untouched (spec §7.3). + usage = job.usage_summary or {} + return { + # `success`/`status` stay at TOP level on every terminal payload -- + # completed, failed and cancelled alike -- so a client branches on + # one key rather than having to know that only the failure shapes + # carry it. They describe the JOB; each extractor keeps its own + # `success` inside its own block. + "success": True, + "status": JobStatus.COMPLETED.lower(), + "extractors": {job.extractor: read_result(job.result_ref)}, + # `total` stays the authoritative billing figure; `by_extractor` + # exists so a multi-extractor job can be attributed. With one + # extractor they are necessarily the same numbers. + "usage_summary": { + "total": usage, + "by_extractor": {job.extractor: usage}, + }, + } + if job.status == JobStatus.FAILED: + return { + "success": False, + "status": "failed", + "error": job.error or "Job failed", + } + return {"success": False, "status": "cancelled"} diff --git a/backend/agent_kv/internal_urls.py b/backend/agent_kv/internal_urls.py new file mode 100644 index 0000000000..2ebffc2b21 --- /dev/null +++ b/backend/agent_kv/internal_urls.py @@ -0,0 +1,31 @@ +from django.urls import path + +from agent_kv.internal_views import ( + FinalizeView, + StageReportView, + SweepView, + TTLCleanupView, +) + +urlpatterns = [ + path( + "jobs//stage/", + StageReportView.as_view(), + name="agent_kv_internal_stage", + ), + path( + "jobs//finalize/", + FinalizeView.as_view(), + name="agent_kv_internal_finalize", + ), + path( + "sweep/", + SweepView.as_view(), + name="agent_kv_internal_sweep", + ), + path( + "ttl-cleanup/", + TTLCleanupView.as_view(), + name="agent_kv_internal_ttl_cleanup", + ), +] diff --git a/backend/agent_kv/internal_views.py b/backend/agent_kv/internal_views.py new file mode 100644 index 0000000000..de4583f5ff --- /dev/null +++ b/backend/agent_kv/internal_views.py @@ -0,0 +1,340 @@ +"""Internal APIs for the agent-kv cloud executor (spec §5.4). + +Mounted under ``/internal/v1/agent-kv/`` and guarded ambiently by +``InternalAPIAuthMiddleware`` -- these views take empty auth/permission +classes. Most require ``org_id`` in the body of every request (they act on +one job belonging to one org); the two maintenance views, ``SweepView`` and +``TTLCleanupView``, are the exception -- they are platform-wide and take no +``org_id`` at all, since they sweep across every org in one call. +""" + +import logging + +from django.db.models import F, JSONField, Value +from django.db.models.expressions import CombinedExpression +from rest_framework.response import Response +from rest_framework.views import APIView + +from agent_kv.maintenance import run_sweep, run_ttl_cleanup +from agent_kv.models import AgentKVJob, JobStatus +from agent_kv.rate_limiter import AgentKVConcurrencyLimiter +from agent_kv.storage import delete_input, delete_result_file, write_result + +logger = logging.getLogger(__name__) + +_VALID_STAGE_STATUSES = frozenset({"running", "done"}) +# `name` is reserved too: `_status_document` builds each stage entry as +# `{"name": name, **stages_json[name]}`, and the spread comes LAST -- so a +# persisted `name` counter overrides it. That renames the stage in every later +# status response and breaks the `if name in stages_json` filter the document +# uses to decide which stages to show, i.e. one counter can make a job's +# progress disappear. +_RESERVED_STAGE_ENTRY_KEYS = frozenset({"status", "seconds", "name"}) +_SCALAR_COUNTER_TYPES = (str, int, float, bool) + + +def _sanitize_counters(counters) -> dict: + """Keep only scalar counters that can't clobber the reserved fields. + + Task-11-review ruling: ``counters`` is untrusted executor payload -- + a key of ``status``/``seconds`` must not be able to override the + endpoint's own reserved fields, and a nested dict/list value must not + be admitted into the persisted stage entry (only scalars are). + """ + if not isinstance(counters, dict): + return {} + return { + key: value + for key, value in counters.items() + if key not in _RESERVED_STAGE_ENTRY_KEYS + and isinstance(value, _SCALAR_COUNTER_TYPES) + } + + +def _stage_merge_expression(stage: str, entry: dict) -> CombinedExpression: + """Build a DB-side jsonb merge expression for one stage's entry. + + Compiles (on Postgres) to ``"stages" || %s``, the jsonb ``||`` + concatenation operator, which does a shallow top-level merge: only the + ``stage`` key in the JSON column is touched, every other key is left + exactly as the current row has it. This is what makes the merge safe + against two concurrent stage reports for two *different* stages racing + each other -- a Python-side ``dict(job.stages or {})`` read, merge, and + full-column ``.update(stages=...)`` would let whichever write lands + second silently clobber the first (task-11-review ruling). + """ + return CombinedExpression( + F("stages"), + "||", + Value({stage: entry}, output_field=JSONField()), + output_field=JSONField(), + ) + + +class StageReportView(APIView): + """Merge one stage's progress into ``job.stages`` (spec §5.4). + + This endpoint is the sole write gate for ``job.stages``: only the + defined shape (``status``, optional ``seconds``, plus sanitized flat + ``counters``) is ever persisted -- unexpected top-level body keys, and + any counter that collides with a reserved key or isn't a scalar, are + dropped rather than stored, so a later status-endpoint spread of a + stage entry can't leak arbitrary payload content (task-9 review). + """ + + authentication_classes: list = [] + permission_classes: list = [] + + def post(self, request, *args, job_id=None, **kwargs): + body = request.data + org_id = body.get("org_id") + if not org_id: + return Response({"detail": "org_id is required"}, status=400) + stage = body.get("stage") + if not stage: + return Response({"detail": "stage is required"}, status=400) + status_value = body.get("status") + if status_value not in _VALID_STAGE_STATUSES: + return Response({"detail": "status must be 'running' or 'done'"}, status=400) + + job_qs = AgentKVJob.objects.filter(id=job_id, organization_id=org_id).exclude( + status__in=list(AgentKVJob.TERMINAL) + ) + job = job_qs.first() + if job is None: + # Two very different situations used to share this response: a late + # report for an already-terminal job (ordinary), and a report for a + # job_id/org_id pair that matches NO row at all (never ordinary). + # + # The second is what an org-slug-vs-FK-pk mix-up looks like, and + # `dispatch.py` documents that exact confusion shipping once + # already. If it recurs, every stage report is a silent 200 no-op, + # every job stays non-terminal and gets reaped as "Job timed out" -- + # a 100% failure rate presenting as timeouts, with nothing in any + # log. So say which one happened. + exists = AgentKVJob.objects.filter(id=job_id, organization_id=org_id).exists() + if not exists: + logger.warning( + "agent-kv stage report for unknown job %s (org=%s); no row " + "matches that job/org pair -- check the org identifier", + job_id, + org_id, + ) + return Response( + { + "ok": True, + "noop": True, + "reason": "already_terminal" if exists else "unknown_job", + } + ) + + # `job.stage` is varchar(32). A longer name was an unhandled 500, and + # `job.stages` could grow unbounded distinct keys from a misbehaving + # or malicious executor. + if not isinstance(stage, str): + # A truthy non-string (e.g. `1`) passed the presence check above and + # then raised TypeError at `len()` -- a 500 for a malformed report. + return Response({"detail": "stage must be a string"}, status=400) + if len(stage) > 32: + return Response({"detail": "stage name too long"}, status=400) + + entry = {"status": status_value} + if "seconds" in body: + # Type-checked rather than taken verbatim: this is untrusted + # executor payload and it is echoed back in the status document, + # so a dict or list here was persisted and returned to the caller. + seconds = body["seconds"] + if not isinstance(seconds, (int, float)) or isinstance(seconds, bool): + return Response({"detail": "seconds must be a number"}, status=400) + entry["seconds"] = seconds + entry.update(_sanitize_counters(body.get("counters"))) + + updates = { + "stages": _stage_merge_expression(stage, entry), + "stage": stage, + } + if job.status in (JobStatus.PENDING, JobStatus.DISPATCHED): + updates["status"] = JobStatus.RUNNING + # Re-guarded at update time: job_qs still excludes TERMINAL, so if + # the job raced to terminal between the read above and this + # UPDATE's WHERE clause, this becomes a silent no-op too. + job_qs.update(**updates) + return Response({"ok": True}) + + +class FinalizeView(APIView): + """Terminalize a job: write its result (on success) then mark it done. + + Idempotent: a job already in a terminal state short-circuits before any + write -- ``write_result`` is only ever called for a job that is still + non-terminal at read time, so a duplicate/late finalize call can never + rewrite an already-written result. The concurrency slot is always + released, in a ``finally``, on every path that actually attempts to + finalize (success, duplicate no-op, or an exception raised while + finalizing) -- but never on a 400 for a malformed body, since nothing + was finalized (and therefore nothing had a slot to release yet). + + On the success path, ``write_result`` unavoidably runs *before* the + ``mark_terminal`` guard (the result has to exist to get a ``result_ref`` + to pass into it) -- so a guard-loss right after the write (a concurrent + cancel/duplicate finalize won instead) leaves a result file on disk with + no ref anywhere pointing at it. That orphan is cleaned up immediately via + ``storage.delete_result_file`` rather than left to leak past even TTL + cleanup (which only ever acts on a job's *stored* ``result_ref``). + + Also deletes the staged input the moment ``mark_terminal`` actually + wins (spec D10: "uploaded document deleted on job completion") -- the + run is over either way, whether it completed or failed, so the input is + no longer needed; the result file this same request may have just + *successfully* written (guard won) is left completely untouched. Never + runs on the duplicate/guard-lost path, since either another writer + already owns cleanup or there's nothing new to terminalize. + """ + + authentication_classes: list = [] + permission_classes: list = [] + + @staticmethod + def _complete(job, job_id, org_id: str, body) -> bool: + """Persist the result and terminalize as COMPLETED. True if the guard won. + + ``write_result`` unavoidably runs *before* the ``mark_terminal`` guard + (the result has to exist to get a ``result_ref`` to pass in), so a + guard-loss right after the write -- a concurrent cancel or duplicate + finalize won instead -- leaves a result file with no ref anywhere + pointing at it. That orphan is deleted here rather than left to leak + past even TTL cleanup, which only ever acts on a job's *stored* + ``result_ref``. Best-effort; never blocks the response. + """ + result_ref = write_result(org_id, str(job_id), body.get("result") or {}) + finalized = AgentKVJob.mark_terminal( + job_id, + org_id, + JobStatus.COMPLETED, + result_ref=result_ref, + usage_summary=body.get("usage_summary"), + ) + if finalized: + job.status = JobStatus.COMPLETED + else: + delete_result_file(result_ref) + return finalized + + @staticmethod + def _fail(job, job_id, org_id: str, body) -> bool: + """Terminalize as FAILED. True if the guard won.""" + finalized = AgentKVJob.mark_terminal( + job_id, org_id, JobStatus.FAILED, error=body.get("error") or "" + ) + if finalized: + job.status = JobStatus.FAILED + return finalized + + @staticmethod + def _drop_staged_input(job, job_id, org_id: str) -> None: + """Delete the staged input now the run is over (spec D10). + + Only ever called when ``mark_terminal`` actually won, so it never runs + on the duplicate/guard-lost path -- there, either another writer owns + cleanup or there is nothing new to terminalize. A CANCELLED job never + reaches here either: cancel goes through JobCancelView/mark_terminal + directly, and a late finalize against an already-CANCELLED job loses + the terminal guard, so a cancelled job's input rides the normal TTL + sweep instead. + + The ref is blanked only when the file is confirmed gone -- otherwise it + is the only handle TTL cleanup could retry from, and blanking it would + orphan the uploaded document. + """ + if delete_input(job): + AgentKVJob.objects.filter(id=job_id, organization_id=org_id).update( + input_ref="" + ) + + def post(self, request, *args, job_id=None, **kwargs): + body = request.data + org_id = body.get("org_id") + if not org_id: + return Response({"detail": "org_id is required"}, status=400) + success = body.get("success") + if not isinstance(success, bool): + # Strict bool check (not truthy/falsy): a malformed/missing + # `success` must 400 rather than silently fall through to the + # failure branch and persist a FAILED job with an empty error. + return Response({"detail": "success must be a boolean"}, status=400) + + job = AgentKVJob.objects.filter(id=job_id, organization_id=org_id).first() + finalized = False + try: + if job is not None and job.status not in AgentKVJob.TERMINAL: + handler = self._complete if success else self._fail + finalized = handler(job, job_id, org_id, body) + if finalized: + self._drop_staged_input(job, job_id, org_id) + finally: + AgentKVConcurrencyLimiter.release(org_id, str(job_id)) + + if job is None: + logger.warning( + "agent-kv finalize for unknown job %s (org=%s); no row matches " + "that job/org pair, so nothing was terminalized -- the job (if " + "it exists under another org identifier) will be reaped as a " + "timeout", + job_id, + org_id, + ) + + return Response( + { + "finalized": finalized, + "webhook_url": job.webhook_url if job else "", + "status": job.status.lower() if job else "", + # Lets the callback worker tell "already terminal" (ordinary) + # from "no such job" (a wiring fault) -- they were byte-identical. + "reason": "ok" if job else "unknown_job", + } + ) + + +class SweepView(APIView): + """Terminalize PENDING-never-dispatched AND stuck jobs (spec §5.4). + + Thin wrapper: the two-phase sweep logic itself -- never-dispatched + PENDING rows past ``AGENT_KV_SWEEP_GRACE_SECONDS`` and stuck + DISPATCHED/RUNNING rows past ``AGENT_KV_STUCK_JOB_GRACE_SECONDS``, each + idempotent and batch-capped -- lives in + :func:`agent_kv.maintenance.run_sweep`, shared with the + ``agent_kv_sweep`` management command (the Kubernetes CronJob mechanism + for the cloud deployment). See that function's docstring for the full + rationale. + + Platform-wide by design (no ``org_id`` in the body, unlike every other + view in this module) -- it is invoked by a periodic maintenance + mechanism (spec §5.4), not the cloud executor acting on one job/org. + """ + + authentication_classes: list = [] + permission_classes: list = [] + + def post(self, request, *args, **kwargs): + return Response(run_sweep()) + + +class TTLCleanupView(APIView): + """Delete staged files for expired jobs and blank their refs (spec §5.4). + + Thin wrapper: the cleanup logic itself -- deleting object-store paths + for expired jobs and blanking their refs, idempotent and batch-capped -- + lives in :func:`agent_kv.maintenance.run_ttl_cleanup`, shared with the + ``agent_kv_ttl_cleanup`` management command (the Kubernetes CronJob + mechanism for the cloud deployment). See that function's docstring for + the full rationale. + + Platform-wide by design, same as :class:`SweepView`. + """ + + authentication_classes: list = [] + permission_classes: list = [] + + def post(self, request, *args, **kwargs): + return Response(run_ttl_cleanup()) diff --git a/backend/agent_kv/key_validator.py b/backend/agent_kv/key_validator.py new file mode 100644 index 0000000000..2daa729cbc --- /dev/null +++ b/backend/agent_kv/key_validator.py @@ -0,0 +1,45 @@ +import logging +import uuid + +from api_v2.api_key_validator import BaseAPIKeyValidator +from api_v2.exceptions import Forbidden + +from agent_kv.models import AgentKVKey + +logger = logging.getLogger(__name__) + +# One message for every rejection reason, deliberately. A caller holding a bad +# key learns only that it is bad -- not whether it was malformed, unknown, +# inactive, or missing an organization -- so the response cannot be used to +# probe which keys exist. Named once so the three raise sites cannot drift +# apart and start leaking that distinction. +_INVALID_KEY = "Invalid api key" + + +class AgentKVKeyValidator(BaseAPIKeyValidator): + @staticmethod + def validate_and_process(self, request, func, api_key, *args, **kwargs): + try: + uuid.UUID(api_key) + except (ValueError, AttributeError): + raise Forbidden(_INVALID_KEY) + try: + key_obj = AgentKVKey.objects.get(key=api_key, is_active=True) + except AgentKVKey.DoesNotExist: + raise Forbidden(_INVALID_KEY) + if key_obj.organization_id is None: + # `organization` is nullable on the model (DefaultOrganizationMixin + # fills it from UserContext at save time), but EVERY downstream use + # of a key is org-scoped: the subscription gate, the concurrency + # limiter, the storage prefix and every job lookup. A key with no + # organization cannot scope any of them, so it is not a usable key. + # Refused once here, at the auth boundary, rather than surfacing as + # a null dereference inside whichever view happens to touch the + # organization first. + logger.error( + "agent-kv key %s has no organization; refusing the request", + key_obj.id, + ) + raise Forbidden(_INVALID_KEY) + kwargs["agent_kv_key"] = key_obj + return func(self, request, *args, **kwargs) diff --git a/backend/agent_kv/maintenance.py b/backend/agent_kv/maintenance.py new file mode 100644 index 0000000000..0a537a9708 --- /dev/null +++ b/backend/agent_kv/maintenance.py @@ -0,0 +1,356 @@ +"""Agent-KV platform-wide maintenance logic (spec §5.4). + +Two independent, idempotent, batch-capped jobs -- the never-dispatched/ +stuck-job sweep and the TTL cleanup of expired staged files -- live here so +there is exactly one implementation shared by both invocation paths: + +* the internal HTTP endpoints (``agent_kv/internal_views.py::SweepView``/ + ``TTLCleanupView``, ``POST /internal/v1/agent-kv/sweep/`` and + ``.../ttl-cleanup/``), driven by the OSS/self-hosted PG-scheduler periodic + task mechanism (``workers/scheduler/agent_kv_tasks.py``); and +* the ``agent_kv_sweep``/``agent_kv_ttl_cleanup`` Django management commands + (``agent_kv/management/commands/``), driven by a Kubernetes CronJob in the + cloud deployment. + +Both callers get the identical dict shape back (``{"swept": N, "timed_out": +M}`` / ``{"cleaned": N}``), and both entrypoints are equally safe to call +more often than needed, or concurrently with each other -- everything below +is either a guarded UPDATE (``AgentKVJob.mark_terminal``) or a targeted +single-row write, never a queryset-wide one. +""" + +import logging +from datetime import timedelta + +from django.conf import settings +from django.db.models import Q +from django.db.models.functions import Coalesce +from django.utils import timezone + +from agent_kv.models import AgentKVJob, JobStatus +from agent_kv.rate_limiter import SLOT_TTL_SECONDS, AgentKVConcurrencyLimiter +from agent_kv.storage import delete_job_files + +logger = logging.getLogger(__name__) + +_MAINTENANCE_BATCH_LIMIT = 500 +# Slots in each TTL-cleanup batch held for rows whose file delete failed before. +# Caps the retry lane (so failures cannot crowd out new expirations) and +# guarantees it (so new expirations cannot crowd out retries) -- review found +# the implementation starving each side in turn when one ordering tried to do +# both. 20% leaves 400 slots for the normal case, where nothing is retrying. +_TTL_RETRY_RESERVE = 100 +_NEVER_DISPATCHED_ERROR = "Job was never dispatched" +_STUCK_JOB_ERROR = "Job timed out" + + +def run_sweep() -> dict: + """Terminalize PENDING-never-dispatched AND stuck jobs (spec §5.4). + + Two independent phases, run every call, each capped and counted + separately: + + **Phase 1 -- never dispatched.** Mirrors ``workflow_manager``'s + undispatched-execution sweep: the submit endpoint commits a PENDING row + before dispatch runs, and an abort in between (client disconnect, worker + crash, pod eviction) can leave it stranded with no owner -- ``PENDING`` + is not a terminal state, and nothing else recovers a job that was never + queued. ``dispatched_at`` is stamped as a positive fact at dispatch + time, so PENDING + older than the grace + ``dispatched_at IS NULL`` + means the dispatch never happened. + + **Phase 2 -- stuck in flight.** A job that *did* dispatch can still + never terminalize: the executor pod can be killed, its callback queue + can be lost, or the cloud engine itself can hang -- none of which + ``mark_terminal`` ever sees, so a DISPATCHED/RUNNING row can sit + forever holding a concurrency slot with no path back to terminal. + ``dispatched_at`` (stamped positively at dispatch, spec §5.3/Fix 1) older + than ``AGENT_KV_STUCK_JOB_GRACE_SECONDS`` is the only signal available -- + there is no heartbeat -- so this phase force-fails anything past that + grace, same as the workflow reaper's stuck-execution recovery. + + Platform-wide by design (no ``org_id`` parameter) -- invoked by a + periodic maintenance mechanism (spec §5.4), not by something acting on + one job/org. + + Idempotent: ``mark_terminal``'s guarded UPDATE only terminalizes a row + still in a non-terminal state, so a job already swept (or one that + legitimately dispatched/finalized/was cancelled since the candidate + query ran) is left alone by a repeat call. ``swept``/``timed_out`` each + count guard successes, not candidates, so a race against a concurrent + finalize/cancel/duplicate sweep is reflected accurately instead of + double-counted. + + Each phase is independently batch-capped at ``_MAINTENANCE_BATCH_LIMIT`` + (oldest-first by its own ordering key) so a large backlog in either + phase -- exactly what a dispatch-path or executor-fleet infra incident + produces, which is also when this sweep matters most -- can't load + unbounded into memory or hold a caller open through a long loop. + Idempotency (above) is what makes this safe to cap: whatever a call + doesn't reach is still there, unchanged, for the next tick. + """ + now = timezone.now() + + never_dispatched_cutoff = now - timedelta( + seconds=settings.AGENT_KV_SWEEP_GRACE_SECONDS + ) + never_dispatched_candidates = AgentKVJob.objects.filter( + status=JobStatus.PENDING, + created_at__lt=never_dispatched_cutoff, + dispatched_at__isnull=True, + ).order_by("created_at")[:_MAINTENANCE_BATCH_LIMIT] + + swept = 0 + for job in never_dispatched_candidates: + # Per-job isolation is structural, not a try/except here: + # mark_terminal is a guarded UPDATE that can't raise on a + # normal outcome, and release() has its own internal + # try/except (rate_limiter.py) -- so one job's failure can't + # abort the loop for the rest of the batch. + org_id = job.organization_id + won = AgentKVJob.mark_terminal( + job.id, + org_id, + JobStatus.FAILED, + error=_NEVER_DISPATCHED_ERROR, + ) + if won: + swept += 1 + # Only a job this call actually terminalized held a slot + # worth releasing here -- one a concurrent finalize/cancel + # won instead already released its own slot on that path. + AgentKVConcurrencyLimiter.release(str(org_id), str(job.id)) + + stuck_cutoff = now - timedelta(seconds=settings.AGENT_KV_STUCK_JOB_GRACE_SECONDS) + # `dispatched_at__lt` OR `dispatched_at IS NULL`, not just the former. + # A DISPATCHED/RUNNING row whose `dispatched_at` is NULL is unreachable by + # a `__lt` filter alone -- SQL `NULL < x` is never true -- so such a row + # would hang in a non-terminal state forever. dispatch.py now stamps + # `dispatched_at` for any non-terminal row precisely so this cannot + # normally happen; this arm is the backstop for the window that remains + # (the worker dying between the enqueue and that stamp), because the cost + # of missing one is a job that never terminalizes at all. Such a row falls + # back to `created_at` for the age test, which is the only timestamp it + # has. + stuck_candidates = ( + AgentKVJob.objects.filter( + status__in=[JobStatus.DISPATCHED, JobStatus.RUNNING], + ) + .filter( + Q(dispatched_at__lt=stuck_cutoff) + | Q(dispatched_at__isnull=True, created_at__lt=stuck_cutoff) + ) + # Coalesce, not a bare `dispatched_at`. The second Q arm exists to + # recover rows whose post-enqueue bookkeeping was lost -- they have + # `dispatched_at IS NULL` -- but Postgres sorts ascending NULLS LAST, + # so with a full batch of non-NULL stuck rows ahead of them those rows + # were never selected. The backstop could not fire in exactly the + # situation it exists for: a backlog. + .order_by(Coalesce("dispatched_at", "created_at"))[:_MAINTENANCE_BATCH_LIMIT] + ) + + timed_out = 0 + for job in stuck_candidates: + org_id = job.organization_id + won = AgentKVJob.mark_terminal( + job.id, + org_id, + JobStatus.FAILED, + error=_STUCK_JOB_ERROR, + ) + if won: + timed_out += 1 + AgentKVConcurrencyLimiter.release(str(org_id), str(job.id)) + + # Phase 3: release slots still held by jobs that were CANCELLED after being + # dispatched. + # + # Cancelling a dispatched job deliberately does NOT release its slot -- the + # executor is still running and still billing, so the slot belongs to the + # finalize callback that will arrive when it finishes. But if that executor + # dies, no callback ever arrives, and neither phase above selects a + # CANCELLED row: phase 1 wants PENDING, phase 2 wants DISPATCHED/RUNNING. + # The slot then sits occupied until Redis expires it. + # + # Bounded on BOTH sides, which is what keeps this from rescanning the same + # ancient rows forever: older than the stuck grace (so a live executor is + # not cut short) and newer than the slot TTL (past which Redis has already + # dropped the entry, so there is nothing left to release). + slot_ttl_floor = now - timedelta(seconds=SLOT_TTL_SECONDS) + abandoned_cancelled = AgentKVJob.objects.filter( + status=JobStatus.CANCELLED, + dispatched_at__isnull=False, + completed_at__lt=stuck_cutoff, + completed_at__gt=slot_ttl_floor, + # NEWEST first, the opposite of the other two phases and deliberate. + # Releasing a slot does not remove its row from this query -- there is + # no "released" marker -- so oldest-first would re-select the same 500 + # rows every sweep whenever the backlog exceeds the batch limit, and + # rows arriving behind them would age out of the window never having + # been looked at. + # + # Newest-first inverts which rows lose: the ones skipped are the oldest, + # i.e. closest to the slot TTL floor, where Redis is about to drop the + # entry anyway and a release buys almost nothing. Every row is seen + # while its release still matters, unless more than + # `_MAINTENANCE_BATCH_LIMIT` become eligible inside one sweep interval + # -- and the TTL remains the backstop for that. + ).order_by("-completed_at")[:_MAINTENANCE_BATCH_LIMIT] + + released = 0 + for job in abandoned_cancelled: + # `release` is an idempotent zrem, so re-releasing a slot the callback + # already freed costs nothing and is not worth a guard. + AgentKVConcurrencyLimiter.release(str(job.organization_id), str(job.id)) + released += 1 + if released: + logger.info( + "agent-kv sweep released %s concurrency slot(s) held by cancelled " + "jobs whose executor never called back", + released, + ) + + # These two counts are the only evidence the sweep ran and the only + # evidence it had to do anything. A sweep that terminalizes a thousand jobs + # as FAILED used to emit nothing at all -- the module has had no logger + # since it was written -- so a backlog of stranded jobs looked identical to + # a quiet, healthy system. + if swept or timed_out: + logger.warning( + "agent-kv sweep terminalized %s never-dispatched and %s stuck job(s); " + "a non-zero count here means jobs were stranded and their " + "concurrency slots held", + swept, + timed_out, + ) + else: + logger.info("agent-kv sweep: nothing to terminalize") + # `released` is reported, not just logged. Phase 3 counted slot releases + # into a local, logged it, and then returned only the first two counts -- + # so `SweepView` (which returns this dict verbatim) and the scheduler task + # (which logs it) could never show that any slot was recovered. That is the + # same "counted, then discarded before the caller sees it" shape as the + # round-1 finding about this function's return value, reintroduced by the + # phase that was added to fix a different half of it. + # + # It matters operationally: a cancelled job whose executor never called + # back holds a slot against the org's cap, and phase 3 is the only thing + # that frees it. Without this key, "the cap is being hit and the sweep is + # recovering slots" and "the cap is being hit and nothing is recovering + # anything" produce identical output. + return {"swept": swept, "timed_out": timed_out, "released": released} + + +def run_ttl_cleanup() -> dict: + """Delete staged files for expired jobs and blank their refs (spec §5.4). + + The job row itself is retained (audit trail) -- only the object-store + paths are dropped, once nothing can read them any more (the status/ + result endpoints already 404 past ``expires_at`` -- Task 9). Blanking a ref + after its file is deleted is what makes a repeat call a no-op: the filters + below only match rows still carrying a non-blank ref, so a job already + cleaned (or one that never staged an input/produced a result) drops out of + the candidate set on the next pass. + + Platform-wide by design, same as :func:`run_sweep`. + + **Two queries, not one ordering.** Retaining a ref whose delete failed is + what makes a retry possible at all, but it also means failed rows compete + with new expirations for a capped batch, and either side can starve the + other. Both directions were observed in review: + + - Oldest-expiry-first alone: 500 permanently-failing rows refill every + batch and nothing newer is ever reached. + - ``cleanup_failed_at NULLS FIRST`` (the first attempt at a fix): the + mirror image -- 500 fresh expirations per tick fill every batch and the + failures are never retried, so their files sit in storage indefinitely. + + One ordering cannot express "neither side starves the other", so the batch + is split instead: retries get up to ``_TTL_RETRY_RESERVE`` slots, fresh + rows get whatever is left. Each side is capped, so each side is guaranteed + capacity whenever it has work. + + The split also buys a cheaper plan. Each query now sorts by ONE column + ascending with a plain equality/IS NULL predicate on the index's leading + column, so the ``(cleanup_failed_at, expires_at)`` index serves both + directly. ``NULLS FIRST`` could not use it at all -- a btree index is + ``NULLS LAST`` by default -- so Postgres had to sort every matching expired + row before applying the limit, which grows with the backlog. + """ + now = timezone.now() + # TERMINAL only. Without this, a job still RUNNING past `expires_at` had + # its staged input deleted out from under the executor -- the TTL is a + # retention policy for finished work, not a kill switch for running work. + # Reachable whenever a job is stuck non-terminal for longer than + # AGENT_KV_RESULT_TTL_DAYS, which is precisely what the sweep's phase 2 + # exists to catch; the sweep terminalizes those first, and then this + # cleans them. + expired = ( + AgentKVJob.objects.filter(expires_at__lt=now) + .filter(status__in=list(AgentKVJob.TERMINAL)) + .filter(Q(input_ref__gt="") | Q(result_ref__gt="")) + ) + + # Retries first, capped at the reserve so they cannot crowd out fresh work. + # Oldest failure first, so failures rotate rather than one row absorbing + # every retry. + retries = list( + expired.filter(cleanup_failed_at__isnull=False).order_by("cleanup_failed_at")[ + :_TTL_RETRY_RESERVE + ] + ) + # Fresh rows take the remaining capacity -- the full batch when there is + # nothing to retry, which is the normal case. + fresh = list( + expired.filter(cleanup_failed_at__isnull=True).order_by("expires_at")[ + : _MAINTENANCE_BATCH_LIMIT - len(retries) + ] + ) + + cleaned = 0 + retained = 0 + for job in retries + fresh: + # Blank only the refs whose files are CONFIRMED gone. Targeted + # single-row update (not a queryset-wide `.update()`) for the same + # reason it always was: one job's refs must never be blanked off the + # back of another job's delete. + # + # A ref left set is the retry handle -- this loop used to blank both + # unconditionally, so a transient object-store failure orphaned the file + # permanently (the candidate filters above only match rows that still + # carry a non-blank ref, so a blanked row can never be reconsidered). + cleared = delete_job_files(job) + fields: dict = dict.fromkeys(cleared, "") + if len(cleared) == 2: + cleaned += 1 + # Fully cleaned rows drop out of the candidate filter anyway (both + # refs blank), so this only matters for a row that failed before and + # succeeded now -- it must not keep a stale failure marker, or it + # would consume a reserve slot it no longer needs. + fields["cleanup_failed_at"] = None + else: + retained += 1 + # Stamped on every failed attempt, not just the first: this is what + # moves the row into the retry lane, and refreshing it rotates the + # retry order among failures instead of letting the earliest-stamped + # row be retried forever. + fields["cleanup_failed_at"] = timezone.now() + AgentKVJob.objects.filter(id=job.id).update(**fields) + # `retained` is reported, not just logged: these rows keep their refs and + # are retried from the reserve lane on a later tick. A `retained` that + # stays high across ticks is the signal that something is wrong with the + # object store rather than with one job -- the split stops either side + # blocking the other, it does not make a persistent fault harmless. + # `retained` is the count this pass could NOT clean -- a file delete that + # failed keeps its ref so the next pass retries it. A `retained` that never + # falls is a stuck object, and stays invisible without this. + if retained: + logger.warning( + "agent-kv TTL cleanup removed %s job(s) and RETAINED %s whose files " + "could not be deleted; those refs are kept for the next pass", + cleaned, + retained, + ) + elif cleaned: + logger.info("agent-kv TTL cleanup removed %s expired job(s)", cleaned) + return {"cleaned": cleaned, "retained": retained} diff --git a/backend/agent_kv/management/__init__.py b/backend/agent_kv/management/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/agent_kv/management/commands/__init__.py b/backend/agent_kv/management/commands/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/agent_kv/management/commands/agent_kv_sweep.py b/backend/agent_kv/management/commands/agent_kv_sweep.py new file mode 100644 index 0000000000..9c091b9403 --- /dev/null +++ b/backend/agent_kv/management/commands/agent_kv_sweep.py @@ -0,0 +1,24 @@ +"""CronJob entrypoint for the Agent-KV never-dispatched/stuck-job sweep. + +Thin wrapper around :func:`agent_kv.maintenance.run_sweep` -- see that +function's docstring for the two-phase sweep logic itself (spec §5.4). This +command is the cloud mechanism for driving it: a Kubernetes CronJob runs +``python manage.py agent_kv_sweep`` on a schedule. The internal +``SweepView``/``POST /internal/v1/agent-kv/sweep/`` endpoint +(``agent_kv/internal_views.py``) delegates to the same function and remains +the OSS/self-hosted mechanism, driven by the PG-scheduler periodic task +(``workers/scheduler/agent_kv_tasks.py``). +""" + +import json + +from django.core.management.base import BaseCommand + + +class Command(BaseCommand): + help = "Terminalize never-dispatched and stuck Agent-KV jobs (CronJob entrypoint)" + + def handle(self, *args, **options): + from agent_kv.maintenance import run_sweep + + self.stdout.write(json.dumps(run_sweep())) diff --git a/backend/agent_kv/management/commands/agent_kv_ttl_cleanup.py b/backend/agent_kv/management/commands/agent_kv_ttl_cleanup.py new file mode 100644 index 0000000000..c8025c40c9 --- /dev/null +++ b/backend/agent_kv/management/commands/agent_kv_ttl_cleanup.py @@ -0,0 +1,24 @@ +"""CronJob entrypoint for the Agent-KV TTL cleanup of expired staged files. + +Thin wrapper around :func:`agent_kv.maintenance.run_ttl_cleanup` -- see that +function's docstring for the cleanup logic itself (spec §5.4). This command +is the cloud mechanism for driving it: a Kubernetes CronJob runs +``python manage.py agent_kv_ttl_cleanup`` on a schedule. The internal +``TTLCleanupView``/``POST /internal/v1/agent-kv/ttl-cleanup/`` endpoint +(``agent_kv/internal_views.py``) delegates to the same function and remains +the OSS/self-hosted mechanism, driven by the PG-scheduler periodic task +(``workers/scheduler/agent_kv_tasks.py``). +""" + +import json + +from django.core.management.base import BaseCommand + + +class Command(BaseCommand): + help = "Delete staged files for expired Agent-KV jobs (CronJob entrypoint)" + + def handle(self, *args, **options): + from agent_kv.maintenance import run_ttl_cleanup + + self.stdout.write(json.dumps(run_ttl_cleanup())) diff --git a/backend/agent_kv/migrations/0001_initial.py b/backend/agent_kv/migrations/0001_initial.py new file mode 100644 index 0000000000..0fa24c9eac --- /dev/null +++ b/backend/agent_kv/migrations/0001_initial.py @@ -0,0 +1,157 @@ +# Generated by Django 4.2.30 on 2026-08-28 12:16 + +import uuid + +import django.db.models.deletion +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + initial = True + + dependencies = [ + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ("account_v2", "0007_organization_restrict_connector_creation"), + ] + + operations = [ + migrations.CreateModel( + name="AgentKVKey", + fields=[ + ("created_at", models.DateTimeField(auto_now_add=True)), + ("modified_at", models.DateTimeField(auto_now=True)), + ( + "id", + models.UUIDField( + default=uuid.uuid4, + editable=False, + primary_key=True, + serialize=False, + ), + ), + ("name", models.CharField(max_length=128)), + ("description", models.CharField(blank=True, default="", max_length=512)), + ("key", models.UUIDField(default=uuid.uuid4, unique=True)), + ("is_active", models.BooleanField(default=True)), + ( + "created_by", + models.ForeignKey( + null=True, + on_delete=django.db.models.deletion.SET_NULL, + related_name="agent_kv_keys_created", + to=settings.AUTH_USER_MODEL, + ), + ), + ( + "modified_by", + models.ForeignKey( + null=True, + on_delete=django.db.models.deletion.SET_NULL, + related_name="+", + to=settings.AUTH_USER_MODEL, + ), + ), + ( + "organization", + models.ForeignKey( + blank=True, + db_comment="Foreign key reference to the Organization model.", + editable=False, + null=True, + on_delete=django.db.models.deletion.CASCADE, + to="account_v2.organization", + ), + ), + ], + options={ + "db_table": "agent_kv_key", + }, + ), + migrations.CreateModel( + name="AgentKVJob", + fields=[ + ("created_at", models.DateTimeField(auto_now_add=True)), + ("modified_at", models.DateTimeField(auto_now=True)), + ( + "id", + models.UUIDField( + default=uuid.uuid4, + editable=False, + primary_key=True, + serialize=False, + ), + ), + ("task_id", models.UUIDField(blank=True, null=True)), + ( + "status", + models.CharField( + choices=[ + ("PENDING", "Pending"), + ("DISPATCHED", "Dispatched"), + ("RUNNING", "Running"), + ("COMPLETED", "Completed"), + ("FAILED", "Failed"), + ("CANCELLED", "Cancelled"), + ], + default="PENDING", + max_length=16, + ), + ), + ("stage", models.CharField(blank=True, default="", max_length=32)), + ("stages", models.JSONField(blank=True, default=dict)), + ("pages_total", models.IntegerField(blank=True, null=True)), + ("input_ref", models.CharField(blank=True, default="", max_length=512)), + ("result_ref", models.CharField(blank=True, default="", max_length=512)), + ("usage_summary", models.JSONField(blank=True, null=True)), + ("error", models.TextField(blank=True, default="")), + ("dispatched_at", models.DateTimeField(blank=True, null=True)), + ("completed_at", models.DateTimeField(blank=True, null=True)), + ("expires_at", models.DateTimeField(blank=True, null=True)), + ("tags", models.JSONField(blank=True, default=list)), + ("custom_data", models.JSONField(blank=True, null=True)), + ("webhook_url", models.URLField(blank=True, default="", max_length=1024)), + ( + "api_key", + models.ForeignKey( + null=True, + on_delete=django.db.models.deletion.SET_NULL, + related_name="jobs", + to="agent_kv.agentkvkey", + ), + ), + ( + "organization", + models.ForeignKey( + blank=True, + db_comment="Foreign key reference to the Organization model.", + editable=False, + null=True, + on_delete=django.db.models.deletion.CASCADE, + to="account_v2.organization", + ), + ), + ], + options={ + "db_table": "agent_kv_job", + }, + ), + migrations.AddConstraint( + model_name="agentkvkey", + constraint=models.UniqueConstraint( + fields=("name", "organization"), name="unique_agent_kv_key_name_per_org" + ), + ), + migrations.AddIndex( + model_name="agentkvjob", + index=models.Index( + fields=["organization", "status"], name="agent_kv_jo_organiz_9b6b97_idx" + ), + ), + migrations.AddIndex( + model_name="agentkvjob", + index=models.Index( + fields=["expires_at"], name="agent_kv_jo_expires_5966b4_idx" + ), + ), + ] diff --git a/backend/agent_kv/migrations/0002_agentkvjob_extractor.py b/backend/agent_kv/migrations/0002_agentkvjob_extractor.py new file mode 100644 index 0000000000..8889c7cd36 --- /dev/null +++ b/backend/agent_kv/migrations/0002_agentkvjob_extractor.py @@ -0,0 +1,21 @@ +"""Record which extractor each job ran. + +Defaulting to "kv" is the historical truth, not a guess: before this column +existed the API accepted exactly one extractor and it was always `kv`. +""" + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("agent_kv", "0001_initial"), + ] + + operations = [ + migrations.AddField( + model_name="agentkvjob", + name="extractor", + field=models.CharField(default="kv", max_length=32), + ), + ] diff --git a/backend/agent_kv/migrations/0003_agentkvjob_cleanup_failed_at.py b/backend/agent_kv/migrations/0003_agentkvjob_cleanup_failed_at.py new file mode 100644 index 0000000000..d67b837ffb --- /dev/null +++ b/backend/agent_kv/migrations/0003_agentkvjob_cleanup_failed_at.py @@ -0,0 +1,41 @@ +"""Record when TTL cleanup last failed on a job, and index the ordering. + +Review follow-up on UN-4044. ``run_ttl_cleanup`` retains a ref whose file +delete failed, because the ref is the only handle a retry can work from. Under +the previous plain ``expires_at`` ordering those rows refilled the capped batch +on every tick, so a persistent object-store fault on 500 rows stalled cleanup +outright and every later expired job kept its files past TTL. + +Ordering is now ``(cleanup_failed_at NULLS FIRST, expires_at)``: a job that has +never failed is always processed ahead of one that has, so failures cannot +block the backlog, and the stamp is refreshed on each failed attempt so +failures rotate rather than one row absorbing every retry. + +A plain (non-CONCURRENT) AddIndex is correct here, unlike the concurrent builds +used elsewhere in this codebase: ``agent_kv_job`` is created by this app's own +``0001_initial`` and has never been deployed, so the table is empty when this +runs and the ACCESS EXCLUSIVE lock has nothing to block. +""" + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("agent_kv", "0002_agentkvjob_extractor"), + ] + + operations = [ + migrations.AddField( + model_name="agentkvjob", + name="cleanup_failed_at", + field=models.DateTimeField(blank=True, null=True), + ), + migrations.AddIndex( + model_name="agentkvjob", + index=models.Index( + fields=["cleanup_failed_at", "expires_at"], + name="agent_kv_jo_cleanup_f38edc_idx", + ), + ), + ] diff --git a/backend/agent_kv/migrations/0004_pg_periodic_tasks.py b/backend/agent_kv/migrations/0004_pg_periodic_tasks.py new file mode 100644 index 0000000000..c010a9b488 --- /dev/null +++ b/backend/agent_kv/migrations/0004_pg_periodic_tasks.py @@ -0,0 +1,116 @@ +"""Remove the Agent-KV maintenance periodics -- the CronJob owns the schedule. + +An earlier revision of this migration SEEDED two ``PgPeriodicTask`` rows +(``agent_kv.sweep`` at ``*/10``, ``agent_kv.ttl_cleanup`` at ``17 3 * * *``) on +the premise that nothing anywhere fired the two maintenance commands -- "no +``PgPeriodicTask``, no beat entry, no CronJob". **That premise was false when it +was written.** ``templates/backend/agent-kv-cronjobs.yaml`` in the cloud chart +already ran the same two management commands, on the same branch, added several +commits earlier. The seeding migration was written in response to a review +finding about the unscheduled sweep without noticing the finding had already +been answered elsewhere in the same change. + +Keeping both would have been worse than either alone: + +* **Two owners for one schedule.** The rows are seeded ``pg_owned: False``, and + the PG scheduler claims only ``WHERE pg_owned AND enabled`` + (``pg_queue/models.py``), so they are inert *today*. But the whole point of + that flag is that a rollout flips it -- and the moment it did, the ``*/10`` + rows and the CronJob's ``*/15`` ticks would both fire, colliding exactly at + ``:30`` and contending on the same job rows. The TTL cadences disagreed by + 24x (daily here, hourly in the chart), so whichever fired would be whichever + owner someone remembered. +* **Only one owner has operational bounds.** The CronJob carries + ``concurrencyPolicy: Forbid`` and ``activeDeadlineSeconds``; a PG periodic has + neither, so a hung sweep under the scheduler has nothing stopping it. +* **The feature only runs where the CronJob is.** This deployment routes + ``table`` only, and the ``agentic_table`` plugin ships in the cloud image. A + pure-OSS install cannot dispatch an Agent-KV job at all, so there is no + stranded row for a sweep to reap there. + +So the CronJob is the single owner, and this migration deletes the rows rather +than merely stopping at not creating them -- any environment that applied the +seeding revision (dev namespaces did) still has them, and leaving them to be +switched on by a future global ``pg_owned`` flip is precisely the hazard being +removed. + +**Self-hosted OSS operators who enable the Agent-KV API must schedule +``manage.py agent_kv_sweep`` and ``manage.py agent_kv_ttl_cleanup`` +themselves** (cron, a systemd timer, or their own ``PgPeriodicTask`` rows with +``pg_owned: True``). Without them a stranded job holds its concurrency slot +until the 6h Redis TTL, and ``AGENT_KV_RESULT_TTL_DAYS`` is advisory -- +staged documents and results stay in the bucket indefinitely, which is a +retention failure rather than a disk-usage one. ``docs/agent-kv-api.md`` §12 +carries this in the deploy checklist. + +Reverse re-creates the rows (still inert) so the migration is reversible and a +``migrate agent_kv 0003`` lands back on the previous tree's state exactly. +""" + +from django.db import migrations + +#: What the seeding revision wrote, kept ONLY so forward can delete exactly +#: those rows by name and reverse can restore them. Not a live schedule -- +#: the cadences that run are in the cloud chart's `backend.agentKvCronJobs`. +PG_PERIODIC_TASKS = [ + { + "name": "agent_kv_sweep", + "task_name": "agent_kv.sweep", + "queue": "scheduler", + "task_args": [], + "task_kwargs": {}, + "cron_string": "*/10 * * * *", + }, + { + "name": "agent_kv_ttl_cleanup", + "task_name": "agent_kv.ttl_cleanup", + "queue": "scheduler", + "task_args": [], + "task_kwargs": {}, + "cron_string": "17 3 * * *", + }, +] + + +def remove_pg_periodic_tasks(apps, schema_editor): + """Delete the seeded rows. Idempotent: a fresh install matches nothing.""" + # snake_case, not the usual `PgPeriodicTask = apps.get_model(...)` Django + # idiom: it is a local variable, and sonar's S117 reads the CamelCase name + # as a naming violation. The historical-model object is the same either way. + periodic_task = apps.get_model("pg_queue", "PgPeriodicTask") + periodic_task.objects.filter( + name__in=[spec["name"] for spec in PG_PERIODIC_TASKS] + ).delete() + + +def create_pg_periodic_tasks(apps, schema_editor): + """Reverse only. Restores the rows the seeding revision created, inert.""" + periodic_task = apps.get_model("pg_queue", "PgPeriodicTask") + for spec in PG_PERIODIC_TASKS: + periodic_task.objects.update_or_create( + name=spec["name"], + defaults={ + "task_name": spec["task_name"], + "queue": spec["queue"], + "task_args": spec["task_args"], + "task_kwargs": spec["task_kwargs"], + "cron_string": spec["cron_string"], + # Platform-wide, not per-org (spec §5.4): both internal + # endpoints sweep across every organization. + "org_id": "", + "enabled": True, + "pg_owned": False, + }, + ) + + +class Migration(migrations.Migration): + dependencies = [ + ("agent_kv", "0003_agentkvjob_cleanup_failed_at"), + # The table this operates on. + ("pg_queue", "0003_pgperiodictask"), + ] + + operations = [ + migrations.RunPython(remove_pg_periodic_tasks, create_pg_periodic_tasks), + ] diff --git a/backend/agent_kv/migrations/0005_agentkvjob_extractor_choices.py b/backend/agent_kv/migrations/0005_agentkvjob_extractor_choices.py new file mode 100644 index 0000000000..7e77da5595 --- /dev/null +++ b/backend/agent_kv/migrations/0005_agentkvjob_extractor_choices.py @@ -0,0 +1,32 @@ +"""Give `extractor` its `choices` and take away its `default`. + +Schema no-op on Postgres: Django keeps both `choices` and a field `default` in +Python only -- no CHECK constraint, no column DEFAULT -- so this `AlterField` +changes model state and emits no DDL that touches data. Existing rows keep +whatever they recorded, including the `kv` rows migration 0002 back-filled. + +The `default="kv"` this removes was correct when 0002 added the column (the +API accepted one extractor and it was always `kv`) and became a trap once +`table` existed: an omitted `extractor=` filed a table job as `kv`, which IS a +valid key in `STAGE_NAMES_BY_EXTRACTOR`, so the status endpoint returned the KV +stage list and dropped `table_extraction` silently. See the field comment on +the model. +""" + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("agent_kv", "0004_pg_periodic_tasks"), + ] + + operations = [ + migrations.AlterField( + model_name="agentkvjob", + name="extractor", + field=models.CharField( + choices=[("kv", "Kv"), ("table", "Table")], max_length=32 + ), + ), + ] diff --git a/backend/agent_kv/migrations/0006_agentkvjob_adapters.py b/backend/agent_kv/migrations/0006_agentkvjob_adapters.py new file mode 100644 index 0000000000..e3a1b0d702 --- /dev/null +++ b/backend/agent_kv/migrations/0006_agentkvjob_adapters.py @@ -0,0 +1,30 @@ +"""Record which platform adapters a job ran on. + +The `table` extractor resolves the CALLER's own adapters rather than operator +env credentials, so adapter choice became a per-request, cost-bearing decision. +Before this column the chosen ids reached `executor_params` and nowhere else -- +not the job row, not the status document, not `usage_summary` -- so the first +question in any billing dispute ("which model did job X use?") needed a join +against `usage_v2` on `run_id` that the API cannot perform and the customer +cannot see. + +Additive and nullable-by-default (`{}`), so existing rows are untouched: they +predate caller-supplied adapters and genuinely ran on env credentials, which +`{}` is the honest representation of. +""" + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("agent_kv", "0005_agentkvjob_extractor_choices"), + ] + + operations = [ + migrations.AddField( + model_name="agentkvjob", + name="adapters", + field=models.JSONField(blank=True, default=dict), + ), + ] diff --git a/backend/agent_kv/migrations/__init__.py b/backend/agent_kv/migrations/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/agent_kv/models.py b/backend/agent_kv/models.py new file mode 100644 index 0000000000..45df84d8dd --- /dev/null +++ b/backend/agent_kv/models.py @@ -0,0 +1,208 @@ +import uuid + +from account_v2.models import User +from django.db import models +from django.utils import timezone +from utils.models.base_model import BaseModel +from utils.models.organization_mixin import DefaultOrganizationMixin + +from agent_kv.constants import TABLE_EXTRACTOR_NAME, V1_EXTRACTOR_NAME + + +class AgentKVKey(DefaultOrganizationMixin, BaseModel): + id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) + name = models.CharField(max_length=128) + description = models.CharField(max_length=512, blank=True, default="") + key = models.UUIDField(default=uuid.uuid4, unique=True) + is_active = models.BooleanField(default=True) + created_by = models.ForeignKey( + User, + on_delete=models.SET_NULL, + null=True, + related_name="agent_kv_keys_created", + ) + modified_by = models.ForeignKey( + User, + on_delete=models.SET_NULL, + null=True, + related_name="+", + ) + + class Meta: + db_table = "agent_kv_key" + constraints = [ + models.UniqueConstraint( + fields=["name", "organization"], + name="unique_agent_kv_key_name_per_org", + ), + ] + + def __str__(self): + return f"{self.name} ({self.organization})" + + +class JobStatus(models.TextChoices): + PENDING = "PENDING" + DISPATCHED = "DISPATCHED" + RUNNING = "RUNNING" + COMPLETED = "COMPLETED" + FAILED = "FAILED" + CANCELLED = "CANCELLED" + + +class JobExtractor(models.TextChoices): + """The extractor names a job row may record. + + `status` has had `choices` since 0001 and `extractor` was the one + stringly-typed field without them -- so nothing but a reader's memory + connected the column to `EXTRACTOR_ROUTES`. + + KV is listed even though this deployment refuses it: the column records + which extractor RAN, and rows written before the carve-out legitimately + say `kv`. Routability is `EXTRACTOR_ROUTES`' job, and these two sets are + deliberately not the same thing -- see `test_table_extractor_routing.py`. + """ + + KV = V1_EXTRACTOR_NAME + TABLE = TABLE_EXTRACTOR_NAME + + +class AgentKVJob(DefaultOrganizationMixin, BaseModel): + TERMINAL = frozenset({JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED}) + + id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) + api_key = models.ForeignKey( + AgentKVKey, + on_delete=models.SET_NULL, + null=True, + related_name="jobs", + ) + task_id = models.UUIDField(null=True, blank=True) + # Which extractor this job ran. v1 dispatches exactly one per job, but + # WHICH one is now a choice, and status/result key their payloads by it -- + # without this column a table job's output would be filed under `kv`. + # + # NO DEFAULT, deliberately. Migration 0002 added this column with + # `default="kv"`, which was the historical truth at the time (the API + # accepted exactly one extractor and it was always `kv`) and is now a + # trap: a creation path that omits `extractor=` files a TABLE job as + # `kv`, and `kv` is a valid key in `STAGE_NAMES_BY_EXTRACTOR`, so + # `_status_document` hands back the KV stage list and silently drops + # `table_extraction` from every status response. The job runs, the caller + # is billed, and the stages array is empty with no warning logged -- + # because nothing is wrong as far as the filter can tell. + # + # With no default, an omission is loud instead: `""` matches no route, so + # `dispatch_job` raises and the job terminalizes as FAILED with an error + # the caller can see, and `_status_document` logs the unknown-extractor + # warning. Which extractor ran is a fact about the job, not something with + # a sensible default. + extractor = models.CharField(max_length=32, choices=JobExtractor.choices) + status = models.CharField( + max_length=16, + choices=JobStatus.choices, + default=JobStatus.PENDING, + ) + stage = models.CharField(max_length=32, blank=True, default="") + stages = models.JSONField(default=dict, blank=True) + pages_total = models.IntegerField(null=True, blank=True) + #: Which platform adapters this run was dispatched with, by role -- + #: ``{"llm": "", "lite_llm": "", "x2text": ""}``. + #: + #: Recorded because adapter choice is now a per-request, caller-controlled, + #: COST-BEARING decision (the `table` extractor runs on the submitter's own + #: adapters, not operator env credentials). Without it the ids reached + #: `executor_params` and nothing else: neither the status document nor + #: `usage_summary` reported them, so "which model did job X use?" -- the + #: first question in any billing dispute -- could only be answered by + #: joining `usage_v2` on `run_id`, which the API cannot do and a customer + #: cannot see at all. + #: + #: Empty for an env-configured extractor (`kv`), which is why it is a dict + #: with a `{}` default rather than three nullable columns: the two + #: credential models coexist, one per extractor. + #: + #: Ids only. Deliberately never `adapter_metadata` -- that is where the + #: provider credentials live, and this column is returned to the caller. + adapters = models.JSONField(default=dict, blank=True) + input_ref = models.CharField(max_length=512, blank=True, default="") + result_ref = models.CharField(max_length=512, blank=True, default="") + usage_summary = models.JSONField(null=True, blank=True) + error = models.TextField(blank=True, default="") + dispatched_at = models.DateTimeField(null=True, blank=True) + completed_at = models.DateTimeField(null=True, blank=True) + expires_at = models.DateTimeField(null=True, blank=True) + # When TTL cleanup last failed to delete one of this job's files. NULL means + # "never attempted, or last attempt succeeded", and it is what separates + # the two lanes run_ttl_cleanup processes: NOT NULL rows are retries, which + # get a reserved slice of each batch, and NULL rows are new expirations, + # which get the rest. Capping both is what stops either side starving the + # other. See run_ttl_cleanup. + cleanup_failed_at = models.DateTimeField(null=True, blank=True) + tags = models.JSONField(default=list, blank=True) + custom_data = models.JSONField(null=True, blank=True) + webhook_url = models.URLField(max_length=1024, blank=True, default="") + + class Meta: + db_table = "agent_kv_job" + indexes = [ + models.Index(fields=["organization", "status"]), + models.Index(fields=["expires_at"]), + # Serves BOTH of run_ttl_cleanup's lanes, each of which sorts by a + # single column ascending behind a predicate on this index's + # leading column: + # retries: WHERE cleanup_failed_at IS NOT NULL ORDER BY cleanup_failed_at + # fresh: WHERE cleanup_failed_at IS NULL ORDER BY expires_at + # Deliberately a plain ascending index. The first version of that + # query asked for `cleanup_failed_at ASC NULLS FIRST`, which a + # btree index cannot serve (btree is NULLS LAST ascending), so + # Postgres sorted every matching expired row before applying the + # 500-row limit -- work that grew with the backlog. Splitting the + # query removed the NULLS FIRST rather than adding a second index + # with a non-default null order. + models.Index(fields=["cleanup_failed_at", "expires_at"]), + ] + + @classmethod + def mark_terminal( + cls, + job_id, + organization_id, + new_status, + *, + error="", + result_ref="", + usage_summary=None, + ) -> bool: + """The ONLY way to reach a terminal state (spec §5.4 write guard). + + Guarded UPDATE: at-least-once callbacks, cancel, and the sweep can all + race; whoever lands first wins and everyone else no-ops. + + The `.exclude(status__in=TERMINAL)` below guards the ROW, not the + argument -- so without the check that opens this method, + `mark_terminal(..., JobStatus.RUNNING)` would stamp + `status=RUNNING, completed_at=now()`: a row that reads as finished to + every TTL/sweep query that keys off `completed_at`, is invisible to + the terminal guard, and can never be terminalized again by anything + that trusts `completed_at`. No caller does this today; the method is + named for the invariant, so it enforces it rather than documenting it. + """ + if new_status not in cls.TERMINAL: + raise ValueError( + f"mark_terminal called with non-terminal status {new_status!r}; " + f"expected one of {sorted(cls.TERMINAL)}" + ) + fields = {"status": new_status, "completed_at": timezone.now()} + if error: + fields["error"] = error + if result_ref: + fields["result_ref"] = result_ref + if usage_summary is not None: + fields["usage_summary"] = usage_summary + updated = ( + cls.objects.filter(id=job_id, organization_id=organization_id) + .exclude(status__in=list(cls.TERMINAL)) + .update(**fields) + ) + return updated == 1 diff --git a/backend/agent_kv/permissions.py b/backend/agent_kv/permissions.py new file mode 100644 index 0000000000..d9ce287517 --- /dev/null +++ b/backend/agent_kv/permissions.py @@ -0,0 +1,8 @@ +from platform_api.permissions import IsOrganizationAdmin as PlatformIsOrganizationAdmin + + +class IsOrganizationAdmin(PlatformIsOrganizationAdmin): + message = "Only organization admins can manage Agent-KV API keys." + + +__all__ = ["IsOrganizationAdmin"] diff --git a/backend/agent_kv/rate_limiter.py b/backend/agent_kv/rate_limiter.py new file mode 100644 index 0000000000..0102e2c1f7 --- /dev/null +++ b/backend/agent_kv/rate_limiter.py @@ -0,0 +1,113 @@ +import logging +import time + +from django.conf import settings +from django_redis import get_redis_connection + +logger = logging.getLogger(__name__) + + +def _limiter_failure_allows_request() -> bool: + """Whether a request proceeds when the limiter backend is unreachable. + + Fails CLOSED unless `AGENT_KV_LIMITER_FAIL_OPEN` is explicitly set. Both + limiters previously returned True on any Redis exception, which is the one + choice that cannot be observed from the outside: a Sentinel failover + silently removed the concurrency ceiling AND the per-key rate ceiling at the + same time, and the API went on accepting billable LLM work as if both still + held. Closed is the safer default for a paid, concurrency-capped API -- a + 429 is recoverable by the caller, an unbounded fan-out is not. + + The waiver remains available because availability-over-accounting is a + legitimate operational choice, but it now has to be made deliberately, in + config, where it is visible -- rather than being the implicit behaviour of + an `except` block. + """ + return bool(getattr(settings, "AGENT_KV_LIMITER_FAIL_OPEN", False)) + + +#: How long a held slot survives in Redis without being released. Public +#: because the sweep needs it: past this age there is nothing left to release, +#: which is what bounds its cancelled-job scan. +SLOT_TTL_SECONDS = 6 * 3600 +_SLOT_TTL_SECONDS = SLOT_TTL_SECONDS + + +def _redis(): + # Same handle acquisition as api_v2.rate_limiter (`redis_cache = + # get_redis_connection("default")`); api_v2 has no reusable helper to + # call into, so the construction line is copied here rather than + # modifying api_v2. + return get_redis_connection("default") + + +class AgentKVConcurrencyLimiter: + @staticmethod + def _key(organization_id: str) -> str: + return f"agent_kv:inflight:{organization_id}" + + # Trim stale slots, count, and claim a slot in ONE atomic server-side + # step. A client-side ``ZCARD`` followed by ``ZADD`` is a check-then-act + # race: N simultaneous submits all observe ``count < limit`` and all get + # accepted (caught live in the Task 13b run: 6 concurrent submits against + # a limit of 5 produced six 202s). Redis runs a script atomically, so at + # most ``limit`` members can ever be added. + _ACQUIRE_SCRIPT = """ +local key = KEYS[1] +redis.call('ZREMRANGEBYSCORE', key, 0, ARGV[3]) +if redis.call('ZCARD', key) >= tonumber(ARGV[4]) then + return 0 +end +redis.call('ZADD', key, ARGV[2], ARGV[1]) +redis.call('EXPIRE', key, ARGV[5]) +return 1 +""" + + @classmethod + def check_and_acquire(cls, organization_id: str, job_id: str) -> bool: + try: + r = _redis() + now = time.time() + key = cls._key(organization_id) + acquired = r.eval( + cls._ACQUIRE_SCRIPT, + 1, + key, + job_id, + now, + now - _SLOT_TTL_SECONDS, + settings.AGENT_KV_CONCURRENT_LIMIT, + _SLOT_TTL_SECONDS, + ) + return bool(int(acquired)) + except Exception: + # Fail CLOSED by default. Both this ceiling and the per-key rate + # ceiling used to `return True` on any Redis error, so a single + # Sentinel failover or pool exhaustion removed BOTH at once while + # the API kept returning 202s for billable LLM work -- with nothing + # but a per-request `logger.warning` to show for it. + # + # `logger.exception` (not warning): the limiter being unavailable is + # an error, and the old level is what let this sit unnoticed. + logger.exception("agent-kv concurrency limiter unavailable") + return _limiter_failure_allows_request() + + @classmethod + def release(cls, organization_id: str, job_id: str) -> None: + try: + _redis().zrem(cls._key(organization_id), job_id) + except Exception: + logger.warning("agent-kv slot release failed", exc_info=True) + + +def check_key_rate(key_id: str) -> bool: + try: + r = _redis() + window = int(time.time() // 60) + key = f"agent_kv:rate:{key_id}:{window}" + count = r.incr(key) + r.expire(key, 120) + return count <= settings.AGENT_KV_KEY_RATE_LIMIT_PER_MINUTE + except Exception: + logger.exception("agent-kv key rate limiter unavailable") + return _limiter_failure_allows_request() diff --git a/backend/agent_kv/serializers.py b/backend/agent_kv/serializers.py new file mode 100644 index 0000000000..081aeabb6b --- /dev/null +++ b/backend/agent_kv/serializers.py @@ -0,0 +1,15 @@ +from rest_framework import serializers + +from agent_kv.models import AgentKVKey + + +class AgentKVKeySerializer(serializers.ModelSerializer): + class Meta: + model = AgentKVKey + fields = ["id", "name", "description", "key", "is_active", "created_at"] + + +class AgentKVKeyWriteSerializer(serializers.ModelSerializer): + class Meta: + model = AgentKVKey + fields = ["name", "description", "is_active"] diff --git a/backend/agent_kv/storage.py b/backend/agent_kv/storage.py new file mode 100644 index 0000000000..b19adce545 --- /dev/null +++ b/backend/agent_kv/storage.py @@ -0,0 +1,165 @@ +"""Object-store staging/results for agent-kv (spec §5.4, §6.4). + +Paths are the contract: {AGENT_KV_STORAGE_DIR_PREFIX}/{org_id}/{job_id}/input{ext} +and .../result.json. No document bytes ever ride the broker. + +The prefix is bucket-rooted (default ``unstract/agent_kv``) like +``WORKFLOW_EXECUTION_DIR_PREFIX``/``API_EXECUTION_DIR_PREFIX``: s3fs/gcsfs read +the first path segment as the *bucket*, so a bucket-less root (the old +``org/{org_id}/...``) makes every write fail with ``NoSuchBucket``. The cloud +executor keys its OCR cache under the same root +(``{prefix}/{org_id}/cache/...``), so both sides must be configured alike. +""" + +import logging +import os +import uuid + +from django.conf import settings + +from unstract.filesystem import FileStorageType, FileSystem + +logger = logging.getLogger(__name__) + + +def _fs(): + return FileSystem(FileStorageType.AGENT_KV).get_file_storage() + + +def _base(org_id: str, job_id: str) -> str: + return f"{settings.AGENT_KV_STORAGE_DIR_PREFIX}/{org_id}/{job_id}" + + +def stage_input(org_id: str, job_id: str, uploaded_file) -> str: + ext = os.path.splitext(uploaded_file.name or "")[1].lower() or ".bin" + ref = f"{_base(org_id, job_id)}/input{ext}" + data = b"".join(uploaded_file.chunks()) + _fs().write(path=ref, mode="wb", data=data) + return ref + + +def write_result( + org_id: str, job_id: str, result: dict, *, nonce: str | None = None +) -> str: + """Write a job result and return its object-store ref. + + The path is UNIQUE per call (``.../result-.json``, nonce defaulting + to a fresh uuid4 hex) rather than a deterministic ``.../result.json``. + ``FinalizeView`` writes the result BEFORE it attempts the terminal-state + guard; on a concurrent duplicate-SUCCESS finalize, the guard loser must + clean up ONLY the orphan it just wrote -- with a shared deterministic path + it would instead delete the exact file the winning row's ``result_ref`` + points at, 500-ing the completed job's result endpoint / losing data + (pre-Greptile critical #3). A unique ref per attempt makes the loser's + ``delete_result_file`` target its own file and nothing else. + + Backward-tolerant: ``read_result``/``delete_job_files``/TTL cleanup all use + the ref STORED on the job row, so any previously written ``result.json`` + ref remains readable and removable unchanged. + """ + ref = f"{_base(org_id, job_id)}/result-{nonce or uuid.uuid4().hex}.json" + _fs().json_dump(path=ref, data=result) + return ref + + +def read_result(result_ref: str) -> dict: + return _fs().json_load(path=result_ref) + + +def delete_job_files(job) -> list[str]: + """Delete the job's staged input and result, returning the ref FIELD names + that are now confirmed to point at nothing. + + The return value exists so a caller can blank only the refs whose files are + actually gone. Blanking unconditionally -- what both callers used to do -- + loses the only handle to a file whose delete FAILED: TTL cleanup selects + candidates by `input_ref > "" OR result_ref > ""`, so a blanked row drops + out of the candidate set permanently and the object sits in the bucket + forever, unreferenced and unbillable-to-anyone (Greptile review #3). + + A ref is reported clear when it was already empty (nothing was ever staged) + or when the delete succeeded. ``FileNotFoundError`` counts as success: the + goal is "the file is not there", and a ref whose object is already gone has + nothing left to retry. Any other exception leaves the ref OUT of the + returned list, so the caller keeps it and the next pass retries. + """ + fh = _fs() + cleared: list[str] = [] + for field in ("input_ref", "result_ref"): + ref = getattr(job, field, "") or "" + if not ref: + cleared.append(field) + continue + try: + fh.rm(path=ref) + except FileNotFoundError: + cleared.append(field) + except Exception: + # Left out of `cleared` deliberately -- the ref is the retry handle. + logger.warning( + "agent-kv cleanup: could not remove %s (ref kept for retry)", + ref, + exc_info=True, + ) + else: + cleared.append(field) + return cleared + + +def delete_result_file(result_ref: str) -> None: + """Best-effort delete of a result file that was just written but never + got a ref persisted onto the job row. + + ``FinalizeView`` writes the result *before* attempting the terminal- + state guard (spec §5.4); if the guard loses the race -- a concurrent + cancel or a duplicate finalize won instead -- the file this call just + wrote is orphaned: nothing anywhere points at it, so it would otherwise + sit in object storage forever, past even TTL cleanup (which only acts on + a job's *stored* ``result_ref``). Tolerant of the file being missing + already, mirroring ``delete_input``/``delete_job_files``. + """ + if not result_ref: + return + try: + _fs().rm(path=result_ref) + except Exception: + logger.warning("agent-kv cleanup: could not remove orphaned %s", result_ref) + + +def delete_input(job) -> bool: + """Delete only the staged input file (spec D10: "uploaded document + deleted on job completion"), returning whether ``input_ref`` is now + confirmed to point at nothing. + + Deliberately narrower than ``delete_job_files``: it never touches + ``result_ref``/the result file. ``FinalizeView`` calls this right after + a job is fully terminalized, in the same request that may have just + written the result -- that file must be left completely alone. + + Same confirmed-clear contract as ``delete_job_files``, and for the same + reason: this swallowed every exception while its caller blanked + ``input_ref`` unconditionally, so a transient object-store error orphaned + the customer's uploaded document permanently -- TTL cleanup only selects + rows that still carry a non-blank ref, so a blanked row can never be + reconsidered. ``delete_job_files`` and its two callers were fixed for this + in the Greptile round; THIS third site was missed, because the fix was + applied by following the one function's callers rather than by grepping for + every site that blanks a ref. + + True when there was nothing to delete or the file is gone (including + already-missing); False when it is still there and the ref must be kept. + """ + if not job.input_ref: + return True + try: + _fs().rm(path=job.input_ref) + except FileNotFoundError: + return True + except Exception: + logger.warning( + "agent-kv cleanup: could not remove %s (ref kept for retry)", + job.input_ref, + exc_info=True, + ) + return False + return True diff --git a/backend/agent_kv/tests/__init__.py b/backend/agent_kv/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/agent_kv/tests/_factories.py b/backend/agent_kv/tests/_factories.py new file mode 100644 index 0000000000..2596355e8f --- /dev/null +++ b/backend/agent_kv/tests/_factories.py @@ -0,0 +1,26 @@ +"""Shared test object factories for the Agent-KV suites.""" + +from account_v2.models import Organization + +from agent_kv.models import AgentKVKey + + +def kv_key(*, slug: str = "acme-slug", org_pk: int = 1, **overrides) -> AgentKVKey: + """A USABLE AgentKVKey — one that carries an organization. + + `AgentKVKey.organization` is nullable on the model, but a key without one is + not usable: every downstream consumer is org-scoped (the subscription gate, + the concurrency limiter, the storage prefix, every job lookup), so + `AgentKVKeyValidator` refuses it. Tests that build a bare + `AgentKVKey(name=...)` are therefore constructing a key the auth layer would + reject, and any behaviour they assert past that point is unreachable in + production. + + `org_pk` and `slug` are deliberately DIFFERENT kinds of value: + `Subscription.organization_id` is a CharField holding the slug, while the + key's own `organization_id` is the FK primary key. Keeping them distinct is + what lets a test catch code that reaches for the wrong one. + """ + key = AgentKVKey(name="k", is_active=True, **overrides) + key.organization = Organization(id=org_pk, organization_id=slug) + return key diff --git a/backend/agent_kv/tests/fixtures/two_page.pdf b/backend/agent_kv/tests/fixtures/two_page.pdf new file mode 100644 index 0000000000..92ee5755ea Binary files /dev/null and b/backend/agent_kv/tests/fixtures/two_page.pdf differ diff --git a/backend/agent_kv/tests/test_adapter_scoping.py b/backend/agent_kv/tests/test_adapter_scoping.py new file mode 100644 index 0000000000..a8984a8cad --- /dev/null +++ b/backend/agent_kv/tests/test_adapter_scoping.py @@ -0,0 +1,323 @@ +"""The adapter lookup must work with NO request-local organization context. + +This is the test whose absence let a P1 through. `_resolved_adapters` was +written with `AdapterInstance.objects`, by analogy to the `organization_id` +filter in `_get_job` -- but the two models have different default managers: + +* `AgentKVJob` uses `BaseModelManager`, which does not auto-filter, so its + explicit `organization_id=` is the only filter. +* `AdapterInstance` sets `AdapterInstanceModelManager`, which inherits + `DefaultOrganizationManagerMixin.get_queryset()` and filters EVERY query by + `UserContext.get_organization()` -- a request-local thread-local set by the + tenant middleware. + +`/agent-kv/` is deliberately whitelisted past that middleware (the org comes +from the Bearer key, not the URL), so the thread-local is unset and the ambient +filter became `organization=None`. Every valid submit was refused with "no such +adapter in this organization". + +Every test in `test_submit_view.py` mocked the lookup, so none of them could +see it. These run against the real manager with the real thread-local unset -- +which is the only arrangement that reproduces the production path. + +DB-backed, so `TestCase` and the integration tier (`backend/conftest.py` +auto-marks on that basis; no manual marker). +""" + +import os +import uuid +from unittest import mock + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from django.test import TestCase # noqa: E402 + +from account_v2.models import Organization # noqa: E402 +from adapter_processor_v2.models import AdapterInstance # noqa: E402 +from rest_framework.exceptions import ValidationError # noqa: E402 + +from agent_kv.execution_views import ( # noqa: E402 + _lookup_adapter, + _resolved_adapters, +) +from utils.user_context import UserContext # noqa: E402 + + +class TestAdapterLookupWithoutRequestContext(TestCase): + def setUp(self): + self.org = Organization.objects.create( + display_name="Acme", organization_id="org_acme", name="acme" + ) + self.other = Organization.objects.create( + display_name="Other", organization_id="org_other", name="other" + ) + self.adapter = AdapterInstance.objects.create( + adapter_name="acme-llm", + adapter_id="anthropic|abc", + adapter_type="LLM", + adapter_metadata={}, + organization=self.org, + ) + # The state the whitelisted route actually runs in: no tenant + # middleware has run, so there is no ambient organization. + UserContext.set_organization_identifier(None) + + def test_the_lookup_finds_the_org_s_own_adapter_with_no_user_context(self): + """The regression. Through `.objects` this returned None and the submit + was a 400 -- for a correctly-configured caller naming their own adapter. + """ + self.assertIsNone( + UserContext.get_organization(), + "precondition: this route has no ambient organization", + ) + + found = _lookup_adapter(self.adapter.id, self.org.id) + + self.assertIsNotNone( + found, + "the caller's own adapter must resolve without a request-local " + "organization; `.objects` filters by one and returns nothing here", + ) + self.assertEqual(found.id, self.adapter.id) + self.assertEqual(found.adapter_type, "LLM") + + def test_the_org_filter_still_excludes_another_org_s_adapter(self): + """Bypassing the AMBIENT filter must not bypass the explicit one. + + `_base_manager` is unfiltered, so the `organization_id` argument is the + only tenant boundary left -- if it stopped working, the fix for the + regression above would have opened the cross-tenant hole it exists to + close. + """ + self.assertIsNone(_lookup_adapter(self.adapter.id, self.other.id)) + + def test_an_unknown_id_resolves_to_none(self): + self.assertIsNone(_lookup_adapter(uuid.uuid4(), self.org.id)) + + def test_the_ambient_manager_really_does_hide_it(self): + """Pins the cause, so the comment on `_lookup_adapter` stays truthful. + + If `AdapterInstance.objects` ever stops auto-filtering by the + thread-local, this fails and `_base_manager` is no longer necessary -- + at which point the explanation should be revisited rather than left + asserting something that is no longer true. + """ + via_objects = AdapterInstance.objects.filter( + id=self.adapter.id, organization_id=self.org.id + ).first() + + self.assertIsNone( + via_objects, + "`.objects` is expected to hide the adapter here -- that is the " + "whole reason `_lookup_adapter` uses `_base_manager`", + ) + + +class TestResolvedAdaptersAgainstRealRows(TestCase): + """`_resolved_adapters` end-to-end against the database, not against mocks. + + `test_submit_view.py` patches `ev._lookup_adapter`, so every property the + gate reads off an adapter -- `adapter_type`, `is_usable`, `is_available` -- + is only ever compared to a value the stand-in chose for itself. That is + enough to test the gate's branching and nothing about whether the branches + match the real columns. These rows get their flags from the model's own + defaults. + + It is also the only place the refusal-oracle property can honestly be + tested: "not yours" and "no such adapter" have to come from two GENUINELY + different database states, which a single mocked lookup cannot produce. + """ + + def setUp(self): + self.org = Organization.objects.create( + display_name="Acme", organization_id="org_acme", name="acme" + ) + self.other = Organization.objects.create( + display_name="Other", organization_id="org_other", name="other" + ) + self.key = mock.Mock(organization_id=self.org.id) + UserContext.set_organization_identifier(None) + + def _adapter(self, *, org=None, adapter_type="LLM", **kwargs): + return AdapterInstance.objects.create( + adapter_name=f"a-{uuid.uuid4().hex[:8]}", + adapter_id="anthropic|abc", + adapter_type=adapter_type, + adapter_metadata={}, + organization=org or self.org, + **kwargs, + ) + + def _entry(self, **roles): + return {"name": "table", "adapters": {r: str(v) for r, v in roles.items()}} + + def _full_entry(self, **overrides): + roles = { + "llm": self._adapter().id, + "lite_llm": self._adapter().id, + "x2text": self._adapter(adapter_type="X2TEXT").id, + } + roles.update(overrides) + return self._entry(**roles) + + def test_a_fully_valid_set_of_the_org_s_own_adapters_resolves(self): + """The happy path, with every flag coming from the DB's own defaults. + + `is_usable` and `is_available` both default True at the model, so this + also pins that the two new refusals do not reject ordinary adapters + nobody has touched. + """ + entry = self._full_entry() + + self.assertEqual(_resolved_adapters(entry, self.key), entry["adapters"]) + + def test_a_wrong_type_is_refused_against_a_real_row(self): + """An X2TEXT id in the `llm` slot. + + Previously asserted only against a mock that was told to report + `adapter_type="X2TEXT"`; this reads the real column, so the comparison + against `AdapterTypes.LLM.value` is pinned to the value the DB stores. + """ + entry = self._full_entry(llm=self._adapter(adapter_type="X2TEXT").id) + + with self.assertRaises(ValidationError) as caught: + _resolved_adapters(entry, self.key) + + self.assertIn("expected 'LLM'", str(caught.exception)) + + def test_an_exhausted_trial_adapter_is_refused(self): + """A1. `is_usable=False` is how billing cuts off a frictionless trial. + + The platform service returns the credentials regardless, so this gate + is the only thing that stops an operator-funded extraction here. + """ + entry = self._full_entry(llm=self._adapter(is_usable=False).id) + + with self.assertRaises(ValidationError) as caught: + _resolved_adapters(entry, self.key) + + self.assertIn("exhausted", str(caught.exception)) + + def test_a_deprecated_adapter_is_refused(self): + entry = self._full_entry(llm=self._adapter(is_available=False).id) + + with self.assertRaises(ValidationError) as caught: + _resolved_adapters(entry, self.key) + + self.assertIn("deprecated", str(caught.exception)) + + def test_another_org_s_adapter_and_a_nonexistent_one_refuse_identically(self): + """The oracle property, from two genuinely different DB states. + + If these two messages ever diverge, `POST /agent-kv/` becomes an oracle + for which adapter UUIDs are real in OTHER organizations: a caller + brute-forcing ids learns "exists but not yours" from "does not exist". + + The mocked version of this test could not fail -- one patched lookup + returns None for both cases, so it compared one branch's message to + itself. Here one id is a real row owned by `self.other` and the other + is a UUID no row has. + """ + foreign = self._adapter(org=self.other).id + nowhere = uuid.uuid4() + + messages = [] + for adapter_id in (foreign, nowhere): + with self.assertRaises(ValidationError) as caught: + _resolved_adapters(self._full_entry(llm=adapter_id), self.key) + messages.append(str(caught.exception)) + + self.assertEqual( + messages[0], + messages[1], + "a foreign adapter and a nonexistent one must be indistinguishable; " + f"got {messages[0]!r} vs {messages[1]!r}", + ) + + def test_no_adapters_requested_is_a_no_op(self): + """The `kv` shape: env-configured extractors send `{}` and must pass.""" + self.assertEqual(_resolved_adapters({"name": "kv", "adapters": {}}, self.key), {}) + + +class TestTheJobRowRecordsWhatItSpent(TestCase): + """The adapters a run was dispatched with must be recoverable AFTERWARDS. + + Adapter choice is per-request, caller-controlled and cost-bearing on the + `table` path. Before the `adapters` column the chosen ids reached + `executor_params` and nothing else -- not the job row, not the status + document, not `usage_summary` -- so "which model did job X use?", the first + question in any billing dispute, needed a `usage_v2` join on `run_id` that + the API cannot perform and the customer cannot see. + + DB-backed because the point is that the value SURVIVES the request. + """ + + def setUp(self): + self.org = Organization.objects.create( + display_name="Acme", organization_id="org_acme", name="acme" + ) + UserContext.set_organization_identifier(None) + + def _job(self, **kwargs): + from agent_kv.models import AgentKVJob, JobExtractor + + job = AgentKVJob( + organization_id=self.org.id, + extractor=JobExtractor.TABLE, + **kwargs, + ) + job.save() + return job + + def test_the_adapters_round_trip_through_the_database(self): + adapters = { + "llm": str(uuid.uuid4()), + "lite_llm": str(uuid.uuid4()), + "x2text": str(uuid.uuid4()), + } + + job = self._job(adapters=adapters) + job.refresh_from_db() + + assert job.adapters == adapters + + def test_the_status_document_reports_them(self): + from agent_kv.execution_views import _status_document + + adapters = {"llm": str(uuid.uuid4())} + doc = _status_document(self._job(adapters=adapters)) + + assert doc["adapters"] == adapters, doc + + def test_a_job_with_no_adapters_reports_an_empty_dict_not_null(self): + """`kv` is env-configured and names none, and rows predating the column + genuinely ran on operator credentials -- `{}` is the honest value for + both, and a stable type is easier for a client than `null | object`. + """ + from agent_kv.execution_views import _status_document + + job = self._job() + + assert job.adapters == {} + assert _status_document(job)["adapters"] == {} + + def test_no_adapter_metadata_is_ever_stored_on_the_job(self): + """The column is returned to the caller, so it must carry ids only. + + `AdapterInstance.adapter_metadata` holds the PROVIDER CREDENTIALS. A + future change that widened this field to "the adapter" rather than "its + id" would publish those to whoever submitted the job. + """ + adapters = {"llm": str(uuid.uuid4()), "x2text": str(uuid.uuid4())} + job = self._job(adapters=adapters) + job.refresh_from_db() + + for role, value in job.adapters.items(): + assert isinstance(value, str), (role, value) + # A bare UUID string, not a serialized adapter. + uuid.UUID(value) diff --git a/backend/agent_kv/tests/test_auth.py b/backend/agent_kv/tests/test_auth.py new file mode 100644 index 0000000000..6006497646 --- /dev/null +++ b/backend/agent_kv/tests/test_auth.py @@ -0,0 +1,130 @@ +import os +import uuid +from unittest import mock + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +import pytest # noqa: E402 +from api_v2.exceptions import Forbidden # noqa: E402 + +from agent_kv.key_validator import AgentKVKeyValidator # noqa: E402 +from agent_kv.models import AgentKVKey # noqa: E402 + + +def _request(auth=None): + r = mock.Mock() + r.headers = {"Authorization": auth} if auth else {} + return r + + +def _wrapped(): + @AgentKVKeyValidator.validate_api_key + def view(self, request, *args, **kwargs): + return kwargs["agent_kv_key"] + + return view + + +def test_missing_key_is_forbidden(): + view, instance, request = _wrapped(), mock.Mock(), _request() + with pytest.raises(Forbidden): + view(instance, request) + + +@mock.patch.object(AgentKVKey, "objects") +def test_unknown_key_is_forbidden(m_objects): + m_objects.get.side_effect = AgentKVKey.DoesNotExist + view, instance = _wrapped(), mock.Mock() + request = _request(f"Bearer {uuid.uuid4()}") + with pytest.raises(Forbidden): + view(instance, request) + + +@mock.patch.object(AgentKVKey, "objects") +def test_valid_key_injected_into_kwargs(m_objects): + from account_v2.models import Organization # noqa: PLC0415 + + key_obj = AgentKVKey(name="k", is_active=True) + # A usable key always has an organization: every downstream use is + # org-scoped, and the validator refuses one without (see the org-less test + # at the bottom of this module). + key_obj.organization = Organization(id=1, organization_id="acme-slug") + m_objects.get.return_value = key_obj + out = _wrapped()(mock.Mock(), _request(f"Bearer {uuid.uuid4()}")) + assert out is key_obj + + +@mock.patch.object(AgentKVKey, "objects") +def test_non_uuid_key_is_forbidden_without_db_hit(m_objects): + view, instance, request = _wrapped(), mock.Mock(), _request("Bearer not-a-uuid") + with pytest.raises(Forbidden): + view(instance, request) + assert not m_objects.get.called + + +def test_prefix_is_whitelisted(): + from django.conf import settings + + assert f"/{settings.AGENT_KV_PATH_PREFIX}" in settings.WHITELISTED_PATHS + + +def test_public_url_wiring_and_decorator_enforced(): + """Regression pin for the two failure modes the decorator-only tests above + can't catch (they exercise the hand-built `_wrapped()` helper, not the + real view or URLconf): + + (a) dropping `include("agent_kv.execution_urls")` from base_urls.py, and + (b) dropping `@AgentKVKeyValidator.validate_api_key` from SubmitView.post. + + This drives the real URL resolution and class-based view dispatch path + instead: (a) breaks `resolve(...)`, and (b) makes the stub answer with + 501 for every request instead of 403 for an unauthenticated one. + """ + from django.conf import settings + from django.urls import resolve + from rest_framework.test import APIRequestFactory + + from agent_kv.execution_views import SubmitView + + resolved = resolve(f"/{settings.AGENT_KV_PATH_PREFIX}/") + assert resolved.func.cls is SubmitView + + request = APIRequestFactory().post(f"/{settings.AGENT_KV_PATH_PREFIX}/") + response = SubmitView.as_view()(request) + assert response.status_code == 403 + + +def test_key_without_an_organization_is_refused(): + """`organization` is nullable on AgentKVKey, but every downstream use of a + key is org-scoped -- the subscription gate, the concurrency limiter, the + storage prefix, every job lookup. Before this guard the first of those to + touch `key.organization` raised AttributeError and the caller got a 500; + found by running a real submit against a stack whose key had no org. + """ + from unittest import mock # noqa: PLC0415 + + from api_v2.exceptions import Forbidden # noqa: PLC0415 + + from agent_kv.key_validator import AgentKVKeyValidator # noqa: PLC0415 + from agent_kv.models import AgentKVKey # noqa: PLC0415 + + orgless = AgentKVKey(name="k", is_active=True) + assert orgless.organization_id is None + with mock.patch.object(AgentKVKey, "objects") as m_objects: + m_objects.get.return_value = orgless + try: + AgentKVKeyValidator.validate_and_process( + object(), + object(), + lambda *a, **k: "reached the view", + "123e4567-e89b-12d3-a456-426614174001", + ) + except Forbidden: + pass + else: + raise AssertionError("an org-less key must not reach the view") diff --git a/backend/agent_kv/tests/test_cross_repo_stage_name.py b/backend/agent_kv/tests/test_cross_repo_stage_name.py new file mode 100644 index 0000000000..fe69212c1b --- /dev/null +++ b/backend/agent_kv/tests/test_cross_repo_stage_name.py @@ -0,0 +1,95 @@ +"""The one constant that drifts silently across the two repos. + +``TABLE_STAGE_NAMES[0]`` here is the stage name ``_status_document`` will +*show*. ``STAGE_TABLE_EXTRACTION`` in the cloud repo's +``workers/plugins/agentic_table/src/api_binding.py`` is the name the executor +actually *sends*. They must be equal, and nothing enforces it: +``StageReportView`` persists whatever name arrives, and ``_status_document`` +then filters a job's recorded stages through the list below. + +So on drift the job completes normally, bills normally, and every status +response reports an empty ``stages`` array. No test on either side can catch +that alone -- each repo is internally consistent -- and from a client's seat it +reads as "the API is broken". This file is one half of the lock; the cloud half +is ``workers/plugins/agentic_table/tests/test_cross_repo_stage_name.py``. + +A literal, not an import: the two repos are separate checkouts and only meet in +the merged tree, so an import would make this test pass vacuously wherever the +other side is absent -- exactly the environments where drift is introduced. +""" + +from agent_kv.constants import TABLE_EXTRACTOR_NAME, TABLE_STAGE_NAMES + +#: Must equal ``STAGE_TABLE_EXTRACTION`` in the cloud repo's +#: ``workers/plugins/agentic_table/src/api_binding.py``. Change one, change +#: both, in the same PR. +CLOUD_STAGE_TABLE_EXTRACTION = "table_extraction" + + +def test_the_status_filter_matches_the_name_the_executor_sends(): + assert TABLE_STAGE_NAMES[0] == CLOUD_STAGE_TABLE_EXTRACTION, ( + "TABLE_STAGE_NAMES[0] no longer equals STAGE_TABLE_EXTRACTION in the " + "cloud repo's workers/plugins/agentic_table/src/api_binding.py. Every " + "table job will complete normally while reporting an empty `stages` " + "array. Update both sides, and the cloud twin of this test, together." + ) + + +def test_the_table_extractor_reports_exactly_one_stage(): + """The table engine exposes no node-level progress hooks, so a second name + here would describe progress the executor cannot substantiate -- it would + be recorded as missing rather than pending, for every job.""" + assert TABLE_STAGE_NAMES == [CLOUD_STAGE_TABLE_EXTRACTION] + assert TABLE_EXTRACTOR_NAME == "table" + + +# --------------------------------------------------------------------------- +# Second cross-repo lock in this pair: the stage-STATUS allowlist. +# +# Same shape as the stage-name lock above and the same class of silent drift, +# but one layer along. `StageReportView` rejects a status outside +# `_VALID_STAGE_STATUSES` with a 400 -- and the cloud reporter +# (`extraction_seams.progress.StageReporter.report`) SWALLOWS 400s by design, +# because its contract is that a failed progress report must never fail a +# paid run. So if the two sets drift, every affected job completes and bills +# normally while its reported stage stays frozen at the last value that +# happened to be accepted. +# +# The cloud side already carried a `!! CROSS-REPO CONSTANT !!` banner on its +# copy with nothing asserting against this one. The cloud twin of this test is +# `workers/plugins/extraction_seams/tests/test_progress.py`. +# --------------------------------------------------------------------------- + +#: Must equal ``VALID_STAGE_STATUSES`` in the cloud repo's +#: ``workers/plugins/extraction_seams/src/progress.py``. Change one, change +#: both, in the same PR. +#: +#: A literal, not an import: the two repos are separate checkouts and only meet +#: in the merged tree, so an import would make this pass vacuously wherever the +#: other side is absent -- exactly the environments where drift is introduced. +CLOUD_VALID_STAGE_STATUSES = frozenset({"running", "done"}) + + +def test_the_accepted_stage_statuses_match_what_the_executor_sends(): + from agent_kv.internal_views import _VALID_STAGE_STATUSES + + assert _VALID_STAGE_STATUSES == CLOUD_VALID_STAGE_STATUSES, ( + "_VALID_STAGE_STATUSES no longer equals VALID_STAGE_STATUSES in the " + "cloud repo's workers/plugins/extraction_seams/src/progress.py. Stage " + "reports will be rejected with a 400 that the cloud reporter swallows, " + "so jobs will complete and bill while their reported stage freezes. " + "Update both sides, and both twins of this test, together." + ) + + +def test_a_terminal_outcome_is_not_an_accepted_stage_status(): + """Why the set is only two values, asserted rather than described. + + Terminal outcomes travel on finalize. `"failed"` as a stage status is the + case that actually shipped once, and accepting it here would let a job + report a terminal stage while its row stayed non-terminal. + """ + from agent_kv.internal_views import _VALID_STAGE_STATUSES + + for terminal in ("failed", "completed", "cancelled", "timed_out"): + assert terminal not in _VALID_STAGE_STATUSES, terminal diff --git a/backend/agent_kv/tests/test_dispatch.py b/backend/agent_kv/tests/test_dispatch.py new file mode 100644 index 0000000000..c63854ae6c --- /dev/null +++ b/backend/agent_kv/tests/test_dispatch.py @@ -0,0 +1,385 @@ +import os +import uuid +from unittest import mock + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +import pytest # noqa: E402 +from django.conf import settings # noqa: E402 + +from agent_kv import dispatch # noqa: E402 +from agent_kv.constants import TABLE_EXTRACTOR_NAME # noqa: E402 +from agent_kv.models import AgentKVJob, JobStatus # noqa: E402 + + +def _fail_only_bookkeeping(m_objects, exc): + """Let the pre-dispatch terminal check succeed, then fail the next query. + + These tests are about a DB error during POST-ENQUEUE bookkeeping, which + must not be reported as a failed dispatch. A blanket `filter.side_effect` + now hits the pre-dispatch terminal check instead -- a different, earlier + failure that correctly DOES abort the dispatch, since a job that cannot be + verified as live must not have money spent on it. + """ + live = mock.MagicMock() + live.exists.return_value = False + m_objects.filter.side_effect = [live, exc] + + +def _filtered_with(m_objects, **expected): + """True if `objects.filter` was ever called with exactly these kwargs. + + Content-based rather than positional: `dispatch_job` makes several filter + calls (the pre-dispatch terminal check, the guarded bookkeeping UPDATE, and + on the slow path a narrowed fallback), and index-pinned assertions broke + every time one was added. + """ + return any(call.kwargs == expected for call in m_objects.filter.call_args_list) + + +def _not_cancelled(m_objects): + """Make the pre-dispatch terminal check report "still live". + + `dispatch_job` re-reads the row immediately before enqueueing, so a cancel + that landed between the submit's save() and the enqueue cannot result in + paid work for a job whose slot was already released. With `objects` mocked + the default `.exists()` is a truthy Mock, which would short-circuit every + dispatch in this file. + """ + m_objects.filter.return_value.exists.return_value = False + + +def _job(): + from account_v2.models import Organization + + j = AgentKVJob(id=uuid.uuid4(), input_ref="org/o/agent_kv/j/input.pdf") + # Unsaved related org with an explicit PK: assigning it caches the + # instance on the job (so ``job.organization`` never hits the DB) and + # sets ``job.organization_id`` to 7; the slug is deliberately different + # from the PK so a test can tell them apart. + j.organization = Organization(id=7, organization_id="org_slug_1") # PK 7 + j.pages_total = 3 + return j + + +@mock.patch.object(dispatch, "_platform_api_key", return_value="pk") +@mock.patch.object(dispatch, "_dispatcher") +@mock.patch.object(AgentKVJob, "objects") +def test_dispatch_success_stamps_job(m_objects, m_disp, m_key): + _not_cancelled(m_objects) + job = _job() + dispatch.dispatch_job( + job, + extractor=TABLE_EXTRACTOR_NAME, + schema={"a": {"description": "d"}}, + options={"number_format": "EU"}, + ) + + ctx = m_disp.return_value.dispatch_with_callback.call_args.args[0] + assert ctx.executor_name == "agentic_table" + assert ctx.operation == "table_extract_api" + assert ctx.run_id == str(job.id) + assert ctx.execution_source == "agent_kv_api" + assert ctx.organization_id == "7" + assert ctx.executor_params["job_id"] == str(job.id) + assert ctx.executor_params["input_ref"] == job.input_ref + assert ctx.executor_params["schema"] == {"a": {"description": "d"}} + assert ctx.executor_params["options"] == {"number_format": "EU"} + assert ctx.executor_params["platform_api_key"] == "pk" + # max_pages is the CAP the engine must enforce, not the measured count + # (job.pages_total, which rides separately and is None for Excel). + assert ctx.executor_params["max_pages"] == settings.AGENT_KV_MAX_PAGES + assert ctx.executor_params["pages_total"] == 3 + + kw = m_disp.return_value.dispatch_with_callback.call_args.kwargs + assert kw["on_success"].task == "agent_kv_complete" + assert kw["on_success"].kwargs == { + "callback_kwargs": {"job_id": str(job.id), "org_id": "7"} + } + assert kw["on_success"].options.get("queue") == "agent_kv_callback" + assert kw["on_error"].task == "agent_kv_error" + assert kw["on_error"].kwargs == { + "callback_kwargs": {"job_id": str(job.id), "org_id": "7"} + } + assert kw["task_id"] == str(job.task_id) + + assert job.status == JobStatus.DISPATCHED + assert job.dispatched_at is not None + # Guarded queryset UPDATE (not job.save()): only a still-PENDING row may + # be advanced to DISPATCHED. + # The pre-dispatch terminal check filters first; the guarded bookkeeping + # UPDATE is the last filter call. + assert _filtered_with(m_objects, id=job.id, status=JobStatus.PENDING) + m_objects.filter.return_value.update.assert_called_once_with( + task_id=job.task_id, + status=JobStatus.DISPATCHED, + dispatched_at=job.dispatched_at, + ) + + +@mock.patch.object(dispatch, "_platform_api_key", return_value="pk") +@mock.patch.object(dispatch, "_dispatcher") +@mock.patch.object(AgentKVJob, "objects") +def test_dispatch_guarded_update_cannot_overwrite_a_terminal_row( + m_objects, m_disp, m_key +): + """Regression for the un-terminalize bug: if the row already raced to a + terminal status (e.g. FAILED, via an instant-failure finalize callback + that beat this post-enqueue bookkeeping), the guarded UPDATE's WHERE + clause (status=PENDING) structurally cannot match it -- 0 rows update, + and the row's real status is left untouched. Proven the same way + ``AgentKVJob.mark_terminal``'s own guard is proven (test_models.py): by + pinning the exact WHERE-clause kwargs and simulating the "no rows + matched" outcome, since this suite runs with no real DB. + """ + _not_cancelled(m_objects) + m_objects.filter.return_value.update.return_value = 0 # simulates a FAILED row + job = _job() + + # Must not raise -- dispatch_job doesn't (and can't meaningfully) act on + # the update's row count; it already told the caller it dispatched. + dispatch.dispatch_job(job, extractor=TABLE_EXTRACTOR_NAME, schema={}, options={}) + + assert _filtered_with(m_objects, id=job.id, status=JobStatus.PENDING) + # The 0-row result now triggers a SECOND guarded write (below), whose + # `.exclude(status__in=TERMINAL)` is what keeps a terminal row untouched + # here -- so the terminal case is still structurally safe. + (_, excl_kwargs) = ( + m_objects.filter.return_value.exclude.call_args.args, + m_objects.filter.return_value.exclude.call_args.kwargs, + ) + assert set(excl_kwargs["status__in"]) == set(AgentKVJob.TERMINAL) + + +# The stranded-job regression. `StageReportView` promotes PENDING -> RUNNING on +# the executor's first stage report, which can land before this post-enqueue +# bookkeeping. The PENDING-guarded UPDATE then matches 0 rows and +# `dispatched_at` stays NULL -- and a non-terminal row with a NULL +# `dispatched_at` is invisible to BOTH sweep phases (phase 1 requires PENDING; +# phase 2's `dispatched_at__lt` can never match a NULL), so the job reports +# `running` forever and `GET result` 409s for the life of the row. +@mock.patch.object(dispatch, "_platform_api_key", return_value="pk") +@mock.patch.object(dispatch, "_dispatcher") +@mock.patch.object(AgentKVJob, "objects") +def test_dispatch_stamps_dispatched_at_when_the_row_already_moved_to_running( + m_objects, m_disp, m_key +): + _not_cancelled(m_objects) + m_objects.filter.return_value.update.return_value = 0 # PENDING guard missed + job = _job() + + dispatch.dispatch_job(job, extractor=TABLE_EXTRACTOR_NAME, schema={}, options={}) + + # Second write is narrowed to a row that still has no dispatch time, and + # excludes terminal rows. + assert _filtered_with(m_objects, id=job.id, dispatched_at__isnull=True) + fallback = m_objects.filter.return_value.exclude.return_value + update_kwargs = fallback.update.call_args.kwargs + assert update_kwargs["dispatched_at"] == job.dispatched_at + assert update_kwargs["task_id"] == job.task_id + # Crucially does NOT write `status`: the row genuinely was dispatched, but + # moving RUNNING back to DISPATCHED would discard the executor's progress. + assert "status" not in update_kwargs + + +@mock.patch.object(dispatch, "_platform_api_key", return_value="pk") +@mock.patch.object(dispatch, "_dispatcher") +@mock.patch.object(AgentKVJob, "objects") +def test_dispatch_does_not_attempt_the_fallback_when_the_pending_guard_won( + m_objects, m_disp, m_key +): + _not_cancelled(m_objects) + m_objects.filter.return_value.update.return_value = 1 # normal path + job = _job() + + dispatch.dispatch_job(job, extractor=TABLE_EXTRACTOR_NAME, schema={}, options={}) + + # Terminal check + the single guarded bookkeeping UPDATE; no fallback. + assert len(m_objects.filter.call_args_list) == 2 + assert not m_objects.filter.return_value.exclude.called + + +@mock.patch.object(dispatch, "_platform_api_key", return_value="pk") +@mock.patch.object(dispatch, "_dispatcher") +def test_enqueue_failure_raises_dispatch_error(m_disp, m_key): + m_disp.return_value.dispatch_with_callback.side_effect = RuntimeError("broker down") + job = _job() + with pytest.raises(dispatch.DispatchError): + dispatch.dispatch_job(job, extractor=TABLE_EXTRACTOR_NAME, schema={}, options={}) + + +@mock.patch.object(dispatch, "_dispatcher") +def test_dispatch_job_uses_platform_api_key_lookup(m_disp): + from platform_settings_v2.platform_auth_service import ( + PlatformAuthenticationService, + ) + + with mock.patch.object( + PlatformAuthenticationService, + "get_active_platform_key", + return_value=mock.Mock(key="the-real-key"), + ): + job = _job() + with mock.patch.object(AgentKVJob, "objects") as m_objects: + # Without this the pre-dispatch terminal check reads a truthy Mock + # and short-circuits, so nothing is ever enqueued. + _not_cancelled(m_objects) + dispatch.dispatch_job( + job, extractor=TABLE_EXTRACTOR_NAME, schema={}, options={} + ) + ctx = m_disp.return_value.dispatch_with_callback.call_args.args[0] + assert ctx.executor_params["platform_api_key"] == "the-real-key" + # The lookup takes the org's public slug, never the row PK (13b F6). + PlatformAuthenticationService.get_active_platform_key.assert_called_once_with( + "org_slug_1" + ) + assert ctx.organization_id == "7" + + +def test_platform_api_key_raises_dispatch_error_when_absent(): + from platform_settings_v2.platform_auth_service import ( + PlatformAuthenticationService, + ) + + with mock.patch.object( + PlatformAuthenticationService, + "get_active_platform_key", + return_value=None, + ): + job = _job() + with pytest.raises(dispatch.DispatchError): + dispatch._platform_api_key(job) + + +@mock.patch.object(dispatch, "_dispatcher") +@mock.patch.object(dispatch, "_platform_api_key") +def test_raw_exception_from_platform_key_lookup_is_wrapped_as_dispatch_error( + m_key, m_disp +): + """Regression: platform-key lookup and context construction must live + inside dispatch_job's try — a raw (non-DispatchError) exception there + (e.g. a transient DB error) must not escape uncaught. + """ + m_key.side_effect = RuntimeError("platform db down") + + job = _job() + with pytest.raises(dispatch.DispatchError): + dispatch.dispatch_job(job, extractor=TABLE_EXTRACTOR_NAME, schema={}, options={}) + + # Never got far enough to enqueue. + assert not m_disp.return_value.dispatch_with_callback.called + + +def test_dispatcher_factory_call_matches_the_live_signature(): + """Guard the OSS seam that a platform change can silently move. + + UN-4046 removed `get_executor_dispatcher(celery_app=...)`; our call site + kept passing it, so every submit raised TypeError inside dispatch_job and + failed the job. Nothing caught it -- the mismatch is invisible to a mocked + dispatcher and only appears when a real request is made. + + Binding our actual call against the real signature fails loudly the next + time that function's parameters change. + """ + from unittest import mock # noqa: PLC0415 + + from agent_kv import dispatch as d # noqa: PLC0415 + + # autospec=True makes the stub enforce the REAL function's signature, so + # this asserts our call site against it rather than against a permissive + # Mock. Re-adding `celery_app=` here raises TypeError, exactly as production + # did. + with mock.patch( + "pg_queue.executor_rpc.get_executor_dispatcher", autospec=True + ) as m_factory: + d._dispatcher() + m_factory.assert_called_once_with() + + +# Post-enqueue bookkeeping must never fail a dispatch that already succeeded. +# `SubmitView` turns a DispatchError into a FAILED job, so a DB hiccup here +# would terminalize a job whose task is on the queue -- the executor then runs, +# calls back, and finds a terminal row it cannot write to, while the caller was +# told nothing was billed for work that did run. The review round added a +# SECOND update after the enqueue, which widened this window. +@mock.patch.object(dispatch, "_platform_api_key", return_value="pk") +@mock.patch.object(dispatch, "_dispatcher") +@mock.patch.object(AgentKVJob, "objects") +def test_bookkeeping_failure_does_not_fail_an_already_queued_dispatch( + m_objects, m_disp, m_key +): + _not_cancelled(m_objects) + _fail_only_bookkeeping(m_objects, OSError("db gone")) + job = _job() + + # Must not raise: raising is what would mark the live job FAILED. + dispatch.dispatch_job(job, extractor=TABLE_EXTRACTOR_NAME, schema={}, options={}) + + # And the enqueue did happen, so the task is genuinely on the queue. + assert m_disp.return_value.dispatch_with_callback.called + + +@mock.patch.object(dispatch, "_platform_api_key", return_value="pk") +@mock.patch.object(dispatch, "_dispatcher") +@mock.patch.object(AgentKVJob, "objects") +def test_bookkeeping_failure_is_logged_at_error_level(m_objects, m_disp, m_key, caplog): + """Swallowing it silently would trade one bad failure mode for another: + the sweep reconciles the row, but nothing would say why it had to. + """ + _not_cancelled(m_objects) + _fail_only_bookkeeping(m_objects, OSError("db gone")) + with caplog.at_level("ERROR", logger=dispatch.logger.name): + dispatch.dispatch_job( + _job(), extractor=TABLE_EXTRACTOR_NAME, schema={}, options={} + ) + assert any(r.levelname == "ERROR" and r.exc_info for r in caplog.records) + + +# --------------------------------------------------------------------------- +# A job cancelled between submit's save() and the enqueue must not be dispatched. +# +# The cancel sees a PENDING, never-dispatched row, so it terminalizes it AND +# releases its concurrency slot -- correctly, nothing had been dispatched yet. +# Enqueueing anyway would then run paid work for a job the caller already +# cancelled, with its slot already handed to the next submit: the concurrency +# ceiling bypassed and the customer billed for a cancelled job. +# +# Reported by Greptile on PR #2317. +# --------------------------------------------------------------------------- +@mock.patch.object(dispatch, "_platform_api_key", return_value="pk") +@mock.patch.object(dispatch, "_dispatcher") +@mock.patch.object(AgentKVJob, "objects") +def test_a_job_cancelled_before_enqueue_is_not_dispatched(m_objects, m_disp, m_key): + job = _job() + m_objects.filter.return_value.exists.return_value = True # already terminal + + dispatch.dispatch_job(job, extractor=TABLE_EXTRACTOR_NAME, schema={}, options={}) + + assert ( + not m_disp.return_value.dispatch_with_callback.called + ), "paid work was enqueued for a job that was already cancelled" + # And no bookkeeping UPDATE either -- there is nothing to advance. + assert not m_objects.filter.return_value.update.called + + +@mock.patch.object(dispatch, "_platform_api_key", return_value="pk") +@mock.patch.object(dispatch, "_dispatcher") +@mock.patch.object(AgentKVJob, "objects") +def test_the_terminal_check_reads_the_row_rather_than_the_in_memory_copy( + m_objects, m_disp, m_key +): + """The in-memory job predates the cancel by construction, so trusting + `job.status` here would never see it. + """ + job = _job() + job.status = JobStatus.PENDING # stale: the row may already be CANCELLED + _not_cancelled(m_objects) + + dispatch.dispatch_job(job, extractor=TABLE_EXTRACTOR_NAME, schema={}, options={}) + + assert _filtered_with(m_objects, id=job.id, status__in=list(AgentKVJob.TERMINAL)) diff --git a/backend/agent_kv/tests/test_internal_views.py b/backend/agent_kv/tests/test_internal_views.py new file mode 100644 index 0000000000..20cba4c677 --- /dev/null +++ b/backend/agent_kv/tests/test_internal_views.py @@ -0,0 +1,836 @@ +import os +import uuid +from unittest import mock + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from rest_framework.test import APIRequestFactory # noqa: E402 + +from agent_kv import internal_views as iv # noqa: E402 +from agent_kv import storage # noqa: E402 +from agent_kv.models import AgentKVJob, JobStatus # noqa: E402 + + +def _post(path, body): + return APIRequestFactory().post(path, body, format="json") + + +def _merge_payload(update_call): + """Unwrap a ``stages=`` jsonb-merge expression's {stage: entry} payload. + + Also pins the expression's shape: a ``||`` (jsonb concatenation) + ``CombinedExpression`` over the ``stages`` column, not a plain dict -- + the whole point of the DB-side merge is that ``.update()`` never + receives a Python-computed full-column dict. + """ + expr = update_call.kwargs["stages"] + assert not isinstance(expr, dict), "stages must be a jsonb-merge expression" + assert expr.connector == "||" + assert expr.lhs.name == "stages" + return expr.rhs.value + + +# --------------------------------------------------------------------------- +# _stage_merge_expression / _sanitize_counters (expression-builder unit tests) +# --------------------------------------------------------------------------- + + +# (u1) the expression builder produces a single-entry {stage: entry} jsonb +# merge payload over the `stages` column via the `||` connector. +def test_stage_merge_expression_builds_single_entry_jsonb_merge(): + expr = iv._stage_merge_expression("extraction", {"status": "running"}) + + assert expr.connector == "||" + assert expr.lhs.name == "stages" + assert expr.rhs.value == {"extraction": {"status": "running"}} + + +# (u2) counters colliding with reserved keys are dropped. +def test_sanitize_counters_drops_reserved_keys(): + assert iv._sanitize_counters({"status": "done", "seconds": 999, "pages": 3}) == { + "pages": 3 + } + + +# (u3) non-scalar counter values (dict/list) are dropped; scalars pass through. +def test_sanitize_counters_drops_non_scalar_values(): + assert iv._sanitize_counters( + {"pages": 3, "nested": {"a": 1}, "arr": [1, 2], "ok": "yes", "flag": True} + ) == {"pages": 3, "ok": "yes", "flag": True} + + +# (u4) a non-dict counters payload (e.g. a string or list) sanitizes to {}. +def test_sanitize_counters_non_dict_input_is_empty(): + assert iv._sanitize_counters("not-a-dict") == {} + assert iv._sanitize_counters(None) == {} + + +# --------------------------------------------------------------------------- +# StageReportView +# --------------------------------------------------------------------------- + + +# (1) org_id missing from body -> 400 (ambient-auth views require it in body). +def test_stage_report_missing_org_id_is_400(): + resp = iv.StageReportView.as_view()( + _post("/x", {"stage": "extraction", "status": "running"}), + job_id=uuid.uuid4(), + ) + assert resp.status_code == 400 + + +# (1b) missing `stage` -> 400. +def test_stage_report_missing_stage_is_400(): + resp = iv.StageReportView.as_view()( + _post("/x", {"org_id": "org1", "status": "running"}), + job_id=uuid.uuid4(), + ) + assert resp.status_code == 400 + + +# (1c) missing/invalid `status` -> 400 (must be exactly "running" or "done"). +def test_stage_report_invalid_status_is_400(): + resp = iv.StageReportView.as_view()( + _post("/x", {"org_id": "org1", "stage": "extraction", "status": "bogus"}), + job_id=uuid.uuid4(), + ) + assert resp.status_code == 400 + + resp_missing = iv.StageReportView.as_view()( + _post("/x", {"org_id": "org1", "stage": "extraction"}), + job_id=uuid.uuid4(), + ) + assert resp_missing.status_code == 400 + + +# (2) first stage report on a PENDING job flips it to RUNNING exactly once, +# and merges the stage entry via a DB-side jsonb `||` expression (not a +# plain dict) targeting only the reported stage's key. +@mock.patch.object(AgentKVJob, "objects") +def test_stage_report_flips_pending_to_running_once(m_jobs): + job = AgentKVJob(status=JobStatus.PENDING, stages={}) + qs = m_jobs.filter.return_value.exclude.return_value + qs.first.return_value = job + + resp = iv.StageReportView.as_view()( + _post( + "/x", + { + "org_id": "org1", + "stage": "extraction", + "status": "running", + }, + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + qs.update.assert_called_once() + call = qs.update.call_args + assert _merge_payload(call) == {"extraction": {"status": "running"}} + assert call.kwargs["stage"] == "extraction" + assert call.kwargs["status"] == JobStatus.RUNNING + + +# (2b) a stage report on an already-RUNNING job does NOT re-flip status -- +# the "once" half of "flips RUNNING once". +@mock.patch.object(AgentKVJob, "objects") +def test_stage_report_on_running_job_does_not_touch_status(m_jobs): + job = AgentKVJob( + status=JobStatus.RUNNING, stages={"extraction": {"status": "running"}} + ) + qs = m_jobs.filter.return_value.exclude.return_value + qs.first.return_value = job + + resp = iv.StageReportView.as_view()( + _post( + "/x", + {"org_id": "org1", "stage": "extraction", "status": "done", "seconds": 1.5}, + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + call = qs.update.call_args + assert "status" not in call.kwargs + assert _merge_payload(call) == {"extraction": {"status": "done", "seconds": 1.5}} + + +# (3) stage report on a CANCELLED (terminal) job -- the update queryset +# excludes TERMINAL, so .first() returns None -- is a 200 no-op, never +# reaching .update(). +@mock.patch.object(AgentKVJob, "objects") +def test_stage_report_on_cancelled_job_is_200_noop(m_jobs): + qs = m_jobs.filter.return_value.exclude.return_value + qs.first.return_value = None + + resp = iv.StageReportView.as_view()( + _post("/x", {"org_id": "org1", "stage": "extraction", "status": "running"}), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + assert not qs.update.called + + +# (4) duplicate stage report (same stage name posted twice) overwrites the +# entry rather than appending. With the DB-side jsonb `||` merge this falls +# out structurally: the merge payload always carries exactly one key (the +# reported stage), so the jsonb `||` operator replaces that key's value on +# the second call instead of stacking a second version alongside it. +@mock.patch.object(AgentKVJob, "objects") +def test_stage_report_duplicate_overwrites_not_appends(m_jobs): + job = AgentKVJob( + status=JobStatus.RUNNING, + stages={"extraction": {"status": "running"}}, + ) + qs = m_jobs.filter.return_value.exclude.return_value + qs.first.return_value = job + + resp = iv.StageReportView.as_view()( + _post( + "/x", + { + "org_id": "org1", + "stage": "extraction", + "status": "done", + "seconds": 2.0, + "counters": {"pages": 3}, + }, + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + call = qs.update.call_args + payload = _merge_payload(call) + assert list(payload.keys()) == ["extraction"] + assert payload == {"extraction": {"status": "done", "seconds": 2.0, "pages": 3}} + assert call.kwargs["stage"] == "extraction" + + +# (5) the stage-report endpoint is the write gate for job.stages -- only the +# defined shape (status, optional seconds, sanitized flat counters) is ever +# persisted; unexpected top-level body keys (e.g. a stray "evil") must not +# leak into the stored entry. +@mock.patch.object(AgentKVJob, "objects") +def test_stage_report_ignores_unexpected_top_level_keys(m_jobs): + job = AgentKVJob(status=JobStatus.RUNNING, stages={}) + qs = m_jobs.filter.return_value.exclude.return_value + qs.first.return_value = job + + resp = iv.StageReportView.as_view()( + _post( + "/x", + { + "org_id": "org1", + "stage": "extraction", + "status": "running", + "evil": "payload", + "another_bogus_key": 123, + }, + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + stored_entry = _merge_payload(qs.update.call_args)["extraction"] + assert stored_entry == {"status": "running"} + + +# (5b) a counter named "status" or "seconds" cannot clobber the endpoint's +# own reserved fields. +@mock.patch.object(AgentKVJob, "objects") +def test_stage_report_counters_cannot_override_reserved_keys(m_jobs): + job = AgentKVJob(status=JobStatus.RUNNING, stages={}) + qs = m_jobs.filter.return_value.exclude.return_value + qs.first.return_value = job + + resp = iv.StageReportView.as_view()( + _post( + "/x", + { + "org_id": "org1", + "stage": "extraction", + "status": "running", + "counters": {"status": "done", "seconds": 999, "pages": 3}, + }, + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + stored_entry = _merge_payload(qs.update.call_args)["extraction"] + assert stored_entry == {"status": "running", "pages": 3} + + +# (5c) a nested-dict (or list) counter value is dropped rather than stored. +@mock.patch.object(AgentKVJob, "objects") +def test_stage_report_counters_drops_nested_values(m_jobs): + job = AgentKVJob(status=JobStatus.RUNNING, stages={}) + qs = m_jobs.filter.return_value.exclude.return_value + qs.first.return_value = job + + resp = iv.StageReportView.as_view()( + _post( + "/x", + { + "org_id": "org1", + "stage": "extraction", + "status": "running", + "counters": {"pages": 3, "nested": {"a": 1}, "arr": [1, 2]}, + }, + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + stored_entry = _merge_payload(qs.update.call_args)["extraction"] + assert stored_entry == {"status": "running", "pages": 3} + + +# --------------------------------------------------------------------------- +# FinalizeView +# --------------------------------------------------------------------------- + + +# (6) org_id missing from body -> 400. +def test_finalize_missing_org_id_is_400(): + resp = iv.FinalizeView.as_view()(_post("/x", {"success": True}), job_id=uuid.uuid4()) + assert resp.status_code == 400 + + +# (6b) missing `success` -> 400, and no slot is released (nothing was +# finalized). +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +def test_finalize_missing_success_is_400(m_release): + resp = iv.FinalizeView.as_view()(_post("/x", {"org_id": "org1"}), job_id=uuid.uuid4()) + assert resp.status_code == 400 + assert not m_release.called + + +# (6c) a non-bool `success` (e.g. a truthy string) -> 400 rather than +# silently falling through to the failure branch and persisting a FAILED +# job with an empty error. +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +def test_finalize_non_bool_success_is_400(m_release, m_mark_terminal): + resp = iv.FinalizeView.as_view()( + _post("/x", {"org_id": "org1", "success": "true"}), job_id=uuid.uuid4() + ) + assert resp.status_code == 400 + assert not m_release.called + assert not m_mark_terminal.called + + +# (7) finalize success: writes the result THEN marks the job terminal (order +# matters -- a duplicate must never rewrite a result nobody asked to change). +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(iv, "write_result") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_success_writes_result_then_marks_terminal( + m_jobs, m_write, m_mark_terminal, m_release +): + job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="https://example.com/hook") + m_jobs.filter.return_value.first.return_value = job + m_write.return_value = "org/o/agent_kv/j/result.json" + m_mark_terminal.return_value = True + + manager = mock.Mock() + manager.attach_mock(m_write, "write_result") + manager.attach_mock(m_mark_terminal, "mark_terminal") + + job_id = uuid.uuid4() + resp = iv.FinalizeView.as_view()( + _post( + "/x", + { + "org_id": "org1", + "success": True, + "result": {"foo": "bar"}, + "usage_summary": {"tokens": 10}, + }, + ), + job_id=job_id, + ) + + assert resp.status_code == 200 + assert resp.data == { + "finalized": True, + "webhook_url": "https://example.com/hook", + "status": "completed", + "reason": "ok", + } + assert [c[0] for c in manager.mock_calls] == ["write_result", "mark_terminal"] + m_write.assert_called_once_with("org1", str(job_id), {"foo": "bar"}) + m_mark_terminal.assert_called_once_with( + job_id, + "org1", + JobStatus.COMPLETED, + result_ref="org/o/agent_kv/j/result.json", + usage_summary={"tokens": 10}, + ) + + +# (7b) finalize success deletes the staged input and blanks input_ref (spec +# D10: "uploaded document deleted on job completion") -- the result file +# itself is untouched by this. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(iv, "delete_input") +@mock.patch.object(iv, "write_result") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_success_deletes_input_and_blanks_ref( + m_jobs, m_write, m_delete_input, m_mark_terminal, m_release +): + job = AgentKVJob( + status=JobStatus.RUNNING, + webhook_url="", + input_ref="org/o/agent_kv/j/input.pdf", + ) + m_jobs.filter.return_value.first.return_value = job + m_write.return_value = "org/o/agent_kv/j/result.json" + + job_id = uuid.uuid4() + resp = iv.FinalizeView.as_view()( + _post("/x", {"org_id": "org1", "success": True, "result": {}}), + job_id=job_id, + ) + + assert resp.status_code == 200 + m_delete_input.assert_called_once_with(job) + m_jobs.filter.return_value.update.assert_called_once_with(input_ref="") + + +# (7c) success finalize whose guard is LOST right after the write (e.g. a +# cancel raced in between the read and mark_terminal's guarded UPDATE) -- +# the just-written result file has no ref anywhere pointing at it, so it's +# cleaned up immediately rather than orphaned forever. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(iv, "delete_result_file") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=False) +@mock.patch.object(iv, "write_result") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_success_guard_loss_deletes_orphaned_result( + m_jobs, m_write, m_mark_terminal, m_delete_result, m_release +): + job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="") + m_jobs.filter.return_value.first.return_value = job + m_write.return_value = "org/o/agent_kv/j/result.json" + + job_id = uuid.uuid4() + resp = iv.FinalizeView.as_view()( + _post( + "/x", + {"org_id": "org1", "success": True, "result": {"foo": "bar"}}, + ), + job_id=job_id, + ) + + assert resp.status_code == 200 + assert resp.data["finalized"] is False + m_delete_result.assert_called_once_with("org/o/agent_kv/j/result.json") + + +# (7d) success finalize that WINS the guard never deletes the result it just +# wrote -- only a guard-LOSS orphan triggers cleanup. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(iv, "delete_result_file") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(iv, "write_result") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_success_guard_win_does_not_delete_result( + m_jobs, m_write, m_mark_terminal, m_delete_result, m_release +): + job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="") + m_jobs.filter.return_value.first.return_value = job + m_write.return_value = "org/o/agent_kv/j/result.json" + + iv.FinalizeView.as_view()( + _post( + "/x", + {"org_id": "org1", "success": True, "result": {"foo": "bar"}}, + ), + job_id=uuid.uuid4(), + ) + + assert not m_delete_result.called + + +# (7e) concurrent duplicate-SUCCESS finalize race, exercised through the REAL +# storage layer (only the object-store FileSystem is faked): two finalize +# attempts land for one job, the first WINS the terminal guard and the second +# LOSES. Because write_result now returns a UNIQUE ref per attempt, the loser +# cleans up only ITS OWN orphaned result file -- the winner's result_ref +# target is left intact (pre-Greptile critical #3: the old deterministic +# result.json path let the loser delete the winner's live result). +@mock.patch.object(storage, "FileSystem") +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", side_effect=[True, False]) +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_duplicate_success_race_loser_deletes_only_its_own_result( + m_jobs, m_mark, m_release, m_fs +): + # Two separate reads that both observed the job as still RUNNING (the race + # window before either guarded UPDATE ran). + winner_job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="", input_ref="") + loser_job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="", input_ref="") + m_jobs.filter.return_value.first.side_effect = [winner_job, loser_job] + + fh = m_fs.return_value.get_file_storage.return_value + written: list[str] = [] + removed: list[str] = [] + fh.json_dump.side_effect = lambda path, data: written.append(path) + fh.rm.side_effect = lambda path: removed.append(path) + + job_id = uuid.uuid4() + body = {"org_id": "org1", "success": True, "result": {"foo": "bar"}} + r1 = iv.FinalizeView.as_view()(_post("/x", body), job_id=job_id) # winner + r2 = iv.FinalizeView.as_view()(_post("/x", body), job_id=job_id) # loser + + assert r1.data["finalized"] is True + assert r2.data["finalized"] is False + + # Both attempts wrote a result; the two refs are DISTINCT. + assert len(written) == 2 + assert written[0] != written[1] + + # The winner's result file (written[0]) is the one mark_terminal stored as + # result_ref -- it must NOT be deleted. Only the loser's orphan is removed. + assert removed == [written[1]] + assert written[0] not in removed + + # The winner stored its OWN specific ref, not a shared deterministic path. + winner_ref = m_mark.call_args_list[0].kwargs["result_ref"] + assert winner_ref == written[0] + + +# (8) duplicate finalize: the job is already terminal (guard excludes it), +# so this is a no-op -- finalized:false and the result is NOT rewritten. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(iv, "write_result") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_duplicate_does_not_rewrite_result( + m_jobs, m_write, m_mark_terminal, m_release +): + job = AgentKVJob(status=JobStatus.COMPLETED, webhook_url="https://example.com/hook") + m_jobs.filter.return_value.first.return_value = job + + resp = iv.FinalizeView.as_view()( + _post("/x", {"org_id": "org1", "success": True, "result": {"foo": "bar"}}), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + assert resp.data == { + "finalized": False, + "webhook_url": "https://example.com/hook", + "status": "completed", + "reason": "ok", + } + assert not m_write.called + assert not m_mark_terminal.called + + +# (8b) duplicate finalize (guard-lost) never calls delete_input either -- +# either another writer already owns cleanup, or there's nothing new to +# terminalize. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(iv, "delete_input") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_duplicate_does_not_call_delete_input( + m_jobs, m_delete_input, m_mark_terminal, m_release +): + job = AgentKVJob( + status=JobStatus.COMPLETED, + webhook_url="", + input_ref="org/o/agent_kv/j/input.pdf", + ) + m_jobs.filter.return_value.first.return_value = job + + resp = iv.FinalizeView.as_view()( + _post("/x", {"org_id": "org1", "success": True, "result": {}}), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + assert not m_delete_input.called + assert not m_jobs.filter.return_value.update.called + + +# (9) finalize failure records the error via mark_terminal(FAILED, ...). +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(iv, "write_result") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_failure_records_error(m_jobs, m_write, m_mark_terminal, m_release): + job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="") + m_jobs.filter.return_value.first.return_value = job + m_mark_terminal.return_value = True + + job_id = uuid.uuid4() + resp = iv.FinalizeView.as_view()( + _post( + "/x", + {"org_id": "org1", "success": False, "error": "LLM provider timed out"}, + ), + job_id=job_id, + ) + + assert resp.status_code == 200 + assert resp.data == { + "finalized": True, + "webhook_url": "", + "status": "failed", + "reason": "ok", + } + assert not m_write.called + m_mark_terminal.assert_called_once_with( + job_id, "org1", JobStatus.FAILED, error="LLM provider timed out" + ) + + +# (9b) finalize failure also deletes the staged input and blanks input_ref +# -- the run is over whether it completed or failed. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(iv, "delete_input") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_failure_deletes_input_and_blanks_ref( + m_jobs, m_delete_input, m_mark_terminal, m_release +): + job = AgentKVJob( + status=JobStatus.RUNNING, + webhook_url="", + input_ref="org/o/agent_kv/j/input.pdf", + ) + m_jobs.filter.return_value.first.return_value = job + + job_id = uuid.uuid4() + resp = iv.FinalizeView.as_view()( + _post("/x", {"org_id": "org1", "success": False, "error": "boom"}), + job_id=job_id, + ) + + assert resp.status_code == 200 + m_delete_input.assert_called_once_with(job) + m_jobs.filter.return_value.update.assert_called_once_with(input_ref="") + + +# (10) the concurrency slot is released on every finalize path -- including +# when the write/finalize work raises -- because release() sits in a +# finally. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(iv, "write_result") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_releases_slot_even_when_write_result_raises( + m_jobs, m_write, m_mark_terminal, m_release +): + job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="") + m_jobs.filter.return_value.first.return_value = job + m_write.side_effect = RuntimeError("storage exploded") + + job_id = uuid.uuid4() + try: + iv.FinalizeView.as_view()( + _post("/x", {"org_id": "org1", "success": True, "result": {}}), + job_id=job_id, + ) + except RuntimeError: + pass + + m_release.assert_called_once_with("org1", str(job_id)) + assert not m_mark_terminal.called + + +# (10b) slot released on the duplicate (already-terminal, no-op) path too. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(iv, "write_result") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_releases_slot_on_duplicate_noop_path( + m_jobs, m_write, m_mark_terminal, m_release +): + job = AgentKVJob(status=JobStatus.CANCELLED, webhook_url="") + m_jobs.filter.return_value.first.return_value = job + + job_id = uuid.uuid4() + iv.FinalizeView.as_view()( + _post("/x", {"org_id": "org1", "success": False, "error": "late"}), + job_id=job_id, + ) + + m_release.assert_called_once_with("org1", str(job_id)) + + +# --------------------------------------------------------------------------- +# URL wiring +# --------------------------------------------------------------------------- + + +# (11) regression pin for the frozen paths (spec Interfaces block): the +# cloud executor calls these exact URLs, so a dropped/renamed include in +# internal_base_urls.py must fail loudly here rather than 404 in prod. +def test_frozen_internal_urls_resolve_to_the_right_views(): + from django.urls import resolve + + job_id = uuid.uuid4() + stage = resolve(f"/internal/v1/agent-kv/jobs/{job_id}/stage/") + assert stage.func.cls is iv.StageReportView + assert stage.kwargs == {"job_id": job_id} + + finalize = resolve(f"/internal/v1/agent-kv/jobs/{job_id}/finalize/") + assert finalize.func.cls is iv.FinalizeView + assert finalize.kwargs == {"job_id": job_id} + + +# --------------------------------------------------------------------------- +# An unknown job/org pair must be distinguishable from an ordinary duplicate. +# +# Both used to return a byte-identical 200 no-op with nothing logged. The +# second is what an org-slug-vs-FK-pk mix-up looks like -- and `dispatch.py` +# documents that exact confusion shipping once already. If it recurs, every +# finalize is a silent no-op, every job stays non-terminal and is reaped as +# "Job timed out": a 100% failure rate presenting as timeouts, invisible. +# +# Reported as 2.12 in the branch review. +# --------------------------------------------------------------------------- +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "objects") +def test_finalize_for_an_unknown_job_says_so_and_logs(m_objects, m_release, caplog): + import logging + + m_objects.filter.return_value.first.return_value = None + + with caplog.at_level(logging.WARNING, logger="agent_kv.internal_views"): + resp = iv.FinalizeView.as_view()( + _post("/x", {"org_id": "org1", "success": True}), job_id=uuid.uuid4() + ) + + assert resp.status_code == 200 + assert resp.data["reason"] == "unknown_job" + assert resp.data["finalized"] is False + assert "unknown job" in caplog.text + + +@mock.patch.object(AgentKVJob, "objects") +def test_stage_report_for_an_unknown_job_says_so_and_logs(m_objects, caplog): + """`already_terminal` is ordinary and stays quiet; `unknown_job` is not.""" + import logging + + m_objects.filter.return_value.exclude.return_value.first.return_value = None + m_objects.filter.return_value.exists.return_value = False + + with caplog.at_level(logging.WARNING, logger="agent_kv.internal_views"): + resp = iv.StageReportView.as_view()( + _post( + "/x", {"org_id": "org1", "stage": "table_extraction", "status": "running"} + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 200 + assert resp.data["reason"] == "unknown_job" + assert "unknown job" in caplog.text + + +@mock.patch.object(AgentKVJob, "objects") +def test_stage_report_for_an_already_terminal_job_is_quiet(m_objects, caplog): + import logging + + m_objects.filter.return_value.exclude.return_value.first.return_value = None + m_objects.filter.return_value.exists.return_value = True + + with caplog.at_level(logging.WARNING, logger="agent_kv.internal_views"): + resp = iv.StageReportView.as_view()( + _post( + "/x", {"org_id": "org1", "stage": "table_extraction", "status": "running"} + ), + job_id=uuid.uuid4(), + ) + + assert resp.data["reason"] == "already_terminal" + assert caplog.records == [], "a late report for a finished job is ordinary" + + +# --------------------------------------------------------------------------- +# Untrusted executor payload must not be able to rewrite the status document. +# Reported as 2.13 in the branch review. +# --------------------------------------------------------------------------- +def test_a_name_counter_cannot_rename_the_stage(): + """`_status_document` builds `{"name": name, **entry}` with the spread + LAST, so a persisted `name` wins -- renaming the stage in every later status + response and breaking the `if name in stages_json` filter that decides + which stages are shown at all. + """ + assert "name" in iv._RESERVED_STAGE_ENTRY_KEYS + assert iv._sanitize_counters({"name": "pwned", "steps": 3}) == {"steps": 3} + + +@mock.patch.object(AgentKVJob, "objects") +def test_a_non_numeric_seconds_is_rejected(m_objects): + """`seconds` is echoed back in the status document, so a dict or list here + was persisted and returned to the caller verbatim. + """ + resp = iv.StageReportView.as_view()( + _post( + "/x", + { + "org_id": "org1", + "stage": "table_extraction", + "status": "running", + "seconds": {"nested": "payload"}, + }, + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 400 + assert "seconds" in str(resp.data) + + +@mock.patch.object(AgentKVJob, "objects") +def test_a_boolean_seconds_is_rejected(m_objects): + """`bool` is a subclass of `int`, so a bare isinstance check would let + `True` through and persist it as a duration. + """ + resp = iv.StageReportView.as_view()( + _post( + "/x", + { + "org_id": "org1", + "stage": "table_extraction", + "status": "running", + "seconds": True, + }, + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 400 + + +@mock.patch.object(AgentKVJob, "objects") +def test_an_overlong_stage_name_is_rejected(m_objects): + """`job.stage` is varchar(32); a longer name was an unhandled 500, and + `job.stages` could grow unbounded distinct keys. + """ + resp = iv.StageReportView.as_view()( + _post( + "/x", + {"org_id": "org1", "stage": "x" * 33, "status": "running"}, + ), + job_id=uuid.uuid4(), + ) + + assert resp.status_code == 400 + assert "too long" in str(resp.data) diff --git a/backend/agent_kv/tests/test_job_views.py b/backend/agent_kv/tests/test_job_views.py new file mode 100644 index 0000000000..48d34b34e6 --- /dev/null +++ b/backend/agent_kv/tests/test_job_views.py @@ -0,0 +1,884 @@ +import os +import uuid +from datetime import timedelta +from unittest import mock + +import django +import pytest +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from django.utils import timezone # noqa: E402 +from rest_framework.test import APIRequestFactory # noqa: E402 + +from agent_kv import execution_views as ev # noqa: E402 +from agent_kv import execution_views_result as evr # noqa: E402 +from agent_kv.constants import V1_EXTRACTOR_NAME # noqa: E402 +from agent_kv.models import AgentKVJob, AgentKVKey, JobStatus # noqa: E402 +from agent_kv.tests._factories import kv_key # noqa: E402 + + +def _authed(method="get", path="/agent-kv/x"): + req = getattr(APIRequestFactory(), method)(path) + req.META["HTTP_AUTHORIZATION"] = "Bearer 123e4567-e89b-12d3-a456-426614174001" + return req + + +# --------------------------------------------------------------------------- +# (1) status for foreign-org job -> 404 (indistinguishable from unknown). +# --------------------------------------------------------------------------- +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_foreign_org_job_is_404(m_keys, m_jobs): + m_keys.get.return_value = kv_key() + m_jobs.get.side_effect = AgentKVJob.DoesNotExist + resp = ev.JobStatusView.as_view()(_authed(), job_id=uuid.uuid4()) + assert resp.status_code == 404 + + +# --------------------------------------------------------------------------- +# (1b) ...and the lookup is actually ORG-SCOPED. +# +# The test above forces `DoesNotExist` and asserts the view 404s, which says +# nothing about the `organization_id` filter: delete +# `organization_id=agent_kv_key.organization_id` from `_get_job` and it stays +# green -- while every org could then read, cancel, delete and fetch the result +# of every other org's job. +# +# The production code IS correctly scoped. What was missing is a test that can +# detect its removal, on the one security-relevant path in the app. Every +# job-scoped view funnels through `_get_job`, so each is covered. +# +# Reported by Greptile on PR #2317. +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("view", "method"), + [ + (ev.JobStatusView, "get"), + (ev.JobResultView, "get"), + (ev.JobCancelView, "post"), + (ev.JobDeleteView, "delete"), + ], +) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_every_job_scoped_view_filters_by_organization(m_keys, m_jobs, view, method): + m_keys.get.return_value = kv_key(org_pk=7) + m_jobs.get.side_effect = AgentKVJob.DoesNotExist + job_id = uuid.uuid4() + + view.as_view()(_authed(method=method), job_id=job_id) + + m_jobs.get.assert_called_once_with(id=job_id, organization_id=7) + + +# --------------------------------------------------------------------------- +# (2) status running -> stages list ordered per STAGE_NAMES, lowercased +# status, only the stages actually present are included. +# --------------------------------------------------------------------------- +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_status_running_builds_ordered_stages_and_lowercases_status(m_keys, m_jobs): + m_keys.get.return_value = kv_key() + job = AgentKVJob( + # Explicit, not the column default: the stage names below are the KV + # list, and `_status_document` filters a job's stages through + # STAGE_NAMES_BY_EXTRACTOR[job.extractor]. This used to ride on + # `extractor` defaulting to `kv`, which made a load-bearing fact + # invisible in the test that depends on it. + extractor=V1_EXTRACTOR_NAME, + status=JobStatus.RUNNING, + stage="extraction", + stages={ + "extraction": {"status": "running"}, + "document_processing": {"status": "done", "seconds": 0.5}, + }, + pages_total=3, + ) + job.created_at = timezone.now() + m_jobs.get.return_value = job + + resp = ev.JobStatusView.as_view()(_authed(), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + # `status` is the JOB's, so it stays top level; stage reporting is + # extractor-scoped (spec §7.2). + assert resp.data["status"] == "running" + kv = resp.data["extractors"]["kv"] + assert kv["stage"] == "extraction" + # STAGE_NAMES order is document_processing, extraction, ... -- "qa" and + # every other configured stage name is absent from job.stages, so only + # these two appear, in that order. + assert [s["name"] for s in kv["stages"]] == [ + "document_processing", + "extraction", + ] + assert resp.data["pages_total"] == 3 + assert "error" not in resp.data + + +# --------------------------------------------------------------------------- +# (3) result before completion (non-terminal, e.g. RUNNING) -> 409 with the +# current status, lowercased (spec §7.3 controller ruling). +# --------------------------------------------------------------------------- +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_result_before_completion_is_409_with_current_status(m_keys, m_jobs): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING) + m_jobs.get.return_value = job + + resp = ev.JobResultView.as_view()(_authed(), job_id=uuid.uuid4()) + + assert resp.status_code == 409 + assert resp.data == {"status": "running"} + + +# --------------------------------------------------------------------------- +# (4) result after expires_at -> 404. +# --------------------------------------------------------------------------- +@mock.patch.object(evr, "read_result") +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_result_after_expiry_is_404(m_keys, m_jobs, m_read): + m_keys.get.return_value = kv_key() + job = AgentKVJob( + status=JobStatus.COMPLETED, + result_ref="org/o/agent_kv/j/result.json", + expires_at=timezone.now() - timedelta(days=1), + ) + m_jobs.get.return_value = job + + resp = ev.JobResultView.as_view()(_authed(), job_id=uuid.uuid4()) + + assert resp.status_code == 404 + assert not m_read.called + + +# --------------------------------------------------------------------------- +# (5) result happy path -> returns read_result payload unchanged. +# --------------------------------------------------------------------------- +@mock.patch.object(evr, "read_result", return_value={"success": True, "fields": {}}) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_result_happy_path_returns_read_result_payload(m_keys, m_jobs, m_read): + m_keys.get.return_value = kv_key() + job = AgentKVJob( + # Explicit: the assertions below namespace the payload under `kv`, + # which is `job.extractor`, not a constant of the response shape. + extractor=V1_EXTRACTOR_NAME, + status=JobStatus.COMPLETED, + result_ref="org/o/agent_kv/j/result.json", + expires_at=timezone.now() + timedelta(days=1), + ) + m_jobs.get.return_value = job + + resp = ev.JobResultView.as_view()(_authed(), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + # The stored blob is the engine's own result; the response namespaces it + # per extractor and adds per-extractor usage attribution (spec §7.3). + # `success`/`status` at top level on every terminal payload, so a client + # branches the same way for completed, failed and cancelled. + assert resp.data["success"] is True + assert resp.data["status"] == "completed" + assert resp.data["extractors"] == {"kv": {"success": True, "fields": {}}} + assert set(resp.data["usage_summary"]) == {"total", "by_extractor"} + assert list(resp.data["usage_summary"]["by_extractor"]) == ["kv"] + m_read.assert_called_once_with(job.result_ref) + + +# --------------------------------------------------------------------------- +# (5b) result for a FAILED job -> 200 with a success:false body carrying the +# job's own (user-safe) error -- spec §7.3: "Failed jobs: {success: false, +# error, timing} with a user-safe error". This also covers the SubmitView +# sync-wait fix: that branch reuses this exact function unconditionally for +# any terminal job, so it now gets a correct 200 body instead of a 404. +# --------------------------------------------------------------------------- +@mock.patch.object(evr, "read_result") +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_result_for_failed_job_is_200_with_success_false_and_error( + m_keys, m_jobs, m_read +): + m_keys.get.return_value = kv_key() + job = AgentKVJob( + status=JobStatus.FAILED, + error="LLM provider timed out", + expires_at=timezone.now() + timedelta(days=1), + ) + m_jobs.get.return_value = job + + resp = ev.JobResultView.as_view()(_authed(), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + assert resp.data == { + "success": False, + "status": "failed", + "error": "LLM provider timed out", + } + assert not m_read.called + + +# --------------------------------------------------------------------------- +# (5c) result for a CANCELLED job -> 200 with a fixed success:false/cancelled +# body (spec §7.3 controller ruling). +# --------------------------------------------------------------------------- +@mock.patch.object(evr, "read_result") +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_result_for_cancelled_job_is_200_with_cancelled_body(m_keys, m_jobs, m_read): + m_keys.get.return_value = kv_key() + job = AgentKVJob( + status=JobStatus.CANCELLED, + expires_at=timezone.now() + timedelta(days=1), + ) + m_jobs.get.return_value = job + + resp = ev.JobResultView.as_view()(_authed(), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + assert resp.data == {"success": False, "status": "cancelled"} + assert not m_read.called + + +# --------------------------------------------------------------------------- +# (5d) result for a COMPLETED job with a blank result_ref (files already +# swept by TTL cleanup, row not yet expired) -> 404, same as the expired +# case -- exercises the new blank-ref branch distinctly from expiry. +# --------------------------------------------------------------------------- +@mock.patch.object(evr, "read_result") +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_result_completed_with_blank_ref_is_404(m_keys, m_jobs, m_read): + m_keys.get.return_value = kv_key() + job = AgentKVJob( + status=JobStatus.COMPLETED, + result_ref="", + expires_at=timezone.now() + timedelta(days=1), + ) + m_jobs.get.return_value = job + + resp = ev.JobResultView.as_view()(_authed(), job_id=uuid.uuid4()) + + assert resp.status_code == 404 + assert not m_read.called + + +# --------------------------------------------------------------------------- +# (6) cancel on RUNNING -> mark_terminal called with CANCELLED, 200. +# --------------------------------------------------------------------------- +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_cancel_on_running_marks_terminal_and_200s( + m_keys, m_jobs, m_mark_terminal, m_release +): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING) + job.organization_id = "org1" + m_jobs.get.return_value = job + + resp = ev.JobCancelView.as_view()(_authed(method="post"), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + assert resp.data == {"status": "cancelled"} + m_mark_terminal.assert_called_once_with( + job.id, job.organization_id, JobStatus.CANCELLED + ) + + +# --------------------------------------------------------------------------- +# (6b) cancel that WINS the terminal guard releases the concurrency slot -- +# a job cancelled BEFORE dispatch gets no finalize callback and the sweep's +# phase-1 only targets PENDING (not CANCELLED), so without this its slot +# would leak until the 6h TTL (pre-Greptile important #4). release() is +# idempotent (zrem), so a later finalize-callback release is harmless. +# --------------------------------------------------------------------------- +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_cancel_win_releases_concurrency_slot(m_keys, m_jobs, m_mark_terminal, m_release): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.PENDING) + job.organization_id = "org1" + m_jobs.get.return_value = job + _dispatched_in_db(m_jobs, False) + + resp = ev.JobCancelView.as_view()(_authed(method="post"), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + m_release.assert_called_once_with("org1", str(job.id)) + + +def _dispatched_in_db(m_jobs, dispatched: bool) -> None: + """Wire the DB re-read `_never_dispatched` performs. + + It no longer reads `job.dispatched_at` off the in-memory row -- that row + predates `mark_terminal`, and a dispatch completing in that gap made the + cancel release a slot whose executor was still running (round-2 C1). It now + asks the database, so these tests have to say what the database would + answer. + """ + m_jobs.filter.return_value.exists.return_value = dispatched + +# --------------------------------------------------------------------------- +# (6c) cancel that LOSES the guard (job already terminal) must NOT release -- +# whoever terminalized it (finalize callback / a prior cancel) owns the slot. +# --------------------------------------------------------------------------- +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=False) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_cancel_loss_does_not_release_slot(m_keys, m_jobs, m_mark_terminal, m_release): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.COMPLETED) + job.organization_id = "org1" + m_jobs.get.return_value = job + + resp = ev.JobCancelView.as_view()(_authed(method="post"), job_id=uuid.uuid4()) + + assert resp.status_code == 409 + assert not m_release.called + + +# --------------------------------------------------------------------------- +# (7) cancel on COMPLETED -> 409 with current status; result untouched +# (the guard lost, so nothing about the stored result is read or written). +# --------------------------------------------------------------------------- +@mock.patch.object(evr, "read_result") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=False) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_cancel_on_completed_is_409_and_result_untouched( + m_keys, m_jobs, m_mark_terminal, m_read +): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.COMPLETED) + job.organization_id = "org1" + m_jobs.get.return_value = job + + resp = ev.JobCancelView.as_view()(_authed(method="post"), job_id=uuid.uuid4()) + + assert resp.status_code == 409 + # Status is lowercased consistently across every endpoint (spec §7.2) -- + # the 409 body used to leak the raw uppercase value (pre-Greptile #5). + assert resp.data == {"status": "completed"} + assert m_mark_terminal.called + assert not m_read.called + + +# --------------------------------------------------------------------------- +# (8) delete calls delete_job_files and blanks both refs. +# --------------------------------------------------------------------------- +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_calls_delete_job_files_and_blanks_refs( + m_keys, m_jobs, m_delete_files, m_save +): + m_keys.get.return_value = kv_key() + job = AgentKVJob( + status=JobStatus.COMPLETED, + input_ref="org/o/agent_kv/j/input.pdf", + result_ref="org/o/agent_kv/j/result.json", + ) + m_jobs.get.return_value = job + m_delete_files.return_value = ["input_ref", "result_ref"] + + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + m_delete_files.assert_called_once_with(job) + assert job.input_ref == "" + assert job.result_ref == "" + # update_fields carries exactly the blanked refs, so a ref left set is not + # written back as "" by a wider save(). + assert m_save.call_args.kwargs["update_fields"] == ["input_ref", "result_ref"] + + +# --------------------------------------------------------------------------- +# (8b) delete on a non-terminal (RUNNING) job cancels it FIRST, before the +# files are deleted -- a still-running job that finalizes late would +# otherwise write a fresh result onto a job the caller just deleted. +# --------------------------------------------------------------------------- +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "delete_job_files") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_on_running_job_cancels_before_deleting_files( + m_keys, m_jobs, m_mark_terminal, m_delete_files, m_save +): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING) + job.organization_id = "org1" + m_jobs.get.return_value = job + m_delete_files.return_value = ["input_ref", "result_ref"] + + manager = mock.Mock() + manager.attach_mock(m_mark_terminal, "mark_terminal") + manager.attach_mock(m_delete_files, "delete_job_files") + + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + m_mark_terminal.assert_called_once_with(job.id, "org1", JobStatus.CANCELLED) + m_delete_files.assert_called_once_with(job) + assert [c[0] for c in manager.mock_calls] == ["mark_terminal", "delete_job_files"] + + +# --------------------------------------------------------------------------- +# (8e) DELETE on an in-flight job RELEASES the concurrency slot, exactly as +# JobCancelView does (Greptile review #2). The slot is taken at submit and +# released by _fail_job_response, the finalize callback and the sweep -- a job +# terminalized here before dispatch hits none of those, and the sweep's phase-1 +# only targets PENDING, never CANCELLED. Without the release each such delete +# leaked a slot until the 6h TTL, and enough of them exhaust the org's +# allowance and start rejecting new submits. +# --------------------------------------------------------------------------- +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "delete_job_files", return_value=["input_ref", "result_ref"]) +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_on_a_running_job_does_NOT_release_the_slot( + m_keys, m_jobs, m_mark_terminal, m_delete_files, m_save, m_release +): + """Inverted: this used to assert the release, which was the bug. + + Nothing revokes a running executor -- `job.task_id` is written at dispatch + and never read again -- so the engine keeps running, keeps calling LLMs and + keeps billing. Handing its slot back immediately let the next submit start + alongside it, so submit-then-delete in a loop ran arbitrarily many + concurrent extractions against a ceiling of AGENT_KV_CONCURRENT_LIMIT, all + paid for. + + The slot is released by `FinalizeView`'s `finally` when the executor's + callback lands. Reported as 2.5 in the branch review. + """ + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING) + job.organization_id = "org1" + job.dispatched_at = timezone.now() + m_jobs.get.return_value = job + + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + assert not m_release.called, ( + "the executor is still running and still billing; its slot belongs to " + "the finalize callback, not to this request" + ) + + +@mock.patch.object(ev, "dispatch_cancelled_webhook") +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "delete_job_files", return_value=["input_ref", "result_ref"]) +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_on_a_never_dispatched_job_DOES_release_the_slot( + m_keys, m_jobs, m_mark_terminal, m_delete_files, m_save, m_release, m_hook +): + """The case the release exists for, and the only one. + + A job cancelled before dispatch gets no finalize callback, and the sweep's + phase-1 only targets PENDING, never CANCELLED -- so without this its slot + leaks until the 6h TTL. + """ + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.PENDING) + job.organization_id = "org1" + job.dispatched_at = None + m_jobs.get.return_value = job + _dispatched_in_db(m_jobs, False) + + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + m_release.assert_called_once_with("org1", str(job.id)) + + +# --------------------------------------------------------------------------- +# (8f) ...but only when THIS request won the terminal-state race. A concurrent +# cancel or finalize that terminalized first is the one accounting for the +# slot; releasing on a lost race would hand a slot back twice. +# --------------------------------------------------------------------------- +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "save") +# A lost race re-reads the row before cleanup (see 8i); this suite touches no +# database, so the read itself is stubbed out here. +@mock.patch.object(AgentKVJob, "refresh_from_db") +@mock.patch.object(ev, "delete_job_files", return_value=["input_ref", "result_ref"]) +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=False) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_does_not_release_the_slot_when_it_loses_the_terminal_race( + m_keys, m_jobs, m_mark_terminal, m_delete_files, m_refresh, m_save, m_release +): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING) + job.organization_id = "org1" + m_jobs.get.return_value = job + + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + assert not m_release.called + + +# --------------------------------------------------------------------------- +# (8g) a ref whose file could NOT be deleted is left set (Greptile review #3). +# It is the only handle TTL cleanup can retry from: its candidate query matches +# rows by `input_ref > "" OR result_ref > ""`, so blanking a ref whose object is +# still in the bucket orphans that object permanently. Still 204 -- the job is +# terminal and the caller's intent is recorded. +# --------------------------------------------------------------------------- +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_keeps_the_ref_whose_file_delete_failed( + m_keys, m_jobs, m_delete_files, m_save +): + m_keys.get.return_value = kv_key() + job = AgentKVJob( + status=JobStatus.COMPLETED, + input_ref="org/o/agent_kv/j/input.pdf", + result_ref="org/o/agent_kv/j/result.json", + ) + m_jobs.get.return_value = job + # Input gone; the result delete raised inside delete_job_files. + m_delete_files.return_value = ["input_ref"] + + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + assert job.input_ref == "" + assert job.result_ref == "org/o/agent_kv/j/result.json" + assert m_save.call_args.kwargs["update_fields"] == ["input_ref"] + + +# --------------------------------------------------------------------------- +# (8c) delete on an already-terminal job never attempts to cancel it again -- +# unchanged behavior for the terminal case. +# --------------------------------------------------------------------------- +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "delete_job_files") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_on_terminal_job_does_not_call_mark_terminal( + m_keys, m_jobs, m_mark_terminal, m_delete_files, m_save +): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.COMPLETED) + job.organization_id = "org1" + m_jobs.get.return_value = job + + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + assert not m_mark_terminal.called + m_delete_files.assert_called_once_with(job) + + +# --------------------------------------------------------------------------- +# (9) every job-scoped endpoint 401s (403, per Forbidden.status_code) +# without a key (spec §6.8 regression). +# +# The brief's shown snippet for this test wraps the call in +# ``pytest.raises(Forbidden)``, mirroring test_auth.py -- but that suite +# calls the ``@validate_api_key``-decorated function directly, bypassing +# DRF's dispatch(). Routed through the real ``.as_view()()`` cycle (as here, +# and as every other view in this module is exercised), DRF's dispatch() +# catches the raised ``Forbidden`` (an APIException) via +# ``drf_standardized_errors``'s exception handler and renders it as a normal +# Response -- exactly like ``test_foreign_org_job_is_404`` above asserts +# ``resp.status_code`` for ``JobNotFound`` rather than expecting a raise. +# Verified empirically (see task-9-report.md); asserting the response here +# instead keeps the same regression coverage without a spurious failure. +# --------------------------------------------------------------------------- +def test_all_job_views_401_without_key(): + for view, method in [ + (ev.JobStatusView, "get"), + (ev.JobResultView, "get"), + (ev.JobCancelView, "post"), + (ev.JobDeleteView, "delete"), + ]: + req = getattr(APIRequestFactory(), method)("/agent-kv/x") + resp = view.as_view()(req, job_id=uuid.uuid4()) + assert resp.status_code == 403 + + +# --------------------------------------------------------------------------- +# (8i) Losing the terminal race must RE-READ the job before touching its files. +# +# The race: DELETE reads a RUNNING job (result_ref ""), the finalize callback +# wins the guarded UPDATE in between and writes COMPLETED + a real result_ref, +# and `mark_terminal` here returns False. Without a refresh the cleanup below +# runs against the stale copy: `delete_job_files` reports an already-empty +# `result_ref` as "cleared" (there was nothing to delete), and the save then +# writes "" OVER the winner's real ref. The job stays COMPLETED, its result +# 404s, and the object is orphaned -- TTL cleanup selects on `result_ref > ""`, +# so a blanked row never comes back. +# +# Reported by Greptile on PR #2317. +# --------------------------------------------------------------------------- +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "delete_job_files") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=False) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_refreshes_the_job_when_it_loses_the_terminal_race( + m_keys, m_jobs, m_mark_terminal, m_delete_files, m_save, m_release +): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING) + job.organization_id = "org1" + m_jobs.get.return_value = job + + # The winner's write, applied by refresh_from_db: COMPLETED + a real ref. + def _win_the_race(): + job.status = JobStatus.COMPLETED + job.result_ref = "org1/job/result.json" + + with mock.patch.object( + AgentKVJob, "refresh_from_db", side_effect=_win_the_race + ) as m_refresh: + # Capture what the cleanup actually SAW. `call_args` would be useless + # here: it records a reference to the same job object the view then + # blanks, so by assertion time it reads "" whether the fix works or not. + seen: list[str] = [] + + def _record(j): + seen.append(j.result_ref) + return ["input_ref", "result_ref"] if j.result_ref else ["input_ref"] + + m_delete_files.side_effect = _record + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + assert m_refresh.called, ( + "a lost terminal race must re-read the job; acting on the stale copy " + "blanks the winner's result_ref and orphans the result file" + ) + # The file helper saw the winner's ref, so the real object is deleted... + assert seen == ["org1/job/result.json"] + # ...and the blanking is of a ref that was genuinely cleared, not a stale "". + assert m_save.call_args.kwargs["update_fields"] == ["input_ref", "result_ref"] + + +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "delete_job_files", return_value=["input_ref", "result_ref"]) +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_does_not_refresh_when_it_wins_the_terminal_race( + m_keys, m_jobs, m_mark_terminal, m_delete_files, m_save, m_release +): + """Winning means nothing else wrote to the row, so the in-memory copy is + current and the extra query would be waste on the common path. + """ + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING) + job.organization_id = "org1" + m_jobs.get.return_value = job + + with mock.patch.object(AgentKVJob, "refresh_from_db") as m_refresh: + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + assert not m_refresh.called + + +# --------------------------------------------------------------------------- +# (8j) Cancellation must deliver the documented terminal webhook. +# +# Cancellation never reaches finalize -- JobCancelView and DELETE terminalize +# the row themselves -- so nothing in the finalize path ever fired a webhook for +# a cancelled job, and a late executor callback cannot either: it loses the +# terminal guard, and `_maybe_webhook` correctly declines a non-fresh finalize. +# A caller who supplied `webhook_url` was simply never told, despite docs §8. +# +# Guarded on `won`, which is what keeps exactly one path owning the +# notification: a cancel that LOST means a finalize won and will send. +# +# Reported by Greptile on PR #2317. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "dispatch_cancelled_webhook") +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_cancel_that_wins_queues_the_terminal_webhook( + m_keys, m_jobs, m_mark_terminal, m_release, m_hook +): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="https://hook.example/x") + job.organization_id = "org1" + m_jobs.get.return_value = job + + resp = ev.JobCancelView.as_view()(_authed(method="post"), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + assert m_hook.called, "a won cancel must notify; nothing else will" + assert m_hook.call_args.args[0] is job + + +@mock.patch.object(ev, "dispatch_cancelled_webhook") +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=False) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_cancel_that_loses_does_not_queue_a_webhook( + m_keys, m_jobs, m_mark_terminal, m_release, m_hook +): + """The finalize that beat it owns the notification. Sending here too would + double-notify the caller for one job. + """ + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="https://hook.example/x") + job.organization_id = "org1" + m_jobs.get.return_value = job + + resp = ev.JobCancelView.as_view()(_authed(method="post"), job_id=uuid.uuid4()) + + assert resp.status_code == 409 + assert not m_hook.called + + +@mock.patch.object(ev, "dispatch_cancelled_webhook") +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "delete_job_files", return_value=["input_ref", "result_ref"]) +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_delete_that_cancels_a_running_job_also_queues_the_webhook( + m_keys, m_jobs, m_mark_terminal, m_delete_files, m_save, m_release, m_hook +): + """DELETE on a non-terminal job cancels it first; that cancellation is just + as terminal, and just as invisible to the caller without this. + """ + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING, webhook_url="https://hook.example/x") + job.organization_id = "org1" + m_jobs.get.return_value = job + + resp = ev.JobDeleteView.as_view()(_authed(method="delete"), job_id=uuid.uuid4()) + + assert resp.status_code == 204 + assert m_hook.called + + +# --------------------------------------------------------------------------- +# (8k) The same narrowing on JobCancelView, which has the identical shape. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "dispatch_cancelled_webhook") +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_cancel_on_a_running_job_does_NOT_release_the_slot( + m_keys, m_jobs, m_mark_terminal, m_release, m_hook +): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.RUNNING) + job.organization_id = "org1" + job.dispatched_at = timezone.now() + m_jobs.get.return_value = job + + resp = ev.JobCancelView.as_view()(_authed(method="post"), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + assert not m_release.called + + +@mock.patch.object(ev, "dispatch_cancelled_webhook") +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_cancel_on_a_never_dispatched_job_DOES_release_the_slot( + m_keys, m_jobs, m_mark_terminal, m_release, m_hook +): + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.PENDING) + job.organization_id = "org1" + job.dispatched_at = None + m_jobs.get.return_value = job + _dispatched_in_db(m_jobs, False) + + resp = ev.JobCancelView.as_view()(_authed(method="post"), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + m_release.assert_called_once_with("org1", str(job.id)) + + +@mock.patch.object(ev, "dispatch_cancelled_webhook") +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +@mock.patch.object(AgentKVKey, "objects") +def test_a_cancel_that_raced_a_dispatch_does_NOT_release_the_slot( + m_keys, m_jobs, m_mark_terminal, m_release, m_hook +): + """Round-2 C1. The in-memory row says undispatched; the database disagrees. + + The interleaving: the cancel reads a PENDING row with `dispatched_at=NULL`, + then `dispatch_job` passes its own terminal re-read, enqueues, and stamps + `dispatched_at` / DISPATCHED. The cancel's `mark_terminal` still wins + (DISPATCHED is not terminal), and the old check -- reading + `job.dispatched_at` off the snapshot -- answered True and released a slot + whose executor was running and billing. + + That is a smaller instance of the very over-release bug the narrowing was + written to fix. The row here is deliberately left looking undispatched, so + the ONLY thing that can produce the right answer is the DB re-read. + """ + m_keys.get.return_value = kv_key() + job = AgentKVJob(status=JobStatus.PENDING) + job.organization_id = "org1" + job.dispatched_at = None # the stale snapshot, as the cancel path sees it + m_jobs.get.return_value = job + _dispatched_in_db(m_jobs, True) # what actually happened meanwhile + + resp = ev.JobCancelView.as_view()(_authed(method="post"), job_id=uuid.uuid4()) + + assert resp.status_code == 200 + assert not m_release.called, ( + "released a slot for a job that HAD been dispatched -- the executor is " + "still running and its finalize callback owns the slot" + ) + # The caller is still told, since cancellation never reaches finalize. + assert m_hook.called + + +@mock.patch.object(AgentKVJob, "objects") +def test_the_undispatched_check_asks_the_database_not_the_row(m_jobs): + """Pins the mechanism, so the C1 regression cannot return quietly. + + If `_never_dispatched` goes back to reading `job.dispatched_at`, this fails: + the row says dispatched (which the old code would honour) while the DB says + it is not, and only the DB answer is correct after a `mark_terminal`. + """ + job = AgentKVJob(status=JobStatus.CANCELLED) + job.dispatched_at = timezone.now() + _dispatched_in_db(m_jobs, False) + + assert ev._never_dispatched(job) is True + + _dispatched_in_db(m_jobs, True) + assert ev._never_dispatched(job) is False diff --git a/backend/agent_kv/tests/test_key_management.py b/backend/agent_kv/tests/test_key_management.py new file mode 100644 index 0000000000..f092f53bdd --- /dev/null +++ b/backend/agent_kv/tests/test_key_management.py @@ -0,0 +1,95 @@ +import os +import uuid +from unittest import mock + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from utils.filters.organization_filter import OrganizationFilterBackend # noqa: E402 + +from agent_kv.models import AgentKVKey # noqa: E402 +from agent_kv.serializers import AgentKVKeyWriteSerializer # noqa: E402 +from agent_kv.views import AgentKVKeyViewSet # noqa: E402 + + +def test_write_serializer_rejects_blank_name(): + s = AgentKVKeyWriteSerializer(data={"name": "", "description": "x"}) + assert not s.is_valid() + assert "name" in s.errors + + +def test_rotate_assigns_fresh_key_and_saves(): + key_obj = AgentKVKey(name="k", key=uuid.uuid4()) + old = key_obj.key + with mock.patch.object(AgentKVKey, "save") as m_save: + view = AgentKVKeyViewSet() + view.get_object = lambda: key_obj + view.format_kwarg = None + view.request = mock.Mock() + resp = view.rotate(view.request, pk=str(key_obj.id)) + assert key_obj.key != old + assert m_save.called + assert resp.status_code == 200 + + +@mock.patch.object(AgentKVKey, "objects") +def test_org_scoping_via_filter_backend_is_engaged(m_objects): + """Org-scoping is enforced by the global filter backend, not the manager. + + Two things must hold for that guard to stay real: (1) the viewset must + keep ``OrganizationFilterBackend`` in its resolved ``filter_backends`` — + so nobody can quietly swap it out later — and (2) ``get_queryset`` must + hand back the *unscoped* base queryset (``AgentKVKey.objects.all()``, + no ``organization=`` kwarg of its own) rather than pre-filtering, which + would keep working (and hide the loss of protection) even if the backend + were ever removed from ``DEFAULT_FILTER_BACKENDS``. + """ + view = AgentKVKeyViewSet() + + assert OrganizationFilterBackend in view.filter_backends + + result = view.get_queryset() + + m_objects.all.assert_called_once_with() + m_objects.filter.assert_not_called() + assert result is m_objects.all.return_value + + +def test_create_returns_id_and_key(): + """POST .../keys/ must echo the read serializer, not the write one. + + A caller who just created a key has no other way to learn its ``id`` or + the raw ``key`` value -- ``AgentKVKeyWriteSerializer`` (used to validate + the request) carries neither. ``serializer.save`` is stubbed rather than + hitting the real ORM: this file has no DB fixture, and ``ModelSerializer + .save()``'s default create path resolves ``AgentKVKey._default_manager`` + independently of any ``mock.patch.object(AgentKVKey, "objects")`` swap + (a different manager reference), so mocking ``.objects`` the way the + other tests in this file do would silently fall through to a real + database call here. + """ + key_obj = AgentKVKey( + id=uuid.uuid4(), name="k", description="", key=uuid.uuid4(), is_active=True + ) + + def _fake_save(self, **kwargs): + self.instance = key_obj + return key_obj + + view = AgentKVKeyViewSet() + view.action = "create" + view.format_kwarg = None + view.kwargs = {} + view.request = mock.Mock(data={"name": "k"}, user=mock.Mock()) + + with mock.patch.object(AgentKVKeyWriteSerializer, "save", _fake_save): + resp = view.create(view.request) + + assert resp.status_code == 201, resp.data + # DRF serializes UUID fields as strings. + assert resp.data["id"] == str(key_obj.id), resp.data + assert resp.data["key"] == str(key_obj.key), resp.data diff --git a/backend/agent_kv/tests/test_maintenance_commands.py b/backend/agent_kv/tests/test_maintenance_commands.py new file mode 100644 index 0000000000..b527b15a6c --- /dev/null +++ b/backend/agent_kv/tests/test_maintenance_commands.py @@ -0,0 +1,54 @@ +"""Management-command entrypoints for the Agent-KV maintenance jobs. + +Same no-test-DB style as the rest of ``agent_kv/tests``: real Django app +registry via ``django.setup()``, no database. ``agent_kv.maintenance``'s two +entrypoints are patched so no query ever runs -- their own behaviour +(candidate queries, guard semantics, batching) is covered where it lives now, +``agent_kv/maintenance.py``, exercised indirectly through the views in +``test_sweeps.py``. What's pinned here is only the command<->function wiring: +each command calls its maintenance function exactly once, with no arguments, +and writes the returned dict to stdout as JSON -- the shape a Kubernetes +CronJob log line needs. +""" + +import json +import os +from io import StringIO +from unittest import mock + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from django.core.management import call_command # noqa: E402 + +from agent_kv import maintenance # noqa: E402 + + +@mock.patch.object(maintenance, "run_sweep") +def test_agent_kv_sweep_command_calls_run_sweep_once_and_prints_its_result( + m_run_sweep, +): + m_run_sweep.return_value = {"swept": 2, "timed_out": 1} + out = StringIO() + + call_command("agent_kv_sweep", stdout=out) + + m_run_sweep.assert_called_once_with() + assert json.loads(out.getvalue()) == {"swept": 2, "timed_out": 1} + + +@mock.patch.object(maintenance, "run_ttl_cleanup") +def test_agent_kv_ttl_cleanup_command_calls_run_ttl_cleanup_once_and_prints_its_result( + m_run_ttl_cleanup, +): + m_run_ttl_cleanup.return_value = {"cleaned": 5} + out = StringIO() + + call_command("agent_kv_ttl_cleanup", stdout=out) + + m_run_ttl_cleanup.assert_called_once_with() + assert json.loads(out.getvalue()) == {"cleaned": 5} diff --git a/backend/agent_kv/tests/test_models.py b/backend/agent_kv/tests/test_models.py new file mode 100644 index 0000000000..56d32a800f --- /dev/null +++ b/backend/agent_kv/tests/test_models.py @@ -0,0 +1,145 @@ +"""Terminal-state write guard: the invariant everything else leans on (spec §5.4).""" + +import os +import uuid +from unittest import mock + +import django +import pytest +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from agent_kv.models import AgentKVJob, JobStatus # noqa: E402 + + +def test_terminal_set_is_exactly_the_three_states(): + assert AgentKVJob.TERMINAL == frozenset( + {JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED} + ) + + +@mock.patch.object(AgentKVJob, "objects") +def test_mark_terminal_excludes_terminal_rows_and_reports_success(m_objects): + m_qs = m_objects.filter.return_value.exclude.return_value + m_qs.update.return_value = 1 + job_id = uuid.uuid4() + ok = AgentKVJob.mark_terminal( + job_id=job_id, + organization_id="org1", + new_status=JobStatus.FAILED, + error="boom", + ) + assert ok is True + filter_kwargs = m_objects.filter.call_args.kwargs + assert filter_kwargs["id"] == job_id + assert filter_kwargs["organization_id"] == "org1" + _, exclude_kwargs = m_objects.filter.return_value.exclude.call_args + assert set(exclude_kwargs["status__in"]) == set(AgentKVJob.TERMINAL) + update_kwargs = m_qs.update.call_args.kwargs + assert update_kwargs["status"] == JobStatus.FAILED + assert update_kwargs["error"] == "boom" + assert "completed_at" in update_kwargs + + +@mock.patch.object(AgentKVJob, "objects") +def test_mark_terminal_on_already_terminal_row_is_noop_false(m_objects): + m_objects.filter.return_value.exclude.return_value.update.return_value = 0 + ok = AgentKVJob.mark_terminal( + job_id=uuid.uuid4(), + organization_id="org1", + new_status=JobStatus.COMPLETED, + ) + assert ok is False + + +# -------------------------------------------------------------------------- +# 2.16: the guarded UPDATE excludes terminal ROWS, not non-terminal ARGUMENTS. +# -------------------------------------------------------------------------- + + +@mock.patch.object(AgentKVJob, "objects") +@pytest.mark.parametrize( + "bad_status", [JobStatus.PENDING, JobStatus.DISPATCHED, JobStatus.RUNNING] +) +def test_mark_terminal_refuses_a_non_terminal_status(m_objects, bad_status): + """`mark_terminal(..., RUNNING)` used to stamp `completed_at=now()`. + + That leaves a row that reads as finished to everything keying off + `completed_at` (the TTL filter, the sweep's cancelled-job phase) while + still being invisible to the terminal guard -- so nothing can ever + terminalize it again. No caller does this today; the method is named for + the invariant, so it enforces it. + """ + with pytest.raises(ValueError, match="non-terminal status"): + AgentKVJob.mark_terminal( + job_id=uuid.uuid4(), + organization_id="org1", + new_status=bad_status, + ) + assert not m_objects.filter.called, "no write may be attempted" + + +@mock.patch.object(AgentKVJob, "objects") +@pytest.mark.parametrize("good_status", sorted(AgentKVJob.TERMINAL)) +def test_mark_terminal_accepts_every_terminal_status(m_objects, good_status): + m_objects.filter.return_value.exclude.return_value.update.return_value = 1 + assert ( + AgentKVJob.mark_terminal( + job_id=uuid.uuid4(), + organization_id="org1", + new_status=good_status, + ) + is True + ) + + +# -------------------------------------------------------------------------- +# 2.17: `extractor` was the one stringly-typed field without `choices`, and it +# defaulted to the extractor this build cannot run. +# -------------------------------------------------------------------------- + + +def test_extractor_declares_its_choices(): + field = AgentKVJob._meta.get_field("extractor") + assert {value for value, _ in field.choices} == {"kv", "table"} + + +def test_extractor_has_no_default(): + """An omitted `extractor=` must be loud, not silently filed as `kv`. + + `kv` IS a valid key in `STAGE_NAMES_BY_EXTRACTOR`, so a table job filed + under it gets the KV stage list and `table_extraction` is dropped from + every status response -- the job runs, the caller is billed, the stages + array comes back empty and nothing logs a warning, because nothing is + wrong as far as the filter can tell. + + With no default the omission produces `""`, which matches no route: the + dispatch raises, the job terminalizes FAILED with a visible error, and + `_status_document` logs the unknown-extractor warning. + """ + field = AgentKVJob._meta.get_field("extractor") + assert not field.has_default(), ( + "migration 0002's `default='kv'` was the historical truth then and is " + "a mis-filing trap now; see the field comment" + ) + assert AgentKVJob(organization_id="o").extractor == "" + + +def test_recordable_extractors_are_a_superset_of_routable_ones(): + """The two sets are deliberately different, so neither is derived. + + `choices` says what a ROW may record -- rows written before the carve-out + legitimately say `kv`. `EXTRACTOR_ROUTES` says what a submit may DISPATCH, + and `kv` is absent from it on purpose. Deriving either from the other + would quietly re-enable the extractor or make old rows unreadable. + """ + from agent_kv.constants import EXTRACTOR_ROUTES + + recordable = {value for value, _ in AgentKVJob._meta.get_field("extractor").choices} + assert set(EXTRACTOR_ROUTES) < recordable, ( + "every routable extractor must be recordable, and `kv` is recordable " + "without being routable" + ) diff --git a/backend/agent_kv/tests/test_periodic_tasks_migration.py b/backend/agent_kv/tests/test_periodic_tasks_migration.py new file mode 100644 index 0000000000..f03c4a7c85 --- /dev/null +++ b/backend/agent_kv/tests/test_periodic_tasks_migration.py @@ -0,0 +1,166 @@ +"""One owner for the two maintenance schedules, and it is the CronJob. + +History, because the test's shape only makes sense with it: an earlier revision +of migration 0004 SEEDED two ``PgPeriodicTask`` rows to fire +``agent_kv.sweep`` and ``agent_kv.ttl_cleanup``, answering a review finding +about the sweep never running. The finding was real; the fix was redundant. +``templates/backend/agent-kv-cronjobs.yaml`` in the cloud chart already ran the +same two management commands, added earlier on the same branch, so the +migration's stated premise -- "no ``PgPeriodicTask``, no beat entry, no +CronJob" -- was false about its own tree. + +The rows were seeded ``pg_owned: False`` and the PG scheduler claims only +``WHERE pg_owned AND enabled``, so nothing double-fired yet. That is the +dangerous kind of correct: the flag exists to be flipped, and a flip would have +put a ``*/10`` sweep against the chart's ``*/15`` -- colliding at ``:30``, +contending on the same job rows -- with TTL cadences disagreeing 24x. + +So 0004 now DELETES those rows and the CronJob is the sole owner. What these +tests defend is that decision, in the two ways it can be undone: someone +re-pointing the forward operation at the seeding function, or someone dropping +the reverse and making the migration one-way. + +The name/queue assertions are kept rather than deleted -- they still guard the +reverse path, and the same wire names are what the CronJob's management +commands resolve to. +""" + +import importlib +import inspect + +import pytest +from django.db import migrations + +# The module name starts with a digit, so it cannot be a plain `import`. +MIGRATION = importlib.import_module("agent_kv.migrations.0004_pg_periodic_tasks") + +#: Wire names registered by `@worker_task(name=...)` in +#: `workers/scheduler/agent_kv_tasks.py`. Duplicated as literals on purpose: +#: the workers package is not importable from the backend's test venv, and an +#: import would make this pass vacuously wherever it is absent -- which is +#: exactly the environment where the two sides drift apart. +WORKER_TASK_NAMES = {"agent_kv.sweep", "agent_kv.ttl_cleanup"} + + +def _specs() -> dict: + return {spec["name"]: spec for spec in MIGRATION.PG_PERIODIC_TASKS} + + +def _run_python_ops() -> list: + return [ + op + for op in MIGRATION.Migration.operations + if isinstance(op, migrations.RunPython) + ] + + +# --------------------------------------------------------------------------- +# The ownership decision. +# --------------------------------------------------------------------------- + + +def test_the_migration_still_has_a_runpython_operation(): + """Guards the empty-operations failure mode. + + Every other assertion in this file reads `PG_PERIODIC_TASKS` or a function's + source, so `operations = []` would leave the whole suite green while the + migration did nothing at all -- applying cleanly and leaving the seeded rows + in place in every environment that already has them. + """ + assert len(_run_python_ops()) == 1, MIGRATION.Migration.operations + + +def test_the_forward_operation_removes_the_rows_rather_than_seeding_them(): + """The decision itself. + + If this fails, someone has re-pointed forward at `create_pg_periodic_tasks` + and there are two owners for one schedule again. Read the migration's + docstring before changing it: the collision is at `:30`, and only the + CronJob side carries `concurrencyPolicy: Forbid` and a deadline. + """ + assert _run_python_ops()[0].code is MIGRATION.remove_pg_periodic_tasks + + +def test_the_migration_is_reversible(): + """`migrate agent_kv 0003` must land back on the previous tree's state. + + An irreversible data migration here would mean a rollback of this release + leaves an environment with no rows and no way to restore them except by + hand -- and the seeding revision is what the previous release shipped. + """ + op = _run_python_ops()[0] + + assert op.reverse_code is MIGRATION.create_pg_periodic_tasks + assert op.reverse_code is not migrations.RunPython.noop + + +def test_removal_is_idempotent_and_targets_only_the_two_seeded_names(): + """A fresh install has no such rows, and re-running must still succeed. + + Equally important: the delete is scoped by name. An unscoped + `PgPeriodicTask.objects.all().delete()` here would wipe the + dashboard_metrics periodics this migration sits alongside. + """ + source = inspect.getsource(MIGRATION.remove_pg_periodic_tasks) + + assert "name__in=" in source, "the delete must be scoped by name" + assert ".all()" not in source + + +# --------------------------------------------------------------------------- +# The reverse path's contents, and the names the CronJob commands share. +# --------------------------------------------------------------------------- + + +def test_both_maintenance_tasks_are_accounted_for(): + assert set(_specs()) == {"agent_kv_sweep", "agent_kv_ttl_cleanup"} + + +def test_every_row_points_at_a_task_the_worker_registers(): + """A row naming a task nothing registers is a schedule that fires into + nothing. Still live for the reverse path, and the same wire names back the + `agent_kv_sweep` / `agent_kv_ttl_cleanup` management commands the CronJob + invokes. + """ + declared = {spec["task_name"] for spec in MIGRATION.PG_PERIODIC_TASKS} + + assert declared == WORKER_TASK_NAMES, ( + f"migration names {sorted(declared)} but workers/scheduler/" + f"agent_kv_tasks.py registers {sorted(WORKER_TASK_NAMES)}" + ) + + +@pytest.mark.parametrize("name", ["agent_kv_sweep", "agent_kv_ttl_cleanup"]) +def test_rows_are_scheduled_onto_the_scheduler_queue(name): + assert _specs()[name]["queue"] == "scheduler" + + +@pytest.mark.parametrize("name", ["agent_kv_sweep", "agent_kv_ttl_cleanup"]) +def test_cron_strings_are_five_field_and_parseable(name): + """A malformed cron string is accepted by the column and then never fires.""" + fields = _specs()[name]["cron_string"].split() + + assert len(fields) == 5, _specs()[name]["cron_string"] + + +def test_the_restored_rows_are_inert(): + """`pg_owned: False` on the reverse path too. + + A reverse that restored them ENABLED would turn a rollback into the + double-fire this migration exists to prevent. + """ + source = inspect.getsource(MIGRATION.create_pg_periodic_tasks) + + assert '"pg_owned": False' in source + assert "update_or_create" in source, "re-running must stay idempotent" + + +def test_restored_rows_are_platform_wide(): + """Both internal endpoints sweep across every organization (spec §5.4), so a + non-empty `org_id` would scope them to one tenant and silently leave every + other org unswept. + """ + source = inspect.getsource(MIGRATION.create_pg_periodic_tasks) + + assert '"org_id": ""' in source + assert '"enabled": True' in source diff --git a/backend/agent_kv/tests/test_queue_wiring_is_derived.py b/backend/agent_kv/tests/test_queue_wiring_is_derived.py new file mode 100644 index 0000000000..e2c9109e82 --- /dev/null +++ b/backend/agent_kv/tests/test_queue_wiring_is_derived.py @@ -0,0 +1,297 @@ +"""Fleet queue wiring, DERIVED from `EXTRACTOR_ROUTES` rather than literal-pinned. + +`workers/tests/test_queue_consumer_wiring.py` guards the same property with +hardcoded queue names (`UNSERVED_KV_QUEUE`, `AGENT_TABLE_QUEUE`). That works +today and catches a real, expensive failure -- but it cannot survive the change +it most needs to survive. Re-enabling the `kv` extractor means adding one entry +to `EXTRACTOR_ROUTES`, at which point `celery_executor_agentic_kv` flips from +"must NOT be advertised" to "must be served", and the literal-pinned assertion +has to be *deleted* to make the suite pass. An assertion whose correct response +to a change is deletion is not protecting anything at that moment. + +This suite derives both directions from the routing table, so the same change +flips the expectation automatically: + +* Every executor in `EXTRACTOR_ROUTES` must have its `celery_executor_` + queue served by every fleet configuration. Unserved means jobs are accepted + and sit in DISPATCHED forever -- `enqueue_task` succeeds, rows land in + `pg_queue_message`, nothing errors at the producer. +* Every executor this module names but deliberately does NOT route + (`EXECUTOR_NAME` = `agentic_kv`, kept dormant by the carve-out) must NOT be + advertised anywhere. Same silent failure from the other side. + +Lives in the backend suite because that is where `EXTRACTOR_ROUTES` is +importable; the fleet configs it reads are plain text. The workers-side suite +keeps its literal assertions as a cheap second opinion -- two independent +statements of the same fact, which is the point. + +**Every path is guarded by `pytest.skip` when absent.** A cross-tree read that +assumed a file existed has reddened CI on this branch once already: the backend +test venv runs from `backend/`, and nothing guarantees a packaging step kept +`docker/` or `workers/` alongside it. +""" + +import re +from pathlib import Path + +import pytest +import yaml + +from agent_kv.constants import ( + EXECUTOR_NAME, + EXTRACTOR_ROUTES, + TABLE_EXECUTOR_NAME, +) + +REPO_ROOT = Path(__file__).resolve().parents[3] +DEV_COMPOSE = REPO_ROOT / "docker" / "docker-compose.yaml" +TEST_COMPOSE = REPO_ROOT / "tests" / "compose" / "docker-compose.test.yaml" +RUN_WORKER = REPO_ROOT / "workers" / "run-worker.sh" +RUN_WORKER_DOCKER = REPO_ROOT / "workers" / "run-worker-docker.sh" + +PG_QUEUE_VAR = "WORKER_PG_QUEUE_CONSUMER_QUEUE" + +#: Every executor name this module knows about. The ones in `EXTRACTOR_ROUTES` +#: must be served; the rest must not be advertised. Derived, so adding a route +#: moves a name from one set to the other with no edit here. +ALL_KNOWN_EXECUTORS = {EXECUTOR_NAME, TABLE_EXECUTOR_NAME} + + +def _queue_name(executor: str) -> str: + """How `pg_queue` derives a queue from an executor name. + + Mirrors the dispatcher's own derivation. If that convention changes, every + assertion here goes stale together rather than one at a time. + """ + return f"celery_executor_{executor}" + + +def _routed_executors() -> set[str]: + return {executor for executor, _operation in EXTRACTOR_ROUTES.values()} + + +def _unrouted_executors() -> set[str]: + return ALL_KNOWN_EXECUTORS - _routed_executors() + + +def _require(path: Path) -> str: + if not path.is_file(): + pytest.skip(f"{path} is not present in this tree") + return path.read_text() + + +def _queues(raw: str) -> set[str]: + """Queue names from a consumer list, ignoring a `${VAR:-default}` wrapper.""" + inner = re.sub(r"^\$\{[^:}]+:-(.*)\}$", r"\1", raw.strip()) + return {q.strip() for q in inner.split(",") if q.strip()} + + +def _compose_executor_queues(path: Path) -> set[str]: + doc = yaml.safe_load(_require(path)) + env = (doc.get("services", {}).get("worker-pg-executor") or {}).get( + "environment" + ) or [] + if isinstance(env, dict): + values = {k: str(v) for k, v in env.items()} + else: + values = {} + for item in env: + key, _, value = str(item).partition("=") + values[key] = value + raw = values.get(PG_QUEUE_VAR) + assert raw is not None, ( + f"{path.name}: worker-pg-executor sets no {PG_QUEUE_VAR}, so the " + f"running pg-queue-consumer falls back to a default of unknown content" + ) + return _queues(raw) + + +def _shell_role_queues(path: Path, pattern: str) -> set[str]: + match = re.search(pattern, _require(path)) + assert match, f"could not find {pattern!r} in {path.name}" + return _queues(match.group(1)) + + +#: Every place a fleet declares which executor queues it drains. The two shell +#: runners are included because a queue missing from a launcher is exactly as +#: undrained as one missing from compose -- and `run-worker-docker.sh` was +#: guarded by nothing at all. +def _all_executor_fleets() -> dict: + return { + "dev compose": lambda: _compose_executor_queues(DEV_COMPOSE), + "e2e test compose": lambda: _compose_executor_queues(TEST_COMPOSE), + "run-worker.sh": lambda: _shell_role_queues( + RUN_WORKER, r'\["\$PG_ROLE_EXECUTOR"\]="executor;([^"]+)"' + ), + } + + +@pytest.mark.parametrize("fleet", sorted(_all_executor_fleets())) +def test_every_routed_executor_has_a_consumer(fleet): + """The forward direction, derived. + + A route with no consumer accepts work and drains nothing, silently -- the + failure that cost this team ~30 hours of firings with zero executions. + """ + served = _all_executor_fleets()[fleet]() + expected = {_queue_name(e) for e in _routed_executors()} + + missing = expected - served + assert not missing, ( + f"{fleet} serves none of {sorted(missing)}, but EXTRACTOR_ROUTES " + f"dispatches there. Jobs will be accepted and sit in DISPATCHED " + f"forever with no error at the producer." + ) + + +@pytest.mark.parametrize("fleet", sorted(_all_executor_fleets())) +def test_no_fleet_advertises_an_unrouted_executor(fleet): + """The inverse direction, derived. + + `agentic_kv` is the live case: this deployment ships no such plugin, so a + fleet listing its queue would accept work nothing can drain. When `kv` is + re-enabled by adding its `EXTRACTOR_ROUTES` entry, this expectation inverts + on its own and `test_every_routed_executor_has_a_consumer` starts requiring + the queue instead -- which is the whole reason this is derived. + """ + served = _all_executor_fleets()[fleet]() + forbidden = {_queue_name(e) for e in _unrouted_executors()} + + advertised = forbidden & served + assert not advertised, ( + f"{fleet} advertises {sorted(advertised)}, but no EXTRACTOR_ROUTES " + f"entry dispatches there and this deployment carries no such plugin. " + f"Remove it, or add the route and ship the plugin." + ) + + +def test_the_derivation_is_not_vacuous(): + """Both sets must be non-empty, or the two tests above assert nothing. + + If `EXTRACTOR_ROUTES` were emptied, `test_every_routed_executor_...` would + pass with an empty expectation -- green while nothing was wired. If every + known executor were routed, the inverse test would pass trivially. This is + the guard on the guards. + """ + assert _routed_executors(), "no routed executors: the forward test is vacuous" + assert _unrouted_executors(), ( + "every known executor is routed, so the inverse test is vacuous. If " + "`kv` was deliberately re-enabled, that is correct -- delete this " + "assertion and say so." + ) + + +# --------------------------------------------------------------------------- +# The callback queue, across BOTH launchers. +# +# `run-worker-docker.sh` carries its own `["ide_callback"]= +# "ide_callback,agent_kv_callback"` map and was guarded by nothing: dropping +# `agent_kv_callback` there passed every test in the tree while a +# docker-launched fleet drained no terminal callbacks, leaving every job +# RUNNING forever. +# --------------------------------------------------------------------------- + +CALLBACK_QUEUE = "agent_kv_callback" + + +def test_the_docker_launcher_drains_the_terminal_callback_queue(): + """The launcher the workers-side suite does not cover.""" + queues = _shell_role_queues( + RUN_WORKER_DOCKER, r'\["ide_callback"\]="([^"]+)"' + ) + + assert CALLBACK_QUEUE in queues, ( + f"run-worker-docker.sh's ide_callback role omits {CALLBACK_QUEUE}; a " + f"docker-launched fleet would run jobs to completion and never " + f"terminalize them -- finalize is what persists the result, deletes " + f"the staged input, releases the slot and fires the webhook." + ) + + +def test_the_host_launcher_drains_the_terminal_callback_queue(): + """Already covered on the workers side; repeated here so the two launchers + are asserted side by side and a reader can see neither is missing. + """ + queues = _shell_role_queues( + RUN_WORKER, r'\["\$PG_ROLE_IDE_CALLBACK"\]="ide_callback;([^"]+)"' + ) + + assert CALLBACK_QUEUE in queues + + +# --------------------------------------------------------------------------- +# The billing backstop, DERIVED from `EXTRACTOR_ROUTES` the same way. +# +# `executor.tasks._LLM_BEARING_OPS` is the set of operations that must log when +# a successful run emits no usage records. That log line is the ONLY signal of +# total billing loss: the cloud `flush()` returns an empty list rather than +# raising, so "the billing chain broke" and "this job legitimately made no LLM +# calls" are otherwise indistinguishable. An op missing from the set is a +# missing ALARM, and the symptom is silence. +# +# `workers/tests/test_llm_bearing_ops.py` guards it against a hand-maintained +# `PAID_OPERATIONS` twin, which enforces symmetry between two literals -- so it +# can only fail when the copies DISAGREE, never when both are wrong together. A +# new paid extractor added to `EXTRACTOR_ROUTES` and declared in neither list +# ships with no billing alarm and nothing goes red. +# +# Every operation this API routes drives LLM calls by construction -- that is +# what an Agent-KV extractor IS -- so the set can be derived rather than +# declared. +# +# Read as TEXT and parsed with `ast`, not imported: `executor.tasks` is a +# workers package and is not importable from the backend's test venv (the +# mirror of why the workers-side suite cannot import `EXTRACTOR_ROUTES`). An +# import would make this pass vacuously in exactly the environment where the +# drift happens. Skipped if the file is absent from this tree. +# --------------------------------------------------------------------------- + +EXECUTOR_TASKS = REPO_ROOT / "workers" / "executor" / "tasks.py" + + +def _llm_bearing_ops_literal() -> set[str]: + """Pull `_LLM_BEARING_OPS`' string members out of the source.""" + import ast + + tree = ast.parse(_require(EXECUTOR_TASKS)) + for node in ast.walk(tree): + if not isinstance(node, ast.Assign): + continue + targets = [t.id for t in node.targets if isinstance(t, ast.Name)] + if "_LLM_BEARING_OPS" not in targets: + continue + return { + element.value + for element in ast.walk(node.value) + if isinstance(element, ast.Constant) and isinstance(element.value, str) + } + pytest.fail(f"no `_LLM_BEARING_OPS` assignment found in {EXECUTOR_TASKS}") + + +def test_every_routed_operation_has_a_billing_backstop(): + """The derived assertion the two literal lists cannot make. + + Add an extractor to `EXTRACTOR_ROUTES` and forget the backstop, and this + fails -- without anyone having had to remember a second list. + """ + declared = _llm_bearing_ops_literal() + routed_ops = {operation for _executor, operation in EXTRACTOR_ROUTES.values()} + + missing = routed_ops - declared + assert not missing, ( + f"{sorted(missing)} are routed by EXTRACTOR_ROUTES but absent from " + f"`_LLM_BEARING_OPS` in workers/executor/tasks.py. Every Agent-KV " + f"extractor drives LLM calls, so a run of one that emits zero usage " + f"records means the billing chain broke -- and without the set " + f"membership, nothing logs it. Add the op there (and to " + f"workers/tests/test_llm_bearing_ops.py's PAID_OPERATIONS)." + ) + + +def test_the_backstop_derivation_is_not_vacuous(): + """If the AST parse silently returned nothing, the test above would pass.""" + assert _llm_bearing_ops_literal(), ( + "parsed no operation names out of `_LLM_BEARING_OPS` -- the assignment " + "shape changed and the derivation above is now asserting against an " + "empty set" + ) diff --git a/backend/agent_kv/tests/test_rate_limiter.py b/backend/agent_kv/tests/test_rate_limiter.py new file mode 100644 index 0000000000..d7d4cdd9a1 --- /dev/null +++ b/backend/agent_kv/tests/test_rate_limiter.py @@ -0,0 +1,192 @@ +import os +import pathlib +from unittest import mock + +import django +import pytest +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from django.conf import settings # noqa: E402 + +from agent_kv import rate_limiter as rl # noqa: E402 + + +@mock.patch.object(rl, "_redis") +@mock.patch.object(rl.time, "time", return_value=1_700_000_000.0) +def test_acquire_under_limit(m_time, m_redis): + """Acquire is ONE atomic server-side script call (trim, count, add, + expire) -- never a client-side ``ZCARD`` followed by ``ZADD``, which is a + check-then-act race that let 6 concurrent submits through a limit of 5 + in the 13b integration run. The e2e concurrency scenario is the real + proof; this pins the script's inputs. + """ + mock_redis = m_redis.return_value + mock_redis.eval.return_value = 1 + + assert rl.AgentKVConcurrencyLimiter.check_and_acquire("org1", "job1") is True + + key = "agent_kv:inflight:org1" + mock_redis.eval.assert_called_once() + args = mock_redis.eval.call_args[0] + script, numkeys, called_key, member, now, ttl_cut, limit, expire = args + assert script is rl.AgentKVConcurrencyLimiter._ACQUIRE_SCRIPT + assert (numkeys, called_key, member) == (1, key, "job1") + assert now == 1_700_000_000.0 + assert ttl_cut == 1_700_000_000.0 - rl._SLOT_TTL_SECONDS + assert limit == settings.AGENT_KV_CONCURRENT_LIMIT + assert expire == rl._SLOT_TTL_SECONDS + # No client-side check-then-act calls remain. + assert not mock_redis.zcard.called + assert not mock_redis.zadd.called + # Self-heal-before-check-before-acquire ordering lives inside the script. + body = script + assert body.index("ZREMRANGEBYSCORE") < body.index("ZCARD") < body.index("ZADD") + + +@mock.patch.object(rl, "_redis") +def test_acquire_at_limit_refused(m_redis): + m_redis.return_value.eval.return_value = 0 + assert rl.AgentKVConcurrencyLimiter.check_and_acquire("org1", "job1") is False + + +@mock.patch.object(rl, "_redis") +def test_release_removes_member(m_redis): + rl.AgentKVConcurrencyLimiter.release("org1", "job1") + m_redis.return_value.zrem.assert_called_once_with("agent_kv:inflight:org1", "job1") + + +@mock.patch.object(rl, "_redis") +def test_redis_error_during_the_script_fails_closed(m_redis): + """Was `test_redis_error_fails_open`, asserting `is True`. + + It pinned the defect as the contract: an error mid-script removed the + concurrency ceiling and the API kept accepting billable work. The failure + now surfaces as a 429 at the caller (`RateLimited`), which is recoverable; + an unbounded fan-out is not. See `AGENT_KV_LIMITER_FAIL_OPEN` for the + deliberate waiver. + """ + m_redis.return_value.eval.side_effect = ConnectionError("down") + with mock.patch.object(settings, "AGENT_KV_LIMITER_FAIL_OPEN", False): + assert rl.AgentKVConcurrencyLimiter.check_and_acquire("org1", "job1") is False + + +@mock.patch.object(rl, "_redis") +@mock.patch.object(rl.time, "time", return_value=1_700_000_000.0) +def test_key_rate_over_limit(m_time, m_redis): + mock_redis = m_redis.return_value + mock_redis.incr.return_value = 61 + + assert rl.check_key_rate("key1") is False + + expected_window = int(1_700_000_000.0 // 60) + expected_key = f"agent_kv:rate:key1:{expected_window}" + mock_redis.incr.assert_called_once_with(expected_key) + mock_redis.expire.assert_called_once_with(expected_key, 120) + + +@mock.patch.object(rl, "_redis") +@mock.patch.object(rl.time, "time", return_value=1_700_000_000.0) +def test_key_rate_under_limit(m_time, m_redis): + mock_redis = m_redis.return_value + mock_redis.incr.return_value = 3 + + assert rl.check_key_rate("key1") is True + + expected_window = int(1_700_000_000.0 // 60) + expected_key = f"agent_kv:rate:key1:{expected_window}" + mock_redis.incr.assert_called_once_with(expected_key) + mock_redis.expire.assert_called_once_with(expected_key, 120) + + +# --------------------------------------------------------------------------- +# Backend-unavailable behaviour. Both limiters used to `return True` on ANY +# Redis exception, which is the one failure mode invisible from outside: a +# Sentinel failover or pool exhaustion removed the concurrency ceiling AND the +# per-key rate ceiling simultaneously, while the API went on returning 202s for +# billable LLM work with only a per-request `logger.warning` to show for it. +# --------------------------------------------------------------------------- + + +@mock.patch.object(rl, "_redis", side_effect=OSError("redis down")) +def test_concurrency_limiter_fails_closed_by_default(m_redis): + with mock.patch.object(settings, "AGENT_KV_LIMITER_FAIL_OPEN", False): + assert rl.AgentKVConcurrencyLimiter.check_and_acquire("org1", "job1") is False + + +@mock.patch.object(rl, "_redis", side_effect=OSError("redis down")) +def test_key_rate_limiter_fails_closed_by_default(m_redis): + with mock.patch.object(settings, "AGENT_KV_LIMITER_FAIL_OPEN", False): + assert rl.check_key_rate("key1") is False + + +@mock.patch.object(rl, "_redis", side_effect=OSError("redis down")) +def test_limiters_fail_open_only_when_the_setting_says_so(m_redis): + """The waiver stays available -- availability-over-accounting is a real + operational choice -- but it has to be made deliberately, in config, rather + than being the implicit behaviour of an `except` block. + """ + with mock.patch.object(settings, "AGENT_KV_LIMITER_FAIL_OPEN", True): + assert rl.AgentKVConcurrencyLimiter.check_and_acquire("org1", "job1") is True + assert rl.check_key_rate("key1") is True + + +@mock.patch.object(rl, "_redis", side_effect=OSError("redis down")) +def test_limiter_unavailability_is_logged_at_error_level(m_redis, caplog): + """`logger.exception`, not `logger.warning`. The old level is part of why + this sat unnoticed -- a limiter being gone is an error, not a warning. + """ + with mock.patch.object(settings, "AGENT_KV_LIMITER_FAIL_OPEN", False): + with caplog.at_level("ERROR", logger=rl.logger.name): + rl.AgentKVConcurrencyLimiter.check_and_acquire("org1", "job1") + rl.check_key_rate("key1") + assert [r.levelname for r in caplog.records] == ["ERROR", "ERROR"] + # The traceback is attached, so the actual Redis failure is diagnosable. + assert all(r.exc_info for r in caplog.records) + + +@mock.patch.object(rl, "_redis", side_effect=OSError("redis down")) +def test_release_still_tolerates_an_unreachable_backend(m_redis): + """Release must stay best-effort regardless: it is called on terminal + paths (finalize, cancel, sweep) where raising would abort the caller's own + work, and a slot it cannot free expires on its own TTL. + """ + rl.AgentKVConcurrencyLimiter.release("org1", "job1") + + +# --------------------------------------------------------------------------- +# The published env table said both limiters "fail open on Redis errors", for +# the whole life of the branch that made them fail CLOSED. An operator reading +# it would size a Redis outage as "requests get through" when the real +# behaviour is "every submit 429s" -- the opposite incident. Docs are the +# contract for a public API, so this is pinned rather than trusted. +# --------------------------------------------------------------------------- + +_DOCS = pathlib.Path(__file__).resolve().parents[3] / "docs" / "agent-kv-api.md" + + +@pytest.mark.parametrize( + "variable", + ["AGENT_KV_CONCURRENT_LIMIT", "AGENT_KV_KEY_RATE_LIMIT_PER_MINUTE"], +) +def test_the_docs_do_not_describe_either_limiter_as_failing_open(variable): + row = next(line for line in _DOCS.read_text().splitlines() if f"`{variable}`" in line) + assert "fails open" not in row.casefold(), row + assert "closed" in row.casefold(), ( + f"the {variable} row must say what a Redis outage does, and it fails " + "closed -- see _limiter_failure_allows_request" + ) + + +def test_the_waiver_flag_is_documented_with_its_real_default(): + text = _DOCS.read_text() + assert "`AGENT_KV_LIMITER_FAIL_OPEN`" in text, ( + "the only way to restore fail-open is undocumented, so the only " + "documented behaviour was the wrong one" + ) + assert not rl._limiter_failure_allows_request() or getattr( + settings, "AGENT_KV_LIMITER_FAIL_OPEN", False + ), "the default must be fail-closed" diff --git a/backend/agent_kv/tests/test_storage.py b/backend/agent_kv/tests/test_storage.py new file mode 100644 index 0000000000..257a179dea --- /dev/null +++ b/backend/agent_kv/tests/test_storage.py @@ -0,0 +1,304 @@ +import os +import re +from unittest import mock + +import django +import pytest +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from django.conf import settings # noqa: E402 +from django.test import override_settings # noqa: E402 + +from agent_kv import storage # noqa: E402 +from agent_kv.models import AgentKVJob # noqa: E402 + + +@mock.patch.object(storage, "FileSystem") +def test_stage_input_path_and_write(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + up = mock.Mock() + up.name = "invoice.PDF" + up.chunks.return_value = [b"a", b"b"] + ref = storage.stage_input("org1", "job1", up) + assert ref == "unstract/agent_kv/org1/job1/input.pdf" + assert fh.write.called + kwargs = fh.write.call_args.kwargs + assert kwargs["path"] == "unstract/agent_kv/org1/job1/input.pdf" + assert kwargs["data"] == b"ab" + + +@mock.patch.object(storage, "FileSystem") +def test_stage_input_defaults_extension_when_missing(m_fs): + up = mock.Mock() + up.name = "noext" + up.chunks.return_value = [b"x"] + ref = storage.stage_input("org1", "job1", up) + assert ref == "unstract/agent_kv/org1/job1/input.bin" + + +@mock.patch.object(storage, "FileSystem") +def test_write_and_read_result_roundtrip_path(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + ref = storage.write_result("org1", "job1", {"success": True}) + # The result path is UNIQUE per finalize attempt (a nonce suffix) so a + # guard-losing concurrent finalize deletes only its own orphan, never the + # winning row's result_ref target (pre-Greptile critical #3). It still + # lives under the deterministic job dir. + assert re.fullmatch( + r"unstract/agent_kv/org1/job1/result-[0-9a-f]{32}\.json", ref + ), ref + fh.json_dump.assert_called_once() + assert fh.json_dump.call_args.kwargs["path"] == ref + assert fh.json_dump.call_args.kwargs["data"] == {"success": True} + + +@mock.patch.object(storage, "FileSystem") +def test_write_result_produces_unique_ref_per_call(m_fs): + # Two writes for the SAME job must yield two DISTINCT refs -- this is what + # makes a concurrent duplicate-success finalize safe: the guard loser's + # delete targets its own ref, not the winner's. + a = storage.write_result("org1", "job1", {"n": 1}) + b = storage.write_result("org1", "job1", {"n": 2}) + assert a != b + assert a.startswith("unstract/agent_kv/org1/job1/result-") + assert b.startswith("unstract/agent_kv/org1/job1/result-") + + +@mock.patch.object(storage, "FileSystem") +def test_write_result_accepts_explicit_nonce(m_fs): + # A caller may pin the nonce (deterministic ref for a given attempt). + ref = storage.write_result("org1", "job1", {"n": 1}, nonce="abc123") + assert ref == "unstract/agent_kv/org1/job1/result-abc123.json" + + +@mock.patch.object(storage, "FileSystem") +def test_read_result_returns_parsed_json(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + fh.json_load.return_value = {"success": True} + out = storage.read_result("unstract/agent_kv/org1/job1/result.json") + assert out == {"success": True} + fh.json_load.assert_called_once_with(path="unstract/agent_kv/org1/job1/result.json") + + +@mock.patch.object(storage, "FileSystem") +def test_delete_job_files_removes_both_refs(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + job = AgentKVJob( + input_ref="org/o/agent_kv/j/input.pdf", + result_ref="org/o/agent_kv/j/result.json", + ) + cleared = storage.delete_job_files(job) + removed = {c.kwargs["path"] for c in fh.rm.call_args_list} + assert removed == {job.input_ref, job.result_ref} + assert cleared == ["input_ref", "result_ref"] + + +@mock.patch.object(storage, "FileSystem") +def test_delete_job_files_tolerates_missing_files(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + fh.rm.side_effect = FileNotFoundError("gone") + job = AgentKVJob( + input_ref="org/o/agent_kv/j/input.pdf", + result_ref="org/o/agent_kv/j/result.json", + ) + # Must not raise, and both count as CLEARED: the goal is "the file is not + # there", and a ref whose object is already gone has nothing to retry. + cleared = storage.delete_job_files(job) + assert fh.rm.call_count == 2 + assert cleared == ["input_ref", "result_ref"] + + +@mock.patch.object(storage, "FileSystem") +def test_delete_job_files_skips_blank_refs(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + job = AgentKVJob(input_ref="", result_ref="") + cleared = storage.delete_job_files(job) + assert not fh.rm.called + # Reported clear: nothing was ever staged, so there is nothing to retry and + # the caller is free to blank (already-blank) refs. + assert cleared == ["input_ref", "result_ref"] + + +# The Greptile-review contract (#3): a ref is reported clear ONLY when its file +# is confirmed gone. A storage error that is not FileNotFoundError leaves the +# ref out, and the ref is what TTL cleanup retries from -- its candidate query +# matches rows by `input_ref > "" OR result_ref > ""`, so a blanked ref whose +# object is still in the bucket orphans that object permanently. +@mock.patch.object(storage, "FileSystem") +def test_delete_job_files_omits_the_ref_whose_delete_raised(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + job = AgentKVJob( + input_ref="org/o/agent_kv/j/input.pdf", + result_ref="org/o/agent_kv/j/result.json", + ) + + def rm(*, path): + if path == job.result_ref: + raise OSError("object store unreachable") + + fh.rm.side_effect = rm + + cleared = storage.delete_job_files(job) + + # Still attempted both -- one failure must not abort the other's cleanup. + assert fh.rm.call_count == 2 + assert cleared == ["input_ref"] + + +@mock.patch.object(storage, "FileSystem") +def test_delete_job_files_reports_nothing_clear_when_both_deletes_raise(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + fh.rm.side_effect = OSError("object store unreachable") + job = AgentKVJob( + input_ref="org/o/agent_kv/j/input.pdf", + result_ref="org/o/agent_kv/j/result.json", + ) + assert storage.delete_job_files(job) == [] + + +# delete_input carries the same confirmed-clear contract as delete_job_files. +# It did not when the Greptile round fixed delete_job_files and its two +# callers: this third ref-blanking site (FinalizeView) was missed, because the +# fix followed one function's callers instead of grepping for every site that +# blanks a ref. A failed delete here orphaned the customer's uploaded document +# permanently -- TTL cleanup only selects rows that still carry a non-blank +# ref, so a blanked row can never be reconsidered. +@mock.patch.object(storage, "FileSystem") +def test_delete_input_reports_clear_on_success(m_fs): + job = AgentKVJob(input_ref="org/o/agent_kv/j/input.pdf") + assert storage.delete_input(job) is True + + +@mock.patch.object(storage, "FileSystem") +def test_delete_input_reports_clear_when_the_file_is_already_gone(m_fs): + m_fs.return_value.get_file_storage.return_value.rm.side_effect = FileNotFoundError + job = AgentKVJob(input_ref="org/o/agent_kv/j/input.pdf") + # Already missing is the goal state -- nothing left to retry. + assert storage.delete_input(job) is True + + +@mock.patch.object(storage, "FileSystem") +def test_delete_input_reports_NOT_clear_when_the_delete_raised(m_fs): + m_fs.return_value.get_file_storage.return_value.rm.side_effect = OSError("down") + job = AgentKVJob(input_ref="org/o/agent_kv/j/input.pdf") + # Must not raise, and must not claim the ref is clear. + assert storage.delete_input(job) is False + + +@mock.patch.object(storage, "FileSystem") +def test_delete_input_reports_clear_for_a_blank_ref(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + assert storage.delete_input(AgentKVJob(input_ref="")) is True + assert not fh.rm.called + + +@mock.patch.object(storage, "FileSystem") +def test_delete_input_removes_only_input_ref(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + job = AgentKVJob( + input_ref="org/o/agent_kv/j/input.pdf", + result_ref="org/o/agent_kv/j/result.json", + ) + storage.delete_input(job) + fh.rm.assert_called_once_with(path=job.input_ref) + + +@mock.patch.object(storage, "FileSystem") +def test_delete_input_tolerates_missing_file(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + fh.rm.side_effect = FileNotFoundError("gone") + job = AgentKVJob(input_ref="org/o/agent_kv/j/input.pdf") + # Must not raise. + storage.delete_input(job) + assert fh.rm.call_count == 1 + + +@mock.patch.object(storage, "FileSystem") +def test_delete_input_skips_blank_ref(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + job = AgentKVJob(input_ref="") + storage.delete_input(job) + assert not fh.rm.called + + +@mock.patch.object(storage, "FileSystem") +def test_delete_result_file_removes_the_given_ref(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + storage.delete_result_file("org/o/agent_kv/j/result.json") + fh.rm.assert_called_once_with(path="org/o/agent_kv/j/result.json") + + +@mock.patch.object(storage, "FileSystem") +def test_delete_result_file_tolerates_missing_file(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + fh.rm.side_effect = FileNotFoundError("gone") + # Must not raise. + storage.delete_result_file("org/o/agent_kv/j/result.json") + assert fh.rm.call_count == 1 + + +@mock.patch.object(storage, "FileSystem") +def test_delete_result_file_skips_blank_ref(m_fs): + fh = m_fs.return_value.get_file_storage.return_value + storage.delete_result_file("") + assert not fh.rm.called + + +def test_default_storage_prefix_is_bucket_rooted(): + """The shipped default roots every path in a real bucket (13b F1). + + s3fs/gcsfs read the first path segment as the bucket, so a prefix without + one (the old ``org/{org_id}/...``) makes every write fail ``NoSuchBucket``. + The default must therefore stay ``unstract/agent_kv`` -- bucket ``unstract`` + (created by the MinIO dev bootstrap), directory ``agent_kv`` -- and must + match the cloud executor's ``AGENT_KV_STORAGE_DIR_PREFIX``. + """ + if "AGENT_KV_STORAGE_DIR_PREFIX" in os.environ: + pytest.skip("AGENT_KV_STORAGE_DIR_PREFIX is overridden in this environment") + assert settings.AGENT_KV_STORAGE_DIR_PREFIX == "unstract/agent_kv" + bucket, _, rest = settings.AGENT_KV_STORAGE_DIR_PREFIX.partition("/") + # Split: a composite assert cannot say WHICH half failed, and these two + # mean different things -- no bucket segment vs. a bucket with no path. + assert bucket + assert rest + + +@mock.patch.object(storage, "FileSystem") +def test_paths_are_rooted_at_the_configured_prefix(m_fs): + up = mock.Mock() + up.name = "invoice.pdf" + up.chunks.return_value = [b"x"] + with override_settings(AGENT_KV_STORAGE_DIR_PREFIX="mybucket/kv-root"): + input_ref = storage.stage_input("org1", "job1", up) + result_ref = storage.write_result("org1", "job1", {"success": True}) + assert input_ref.startswith("mybucket/kv-root/") + assert result_ref.startswith("mybucket/kv-root/") + assert input_ref == "mybucket/kv-root/org1/job1/input.pdf" + assert re.fullmatch( + r"mybucket/kv-root/org1/job1/result-[0-9a-f]{32}\.json", result_ref + ), result_ref + + +def test_storage_prefix_normalisation_matches_executor(monkeypatch): + """A sloppy operator value must resolve to the same root on both sides: + the cloud executor strips whitespace and edge slashes before keying its + OCR cache, so the backend must too, else inputs and cache split roots. + """ + import importlib + + monkeypatch.setenv("AGENT_KV_STORAGE_DIR_PREFIX", " /unstract/agent_kv/ ") + from backend.settings import base + + importlib.reload(base) + try: + assert base.AGENT_KV_STORAGE_DIR_PREFIX == "unstract/agent_kv" + monkeypatch.setenv("AGENT_KV_STORAGE_DIR_PREFIX", " ") + importlib.reload(base) + assert base.AGENT_KV_STORAGE_DIR_PREFIX == "unstract/agent_kv" + finally: + monkeypatch.delenv("AGENT_KV_STORAGE_DIR_PREFIX", raising=False) + importlib.reload(base) diff --git a/backend/agent_kv/tests/test_submit_serializer.py b/backend/agent_kv/tests/test_submit_serializer.py new file mode 100644 index 0000000000..a94d77dcfb --- /dev/null +++ b/backend/agent_kv/tests/test_submit_serializer.py @@ -0,0 +1,552 @@ +"""Submit-serializer rules for the table extractor. + +This deployment routes `table` and nothing else (`EXTRACTOR_ROUTES`), so every +submit built here names `table`. The request-level rules -- file type, size, +page cap, page range, timeout, the extractor-scoped wire format -- are the same +for any extractor and are exercised through that one. + +`KVOptionsSerializer` is still in the tree and still tested, but it is tested +*directly* rather than through a submit: with `kv` out of `EXTRACTOR_ROUTES` a +`kv` entry is refused at `validate_name` before any options validator runs, so +driving those rules through `SubmitSerializer` would assert nothing about them. +""" + +import json +import os +from unittest import mock + +import django +import pytest +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from django.core.files.uploadedfile import SimpleUploadedFile # noqa: E402 +from rest_framework import serializers # noqa: E402 + +from agent_kv.execution_serializers import ( # noqa: E402 + ExtractorSerializer, + KVOptionsSerializer, + SubmitSerializer, +) + +FIXTURES = os.path.join(os.path.dirname(__file__), "fixtures") + +#: The table extractor's `keys`: the thing being asked for is a table, named. +TABLE_KEYS = {"target_table": "Rent roll"} + + +def _pdf_upload(name="doc.pdf"): + with open(os.path.join(FIXTURES, "two_page.pdf"), "rb") as f: + return SimpleUploadedFile(name, f.read(), content_type="application/pdf") + + +# Fields that describe the REQUEST rather than an extractor (spec §7.1). +# Anything else passed to _data() is routed into the table extractor's options, +# so each test still reads as "submit with this one thing changed". +_JOB_LEVEL = { + "file", + "extractors", + "page_start", + "page_end", + "timeout", + "tags", + "custom_data", + "webhook_url", +} + + +#: The platform adapters a `table` submit must name. All three are REQUIRED -- +#: the engine needs two LLMs and an OCR source, and this API resolves them from +#: the caller's own adapter instances rather than from operator env vars. +#: Ownership and type are checked in the VIEW against the Bearer key's org +#: (`execution_views._resolved_adapters`); the serializer only checks shape, so +#: these can be any UUIDs here. +TABLE_ADAPTERS = { + "llm": "11111111-1111-1111-1111-111111111111", + "lite_llm": "22222222-2222-2222-2222-222222222222", + "x2text": "33333333-3333-3333-3333-333333333333", +} + + +def _data(**over): + """Build a submit payload in the extractor-scoped wire format (§7.0).""" + keys = over.pop("keys", TABLE_KEYS) + adapters = over.pop("adapters", TABLE_ADAPTERS) + options = {k: over.pop(k) for k in list(over) if k not in _JOB_LEVEL} + d = { + "file": _pdf_upload(), + "extractors": json.dumps( + [ + { + "name": "table", + "keys": keys, + "adapters": adapters, + "options": options, + } + ] + ), + } + d.update(over) + return d + + +def _errs(s): + """All validation errors as one string. + + Per-extractor failures surface nested under `extractors`, so asserting on a + specific top-level key would just be asserting on DRF's nesting shape rather + than on the rejection actually happening. + """ + return str(s.errors) + + +def _defaults(m): + """Set every AGENT_KV_* attribute the serializer reads to its production default.""" + m.AGENT_KV_MAX_FILE_SIZE_MB = 50 + m.AGENT_KV_MAX_PAGES = 100 + m.AGENT_KV_MAX_SCHEMA_BYTES = 262_144 + m.AGENT_KV_MAX_CALCULATIONS_BYTES = 20_000 + m.AGENT_KV_MAX_TIMEOUT_SECONDS = 300 + m.AGENT_KV_CALCULATIONS_ENABLED = False + m.AGENT_KV_STRUCTURED_OUTPUT_ENABLED = False + + +# --------------------------------------------------------------------------- +# The routing this deployment actually ships +# --------------------------------------------------------------------------- +def test_a_kv_submit_is_refused_at_the_serializer(): + """The guard on the single most important line in this carve-out. + + `agentic_kv` is not deployed here, so nothing consumes + `celery_executor_agentic_kv`. If `kv` were routable, this submit would + return 202, dispatch into a queue with no consumer, and sit in DISPATCHED + forever -- no error at the producer, nothing in any log to find. A 400 at + the serializer is the whole difference between a clear refusal and a job + that silently never runs. + """ + s = SubmitSerializer( + data=_data( + extractors=json.dumps( + [{"name": "kv", "keys": {"total": {"description": "Grand total"}}}] + ) + ) + ) + assert not s.is_valid() + assert "unknown extractor" in _errs(s) + assert "'kv'" in _errs(s) + + +def test_table_is_the_only_supported_extractor(): + from agent_kv.execution_serializers import SUPPORTED_EXTRACTORS + + assert SUPPORTED_EXTRACTORS == ("table",) + + +# --------------------------------------------------------------------------- +# Request-level rules, exercised through the extractor this deployment serves +# --------------------------------------------------------------------------- +def test_valid_submit_counts_pages_and_defaults_its_options(): + s = SubmitSerializer(data=_data()) + assert s.is_valid(), s.errors + assert s.pages_total == 2 + entry = s.validated_data["extractors"][0] + assert entry["name"] == "table" + assert entry["keys"]["target_table"] == "Rent roll" + assert entry["options"]["number_format"] == "US" + assert entry["options"]["enable_header_mapping"] is False + + +def test_a_table_entry_without_a_target_table_is_rejected(): + """`target_table` is the engine's one required extraction parameter; the + caller's actual mistake is not naming the table. + """ + s = SubmitSerializer(data=_data(keys={})) + assert not s.is_valid() + assert "target_table" in _errs(s) + + +def test_disallowed_extension_rejected(): + bad = SimpleUploadedFile("doc.exe", b"MZ", content_type="application/x-dos") + s = SubmitSerializer(data=_data(file=bad)) + assert not s.is_valid() + assert "file" in s.errors + + +def test_oversize_file_rejected(): + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + m.AGENT_KV_MAX_FILE_SIZE_MB = 0 + s = SubmitSerializer(data=_data()) + assert not s.is_valid() + assert "file" in s.errors + + +def test_page_cap_rejected(): + """No range given, so the whole document is the selection.""" + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + m.AGENT_KV_MAX_PAGES = 1 + s = SubmitSerializer(data=_data()) + assert not s.is_valid() + assert "pages" in str(s.errors).lower() + + +def test_page_cap_counts_the_selected_range_not_the_whole_document(): + """The cap bounds the work the job will DO, not the size of the file. + + A caller asking for one page of a two-page document is requesting one page + of OCR and extraction. Counting the whole document rejected requests that + were inside the documented limit. Reported by Greptile on #2317. + """ + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + m.AGENT_KV_MAX_PAGES = 1 + s = SubmitSerializer(data=_data(page_start=1, page_end=1)) + assert s.is_valid(), s.errors + assert s.pages_selected == 1 + # The measured document count is unchanged -- it is what metering and + # the status document report, and only the cap comparison moved. + assert s.pages_total == 2 + + +def test_an_open_ended_range_is_capped_at_the_last_page(): + """`page_end` past the end selects to the end, it does not inflate the count.""" + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + m.AGENT_KV_MAX_PAGES = 2 + s = SubmitSerializer(data=_data(page_start=2, page_end=999)) + assert s.is_valid(), s.errors + assert s.pages_selected == 1 + + +def test_a_selected_range_over_the_cap_is_still_rejected(): + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + m.AGENT_KV_MAX_PAGES = 1 + s = SubmitSerializer(data=_data(page_start=1, page_end=2)) + assert not s.is_valid() + assert "Requested 2 pages" in _errs(s) + + +def test_page_start_past_the_end_of_the_document_is_rejected(): + """Otherwise the selection is empty and the job runs over nothing, billed.""" + s = SubmitSerializer(data=_data(page_start=5)) + assert not s.is_valid() + assert "past the end" in _errs(s) + + +def test_extractors_not_json_rejected(): + s = SubmitSerializer(data=_data(extractors="{not json")) + assert not s.is_valid() + assert "extractors" in s.errors + + +def test_timeout_bounds(): + s = SubmitSerializer(data=_data(timeout=301)) + assert not s.is_valid() + s2 = SubmitSerializer(data=_data(timeout=0)) + assert s2.is_valid(), s2.errors + + +def test_page_range_validation(): + s = SubmitSerializer(data=_data(page_start=5, page_end=2)) + assert not s.is_valid() + + +def test_unreadable_pdf_is_field_error_and_stream_is_rewound(): + bad = SimpleUploadedFile("doc.pdf", b"not a pdf", content_type="application/pdf") + s = SubmitSerializer(data=_data(file=bad)) + assert not s.is_valid() + assert "file" in s.errors + # The `finally: f.seek(0)` in validate() must run on the error path too, + # so a downstream reader (e.g. the view persisting the upload) still sees + # the whole stream from the start. + assert bad.tell() == 0 + assert bad.read() == b"not a pdf" + + +# --------------------------------------------------------------------------- +# Extractor-scoped wire format (spec §7.0/§7.1). These rules exist so a caller +# learns immediately that something is unsupported, rather than having the +# request quietly run as something other than what they asked for. +# --------------------------------------------------------------------------- +def test_unknown_extractor_name_rejected(): + s = SubmitSerializer( + data=_data(extractors=json.dumps([{"name": "bogus", "keys": TABLE_KEYS}])) + ) + assert not s.is_valid() + assert "unknown extractor" in _errs(s) + + +def test_multiple_extractors_rejected_until_fan_out_exists(): + """The FORMAT is fixed before launch; the fan-out execution is not built. + + Accepting two entries and running only the first would be the silent kind + of wrong -- the caller is billed for a job that ignored half the request. + """ + two = [ + {"name": "table", "keys": TABLE_KEYS}, + {"name": "table", "keys": TABLE_KEYS}, + ] + s = SubmitSerializer(data=_data(extractors=json.dumps(two))) + assert not s.is_valid() + assert "multiple extractors are not supported yet" in _errs(s) + + +def test_empty_or_non_list_extractors_rejected(): + for bad in ("[]", '{"name": "table"}', '"table"'): + s = SubmitSerializer(data=_data(extractors=bad)) + assert not s.is_valid(), bad + assert "non-empty JSON array" in _errs(s) + + +def test_unknown_option_is_rejected_not_silently_dropped(): + """DRF drops unknown fields by default; for per-extractor options that would + mean a typo'd or misaddressed knob silently changing what the job runs. + """ + s = SubmitSerializer(data=_data(number_formatt="EU")) # typo for `number_format` + assert not s.is_valid() + assert "unknown options for extractor 'table'" in _errs(s) + assert "number_formatt" in _errs(s) + + +def test_a_kv_option_on_a_table_entry_is_rejected(): + """The cross-extractor case the per-extractor serializers exist for: `qa` + is a KV pipeline knob and means nothing to the table engine. Dropped + silently, the caller would believe they had turned something off. + """ + s = SubmitSerializer(data=_data(qa=False)) + assert not s.is_valid() + assert "unknown options for extractor 'table'" in _errs(s) + assert "qa" in _errs(s) + + +def test_old_flat_format_is_no_longer_accepted(): + """Hard switch (§7.1): the pre-§7.0 shape has no alias. A caller still + sending the flat form must get a clear 400, not a job that silently ran + with default options. + """ + s = SubmitSerializer( + data={ + "file": _pdf_upload(), + "keys": json.dumps(TABLE_KEYS), + "number_format": "EU", + } + ) + assert not s.is_valid() + assert "extractors" in s.errors # the now-required field is missing + + +def test_unknown_key_on_the_extractor_entry_is_rejected(): + """`options` rejects unknowns; the entry itself must too. + + DRF drops unrecognised keys, so a near-miss like `option` (singular) for + `options` would be discarded whole: `options` defaults to {}, the job runs + with every default, and the caller gets a 202 with no sign their + configuration was ignored. + """ + s = SubmitSerializer( + data=_data( + extractors=json.dumps( + [{"name": "table", "keys": TABLE_KEYS, "option": {"number_format": "EU"}}] + ) + ) + ) + assert not s.is_valid() + assert "unknown keys on extractor entry" in _errs(s) + assert "option" in _errs(s) + + +# --------------------------------------------------------------------------- +# KVOptionsSerializer, tested directly. +# +# The class is dormant on this deployment -- `kv` is not routable, so nothing +# reaches it through a submit -- but it stays in the tree so the branch that +# carries the KV engine re-enables the extractor with one line rather than a +# content merge. Dormant and untested is how a one-line re-enable turns into a +# regression, so its rules are asserted here against the serializer itself. +# --------------------------------------------------------------------------- +def _kv_options(**over): + s = KVOptionsSerializer(data=over) + return s + + +def test_kv_options_default_qa_and_challenge_on(): + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + s = _kv_options() + assert s.is_valid(), s.errors + assert s.validated_data["qa"] is True + assert s.validated_data["challenge"] is True + assert s.validated_data["extraction_mode"] == "whole-doc" + + +def test_kv_calculations_rejected_when_disabled(): + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + m.AGENT_KV_CALCULATIONS_ENABLED = False + s = _kv_options(calculations="annualize rent") + assert not s.is_valid() + assert "not available" in str(s.errors) + + +def test_kv_calculations_accepted_when_enabled(): + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + m.AGENT_KV_CALCULATIONS_ENABLED = True + assert _kv_options(calculations="annualize rent").is_valid() + + +def test_kv_calculations_cap(): + # AGENT_KV_CALCULATIONS_ENABLED must be mocked True here: real settings + # default it False, so an unmocked run hits the feature-gate branch (a + # different error) instead of the byte-size cap this test is named for. + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + m.AGENT_KV_CALCULATIONS_ENABLED = True + s = _kv_options(calculations="x" * 30_000) + assert not s.is_valid() + assert "calculations" in str(s.errors) + assert "20000 bytes" in str(s.errors) + + +def test_kv_structured_output_rejected_when_disabled(): + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + m.AGENT_KV_STRUCTURED_OUTPUT_ENABLED = False + s = _kv_options(structured_output=True) + assert not s.is_valid() + assert "structured_output" in str(s.errors) + + +def test_kv_empty_calculations_and_false_structured_output_pass_when_disabled(): + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + assert _kv_options(calculations="", structured_output=False).is_valid() + + +def test_kv_unknown_option_is_rejected(): + with mock.patch("agent_kv.execution_serializers.settings") as m: + _defaults(m) + s = _kv_options(qaa=True) # typo for `qa` + assert not s.is_valid() + assert "unknown options for extractor 'kv'" in str(s.errors) + + +# --------------------------------------------------------------------------- +# Extractor identity must be decided ONCE. +# +# `validate_keys` used to branch on raw `initial_data["name"]` while +# `validate_name` saw the value DRF had already trimmed +# (`CharField.trim_whitespace` defaults True). So `" table "` passed the name +# check as `table` and then took the **kv** branch for keys, never running +# `TableKeysSerializer`. +# +# A kv-shaped payload compiles cleanly as a KV schema, so the submit returned +# 202 and dispatched to `agentic_table` with a `target_table` that is a dict +# rather than the string the binding requires -- staged, billed, then failed at +# the executor. +# +# Reported as 2.1 in the branch review. +# --------------------------------------------------------------------------- +def test_a_padded_extractor_name_still_validates_keys_as_that_extractor(): + s = SubmitSerializer( + data=_data( + extractors=json.dumps( + [ + { + # Trimmed to "table" by validate_name... + "name": " table ", + # ...but this is a KV-shaped schema, not a table one. It + # compiles fine as KV, which is how it used to reach 202. + "keys": {"target_table": {"description": "Grand total"}}, + } + ] + ) + ) + ) + + assert not s.is_valid(), ( + "a kv-shaped keys payload was accepted for the table extractor; the " + "job would stage, bill and then fail at the executor on target_table" + ) + assert "target_table" in _errs(s) + + +def test_a_padded_name_is_still_normalised_for_routing(): + """The trim itself is fine and worth keeping -- the bug was deciding + identity twice, not accepting whitespace. + """ + s = SubmitSerializer( + data=_data( + extractors=json.dumps( + [{"name": " table ", "keys": TABLE_KEYS, "adapters": TABLE_ADAPTERS}] + ) + ) + ) + + assert s.is_valid(), s.errors + assert s.validated_data["extractors"][0]["name"] == "table" + + +def test_table_keys_are_validated_through_the_keys_serializer(): + """Reaching TableKeysSerializer is what rejects unknown members; the kv + fallback would have compiled them as a schema instead. + """ + s = SubmitSerializer( + data=_data( + extractors=json.dumps( + [{"name": "table", "keys": {"target_table": "T", "bogus": 1}}] + ) + ) + ) + + assert not s.is_valid() + assert "unknown keys for extractor 'table'" in _errs(s) + + +# --------------------------------------------------------------------------- +# 2.20: `compile_schema` is called in the no-keys-serializer fall-through for +# what it REFUSES, not for what it returns. +# +# The `CompiledSchema` used to be stashed on `ExtractorSerializer.compiled` and +# collected into `SubmitSerializer.compiled`, which no non-test code read -- +# `dispatch_job` sends the raw `keys` dict and the engine recompiles. Deleting +# the attribute must not quietly delete the 400, which is the only reason the +# call is there. +# +# The branch is DORMANT on this deployment: `kv` is the only extractor without +# a keys serializer and `validate_name` refuses it before `_validated_keys` +# runs, so it is reached by calling it directly. That is the point -- it is the +# contract for the next extractor added without one, and nothing else covers it. +# --------------------------------------------------------------------------- + + +def test_the_schema_compiler_fallthrough_still_rejects_a_bad_schema(): + ser = ExtractorSerializer() + + with pytest.raises(serializers.ValidationError) as exc: + # A leaf with no `description` is refused by the ported compiler. + ser._validated_keys("some_future_extractor", {"total": {}}) + + assert "keys" in exc.value.detail + + +def test_the_schema_compiler_fallthrough_returns_the_raw_spec(): + """What reaches `dispatch_job` is the submitted dict, unchanged.""" + ser = ExtractorSerializer() + spec = {"total": {"description": "The grand total", "format": "currency"}} + + assert ser._validated_keys("some_future_extractor", spec) is spec + + +def test_neither_serializer_still_carries_a_compiled_attribute(): + """A `compiled` that is populated and never read reads as plumbing.""" + for cls in (ExtractorSerializer, SubmitSerializer): + assert not hasattr(cls, "compiled"), ( + f"{cls.__name__}.compiled was dead: `dispatch_job` sends the raw " + "keys dict, and the compiled form cannot cross the queue anyway" + ) diff --git a/backend/agent_kv/tests/test_submit_view.py b/backend/agent_kv/tests/test_submit_view.py new file mode 100644 index 0000000000..e3c285830c --- /dev/null +++ b/backend/agent_kv/tests/test_submit_view.py @@ -0,0 +1,1146 @@ +import os +import uuid +from unittest import mock + +import django +import pytest +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from django.conf import settings # noqa: E402 +from django.utils import timezone # noqa: E402 +from rest_framework.test import APIRequestFactory # noqa: E402 + +from agent_kv import execution_views as ev # noqa: E402 +from agent_kv.models import AgentKVJob, AgentKVKey, JobStatus # noqa: E402 +from agent_kv.tests._factories import kv_key # noqa: E402 + + +def _post(data=None): + return APIRequestFactory().post("/agent-kv/", data or {}, format="multipart") + + +def _plugin_with_gate(denial=None): + """A plugin dict shaped like the cloud build's, with a gate that admits. + + Every submit test needs one now: a plugin WITHOUT `service_class` means the + deployment cannot check entitlement, and the view refuses rather than + dispatching billable work unmetered (see + `test_plugin_without_a_service_class_is_refused`). + """ + gate = mock.Mock() + gate.check.return_value = denial + return {"module": object(), "service_class": lambda: gate} + + +def _authed_post(data=None): + req = _post(data) + req.META["HTTP_AUTHORIZATION"] = "Bearer 123e4567-e89b-12d3-a456-426614174001" + return req + + +# Fields describing the REQUEST rather than an extractor (spec §7.1); anything +# else an override names is routed into the kv extractor's options, so each test +# still reads as "submit with this one thing changed". +_JOB_LEVEL = { + "file", + "extractors", + "page_start", + "page_end", + "timeout", + "tags", + "custom_data", + "webhook_url", +} + + +def _valid_validated_data(**overrides): + """`SubmitSerializer.validated_data` in the extractor-scoped shape (§7.0).""" + # `table`, not `kv`. The carve-out re-pointed the real-serializer tests and + # missed this mocked payload, so every view test driven by it called + # `dispatch_job(extractor="kv")` -- whose FIRST statement is + # `EXTRACTOR_ROUTES["kv"]`, a KeyError raised OUTSIDE the try. That landed + # in SubmitView's belt-and-braces `except Exception` and produced the exact + # 500 some of these tests assert, so they passed while never reaching the + # code they name. It also meant the whole file exercised the view with an + # `extractors` payload the real serializer rejects with a 400. + keys = overrides.pop("keys", {"target_table": "Rent roll"}) + # EXPLICIT, and empty by default. `_resolved_adapters` reads + # `entry.get("adapters")` and returns early when it is falsy -- a MagicMock + # auto-attribute here would be TRUTHY, so the gate would fall through to a + # real `AdapterInstance` query and every view test would die on "Database + # access not allowed". Tests that exercise the gate set it and mock the + # lookup; the rest leave it empty and never touch the ORM. + adapters = overrides.pop("adapters", {}) + options = { + "instructions": "", + "json_structure": "", + "enable_header_mapping": False, + "correct_number_separators": False, + "number_format": "US", + } + options.update({k: overrides.pop(k) for k in list(overrides) if k not in _JOB_LEVEL}) + data = { + "file": mock.Mock(name="uploaded_file"), + "extractors": [ + {"name": "table", "keys": keys, "adapters": adapters, "options": options} + ], + "page_start": 1, + "page_end": None, + "timeout": 0, + "tags": [], + "custom_data": None, + "webhook_url": "", + } + data.update(overrides) + return data + + +def _mock_serializer(m_cls, **overrides): + instance = m_cls.return_value + instance.is_valid.return_value = True + instance.validated_data = _valid_validated_data(**overrides) + instance.pages_total = 3 + return instance + + +#: A `table` submit's three required adapter roles. +_TABLE_ADAPTERS = { + "llm": "11111111-1111-1111-1111-111111111111", + "lite_llm": "22222222-2222-2222-2222-222222222222", + "x2text": "33333333-3333-3333-3333-333333333333", +} + +#: Which `AdapterTypes` value each role must resolve to. +_ROLE_TYPES = {"llm": "LLM", "lite_llm": "LLM", "x2text": "X2TEXT"} + + +def _adapters_owned_by( + organization_id, *, types=None, missing=(), unusable=(), unavailable=() +): + """Patch the single-adapter lookup `_resolved_adapters` performs. + + Patches `ev._lookup_adapter`, the view's own seam, rather than the ORM. + Two reasons, both learned the hard way: `AdapterInstance._base_manager` is + a read-only property and cannot be patched at all, and patching `.objects` + is what hid the defect where the gate refused every valid submit (that + manager auto-filters by a request-local org this route never sets). + + So these tests assert the GATE's logic -- org match, type match, identical + refusals -- and `tests/test_adapter_scoping.py` asserts that the real + lookup finds a real adapter with no request context. Neither covers the + other. + + An id in `missing` resolves to None (not this org's, or nonexistent -- the + gate cannot and must not tell those apart); everything else resolves to an + adapter of the type `types` gives, defaulting to the right one per role. + `unusable` / `unavailable` flip `is_usable` / `is_available` for the ids + named, which is how the trial-exhaustion and deprecation refusals are + reached here. + + Both flags are set EXPLICITLY on every stand-in, never left to `Mock`'s + auto-attribute. A bare `mock.Mock()` returns a truthy child for any + attribute, so a gate that reads `adapter.is_usable` would be satisfied by a + mock that was never asked to model it -- the gate could be deleted and + these tests would still pass. The real values live in + `tests/test_adapter_scoping.py`, against rows the DB defaulted itself. + """ + type_by_id = { + aid: (types or {}).get(role, _ROLE_TYPES[role]) + for role, aid in _TABLE_ADAPTERS.items() + } + missing_ids = {str(m) for m in missing} + unusable_ids = {str(m) for m in unusable} + unavailable_ids = {str(m) for m in unavailable} + + def _lookup(adapter_id, org_id): + aid = str(adapter_id) + if aid in missing_ids or org_id != organization_id: + return None + return mock.Mock( + id=aid, + adapter_type=type_by_id.get(aid, "LLM"), + is_usable=aid not in unusable_ids, + is_available=aid not in unavailable_ids, + ) + + return mock.patch.object(ev, "_lookup_adapter", side_effect=_lookup) + + +def _stamp_created_at(job, *args, **kwargs): + """save() side effect mimicking auto_now_add for a fully-mocked save(). + + A plain ``mock.patch.object(AgentKVJob, "save")`` never runs Django's real + save() machinery, so ``created_at`` (auto_now_add) is never stamped. The + view reads ``job.created_at.isoformat()`` on the 202 path, so tests that + reach it need ``autospec=True`` (to get ``self``) plus this side effect. + """ + if job.created_at is None: + job.created_at = timezone.now() + + +# --------------------------------------------------------------------------- +# 501 before anything: the agent-kv plugin probe fails first, ahead of any +# staging/dispatch/DB work. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "get_plugin", return_value=None) +@mock.patch.object(AgentKVKey, "objects") +def test_absent_plugin_501s_before_anything(m_keys, m_plugin): + m_keys.get.return_value = kv_key() + resp = ev.SubmitView.as_view()(_authed_post()) + assert resp.status_code == 501 + + +# --------------------------------------------------------------------------- +# 429: per-key rate limit refusal. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "check_key_rate", return_value=False) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_key_rate_limited_429s(m_keys, m_plugin, m_rate): + m_keys.get.return_value = kv_key() + resp = ev.SubmitView.as_view()(_authed_post()) + assert resp.status_code == 429 + assert m_rate.called + + +# --------------------------------------------------------------------------- +# 429: concurrency-slot refusal. No job row must be persisted (save() never +# called) and staging must never run. +# --------------------------------------------------------------------------- +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "stage_input") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_concurrency_limited_429s_with_no_job_row( + m_keys, m_plugin, m_rate, m_serializer_cls, m_limiter, m_stage, m_save +): + m_keys.get.return_value = kv_key() + _mock_serializer(m_serializer_cls) + m_limiter.check_and_acquire.return_value = False + + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 429 + assert not m_save.called + assert not m_stage.called + + +# --------------------------------------------------------------------------- +# Staging/save failure: the concurrency slot must be released and the +# client gets a user-safe 500 — no leaked exception text, no stuck slot. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "stage_input") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_stage_input_failure_releases_slot_and_500s_with_safe_body( + m_keys, m_plugin, m_rate, m_serializer_cls, m_stage, m_limiter +): + m_keys.get.return_value = kv_key() + _mock_serializer(m_serializer_cls) + m_limiter.check_and_acquire.return_value = True + m_stage.side_effect = OSError("object store unreachable: leaked-secret-bucket-key") + + with mock.patch.object(AgentKVJob, "mark_terminal") as m_mark_terminal: + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 500 + # The internal exception text must never reach the client. + assert "leaked-secret-bucket-key" not in str(resp.data) + assert resp.data["status"] == JobStatus.FAILED.lower() + assert resp.data["error"] == "Job could not be accepted; nothing was billed." + assert "job_id" in resp.data + + # stage_input raised before job.save() ever ran: no row exists, so + # mark_terminal must not be called against a nonexistent job. + assert not m_mark_terminal.called + assert m_limiter.release.called + released_job_id = m_limiter.release.call_args.args[1] + assert released_job_id == resp.data["job_id"] + + +# --------------------------------------------------------------------------- +# Happy path: 202 with job_id/status/status_url; staging + dispatch happen; +# job is persisted. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save", autospec=True) +@mock.patch.object(ev, "stage_input", return_value="org/o/agent_kv/j/input.pdf") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_happy_path_returns_202_with_job_id_status_and_status_url( + m_keys, + m_plugin, + m_rate, + m_serializer_cls, + m_limiter, + m_stage, + m_save, + m_dispatch, +): + m_keys.get.return_value = kv_key() + _mock_serializer(m_serializer_cls) + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + + def _side_effect(job, *, extractor, schema, options, adapters=None): + job.status = JobStatus.DISPATCHED + job.dispatched_at = timezone.now() + + m_dispatch.side_effect = _side_effect + + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 202 + assert {"job_id", "status", "status_url"} <= set(resp.data.keys()) + # job_id must be a real UUID string. + uuid.UUID(resp.data["job_id"]) + assert ( + resp.data["status_url"] + == f"/{settings.AGENT_KV_PATH_PREFIX}/{resp.data['job_id']}" + ) + assert resp.data["status"] == JobStatus.DISPATCHED.lower() + + assert m_limiter.check_and_acquire.called + assert m_stage.called + assert m_save.called + assert m_dispatch.called + + +# --------------------------------------------------------------------------- +# The 9-key options dict SubmitView.post builds and forwards to dispatch_job +# must map every field correctly -- a dropped/typo'd key here would otherwise +# pass silently since nothing else asserts on dispatch_job's call args. +# +# It carries the RAW `keys` dict, not a `CompiledSchema`. The serializer used +# to stash the compiled form and nothing read it; the engine recompiles from +# the raw dict by design (see `compile.py`'s docstring), and the compiled form +# could not cross the queue anyway. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save", autospec=True) +@mock.patch.object(ev, "stage_input", return_value="org/o/agent_kv/j/input.pdf") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_dispatch_job_called_with_expected_options_and_schema( + m_keys, + m_plugin, + m_rate, + m_serializer_cls, + m_limiter, + m_stage, + m_save, + m_dispatch, +): + m_keys.get.return_value = kv_key() + keys_schema = {"target_table": "Rent roll"} + _mock_serializer( + m_serializer_cls, + keys=keys_schema, + number_format="EU", + enable_header_mapping=True, + instructions="skip the header row", + page_start=2, + page_end=5, + ) + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + + ev.SubmitView.as_view()(_authed_post()) + + assert m_dispatch.called + kwargs = m_dispatch.call_args.kwargs + # The extractor name itself was never asserted, so this test could not tell + # a `table` dispatch from a `kv` one -- which is how the mocked payload + # stayed on `kv` right through the carve-out. + assert kwargs["extractor"] == "table" + assert kwargs["schema"] == keys_schema + assert kwargs["options"] == { + "instructions": "skip the header row", + "json_structure": "", + "enable_header_mapping": True, + "correct_number_separators": False, + "number_format": "EU", + "page_start": 2, + "page_end": 5, + } + + +# --------------------------------------------------------------------------- +# Dispatch failure: job marked FAILED via mark_terminal, concurrency slot +# released, 500 response with a user-safe (non-leaking) error message. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "stage_input", return_value="org/o/agent_kv/j/input.pdf") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_dispatch_failure_marks_job_failed_releases_slot_and_500s( + m_keys, + m_plugin, + m_rate, + m_serializer_cls, + m_stage, + m_save, + m_dispatch, + m_mark_terminal, + m_limiter, +): + m_keys.get.return_value = kv_key() + _mock_serializer(m_serializer_cls) + m_limiter.check_and_acquire.return_value = True + m_dispatch.side_effect = ev.DispatchError("broker credentials: super-secret-token") + + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 500 + # The internal exception text must never reach the client. + assert "super-secret-token" not in str(resp.data) + assert resp.data["status"] == JobStatus.FAILED.lower() + assert "job_id" in resp.data + + assert m_mark_terminal.called + call = m_mark_terminal.call_args + assert call.args[2] == JobStatus.FAILED + assert call.kwargs["error"] == "Job could not be dispatched; nothing was billed." + + assert m_limiter.release.called + released_job_id = m_limiter.release.call_args.args[1] + assert released_job_id == str(call.args[0]) + + +# --------------------------------------------------------------------------- +# Belt-and-braces: a raw (non-DispatchError) exception out of dispatch_job +# must get IDENTICAL cleanup to a DispatchError — nothing may escape +# unhandled. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "stage_input", return_value="org/o/agent_kv/j/input.pdf") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_dispatch_job_raising_non_dispatch_error_is_still_caught_and_cleaned_up( + m_keys, + m_plugin, + m_rate, + m_serializer_cls, + m_stage, + m_save, + m_dispatch, + m_mark_terminal, + m_limiter, +): + m_keys.get.return_value = kv_key() + _mock_serializer(m_serializer_cls) + m_limiter.check_and_acquire.return_value = True + m_dispatch.side_effect = RuntimeError("unexpected: leaked-secret-abc") + + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 500 + assert "leaked-secret-abc" not in str(resp.data) + assert resp.data["status"] == JobStatus.FAILED.lower() + assert resp.data["error"] == "Job could not be dispatched; nothing was billed." + + assert m_mark_terminal.called + assert m_mark_terminal.call_args.args[2] == JobStatus.FAILED + assert m_limiter.release.called + + +# --------------------------------------------------------------------------- +# End-to-end regression for the widened dispatch.py try: a raw RuntimeError +# from the platform-key lookup deep inside the REAL dispatch_job must still +# result in mark_terminal + release + a safe 500 at the view layer. +# --------------------------------------------------------------------------- +@mock.patch("agent_kv.dispatch._platform_api_key") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(AgentKVJob, "save", autospec=True) +@mock.patch.object(ev, "stage_input", return_value="org/o/agent_kv/j/input.pdf") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_platform_key_lookup_failure_inside_real_dispatch_is_caught_end_to_end( + m_keys, + m_plugin, + m_rate, + m_serializer_cls, + m_stage, + m_save, + m_limiter, + m_mark_terminal, + m_platform_key, +): + # ev.dispatch_job is intentionally left real here — only the platform-key + # lookup deep inside it is mocked to raise, proving the widened + # dispatch.py try (Fix 2a) plus the view's cleanup (Fix 2b) work together. + m_keys.get.return_value = kv_key() + _mock_serializer(m_serializer_cls) + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + m_platform_key.side_effect = RuntimeError("platform db down: leaked-secret-xyz") + + resp = ev.SubmitView.as_view()(_authed_post()) + + # The assertion this test was missing. Without it it passed while never + # reaching the code it names: the mocked payload said `kv`, so + # `dispatch_job`'s first statement -- `EXTRACTOR_ROUTES["kv"]` -- raised a + # KeyError OUTSIDE the try, the view's belt-and-braces handler produced the + # same 500, and `_platform_api_key` was never called. The + # "leaked-secret-xyz" check below was vacuous: that string was never + # produced. + assert m_platform_key.called, ( + "the platform-key lookup was never reached, so this test proves " + "nothing about the widened dispatch.py try it exists to cover" + ) + + assert resp.status_code == 500 + assert "leaked-secret-xyz" not in str(resp.data) + assert resp.data["status"] == JobStatus.FAILED.lower() + assert resp.data["error"] == "Job could not be dispatched; nothing was billed." + + assert m_mark_terminal.called + assert m_limiter.release.called + + +# --------------------------------------------------------------------------- +# timeout=0: the wait branch must never run (no sleep, no polling, no +# attempt to import the Task-9 result_payload module). +# --------------------------------------------------------------------------- +@mock.patch("agent_kv.execution_views.time.sleep") +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save", autospec=True) +@mock.patch.object(ev, "stage_input", return_value="org/o/agent_kv/j/input.pdf") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_timeout_zero_returns_immediately_without_polling( + m_keys, + m_plugin, + m_rate, + m_serializer_cls, + m_limiter, + m_stage, + m_save, + m_dispatch, + m_sleep, +): + m_keys.get.return_value = kv_key() + _mock_serializer(m_serializer_cls, timeout=0) + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + + with mock.patch.object(AgentKVJob, "refresh_from_db") as m_refresh: + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 202 + assert not m_sleep.called + assert not m_refresh.called + + +# --------------------------------------------------------------------------- +# Sync-wait regression (spec §7.3 controller ruling on task-9-report.md +# concern 3): if the job fails during the wait window, the inline +# synchronous response must still be a 200 carrying the failure body -- not +# a 404. Before the fix, ``result_payload`` raised ``JobNotFound`` for any +# terminal-but-not-COMPLETED job (blank ``result_ref``), which escaped this +# view as an unhandled 404. Uses the real (unmocked) ``result_payload`` so +# the fix is exercised end-to-end, not just at the unit level. +# --------------------------------------------------------------------------- +@mock.patch("agent_kv.execution_views.time.sleep") +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save", autospec=True) +@mock.patch.object(ev, "stage_input", return_value="org/o/agent_kv/j/input.pdf") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_sync_wait_returns_200_with_failure_body_when_job_fails_mid_wait( + m_keys, + m_plugin, + m_rate, + m_serializer_cls, + m_limiter, + m_stage, + m_save, + m_dispatch, + m_sleep, +): + m_keys.get.return_value = kv_key() + _mock_serializer(m_serializer_cls, timeout=5) + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + + def _refresh_side_effect(job, *args, **kwargs): + job.status = JobStatus.FAILED + job.error = "LLM provider timed out" + + with mock.patch.object( + AgentKVJob, "refresh_from_db", autospec=True, side_effect=_refresh_side_effect + ): + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 200 + assert resp.data == { + "success": False, + "status": "failed", + "error": "LLM provider timed out", + } + assert not m_sleep.called + + +# --------------------------------------------------------------------------- +# Subscription admission (§6.6, following the API deployment path). +# +# Deployments are billed-gated by cloud's SubscriptionMiddleware, which resolves +# the org from the URL (/deployment/api/{org_name}/...). Agent-KV's URL carries +# no org segment -- the org lives in the Bearer key -- so that middleware +# resolves org_id=None for these requests and lets every one of them through. +# The gate is therefore invoked here, after key validation. +# --------------------------------------------------------------------------- + + +def _key_with_org(slug="acme-slug", pk=42): + """A key whose FK pk and org slug are deliberately DIFFERENT values. + + `Subscription.organization_id` is a CharField holding the slug; the FK pk is + an int. Making them differ is what lets the tests below detect the wrong one + being passed -- with a single shared value the assertion would pass either + way and the gate could silently never match a subscription row. + """ + from account_v2.models import Organization # noqa: PLC0415 + + key = AgentKVKey(name="k", is_active=True) + key.organization = Organization(id=pk, organization_id=slug) + return key + + +@mock.patch.object(AgentKVJob, "save") +@mock.patch.object(ev, "stage_input") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(AgentKVKey, "objects") +def test_subscription_denial_is_returned_verbatim_and_starts_no_work( + m_keys, m_rate, m_serializer_cls, m_limiter, m_stage, m_save +): + """A 402 from the gate must reach the client unchanged -- same status and + body an API deployment returns for the same subscription state -- and must + stop the request before any billable work, slot or row. + """ + from django.http import HttpResponse # noqa: PLC0415 + + m_keys.get.return_value = _key_with_org() + _mock_serializer(m_serializer_cls) + denial = HttpResponse(b'{"errors": "Trial period expired."}', status=402) + gate = mock.Mock() + gate.check.return_value = denial + + with mock.patch.object( + ev, "get_plugin", return_value={"module": object(), "service_class": lambda: gate} + ): + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 402 + assert not m_save.called + assert not m_stage.called + assert not m_limiter.check_and_acquire.called + + +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(AgentKVKey, "objects") +def test_subscription_gate_is_passed_the_org_slug_not_the_fk_pk( + m_keys, m_rate, m_serializer_cls, m_limiter +): + """The gate must receive `key.organization.organization_id` (the slug that + `Subscription.organization_id` is keyed on), NOT `key.organization_id` (the + Organization FK primary key). + + Passing the pk matches no subscription row, and the shared policy reads "no + row" as "nothing to enforce" -- so the gate would admit every request while + looking fully wired. This test is the only thing standing between that + one-attribute slip and a billing gate that never fires. + """ + m_keys.get.return_value = _key_with_org(slug="acme-slug", pk=42) + _mock_serializer(m_serializer_cls) + m_limiter.check_and_acquire.return_value = False # stop early; gate already ran + gate = mock.Mock() + gate.check.return_value = None + + with mock.patch.object( + ev, "get_plugin", return_value={"module": object(), "service_class": lambda: gate} + ): + ev.SubmitView.as_view()(_authed_post()) + + assert gate.check.called + passed_org = gate.check.call_args.args[0] + assert passed_org == "acme-slug", f"gate got {passed_org!r}, expected the org slug" + assert passed_org != 42 + + +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value={"module": object()}) +@mock.patch.object(AgentKVKey, "objects") +def test_plugin_without_a_service_class_is_refused( + m_keys, m_plugin, m_rate, m_serializer_cls, m_limiter +): + """A plugin that exposes no `service_class` cannot check entitlement, so + the submit is REFUSED rather than admitted. + + This used to degrade to "proceed", to tolerate a cloud image predating the + gate. But the admitted request dispatches billable LLM and OCR work, and + this route's URL carries no org segment, so `SubscriptionMiddleware` cannot + catch it downstream either -- a mixed deploy would run unmetered paid work + with nothing anywhere enforcing entitlement. Reported by Greptile on #2317. + + 503, not 402: the subscription was never evaluated, so reporting it as + denied would send an operator to the billing system for what is an + image-pairing problem. + """ + m_keys.get.return_value = _key_with_org() + _mock_serializer(m_serializer_cls) + m_limiter.check_and_acquire.return_value = False + + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 503, resp.data + # Refused BEFORE a slot was taken or anything was staged. + assert not m_limiter.check_and_acquire.called + + +# --------------------------------------------------------------------------- +# Integration: the REAL serializer through the REAL view. +# +# Every other test in this module patches SubmitSerializer and feeds the view a +# `validated_data` built by _valid_validated_data() -- so the view is asserted +# against a shape this file makes up, not the one the serializer actually +# emits. The serializer suite has the mirror-image blind spot: it validates +# payloads and never runs a view. Both pass even if the two disagree, which is +# precisely the seam an extractor-scoped wire format (§7.0) moves. +# +# This test builds a real multipart upload, runs it through the real serializer +# and the real view, and asserts what reaches dispatch_job -- the frozen +# OSS<->cloud contract on the far side. +# --------------------------------------------------------------------------- +def _real_multipart_post(extractors, **job_level): + import json as _json # noqa: PLC0415 + import os as _os # noqa: PLC0415 + + from django.core.files.uploadedfile import SimpleUploadedFile # noqa: PLC0415 + + fixture = _os.path.join(_os.path.dirname(__file__), "fixtures", "two_page.pdf") + with open(fixture, "rb") as fh: + upload = SimpleUploadedFile("doc.pdf", fh.read(), content_type="application/pdf") + payload = {"file": upload, "extractors": _json.dumps(extractors), **job_level} + req = APIRequestFactory().post("/agent-kv/", payload, format="multipart") + req.META["HTTP_AUTHORIZATION"] = "Bearer 123e4567-e89b-12d3-a456-426614174001" + return req + + +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save", autospec=True) # autospec: see _stamp_created_at +@mock.patch.object(ev, "stage_input") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_real_serializer_through_real_view_reaches_dispatch_intact_for_table( + m_keys, m_plugin, m_rate, m_limiter, m_stage, m_save, m_dispatch +): + """The table path's version of the test above -- and the exact contract + that already broke once: an earlier version had the consumer reading + `target_table` off the top level of `executor_params` while the producer + nested it under `schema`, which would have failed every table job, and it + survived two rounds of unit tests on both sides because each side was + internally consistent on its own. Only a test that runs the real + serializer AND the real view together, and inspects what lands at + `dispatch_job`, can catch that kind of drift. + """ + m_keys.get.return_value = kv_key() + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + schema = {"target_table": "Rent rolls"} + + with _adapters_owned_by(kv_key().organization_id): + resp = ev.SubmitView.as_view()( + _real_multipart_post( + [ + { + "name": "table", + "keys": schema, + "adapters": _TABLE_ADAPTERS, + "options": {"instructions": "skip totals"}, + } + ], + ) + ) + + assert resp.status_code == 202, resp.data + kwargs = m_dispatch.call_args.kwargs + assert kwargs["extractor"] == "table" + # The caller's own adapter instances reach the executor, by role. This is + # the whole point of the adapter wire format: the engine resolves these + # through the platform service instead of reading operator env vars. + assert kwargs["adapters"] == _TABLE_ADAPTERS + # Nested under `schema`, matching exactly what the cloud executor reads + # (`params["schema"]["target_table"]") -- NOT hoisted to the top level. + assert kwargs["schema"] == {"target_table": "Rent rolls"} + assert kwargs["options"]["instructions"] == "skip totals" + + # The job row records which extractor it ran (execution_views.py:142) -- + # asserted here since nothing else in this suite reads it back. + job = m_dispatch.call_args.args[0] + assert job.extractor == "table" + + +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_real_serializer_rejects_the_old_flat_shape_with_400(m_keys, m_plugin, m_rate): + """End-to-end proof of the hard switch: a caller on the pre-§7.0 format gets + a 400 from the real stack, not a job that quietly ran with defaults. + """ + import json as _json # noqa: PLC0415 + import os as _os # noqa: PLC0415 + + from django.core.files.uploadedfile import SimpleUploadedFile # noqa: PLC0415 + + m_keys.get.return_value = kv_key() + fixture = _os.path.join(_os.path.dirname(__file__), "fixtures", "two_page.pdf") + with open(fixture, "rb") as fh: + upload = SimpleUploadedFile("doc.pdf", fh.read(), content_type="application/pdf") + req = APIRequestFactory().post( + "/agent-kv/", + { + "file": upload, + "keys": _json.dumps({"total": {"description": "T"}}), + "qa": "false", + }, + format="multipart", + ) + req.META["HTTP_AUTHORIZATION"] = "Bearer 123e4567-e89b-12d3-a456-426614174001" + + resp = ev.SubmitView.as_view()(req) + + assert resp.status_code == 400 + assert "extractors" in str(resp.data) + + +# --------------------------------------------------------------------------- +# (9) The `extractors` file part is read under a bound, and decoded strictly. +# +# The size cap lives in `validate_extractors`, i.e. AFTER the part was +# materialised, and `DATA_UPLOAD_MAX_MEMORY_SIZE` excludes file-typed parts -- +# so a 500 MB part was read in full before being rejected at 256 KiB. One such +# request per worker process OOMs the pod: a cheap denial of service. +# +# Reported as 2.7 in the branch review. +# --------------------------------------------------------------------------- +from django.core.files.uploadedfile import SimpleUploadedFile # noqa: E402 + + +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_an_oversized_extractors_part_is_rejected_by_size(m_keys, m_plugin, m_rate): + """The cap is now consulted BEFORE the part is materialised. + + Note what this does and does not prove: it pins the 400, which is the + behaviour a caller sees. It cannot observe the read size, because the + multipart parser builds its own file object from the wire and the instance + constructed here never reaches the view. The bounded read itself is visible + in `SubmitView.post` as `part.read(limit + 1)`; a test that claimed to + verify it through this path would be asserting on an object the view never + touched. + """ + m_keys.get.return_value = _key_with_org() + cap = 64 + part = SimpleUploadedFile("extractors.json", b"x" * (cap * 50)) + + req = APIRequestFactory().post("/agent-kv/", {"extractors": part}, format="multipart") + req.META["HTTP_AUTHORIZATION"] = "Bearer 123e4567-e89b-12d3-a456-426614174001" + with mock.patch.object(ev.settings, "AGENT_KV_MAX_SCHEMA_BYTES", cap): + resp = ev.SubmitView.as_view()(req) + + assert resp.status_code == 400 + assert "exceeds" in str(resp.data) + + +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_a_non_utf8_extractors_part_is_rejected(m_keys, m_plugin, m_rate): + """`errors="replace"` let a latin-1 key name decode to U+FFFD and then + compile cleanly -- a malformed payload became a job that ran against a + schema the caller never wrote. + """ + m_keys.get.return_value = _key_with_org() + + class _Latin1Part: + def read(self, n=-1): + return b'[{"name": "table", "keys": {"target_table": "\xe9"}}]' + + req = APIRequestFactory().post( + "/agent-kv/", {"extractors": _Latin1Part()}, format="multipart" + ) + req.META["HTTP_AUTHORIZATION"] = "Bearer 123e4567-e89b-12d3-a456-426614174001" + resp = ev.SubmitView.as_view()(req) + + assert resp.status_code == 400 + assert "UTF-8" in str(resp.data) + + +# --------------------------------------------------------------------------- +# (10) A DB failure after staging must not orphan the upload. +# +# `run_ttl_cleanup` selects candidates from AgentKVJob rows, so an object whose +# row was never saved is unreachable by every cleanup path there is -- customer +# data sitting in the bucket that nobody can find or delete. +# +# Reported as 2.8 in the branch review. +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "delete_job_files") +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(ev.AgentKVConcurrencyLimiter, "check_and_acquire", return_value=True) +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "save", side_effect=RuntimeError("db down")) +@mock.patch.object(ev, "stage_input", return_value="org1/job/in.pdf") +@mock.patch.object(ev, "SubmitSerializer") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_a_save_failure_after_staging_removes_the_staged_object( + m_keys, + m_plugin, + m_rate, + m_serializer_cls, + m_stage, + m_save, + m_mark_terminal, + m_acquire, + m_release, + m_delete, +): + m_keys.get.return_value = _key_with_org() + _mock_serializer(m_serializer_cls) + + resp = ev.SubmitView.as_view()(_authed_post()) + + assert resp.status_code == 500 + assert m_delete.called, ( + "the staged upload was left in the bucket with no job row carrying its " + "ref; TTL cleanup selects from job rows, so it can never be reached" + ) + + +# --------------------------------------------------------------------------- +# The adapter tenancy gate. +# +# `POST /agent-kv/` names platform adapters by id. The platform service DOES +# re-scope its lookups by organization (`WHERE id=%s and organization_id=%s`, +# against the org of the job's own platform key), so this gate is defence in +# depth rather than the only thing standing between a caller and another +# tenant's credential -- the comment here previously claimed the latter. +# +# It is still load-bearing for what the platform service does not do: produce a +# 400 at submit naming the role (instead of an `SdkError` mid-run, after a slot +# and staging are spent), and refuse an exhausted trial (`is_usable`) or a +# deprecated adapter (`is_available`), neither of which the platform service +# checks before handing the credentials back. +# --------------------------------------------------------------------------- + + +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save", autospec=True) +@mock.patch.object(ev, "stage_input") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +@pytest.mark.parametrize("stolen_role", ["llm", "lite_llm", "x2text"]) +def test_an_adapter_from_another_org_is_refused( + m_keys, m_plugin, m_rate, m_limiter, m_stage, m_save, m_dispatch, stolen_role +): + """Parametrised across all three roles: one unguarded slot is enough.""" + key = kv_key() + m_keys.get.return_value = key + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + + with _adapters_owned_by( + key.organization_id, missing=[_TABLE_ADAPTERS[stolen_role]] + ): + resp = ev.SubmitView.as_view()( + _real_multipart_post( + [ + { + "name": "table", + "keys": {"target_table": "Rent rolls"}, + "adapters": _TABLE_ADAPTERS, + "options": {}, + } + ], + ) + ) + + assert resp.status_code == 400, resp.data + assert "no such adapter in this organization" in str(resp.data) + assert stolen_role in str(resp.data) + assert not m_dispatch.called, "an unowned adapter must never reach the executor" + assert not m_stage.called, "and nothing may be staged or billed first" + + +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save", autospec=True) +@mock.patch.object(ev, "stage_input") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_an_adapter_of_the_wrong_type_is_refused( + m_keys, m_plugin, m_rate, m_limiter, m_stage, m_save, m_dispatch +): + """An X2TEXT id in the `llm` slot resolves fine and then fails deep inside + the engine as a provider error -- which reads like a broken model rather + than two swapped UUIDs. Caught at submit instead. + """ + key = kv_key() + m_keys.get.return_value = key + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + + with _adapters_owned_by(key.organization_id, types={"llm": "X2TEXT"}): + resp = ev.SubmitView.as_view()( + _real_multipart_post( + [ + { + "name": "table", + "keys": {"target_table": "Rent rolls"}, + "adapters": _TABLE_ADAPTERS, + "options": {}, + } + ], + ) + ) + + assert resp.status_code == 400, resp.data + assert "expected 'LLM'" in str(resp.data) + assert not m_dispatch.called + + +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save", autospec=True) +@mock.patch.object(ev, "stage_input") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_an_exhausted_trial_adapter_is_refused_before_anything_is_billed( + m_keys, m_plugin, m_rate, m_limiter, m_stage, m_save, m_dispatch +): + """`is_usable=False` must 400, not 202. + + A frictionlessly onboarded org runs on operator-funded sample credentials; + billing flips `is_usable` when the free allowance is gone, and the IDE, + workflows and Prompt Studio all refuse from that moment. The platform + service does NOT check the flag -- it hands the credentials back -- so if + this endpoint does not check it either, the submit is accepted and the + OPERATOR pays for the caller's extraction. + """ + key = kv_key() + m_keys.get.return_value = key + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + + with _adapters_owned_by(key.organization_id, unusable=[_TABLE_ADAPTERS["llm"]]): + resp = ev.SubmitView.as_view()( + _real_multipart_post( + [ + { + "name": "table", + "keys": {"target_table": "Rent rolls"}, + "adapters": _TABLE_ADAPTERS, + "options": {}, + } + ], + ) + ) + + assert resp.status_code == 400 + assert "exhausted" in str(resp.data), resp.data + assert not m_dispatch.called, "an exhausted trial must not reach the executor" + assert not m_stage.called, "nor be charged for staging the upload" + + +@mock.patch.object(ev, "dispatch_job") +@mock.patch.object(AgentKVJob, "save", autospec=True) +@mock.patch.object(ev, "stage_input") +@mock.patch.object(ev, "AgentKVConcurrencyLimiter") +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(ev, "get_plugin", return_value=_plugin_with_gate()) +@mock.patch.object(AgentKVKey, "objects") +def test_a_deprecated_adapter_is_refused_at_submit( + m_keys, m_plugin, m_rate, m_limiter, m_stage, m_save, m_dispatch +): + """`is_available=False` means the SDK registry no longer carries it. + + Left to the executor this raises `InValidAdapterId` deep in the engine and + reaches the caller as a mid-run extraction failure, on a job that already + took a slot and billed for staging. + """ + key = kv_key() + m_keys.get.return_value = key + m_limiter.check_and_acquire.return_value = True + m_save.side_effect = _stamp_created_at + + with _adapters_owned_by( + key.organization_id, unavailable=[_TABLE_ADAPTERS["x2text"]] + ): + resp = ev.SubmitView.as_view()( + _real_multipart_post( + [ + { + "name": "table", + "keys": {"target_table": "Rent rolls"}, + "adapters": _TABLE_ADAPTERS, + "options": {}, + } + ], + ) + ) + + assert resp.status_code == 400 + assert "deprecated" in str(resp.data), resp.data + assert not m_dispatch.called + assert not m_stage.called diff --git a/backend/agent_kv/tests/test_sweeps.py b/backend/agent_kv/tests/test_sweeps.py new file mode 100644 index 0000000000..94001448ef --- /dev/null +++ b/backend/agent_kv/tests/test_sweeps.py @@ -0,0 +1,844 @@ +"""Agent-KV never-dispatched sweep and TTL cleanup internal endpoints +(spec §5.4, task-14-brief.md). + +Same mock-based style as test_internal_views.py: no real DB, every +``AgentKVJob.objects`` (and, for TTLCleanupView, ``delete_job_files``) call +is mocked and its arguments/ordering are asserted directly. That is also +the mechanism for the two predicate-shaped guarantees in the brief that a +mock can't literally execute against a database: + +* "non-expired jobs untouched" -- proven by asserting the exact + ``expires_at__lt`` filter kwarg the candidate query is built with. +* "blank-ref rows excluded / a second run over the same set is a no-op" -- + proven by asserting the exact ``Q(input_ref__gt="") | Q(result_ref__gt="")`` + filter the candidate query is built with: a row TTLCleanupView just + blanked no longer satisfies that predicate, so it drops out of the next + call's candidate set. + +These tests exercise the two-phase sweep and TTL-cleanup logic through the +(now-thin) ``SweepView``/``TTLCleanupView`` -- the logic itself lives in +``agent_kv.maintenance`` (moved there so the ``agent_kv_sweep``/ +``agent_kv_ttl_cleanup`` management commands can share it), which is why +``delete_job_files`` is patched on the ``maintenance`` module below rather +than on ``internal_views``. +""" + +import logging +import os +import uuid +from datetime import timedelta +from unittest import mock + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from django.conf import settings # noqa: E402 +from django.db.models import Q # noqa: E402 +from django.db.models.functions import Coalesce # noqa: E402 +from django.utils import timezone # noqa: E402 +from rest_framework.test import APIRequestFactory # noqa: E402 + +from agent_kv import internal_views as iv # noqa: E402 +from agent_kv import maintenance # noqa: E402 +from agent_kv.models import AgentKVJob, JobStatus # noqa: E402 + + +def _post(path, body=None): + return APIRequestFactory().post(path, body or {}, format="json") + + +# --------------------------------------------------------------------------- +# SweepView +# --------------------------------------------------------------------------- +# +# SweepView runs two independent phases per call -- never-dispatched PENDING +# jobs, then stuck DISPATCHED/RUNNING jobs -- each its own +# filter().order_by()[:500] chain against the same (mocked) AgentKVJob.objects +# manager. ``_wire_sweep_phases`` gives each phase call its own Mock object +# (via `side_effect`, keyed on call ORDER: phase 1 first, then phase 2) so a +# test can assert on -- and control the candidates of -- one phase without the +# other phase's identical-shaped chain aliasing it. + + +def _wire_sweep_phases(m_objects, never_dispatched=(), stuck=(), cancelled=()): + """Wire the sweep's three phase queries. + + Phase 2 is `filter(status__in=...).filter(Q(...) | Q(...))` -- two levels, + because the age test is an OR: `dispatched_at < cutoff` OR + `dispatched_at IS NULL AND created_at < cutoff`. A `__lt` filter alone can + never match a NULL (SQL `NULL < x` is not true), so without that second arm + a DISPATCHED/RUNNING row with a NULL `dispatched_at` hangs forever. + + Phase 3 releases slots still held by jobs CANCELLED after dispatch, whose + executor died without calling back -- neither phase above selects a + CANCELLED row, so the slot would sit occupied until Redis expired it. + """ + phase1_qs = mock.MagicMock() + phase1_qs.order_by.return_value.__getitem__.return_value = list(never_dispatched) + phase2_status_qs = mock.MagicMock() + phase2_qs = phase2_status_qs.filter.return_value + phase2_qs.order_by.return_value.__getitem__.return_value = list(stuck) + phase3_qs = mock.MagicMock() + phase3_qs.order_by.return_value.__getitem__.return_value = list(cancelled) + m_objects.filter.side_effect = [phase1_qs, phase2_status_qs, phase3_qs] + return phase1_qs, phase2_qs, phase2_status_qs + + +# (1) the never-dispatched-phase candidate query is exactly PENDING + older +# than the grace + dispatched_at IS NULL -- assert the filter kwargs directly +# (this IS the proof that "only PENDING+old+undispatched" are swept; nothing +# else does a real DB round trip in this suite) -- and, per the task-14-review +# ruling, oldest-created-first and capped at 500: an unbounded queryset would +# load a whole infra-incident backlog into memory and hold the request open +# through it, exactly when the sweep matters most. +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_sweep_queries_pending_older_than_grace_and_undispatched( + m_objects, m_mark_terminal +): + frozen_now = timezone.now() + phase1_qs, phase2_qs, phase2_status_qs = _wire_sweep_phases(m_objects) + + with mock.patch.object(timezone, "now", return_value=frozen_now): + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"swept": 0, "timed_out": 0, "released": 0} + filter_kwargs = m_objects.filter.call_args_list[0].kwargs + assert filter_kwargs["status"] == JobStatus.PENDING + assert filter_kwargs["dispatched_at__isnull"] is True + assert filter_kwargs["created_at__lt"] == frozen_now - timedelta( + seconds=settings.AGENT_KV_SWEEP_GRACE_SECONDS + ) + phase1_qs.order_by.assert_called_once_with("created_at") + phase1_qs.order_by.return_value.__getitem__.assert_called_once_with( + slice(None, 500, None) + ) + assert not m_mark_terminal.called + + +# (2) each never-dispatched candidate is terminalized via +# mark_terminal(FAILED, "Job was never dispatched") -- the guarded write, not +# a raw .update(). +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_sweep_terminalizes_each_candidate_as_failed_never_dispatched( + m_objects, m_mark_terminal, m_release +): + job = AgentKVJob(id=uuid.uuid4(), organization_id="org1") + _wire_sweep_phases(m_objects, never_dispatched=[job]) + m_mark_terminal.return_value = True + + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + m_mark_terminal.assert_called_once_with( + job.id, "org1", JobStatus.FAILED, error="Job was never dispatched" + ) + + +# (3) a job the guard actually wins gets its concurrency slot released, with +# its own org id and job id. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_sweep_releases_the_concurrency_slot_of_each_swept_job( + m_objects, m_mark_terminal, m_release +): + job = AgentKVJob(id=uuid.uuid4(), organization_id="org7") + _wire_sweep_phases(m_objects, never_dispatched=[job]) + m_mark_terminal.return_value = True + + iv.SweepView.as_view()(_post("/x")) + + m_release.assert_called_once_with("org7", str(job.id)) + + +# (4) a candidate that LOSES the mark_terminal guard (raced to terminal by +# a concurrent finalize/cancel/duplicate sweep between the candidate read +# and the guarded write) is not counted as swept and its slot is not +# released here -- whichever path won the race already released it. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_sweep_count_reflects_guard_outcomes_not_candidate_count( + m_objects, m_mark_terminal, m_release +): + won_job = AgentKVJob(id=uuid.uuid4(), organization_id="org1") + lost_job = AgentKVJob(id=uuid.uuid4(), organization_id="org1") + _wire_sweep_phases(m_objects, never_dispatched=[won_job, lost_job]) + m_mark_terminal.side_effect = [True, False] + + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"swept": 1, "timed_out": 0, "released": 0} + m_release.assert_called_once_with("org1", str(won_job.id)) + + +# (5) no candidates -> {"swept": 0, "timed_out": 0}, and no terminalize/ +# release side effects. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_sweep_with_no_candidates_is_a_pure_noop(m_objects, m_mark_terminal, m_release): + _wire_sweep_phases(m_objects) + + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"swept": 0, "timed_out": 0, "released": 0} + assert not m_mark_terminal.called + assert not m_release.called + + +# --------------------------------------------------------------------------- +# SweepView -- stuck-job (phase 2) terminalizer (Fix 8) +# --------------------------------------------------------------------------- + + +# (5a) the stuck-job-phase candidate query is exactly +# DISPATCHED/RUNNING + dispatched_at older than the stuck grace, ordered +# oldest-dispatched-first and capped at 500 -- same batch-safety rationale as +# phase 1. +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_stuck_sweep_queries_dispatched_and_running_older_than_stuck_grace( + m_objects, m_mark_terminal +): + frozen_now = timezone.now() + phase1_qs, phase2_qs, phase2_status_qs = _wire_sweep_phases(m_objects) + + with mock.patch.object(timezone, "now", return_value=frozen_now): + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"swept": 0, "timed_out": 0, "released": 0} + filter_kwargs = m_objects.filter.call_args_list[1].kwargs + assert set(filter_kwargs["status__in"]) == {JobStatus.DISPATCHED, JobStatus.RUNNING} + cutoff = frozen_now - timedelta(seconds=settings.AGENT_KV_STUCK_JOB_GRACE_SECONDS) + # The age test is an OR, and the second arm is load-bearing: a `__lt` + # filter alone can never match a NULL `dispatched_at`, so such a row would + # be invisible to this phase AND to phase 1 (which requires PENDING), and + # would hang non-terminal forever. A row with no dispatch time falls back + # to `created_at` -- the only timestamp it has. + (age_q,), age_kwargs = phase2_status_qs.filter.call_args + assert age_kwargs == {} + assert age_q == ( + Q(dispatched_at__lt=cutoff) | Q(dispatched_at__isnull=True, created_at__lt=cutoff) + ) + # Coalesce, NOT a bare "dispatched_at". This assertion used to pin the + # bare column, which is the bug: Postgres sorts ascending NULLS LAST, so + # with a full batch of non-NULL stuck rows ahead of them, the + # `dispatched_at IS NULL` rows the second Q arm exists to recover were + # never selected. The backstop could not fire in exactly the situation it + # exists for -- a backlog. (2.14 in the branch review.) + phase2_qs.order_by.assert_called_once_with(Coalesce("dispatched_at", "created_at")) + phase2_qs.order_by.return_value.__getitem__.assert_called_once_with( + slice(None, 500, None) + ) + assert not m_mark_terminal.called + + +# (5b) each stuck candidate is terminalized via mark_terminal(FAILED, +# "Job timed out") -- distinct error text from the never-dispatched phase. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_stuck_sweep_terminalizes_each_candidate_as_failed_timed_out( + m_objects, m_mark_terminal, m_release +): + job = AgentKVJob(id=uuid.uuid4(), organization_id="org1", status=JobStatus.RUNNING) + _wire_sweep_phases(m_objects, stuck=[job]) + m_mark_terminal.return_value = True + + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"swept": 0, "timed_out": 1, "released": 0} + m_mark_terminal.assert_called_once_with( + job.id, "org1", JobStatus.FAILED, error="Job timed out" + ) + + +# (5c) a stuck job the guard wins gets its concurrency slot released. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_stuck_sweep_releases_the_concurrency_slot_of_each_timed_out_job( + m_objects, m_mark_terminal, m_release +): + job = AgentKVJob(id=uuid.uuid4(), organization_id="org9", status=JobStatus.DISPATCHED) + _wire_sweep_phases(m_objects, stuck=[job]) + m_mark_terminal.return_value = True + + iv.SweepView.as_view()(_post("/x")) + + m_release.assert_called_once_with("org9", str(job.id)) + + +# (5d) a stuck candidate that LOSES the guard (raced to terminal by a +# concurrent finalize/cancel/duplicate sweep) is not counted as timed_out and +# its slot is not released here. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_stuck_sweep_count_reflects_guard_outcomes_not_candidate_count( + m_objects, m_mark_terminal, m_release +): + won_job = AgentKVJob(id=uuid.uuid4(), organization_id="org1") + lost_job = AgentKVJob(id=uuid.uuid4(), organization_id="org1") + _wire_sweep_phases(m_objects, stuck=[won_job, lost_job]) + m_mark_terminal.side_effect = [True, False] + + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"swept": 0, "timed_out": 1, "released": 0} + m_release.assert_called_once_with("org1", str(won_job.id)) + + +# (5e) the two phases' counts are independent -- a hit in one phase doesn't +# affect the other's count, and both run on the same call. +@mock.patch.object(iv.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal") +@mock.patch.object(AgentKVJob, "objects") +def test_sweep_reports_both_phase_counts_independently( + m_objects, m_mark_terminal, m_release +): + never_dispatched_job = AgentKVJob(id=uuid.uuid4(), organization_id="org1") + stuck_job = AgentKVJob(id=uuid.uuid4(), organization_id="org1") + _wire_sweep_phases( + m_objects, never_dispatched=[never_dispatched_job], stuck=[stuck_job] + ) + m_mark_terminal.return_value = True + + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"swept": 1, "timed_out": 1, "released": 0} + + +# --------------------------------------------------------------------------- +# TTLCleanupView +# --------------------------------------------------------------------------- + + +class _Lane: + """One of run_ttl_cleanup's two candidate queries, recording how it was used. + + A real object rather than a ``Mock``: the lane is sorted and then SLICED, + and a Mock whose ``__getitem__`` is wrapped to record the slice ends up + calling itself (the wrapper re-enters the same mock), which is how the + first version of this helper recursed instead of asserting. + """ + + def __init__(self, rows): + self.rows = list(rows) + self.order_by_args = None + self.slice = None + + def order_by(self, *args): + self.order_by_args = args + return self + + def __getitem__(self, sl): + self.slice = sl + return self.rows[sl] + + +class _Lanes: + """Stand-in for the narrowed expired queryset, dispatching to either lane. + + run_ttl_cleanup narrows the expired set once and queries it TWICE -- + ``cleanup_failed_at__isnull=False`` (retries, a reserved slice) and + ``...=True`` (new expirations, the remainder). A plain ``Mock`` hands back + the same child for both calls regardless of arguments, so the two lanes + would yield identical rows and every job would be processed twice. + """ + + def __init__(self, retries, fresh): + self.retry = _Lane(retries) + self.fresh = _Lane(fresh) + self.filter_kwargs: list[dict] = [] + + def filter(self, *_args, **kwargs): + self.filter_kwargs.append(kwargs) + is_null = kwargs.get("cleanup_failed_at__isnull") + return self.retry if is_null is False else self.fresh + + +def _ttl_lanes(m_objects, *, retries=(), fresh=()): + """Wire a mocked ``AgentKVJob.objects`` to both TTL-cleanup lanes. + + Three chained filters now: `expires_at < now`, then TERMINAL-only (2.15 -- + a running job must not have its input deleted out from under it), then the + non-blank-ref Q. + """ + lanes = _Lanes(retries, fresh) + m_objects.filter.return_value.filter.return_value.filter.return_value = lanes + return lanes + + +# (6) the candidate query is exactly expires_at < now AND (non-blank +# input_ref OR non-blank result_ref), capped at 500 across both lanes. +# This is what proves BOTH "non-expired untouched" (the expires_at__lt +# half) and "blank-ref rows excluded / second run is a no-op" (the Q half: +# a row TTLCleanupView just blanked no longer satisfies `__gt=""`). +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_queries_expired_jobs_with_a_nonblank_ref(m_objects, m_delete): + frozen_now = timezone.now() + lanes = _ttl_lanes(m_objects) + + with mock.patch.object(timezone, "now", return_value=frozen_now): + resp = iv.TTLCleanupView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"cleaned": 0, "retained": 0} + assert m_objects.filter.call_args_list[0].kwargs == {"expires_at__lt": frozen_now} + # TERMINAL-only: a job still RUNNING past its TTL must keep its staged + # input. The retention policy covers finished work, not running work. + status_filter = m_objects.filter.return_value.filter.call_args + assert set(status_filter.kwargs["status__in"]) == set(AgentKVJob.TERMINAL) + (q_arg,), q_kwargs = ( + m_objects.filter.return_value.filter.return_value.filter.call_args + ) + assert q_kwargs == {} + assert q_arg == (Q(input_ref__gt="") | Q(result_ref__gt="")) + # Two lanes off that one narrowed set: retries first, then new expirations. + assert lanes.filter_kwargs == [ + {"cleanup_failed_at__isnull": False}, + {"cleanup_failed_at__isnull": True}, + ] + assert not m_delete.called + + +# (6b) each lane sorts by a SINGLE named column ascending -- never +# `cleanup_failed_at NULLS FIRST`. A btree index is NULLS LAST ascending, so +# that ordering could not use the (cleanup_failed_at, expires_at) index at all +# and Postgres had to sort every matching expired row before applying the +# limit, work that grows with the backlog. Splitting the query removed it. +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_lanes_sort_by_a_plain_ascending_column(m_objects, m_delete): + lanes = _ttl_lanes(m_objects) + + iv.TTLCleanupView.as_view()(_post("/x")) + + assert lanes.retry.order_by_args == ("cleanup_failed_at",) + assert lanes.fresh.order_by_args == ("expires_at",) + + +# (6c) with nothing to retry, the fresh lane still gets the WHOLE batch -- the +# reserve is a cap on the retry lane, not a permanent tax on the normal case. +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_fresh_lane_gets_the_whole_batch_when_nothing_is_retrying( + m_objects, m_delete +): + lanes = _ttl_lanes(m_objects) + + iv.TTLCleanupView.as_view()(_post("/x")) + + assert lanes.retry.slice == slice(None, 100, None) + assert lanes.fresh.slice == slice(None, 500, None) + + +# (6d) THE fix for the second starvation direction. With the retry lane full, +# the fresh lane shrinks by exactly that many, so the two together never exceed +# the batch cap -- and crucially the retry lane is served FIRST, so a steady +# stream of new expirations can no longer push retries out of every batch and +# leave their files in storage indefinitely. Both orderings tried before this +# starved one side: oldest-expiry-first starved fresh work, NULLS FIRST starved +# retries. +@mock.patch.object(maintenance, "delete_job_files", return_value=["input_ref"]) +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_reserves_capacity_for_retries_under_a_flood_of_new_expiries( + m_objects, m_delete +): + retries = [ + AgentKVJob(id=uuid.uuid4(), input_ref=f"r{i}", cleanup_failed_at=timezone.now()) + for i in range(100) + ] + fresh = [AgentKVJob(id=uuid.uuid4(), input_ref=f"f{i}") for i in range(500)] + lanes = _ttl_lanes(m_objects, retries=retries, fresh=fresh) + + resp = iv.TTLCleanupView.as_view()(_post("/x")) + + assert lanes.retry.slice == slice(None, 100, None) + assert lanes.fresh.slice == slice(None, 400, None) + # Every retry was attempted even though 500 fresh rows were queued behind + # them, and the batch cap still held: 100 + 400 == 500. + assert m_delete.call_count == 500 + assert resp.data == {"cleaned": 0, "retained": 500} + + +# (7) files are deleted BEFORE the refs are blanked -- order matters: a +# delete failure must not blank a ref pointing at a file that's still there. +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_deletes_files_before_blanking_refs(m_objects, m_delete): + job_id = uuid.uuid4() + job = AgentKVJob( + id=job_id, input_ref="org/o/j/input.pdf", result_ref="org/o/j/result.json" + ) + m_delete.return_value = ["input_ref", "result_ref"] + m_qs = m_objects.filter.return_value + _ttl_lanes(m_objects, fresh=[job]) + + manager = mock.Mock() + manager.attach_mock(m_delete, "delete_job_files") + manager.attach_mock(m_qs.update, "update") + + resp = iv.TTLCleanupView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"cleaned": 1, "retained": 0} + assert [c[0] for c in manager.mock_calls] == ["delete_job_files", "update"] + m_delete.assert_called_once_with(job) + + +# (8) both refs are blanked via `.update()` (not `job.save()`), targeting +# exactly this job's row, and the row itself is left in place (no .delete() +# call is ever made on the queryset). +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_blanks_both_refs_for_the_job_row(m_objects, m_delete): + job_id = uuid.uuid4() + job = AgentKVJob(id=job_id, input_ref="org/o/j/input.pdf", result_ref="") + m_delete.return_value = ["input_ref", "result_ref"] + m_qs = m_objects.filter.return_value + _ttl_lanes(m_objects, fresh=[job]) + + iv.TTLCleanupView.as_view()(_post("/x")) + + assert m_objects.filter.call_args_list[-1].kwargs == {"id": job_id} + m_qs.update.assert_called_once_with( + input_ref="", result_ref="", cleanup_failed_at=None + ) + assert not m_qs.delete.called + + +# (9) `cleaned` counts jobs actually processed this call, across multiple +# candidates. +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_returns_count_of_jobs_cleaned(m_objects, m_delete): + job1 = AgentKVJob(id=uuid.uuid4(), input_ref="a", result_ref="") + job2 = AgentKVJob(id=uuid.uuid4(), input_ref="", result_ref="b") + m_delete.return_value = ["input_ref", "result_ref"] + m_qs = m_objects.filter.return_value + _ttl_lanes( + m_objects, + fresh=[ + job1, + job2, + ], + ) + + resp = iv.TTLCleanupView.as_view()(_post("/x")) + + assert resp.data == {"cleaned": 2, "retained": 0} + assert m_delete.call_count == 2 + assert m_qs.update.call_count == 2 + + +# (10) no candidates -> {"cleaned": 0}, nothing deleted or updated. This is +# also the concrete shape of a "second run over the same set" once every +# candidate from the first run has had its refs blanked: the (mocked) +# candidate query simply returns nothing. +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_with_no_candidates_is_a_pure_noop(m_objects, m_delete): + m_qs = m_objects.filter.return_value + _ttl_lanes(m_objects, fresh=[]) + + resp = iv.TTLCleanupView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data == {"cleaned": 0, "retained": 0} + assert not m_delete.called + assert not m_qs.update.called + + +# (10b) THE case test (7)'s comment claims ("a delete failure must not blank a +# ref pointing at a file that's still there") but which nothing actually +# asserted until the Greptile review: test (7) only pinned the ORDER of the two +# calls, and ordering is irrelevant when the update blanks both refs regardless +# of what the delete returned. Before the fix this test fails -- `update` is +# called with both refs blanked and the still-present result file loses its only +# handle, since the candidate query below matches only rows with a non-blank +# ref. +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_keeps_the_ref_whose_file_delete_failed(m_objects, m_delete): + job_id = uuid.uuid4() + job = AgentKVJob( + id=job_id, input_ref="org/o/j/input.pdf", result_ref="org/o/j/result.json" + ) + # Input gone, result delete raised inside delete_job_files. + m_delete.return_value = ["input_ref"] + m_qs = m_objects.filter.return_value + _ttl_lanes(m_objects, fresh=[job]) + + resp = iv.TTLCleanupView.as_view()(_post("/x")) + + # Only the confirmed-gone ref is blanked; result_ref is NOT in the update, + # so the row still matches `result_ref__gt=""` and the next pass retries. + (), kwargs = m_qs.update.call_args + assert kwargs["input_ref"] == "" + assert "result_ref" not in kwargs + # Stamped so the nulls-first ordering pushes this row behind all never-failed + # work on the next tick -- retried, but unable to block the backlog. + assert kwargs["cleanup_failed_at"] is not None + # Not counted as cleaned -- it is unfinished, and `retained` is what makes a + # permanently-failing backlog visible instead of silently draining to zero. + assert resp.data == {"cleaned": 0, "retained": 1} + + +# (10c) nothing confirmed gone -> no ref is blanked, but the failure is still +# recorded. Skipping the write entirely would leave cleanup_failed_at NULL, and +# a NULL sorts FIRST under the candidate ordering -- so the row would hold the +# head of every batch indefinitely, which is the starvation this ordering +# exists to prevent. +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_blanks_no_ref_but_still_stamps_the_failure(m_objects, m_delete): + job = AgentKVJob( + id=uuid.uuid4(), input_ref="org/o/j/input.pdf", result_ref="org/o/j/r.json" + ) + m_delete.return_value = [] + m_qs = m_objects.filter.return_value + _ttl_lanes(m_objects, fresh=[job]) + + resp = iv.TTLCleanupView.as_view()(_post("/x")) + + # No ref is blanked -- but the failure IS stamped, or the row would sort + # nulls-first forever and keep its place at the head of every batch. + (), kwargs = m_qs.update.call_args + assert kwargs == {"cleanup_failed_at": mock.ANY} + assert kwargs["cleanup_failed_at"] is not None + assert resp.data == {"cleaned": 0, "retained": 1} + + +# (10d) a row that failed before and succeeds now must have its stale failure +# marker cleared. It drops out of the candidate filter anyway (both refs blank), +# so this matters for the audit trail and for anything reading +# cleanup_failed_at as "currently failing" rather than "failed once". +@mock.patch.object(maintenance, "delete_job_files") +@mock.patch.object(AgentKVJob, "objects") +def test_ttl_cleanup_clears_a_stale_failure_marker_on_success(m_objects, m_delete): + job = AgentKVJob( + id=uuid.uuid4(), + input_ref="org/o/j/input.pdf", + result_ref="org/o/j/result.json", + cleanup_failed_at=timezone.now(), + ) + m_delete.return_value = ["input_ref", "result_ref"] + m_qs = m_objects.filter.return_value + _ttl_lanes(m_objects, fresh=[job]) + + resp = iv.TTLCleanupView.as_view()(_post("/x")) + + m_qs.update.assert_called_once_with( + input_ref="", result_ref="", cleanup_failed_at=None + ) + assert resp.data == {"cleaned": 1, "retained": 0} + + +# --------------------------------------------------------------------------- +# URL wiring +# --------------------------------------------------------------------------- + + +# (11) regression pin for the frozen paths (spec Interfaces block): the +# PG-scheduler/reaper periodic mechanism calls these exact URLs, so a +# dropped/renamed include in internal_base_urls.py must fail loudly here +# rather than 404 in prod. +def test_frozen_sweep_and_ttl_cleanup_urls_resolve_to_the_right_views(): + from django.urls import resolve + + sweep = resolve("/internal/v1/agent-kv/sweep/") + assert sweep.func.cls is iv.SweepView + + ttl_cleanup = resolve("/internal/v1/agent-kv/ttl-cleanup/") + assert ttl_cleanup.func.cls is iv.TTLCleanupView + + +# --------------------------------------------------------------------------- +# Phase 3: slots held by jobs CANCELLED after dispatch. +# +# Cancelling a dispatched job deliberately does not release its slot -- the +# executor is still running and still billing, so the slot belongs to the +# finalize callback. But if that executor dies no callback arrives, and neither +# phase 1 (PENDING) nor phase 2 (DISPATCHED/RUNNING) selects a CANCELLED row, +# so the slot sat occupied until Redis expired it six hours later. +# +# Reported by Greptile on PR #2317. +# --------------------------------------------------------------------------- +@mock.patch.object(maintenance.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +def test_sweep_releases_slots_held_by_abandoned_cancelled_jobs( + m_objects, m_mark_terminal, m_release +): + job = AgentKVJob(status=JobStatus.CANCELLED) + job.organization_id = "org1" + _wire_sweep_phases(m_objects, cancelled=[job]) + + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + m_release.assert_called_once_with("org1", str(job.id)) + # The job is ALREADY terminal; phase 3 only frees the slot. + assert not m_mark_terminal.called + + +@mock.patch.object(maintenance.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +def test_phase_three_s_release_count_reaches_the_caller( + m_objects, m_mark_terminal, m_release +): + """`released` must be in the RESPONSE, not only in a log line. + + Phase 3 counted its releases into a local, logged them, and returned only + `{"swept", "timed_out"}`. `SweepView` passes that dict through verbatim and + the scheduler task logs it, so no caller could ever see that a slot had + been recovered -- the same "counted then discarded" shape as the round-1 + finding about this function's return value, reintroduced by the phase added + to fix another part of it. + + Operationally: a cancelled job whose executor never called back holds a + slot against the org's cap, and phase 3 is the only thing that frees it. + Without this key, "the cap is being hit and the sweep is recovering slots" + and "the cap is being hit and nothing is recovering anything" look the same + from outside. + """ + jobs = [] + for _ in range(3): + job = AgentKVJob(status=JobStatus.CANCELLED) + job.organization_id = "org1" + jobs.append(job) + _wire_sweep_phases(m_objects, cancelled=jobs) + + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.status_code == 200 + assert resp.data["released"] == 3, resp.data + assert m_release.call_count == 3 + + +@mock.patch.object(maintenance.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +def test_a_sweep_that_strands_jobs_says_so_at_warning( + m_objects, m_mark_terminal, m_release, caplog +): + """The logging added for round-1 finding 2, pinned. + + Nothing anywhere used `caplog` on this module, so all five log lines could + be deleted and the suite stayed green -- `test_agent_kv_scheduler_tasks.py` + asserts the returned dict, which passed before the logging existed too. + + WARNING specifically, not INFO: a non-zero count means jobs were stranded + and their concurrency slots held, which is the condition an operator needs + surfaced. Before this, a sweep terminalizing a thousand jobs as FAILED + emitted nothing at all, so a backlog of stranded jobs was indistinguishable + from a quiet, healthy system. + """ + job = AgentKVJob(id=uuid.uuid4(), organization_id="org1") + _wire_sweep_phases(m_objects, never_dispatched=[job]) + + with caplog.at_level(logging.INFO, logger="agent_kv.maintenance"): + resp = iv.SweepView.as_view()(_post("/x")) + + assert resp.data["swept"] == 1 + terminalized = [ + r + for r in caplog.records + if r.levelno >= logging.WARNING and "terminalized" in r.getMessage() + ] + assert len(terminalized) == 1, [r.getMessage() for r in caplog.records] + assert "1 never-dispatched" in terminalized[0].getMessage() + + +@mock.patch.object(maintenance.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "mark_terminal", return_value=True) +@mock.patch.object(AgentKVJob, "objects") +def test_a_quiet_sweep_still_records_that_it_RAN( + m_objects, m_mark_terminal, m_release, caplog +): + """The other half, and the one that matters more. + + "Ran and found nothing" and "never ran at all" produce the same counts, and + distinguishing them is the whole reason this line exists -- the sweep + shipped unscheduled and the only symptom was an absence. So the quiet path + must log too, at INFO. + """ + _wire_sweep_phases(m_objects) + + with caplog.at_level(logging.INFO, logger="agent_kv.maintenance"): + iv.SweepView.as_view()(_post("/x")) + + assert any( + "nothing to terminalize" in r.getMessage() for r in caplog.records + ), [r.getMessage() for r in caplog.records] + + +@mock.patch.object(maintenance.AgentKVConcurrencyLimiter, "release") +@mock.patch.object(AgentKVJob, "objects") +def test_phase_three_is_bounded_on_both_sides(m_objects, m_release): + """Older than the stuck grace, so a live executor is not cut short; and + newer than the slot TTL, past which Redis has already dropped the entry and + there is nothing left to release. Without the second bound this would + rescan every cancelled job ever, forever. + """ + frozen_now = timezone.now() + _wire_sweep_phases(m_objects) + + with mock.patch.object(timezone, "now", return_value=frozen_now): + iv.SweepView.as_view()(_post("/x")) + + kwargs = m_objects.filter.call_args_list[2].kwargs + assert kwargs["status"] == JobStatus.CANCELLED + assert kwargs["dispatched_at__isnull"] is False + stuck_cutoff = frozen_now - timedelta( + seconds=settings.AGENT_KV_STUCK_JOB_GRACE_SECONDS + ) + assert kwargs["completed_at__lt"] == stuck_cutoff + assert kwargs["completed_at__gt"] == frozen_now - timedelta( + seconds=maintenance.SLOT_TTL_SECONDS + ) + + +def test_phase_three_takes_the_newest_eligible_rows_first(): + """Releasing a slot does not remove its row from phase 3's query -- there + is no "released" marker -- so oldest-first would re-select the same batch + every sweep under backlog, and rows arriving behind it would age out of the + window unprocessed. + + Newest-first inverts which rows lose: the skipped ones are the oldest, i.e. + closest to the slot TTL floor where Redis is about to drop the entry anyway + and a release buys almost nothing. + + Reported by Greptile on PR #2317. + """ + phase1 = mock.MagicMock() + phase1.order_by.return_value.__getitem__.return_value = [] + phase2_status = mock.MagicMock() + phase2_status.filter.return_value.order_by.return_value.__getitem__.return_value = [] + phase3 = mock.MagicMock() + phase3.order_by.return_value.__getitem__.return_value = [] + + with mock.patch.object(AgentKVJob, "objects") as m_objects: + m_objects.filter.side_effect = [phase1, phase2_status, phase3] + maintenance.run_sweep() + + phase3.order_by.assert_called_once_with("-completed_at") diff --git a/backend/agent_kv/tests/test_table_extractor_routing.py b/backend/agent_kv/tests/test_table_extractor_routing.py new file mode 100644 index 0000000000..696afb9573 --- /dev/null +++ b/backend/agent_kv/tests/test_table_extractor_routing.py @@ -0,0 +1,291 @@ +"""Routing a `table` extractor entry to the table executor. + +The wire format is extractor-scoped (§7.0): `extractors: [{name, keys, options}]`. +v1 accepts exactly one entry, but WHICH one is now a choice, so the job row has +to record it -- status and result used to key everything under the hardcoded +`kv` name, which would have filed a table job's output under the wrong +extractor. +""" + +import json +import os +import uuid +from unittest import mock + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from agent_kv.constants import ( # noqa: E402 + EXTRACTOR_ROUTES, + STAGE_NAMES_BY_EXTRACTOR, + TABLE_EXTRACTOR_NAME, + V1_EXTRACTOR_NAME, +) +from agent_kv.execution_serializers import ( # noqa: E402 + _OPTIONS_SERIALIZERS, + SUPPORTED_EXTRACTORS, + ExtractorSerializer, +) +from agent_kv.models import AgentKVJob, JobStatus # noqa: E402 + + +def test_table_is_the_only_supported_extractor(): + """`kv` is out of EXTRACTOR_ROUTES because this deployment ships no + `agentic_kv` plugin, so nothing drains `celery_executor_agentic_kv`. + SUPPORTED_EXTRACTORS is derived from that table, so the omission is what + turns a `kv` submit into a 400 rather than a 202 for a job that never runs. + """ + assert set(SUPPORTED_EXTRACTORS) == {TABLE_EXTRACTOR_NAME} + assert V1_EXTRACTOR_NAME not in EXTRACTOR_ROUTES + + +def test_every_supported_extractor_has_a_route_stage_list_and_options_serializer(): + """A supported extractor with no route dispatches nowhere; with no stage + list its progress is recorded and then filtered out of the status + document; with no options serializer, a submit for it raises an + uncaught KeyError at `_OPTIONS_SERIALIZERS[data["name"]]` instead of a + validation error -- `validate_name` already passed it, since + `SUPPORTED_EXTRACTORS` is derived from `EXTRACTOR_ROUTES`. + """ + for name in SUPPORTED_EXTRACTORS: + assert name in EXTRACTOR_ROUTES, name + assert name in STAGE_NAMES_BY_EXTRACTOR, name + assert name in _OPTIONS_SERIALIZERS, name + + +def test_the_table_route_targets_the_existing_executor(): + """R1: the queue is derived from the executor name, and + celery_executor_agentic_table is already wired. A new executor name here + would be accepted, dispatch silently, and never drain. + """ + executor, operation = EXTRACTOR_ROUTES[TABLE_EXTRACTOR_NAME] + assert executor == "agentic_table" + assert operation == "table_extract_api" + + +def test_a_table_entry_validates_with_its_own_options(): + data = { + "name": "table", + "keys": {"target_table": "Rent rolls"}, + "options": {"instructions": "skip totals rows"}, + } + s = ExtractorSerializer(data=data) + assert s.is_valid(), s.errors + assert s.validated_data["options"]["instructions"] == "skip totals rows" + + +def test_a_table_entry_requires_a_target_table(): + s = ExtractorSerializer(data={"name": "table", "keys": {}, "options": {}}) + assert not s.is_valid() + assert "target_table" in json.dumps(s.errors) + + +def test_kv_options_are_rejected_on_a_table_entry(): + """The whole point of per-extractor options: an option aimed at the wrong + extractor must not be silently dropped. + """ + s = ExtractorSerializer( + data={"name": "table", "keys": {"target_table": "T"}, "options": {"qa": True}} + ) + assert not s.is_valid() + assert "qa" in json.dumps(s.errors) + + +def test_a_kv_entry_is_refused_by_name(): + """Options scoping is moot for `kv` here: `validate_name` refuses the entry + before any options validator runs. What matters is that the refusal names + the extractor, so a caller sending a KV payload to this deployment learns + why rather than seeing a schema complaint. + """ + s = ExtractorSerializer( + data={ + "name": "kv", + "keys": {"total": {"description": "d"}}, + "options": {"qa": False}, + } + ) + assert not s.is_valid() + assert "unknown extractor" in json.dumps(s.errors) + assert "kv" in json.dumps(s.errors) + + +def _job(extractor=TABLE_EXTRACTOR_NAME, **overrides): + """An unsaved job with an unsaved org -- this suite's established pattern. + + `test_dispatch.py` builds jobs exactly this way and the whole `agent_kv` + suite touches no database (there is no `django_db` marker anywhere in it). + Assigning an unsaved `Organization` with an explicit PK caches it on the + job, so `job.organization` never hits the DB, and the slug is deliberately + different from the PK so a test can tell which one the code reached for. + """ + from account_v2.models import Organization + + job = AgentKVJob(id=uuid.uuid4(), input_ref="org/o/agent_kv/j/input.pdf") + job.organization = Organization(id=7, organization_id="org_slug_1") + job.extractor = extractor + job.pages_total = 3 + for key, value in overrides.items(): + setattr(job, key, value) + return job + + +def test_the_kv_route_constants_survive_for_re_enablement(): + """`kv` is unroutable, not deleted. The executor/operation pair, the stage + list and the options serializer all stay so the branch carrying the engine + re-enables the extractor by restoring one dict entry -- not by merging + content back into the files it rewrites most. + """ + from agent_kv.constants import ( + EXECUTOR_NAME, + OPERATION_KV_EXTRACT, + STAGE_NAMES_BY_EXTRACTOR, + ) + from agent_kv.execution_serializers import _OPTIONS_SERIALIZERS + + assert EXECUTOR_NAME == "agentic_kv" + assert OPERATION_KV_EXTRACT == "kv_extract" + assert V1_EXTRACTOR_NAME in STAGE_NAMES_BY_EXTRACTOR + assert V1_EXTRACTOR_NAME in _OPTIONS_SERIALIZERS + + +def test_the_status_document_keys_stages_by_the_extractor_that_ran(): + from agent_kv.execution_views import _status_document + + job = _job( + stage="table_extraction", + stages={"table_extraction": {"status": "done", "seconds": 12.5}}, + ) + doc = _status_document(job) + + assert TABLE_EXTRACTOR_NAME in doc["extractors"] + assert V1_EXTRACTOR_NAME not in doc["extractors"] + # R7: recorded AND visible. StageReportView persists any stage name the + # executor sends, but _status_document filters through the extractor's own + # list -- with the KV list this array would come back empty. + assert doc["extractors"][TABLE_EXTRACTOR_NAME]["stages"] == [ + {"name": "table_extraction", "status": "done", "seconds": 12.5} + ] + + +def test_a_kv_job_still_reports_its_kv_stages(): + from agent_kv.execution_views import _status_document + + job = _job( + extractor=V1_EXTRACTOR_NAME, + stage="extraction", + stages={ + "document_processing": {"status": "done"}, + "extraction": {"status": "running"}, + }, + ) + doc = _status_document(job) + + names = [s["name"] for s in doc["extractors"][V1_EXTRACTOR_NAME]["stages"]] + assert names == ["document_processing", "extraction"] + + +def test_the_result_payload_keys_by_the_extractor_that_ran(): + from agent_kv.execution_views_result import result_payload + + job = _job( + status=JobStatus.COMPLETED, + result_ref="ref.json", + usage_summary={"pages": 3}, + ) + with mock.patch( + "agent_kv.execution_views_result.read_result", + return_value={"tables": [{"unit": "A1"}]}, + ): + payload = result_payload(job) + + assert payload["success"] is True + assert payload["status"] == "completed" + assert payload["extractors"][TABLE_EXTRACTOR_NAME] == {"tables": [{"unit": "A1"}]} + assert payload["usage_summary"]["by_extractor"][TABLE_EXTRACTOR_NAME] == {"pages": 3} + assert V1_EXTRACTOR_NAME not in payload["extractors"] + + +def test_the_extractor_column_has_no_default(): + """Was `test_the_extractor_column_defaults_to_kv`, asserting `== "kv"`. + + It pinned the defect as the contract. The default WAS the historical truth + when migration 0002 added the column -- the API accepted exactly one + extractor and it was always `kv` -- and it became a mis-filing trap the + moment `table` existed: `kv` is a valid key in `STAGE_NAMES_BY_EXTRACTOR`, + so a table job filed under it gets the KV stage list and + `table_extraction` is silently dropped from every status response, with + nothing logged because nothing looks wrong to the filter. + + `""` matches no route, so an omission now fails the dispatch visibly. The + row-level truth 0002 back-filled is untouched: existing `kv` rows keep + saying `kv`, which is what `test_a_job_whose_extractor_has_no_stage_list_ + reports_no_stages` below depends on. + """ + assert AgentKVJob().extractor == "" + assert V1_EXTRACTOR_NAME in { + value for value, _ in AgentKVJob._meta.get_field("extractor").choices + }, "a pre-carve-out row must still be a legal value to READ back" + + +def test_a_job_whose_extractor_has_no_stage_list_reports_no_stages(caplog): + """A retired extractor name with surviving rows must not 500 `GET status`. + + `_status_document` used to subscript `STAGE_NAMES_BY_EXTRACTOR`, so such a + row raised `KeyError` on status while `GET result` kept working -- the + result payload keys by `job.extractor` without consulting that table. An + empty stage array is the honest answer for an extractor this build no + longer knows how to describe. + """ + from agent_kv.execution_views import _status_document + + job = _job( + extractor="retired_extractor", + stage="something", + stages={"something": {"status": "done"}}, + ) + + with caplog.at_level("WARNING", logger="agent_kv.execution_views"): + doc = _status_document(job) + + assert doc["extractors"]["retired_extractor"]["stages"] == [] + assert doc["status"] == job.status.lower() + assert "retired_extractor" in caplog.text + + +# --------------------------------------------------------------------------- +# `ValidateView` is part of the carve-out: the class stays in the tree and the +# route does not. Both halves are asserted here, including that the class says +# so itself -- a reader of `execution_views.py` sees a complete, decorated, +# live-looking endpoint, and the only thing that decided otherwise was a file +# they may never open. +# --------------------------------------------------------------------------- + + +def test_validate_is_not_a_public_route_on_this_deployment(): + from agent_kv import execution_urls + + names = {p.name for p in execution_urls.urlpatterns} + assert names == { + "agent_kv_submit", + "agent_kv_status", + "agent_kv_result", + "agent_kv_cancel", + }, ( + "the public route set changed; `/validate` compiles a `kv` schema and " + "`kv` is the one extractor this build refuses, so routing it would " + "advertise validation for an extractor every submit 400s" + ) + + +def test_validate_view_documents_that_it_is_unrouted(): + from agent_kv.execution_views import ValidateView + + doc = ValidateView.__doc__ or "" + assert "NOT ROUTED" in doc, ( + "ValidateView is a complete, decorated, live-looking endpoint that " + "nothing reaches. Say so on the class, not only in execution_urls.py" + ) diff --git a/backend/agent_kv/tests/test_validate_view.py b/backend/agent_kv/tests/test_validate_view.py new file mode 100644 index 0000000000..3d6b5ffcba --- /dev/null +++ b/backend/agent_kv/tests/test_validate_view.py @@ -0,0 +1,130 @@ +import json +import os +import uuid +from unittest import mock + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from rest_framework.test import APIRequestFactory # noqa: E402 + +from agent_kv import execution_views as ev # noqa: E402 +from agent_kv.models import AgentKVKey # noqa: E402 +from agent_kv.tests._factories import kv_key # noqa: E402 + + +def _authed(method="post", path="/agent-kv/validate", data=None): + if data is None: + data = {} + req = getattr(APIRequestFactory(), method)( + path, + json.dumps(data) if data else "", + content_type="application/json", + ) + req.META["HTTP_AUTHORIZATION"] = "Bearer 123e4567-e89b-12d3-a456-426614174001" + return req + + +# --------------------------------------------------------------------------- +# (1) Valid schema -> returns valid:true with correct counts +# --------------------------------------------------------------------------- +# These three mock `check_key_rate` explicitly, like test_submit_view.py does. +# They previously relied on the limiter's ambient behaviour: it failed OPEN on +# any Redis error, so with no Redis in the unit lane it returned True and the +# tests passed for a reason unrelated to what they assert. The limiter now +# fails CLOSED (a Sentinel failover used to remove the concurrency and rate +# ceilings at once while the API kept accepting billable work), so an +# unreachable Redis is a 429 and these became 429 == 200. The mock is the right +# fix either way: these are schema-validation tests, and the rate limiter +# should not be in their path at all. +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(AgentKVKey, "objects") +def test_validate_valid_schema_returns_counts(m_keys, m_rate): + m_keys.get.return_value = kv_key() + schema = { + "quotation_number": {"description": "The quote number", "required": True}, + "customer": {"name": {"description": "Bill-to name"}}, + "line_items": { + "description": "One row per line", + "_key": "sku", + "_array": { + "sku": {"description": "SKU"}, + "total": {"description": "Line total", "format": "currency"}, + }, + }, + "_constraints": ["count('line_items') >= 1"], + } + req = _authed(data={"keys": schema}) + + resp = ev.ValidateView.as_view()(req) + + assert resp.status_code == 200 + assert resp.data["valid"] is True + assert "leaves" in resp.data + assert "arrays" in resp.data + assert "constraints" in resp.data + assert isinstance(resp.data["leaves"], int) + assert isinstance(resp.data["arrays"], int) + assert isinstance(resp.data["constraints"], int) + + +# --------------------------------------------------------------------------- +# (2) Invalid schema -> returns valid:false with error message verbatim +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(AgentKVKey, "objects") +def test_validate_invalid_schema_returns_error(m_keys, m_rate): + m_keys.get.return_value = kv_key() + # Pass a non-dict as top-level schema + req = _authed(data={"keys": "not a dict"}) + + resp = ev.ValidateView.as_view()(req) + + assert resp.status_code == 200 + assert resp.data["valid"] is False + assert "error" in resp.data + assert resp.data["error"] == "Top-level key schema must be a JSON object" + + +# --------------------------------------------------------------------------- +# (3) No key (no Authorization header) -> 403 Forbidden +# --------------------------------------------------------------------------- +def test_validate_no_key_returns_403(): + req = APIRequestFactory().post("/agent-kv/validate", data={}) + resp = ev.ValidateView.as_view()(req) + assert resp.status_code == 403 + + +# --------------------------------------------------------------------------- +# (4) Over rate limit -> 429 +# --------------------------------------------------------------------------- +@mock.patch("agent_kv.execution_views.check_key_rate", return_value=False) +@mock.patch.object(AgentKVKey, "objects") +def test_validate_over_rate_limit_returns_429(m_keys, m_check_rate): + # A concrete id, because the rate limiter is keyed on it. + key = kv_key(id=uuid.uuid4()) + m_keys.get.return_value = key + req = _authed(data={"keys": {}}) + + resp = ev.ValidateView.as_view()(req) + + assert resp.status_code == 429 + + +# --------------------------------------------------------------------------- +# (5) Missing 'keys' in request body -> 400 +# --------------------------------------------------------------------------- +@mock.patch.object(ev, "check_key_rate", return_value=True) +@mock.patch.object(AgentKVKey, "objects") +def test_validate_missing_keys_returns_400(m_keys, m_rate): + m_keys.get.return_value = kv_key() + req = _authed(data={}) + + resp = ev.ValidateView.as_view()(req) + + assert resp.status_code == 400 + assert resp.data["detail"] == "body must include 'keys'" diff --git a/backend/agent_kv/urls.py b/backend/agent_kv/urls.py new file mode 100644 index 0000000000..206b7f0d95 --- /dev/null +++ b/backend/agent_kv/urls.py @@ -0,0 +1,23 @@ +from django.urls import path + +from agent_kv.views import AgentKVKeyViewSet + +urlpatterns = [ + path( + "keys/", + AgentKVKeyViewSet.as_view({"get": "list", "post": "create"}), + name="agent_kv_key_list", + ), + path( + "keys//", + AgentKVKeyViewSet.as_view( + {"get": "retrieve", "patch": "partial_update", "delete": "destroy"} + ), + name="agent_kv_key_detail", + ), + path( + "keys//rotate/", + AgentKVKeyViewSet.as_view({"post": "rotate"}), + name="agent_kv_key_rotate", + ), +] diff --git a/backend/agent_kv/views.py b/backend/agent_kv/views.py new file mode 100644 index 0000000000..466a4c7eee --- /dev/null +++ b/backend/agent_kv/views.py @@ -0,0 +1,55 @@ +import uuid + +from rest_framework import status, viewsets +from rest_framework.decorators import action +from rest_framework.permissions import IsAuthenticated +from rest_framework.response import Response +from rest_framework.serializers import BaseSerializer + +from agent_kv.models import AgentKVKey +from agent_kv.permissions import IsOrganizationAdmin +from agent_kv.serializers import AgentKVKeySerializer, AgentKVKeyWriteSerializer + + +class AgentKVKeyViewSet(viewsets.ModelViewSet): + permission_classes = [IsAuthenticated, IsOrganizationAdmin] + + def get_queryset(self): + # No explicit org filter here — the global ``OrganizationFilterBackend`` + # (in DEFAULT_FILTER_BACKENDS) scopes every DRF operation by the + # current org's ``organization`` FK on this model. + return AgentKVKey.objects.all() + + def get_serializer_class(self): + if self.action in ("create", "update", "partial_update"): + return AgentKVKeyWriteSerializer + return AgentKVKeySerializer + + def perform_create(self, serializer: BaseSerializer) -> None: + serializer.save(created_by=self.request.user) + + def create(self, request, *args, **kwargs): + """Respond with the read serializer, not the write one (DRF's default). + + ``AgentKVKeyWriteSerializer`` (used to validate the request, per + ``get_serializer_class``) only carries ``name``/``description``/ + ``is_active`` -- the standard ``CreateModelMixin.create`` would echo + that same serializer back, meaning the caller who just created a key + would never see ``id`` or the raw ``key`` value itself. Both are + server-generated and this is the only response that will ever carry + the plaintext key (list/retrieve return it too today, but rotate is + the only other place a caller can *see* a fresh one) -- mirrors + ``GlobalApiDeploymentKeyViewSet.create``. + """ + serializer = self.get_serializer(data=request.data) + serializer.is_valid(raise_exception=True) + self.perform_create(serializer) + response_serializer = AgentKVKeySerializer(serializer.instance) + return Response(response_serializer.data, status=status.HTTP_201_CREATED) + + @action(detail=True, methods=["post"]) + def rotate(self, request, pk=None): + key_obj = self.get_object() + key_obj.key = uuid.uuid4() + key_obj.save(update_fields=["key", "modified_at"]) + return Response(AgentKVKeySerializer(key_obj).data) diff --git a/backend/backend/base_urls.py b/backend/backend/base_urls.py index 0b763e3c59..30d5d790da 100644 --- a/backend/backend/base_urls.py +++ b/backend/backend/base_urls.py @@ -37,6 +37,8 @@ f"{settings.API_DEPLOYMENT_PATH_PREFIX}/pipeline/", include("pipeline_v2.public_api_urls"), ), + # Agent-KV product API (views own auth; prefix is middleware-whitelisted) + path(f"{settings.AGENT_KV_PATH_PREFIX}/", include("agent_kv.execution_urls")), path("", include("health.urls")), # Internal API for worker communication path("internal/", include("backend.internal_base_urls")), diff --git a/backend/backend/internal_base_urls.py b/backend/backend/internal_base_urls.py index 30d4c8df22..ab184c8fb0 100644 --- a/backend/backend/internal_base_urls.py +++ b/backend/backend/internal_base_urls.py @@ -246,6 +246,12 @@ def test_middleware_debug(request): include("api_v2.internal_urls"), name="api_deployments_internal", ), + # Agent-KV stage-report and finalize APIs for the cloud executor + path( + "v1/agent-kv/", + include("agent_kv.internal_urls"), + name="agent_kv_internal", + ), # Platform configuration and settings APIs path( "v1/platform-settings/", diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index aa03b577ae..f90416969c 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -204,6 +204,48 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | os.environ.get("API_DEPLOYMENT_RATE_LIMIT_LOCK_BLOCKING_TIMEOUT", 5) ) +# Agent-KV API +AGENT_KV_PATH_PREFIX = os.environ.get("AGENT_KV_PATH_PREFIX", "agent-kv") +# Bucket-rooted object-store root for staged inputs/results, mirroring +# WORKFLOW_EXECUTION_DIR_PREFIX / API_EXECUTION_DIR_PREFIX: the FIRST segment is +# the bucket (s3fs/gcsfs treat it that way), so it must already exist. MUST match +# the cloud executor's AGENT_KV_STORAGE_DIR_PREFIX -- the executor keys its OCR +# cache under the same root. +AGENT_KV_STORAGE_DIR_PREFIX = ( + os.environ.get("AGENT_KV_STORAGE_DIR_PREFIX", "unstract/agent_kv").strip().strip("/") + or "unstract/agent_kv" +) # same normalisation as the cloud executor's reader (whitespace, edge slashes) +AGENT_KV_MAX_FILE_SIZE_MB = int(os.environ.get("AGENT_KV_MAX_FILE_SIZE_MB", 50)) +AGENT_KV_MAX_PAGES = int(os.environ.get("AGENT_KV_MAX_PAGES", 100)) +AGENT_KV_MAX_CALCULATIONS_BYTES = int( + os.environ.get("AGENT_KV_MAX_CALCULATIONS_BYTES", 20_000) +) +AGENT_KV_MAX_SCHEMA_BYTES = int(os.environ.get("AGENT_KV_MAX_SCHEMA_BYTES", 262_144)) +AGENT_KV_RESULT_TTL_DAYS = int(os.environ.get("AGENT_KV_RESULT_TTL_DAYS", 7)) +AGENT_KV_MAX_TIMEOUT_SECONDS = int(os.environ.get("AGENT_KV_MAX_TIMEOUT_SECONDS", 300)) +AGENT_KV_CONCURRENT_LIMIT = int(os.environ.get("AGENT_KV_CONCURRENT_LIMIT", 5)) +AGENT_KV_KEY_RATE_LIMIT_PER_MINUTE = int( + os.environ.get("AGENT_KV_KEY_RATE_LIMIT_PER_MINUTE", 60) +) +# When Redis is unreachable, do requests proceed? Default False = fail CLOSED +# (429). Both Agent-KV limiters used to fail open implicitly, inside an `except` +# block, so a Sentinel failover removed the concurrency ceiling and the per-key +# rate ceiling at once while the API kept accepting billable LLM work. Fail-open +# remains available as a deliberate, visible operational choice. +AGENT_KV_LIMITER_FAIL_OPEN = ( + os.environ.get("AGENT_KV_LIMITER_FAIL_OPEN", "false").lower() == "true" +) +AGENT_KV_SWEEP_GRACE_SECONDS = int(os.environ.get("AGENT_KV_SWEEP_GRACE_SECONDS", 3600)) +AGENT_KV_STUCK_JOB_GRACE_SECONDS = int( + os.environ.get("AGENT_KV_STUCK_JOB_GRACE_SECONDS", 21600) +) +AGENT_KV_CALCULATIONS_ENABLED = ( + os.environ.get("AGENT_KV_CALCULATIONS_ENABLED", "false").lower() == "true" +) +AGENT_KV_STRUCTURED_OUTPUT_ENABLED = ( + os.environ.get("AGENT_KV_STRUCTURED_OUTPUT_ENABLED", "false").lower() == "true" +) + DB_NAME = os.environ.get("DB_NAME", "unstract_db") DB_USER = os.environ.get("DB_USER", "unstract_dev") @@ -454,6 +496,7 @@ def filter(self, record): "dashboard_metrics", "platform_api", "global_api_deployment_key", + "agent_kv", ) TENANT_APPS = [] @@ -978,10 +1021,16 @@ def filter(self, record): "/static", ] WHITELISTED_PATHS = [f"/{PATH_PREFIX}{PATH}" for PATH in WHITELISTED_PATHS_LIST] -# White lists workflow-api-deployment path. This also covers the deployment MCP -# server, which hangs off the same URL and authenticates with the deployment's -# own API key rather than a session. -WHITELISTED_PATHS.append(f"/{API_DEPLOYMENT_PATH_PREFIX}") +WHITELISTED_PATHS.extend( + [ + # workflow-api-deployment. Also covers the deployment MCP server, which + # hangs off the same URL and authenticates with the deployment's own API + # key rather than a session. + f"/{API_DEPLOYMENT_PATH_PREFIX}", + # Agent-KV public API: bearer-key auth happens in the views. + f"/{AGENT_KV_PATH_PREFIX}", + ] +) # Whitelisting health check API WHITELISTED_PATHS.append("/health") diff --git a/backend/backend/urls_v2.py b/backend/backend/urls_v2.py index 4b9636a35c..6e66e9561a 100644 --- a/backend/backend/urls_v2.py +++ b/backend/backend/urls_v2.py @@ -32,6 +32,7 @@ path("", include("feature_flag.urls")), path("workflow/", include("workflow_manager.urls")), path("platform/", include("platform_settings_v2.urls")), + path("agent-kv/", include("agent_kv.urls")), path("api/", include("api_v2.urls")), path("usage/", include("usage_v2.urls")), path("notifications/", include("notification_v2.urls")), diff --git a/backend/pyproject.toml b/backend/pyproject.toml index c669d77c00..f99750f4d0 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -42,6 +42,7 @@ dependencies = [ "django-filter>=24.3", "httpx>=0.27.0", # Hence required to add all indirect local dependencies too here. + "unstract-agent-kv-schema", "unstract-connectors", "unstract-core", "unstract-filesystem", @@ -90,6 +91,7 @@ deploy = [ ] [tool.uv.sources] +unstract-agent-kv-schema = { path = "../unstract/agent-kv-schema", editable = true } unstract-connectors = { path = "../unstract/connectors", editable = true } unstract-core = { path = "../unstract/core", editable = true } unstract-filesystem = { path = "../unstract/filesystem", editable = true } diff --git a/backend/sample.env b/backend/sample.env index ca5456167d..78b526f1d1 100644 --- a/backend/sample.env +++ b/backend/sample.env @@ -186,6 +186,34 @@ API_DEPLOYMENT_RATE_LIMIT_LOCK_TIMEOUT=2 # Redis lock blocking timeout (in seconds) - how long to wait to acquire lock API_DEPLOYMENT_RATE_LIMIT_LOCK_BLOCKING_TIMEOUT=5 +# Agent-KV API (backend/backend/settings/base.py) — see docs/agent-kv-api.md +# --------------------------------------------------------------------------- +AGENT_KV_PATH_PREFIX=agent-kv # Public URL prefix (whitelisted past tenant middleware) +AGENT_KV_MAX_FILE_SIZE_MB=50 # Submit-time file size cap (MB) +AGENT_KV_MAX_PAGES=100 # Pre-OCR page cap for PDFs/images; also the post-OCR cap the engine enforces +AGENT_KV_MAX_CALCULATIONS_BYTES=20000 # Byte cap on the optional `calculations` field +AGENT_KV_MAX_SCHEMA_BYTES=262144 # Byte cap on the raw `keys` JSON schema (256 KiB) +AGENT_KV_RESULT_TTL_DAYS=7 # Result/input retention window, stamped at submit time +AGENT_KV_MAX_TIMEOUT_SECONDS=300 # Upper bound on the submit `timeout` (sync-wait) field +AGENT_KV_CONCURRENT_LIMIT=5 # Per-org concurrent in-flight job cap +AGENT_KV_KEY_RATE_LIMIT_PER_MINUTE=60 # Per-key request rate limit (submit + validate) +AGENT_KV_LIMITER_FAIL_OPEN=false # true = both limiters above allow the request when Redis is down (default: refuse with 429) +AGENT_KV_SWEEP_GRACE_SECONDS=3600 # Age before a never-dispatched PENDING job is swept +AGENT_KV_STUCK_JOB_GRACE_SECONDS=21600 # Age (from dispatch) before a stuck DISPATCHED/RUNNING job is force-failed (6h) +AGENT_KV_CALCULATIONS_ENABLED=false # Gates the submit `calculations` field; engine can't run it yet +AGENT_KV_STRUCTURED_OUTPUT_ENABLED=false # Gates the submit `structured_output` field; engine can't run it yet + +# Bucket-rooted object-store root for staged inputs/results +# ({AGENT_KV_STORAGE_DIR_PREFIX}/{org_id}/{job_id}/…), like WORKFLOW_EXECUTION_DIR_PREFIX. +# The FIRST segment is the BUCKET (s3fs/gcsfs treat it so) and must already exist — +# the MinIO dev bootstrap creates `unstract`. MUST match the cloud executor's +# AGENT_KV_STORAGE_DIR_PREFIX: the executor keys its OCR cache under the same root. +AGENT_KV_STORAGE_DIR_PREFIX="unstract/agent_kv" + +# Object storage credentials for those files — same shape as +# WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS; required, no default. +AGENT_KV_FILE_STORAGE_CREDENTIALS='{"provider": "minio", "credentials": {"endpoint_url": "http://unstract-minio:9000", "key": "minio", "secret": "minio123"}}' + # Default user auth credentials DEFAULT_AUTH_USERNAME= DEFAULT_AUTH_PASSWORD= diff --git a/backend/uv.lock b/backend/uv.lock index e13d83abfe..be8e962aa4 100644 --- a/backend/uv.lock +++ b/backend/uv.lock @@ -3698,6 +3698,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8f/13/b688a905653871b10b4ff0403c2ff562c17a0bd50be0d44324f3c85ca48f/ujson-5.12.1-graalpy312-graalpy250_312_native-win_amd64.whl", hash = "sha256:4ea0c490c702c20495e97345acfcf0c2f3153e658ef537ff111929c48b89e10a", size = 45988, upload-time = "2026-05-05T22:04:43.36Z" }, ] +[[package]] +name = "unstract-agent-kv-schema" +version = "0.1.0" +source = { editable = "../unstract/agent-kv-schema" } + [[package]] name = "unstract-backend" version = "0.0.1" @@ -3731,6 +3736,7 @@ dependencies = [ { name = "social-auth-app-django" }, { name = "social-auth-core" }, { name = "stripe" }, + { name = "unstract-agent-kv-schema" }, { name = "unstract-connectors" }, { name = "unstract-core" }, { name = "unstract-filesystem" }, @@ -3798,6 +3804,7 @@ requires-dist = [ { name = "social-auth-app-django", specifier = "==5.3.0" }, { name = "social-auth-core", specifier = "==4.4.2" }, { name = "stripe", specifier = ">=12.5.1" }, + { name = "unstract-agent-kv-schema", editable = "../unstract/agent-kv-schema" }, { name = "unstract-connectors", editable = "../unstract/connectors" }, { name = "unstract-core", editable = "../unstract/core" }, { name = "unstract-filesystem", editable = "../unstract/filesystem" }, diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index e89ca0e19d..eedfdf7a39 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -71,9 +71,10 @@ services: - PG_SCHEDULER_ADOPT_PERIODICS=true labels: - traefik.enable=true - - traefik.http.routers.backend.rule=Host(`frontend.unstract.localhost`) && (PathPrefix(`/api/v1`) || PathPrefix(`/deployment`) || PathPrefix(`/public`)) + - traefik.http.routers.backend.rule=Host(`frontend.unstract.localhost`) && (PathPrefix(`/api/v1`) || PathPrefix(`/deployment`) || PathPrefix(`/public`) || PathPrefix(`/agent-kv`)) - traefik.http.services.backend.loadbalancer.server.port=8000 + # Celery Flower celery-flower: <<: *host_gateway image: unstract/backend:${VERSION} @@ -112,7 +113,7 @@ services: - UNSTRACT_APPS_VERSION=${VERSION} labels: - traefik.enable=true - - traefik.http.routers.frontend.rule=Host(`frontend.unstract.localhost`) && !PathPrefix(`/api/v1`) && !PathPrefix(`/deployment`) && !PathPrefix(`/public`) + - traefik.http.routers.frontend.rule=Host(`frontend.unstract.localhost`) && !PathPrefix(`/api/v1`) && !PathPrefix(`/deployment`) && !PathPrefix(`/public`) && !PathPrefix(`/agent-kv`) - traefik.http.services.frontend.loadbalancer.server.port=80 platform-service: @@ -543,7 +544,13 @@ services: - APPLICATION_NAME=unstract-worker-pg-ide-callback - WORKER_BARRIER_BACKEND=pg - WORKER_PG_QUEUE_CONSUMER_WORKER_TYPE=ide_callback - - WORKER_PG_QUEUE_CONSUMER_QUEUE=ide_callback + # agent_kv_callback carries the Agent-KV terminal callbacks (finalize: + # persist the result, delete the staged input, release the concurrency + # slot, fire the webhook). Without it a job runs to success and then + # sits in RUNNING forever -- the executor's continuation is enqueued + # and nothing drains it. The Celery ide_callback worker has carried + # both queues since the callbacks landed; its PG twin did not. + - WORKER_PG_QUEUE_CONSUMER_QUEUE=ide_callback,agent_kv_callback - WORKER_PG_QUEUE_CONSUMER_HEALTH_PORT=8090 - WORKER_PG_QUEUE_CONSUMER_CONCURRENCY=${PG_IDE_CALLBACK_CONCURRENCY:-2} # Log publisher — see the note on the backend service (UN-4046). diff --git a/docker/sample.env b/docker/sample.env index 47905773da..de64df726f 100644 --- a/docker/sample.env +++ b/docker/sample.env @@ -99,6 +99,56 @@ SCHEDULER_TASK_TIME_LIMIT=1800 # Scheduler timeout (30 minutes) LOG_CONSUMER_AUTOSCALE=2,1 # Max 2, min 1 workers LOG_CONSUMER_TASK_TIME_LIMIT=600 # Log processing timeout (10 minutes) +# IDE Callback Worker (worker-ide-callback) — Agent-KV terminal callbacks +# --------------------------------------------------------------------------- +# The ide_callback worker also owns agent_kv_complete/agent_kv_error +# (Agent-KV job finalize callbacks, spec §5.3) on their own queue, +# `agent_kv_callback`, alongside its existing `ide_callback` queue. +# unstract/backend's agent_kv.dispatch dispatches those tasks straight to +# that queue name, so worker-ide-callback must consume both queues or +# Agent-KV jobs will dispatch but never finalize. +# +# NOTE: unlike CALLBACK/FILE_PROCESSING/etc. above, this worker's queue list +# is NOT overridable from this file — there is no CELERY_QUEUES_IDE_CALLBACK +# var wired into workers/run-worker-docker.sh's queue resolution. The queue +# list is a hardcoded default in that script (WORKER_QUEUES["ide_callback"]), +# already set to "ide_callback,agent_kv_callback" there. Nothing to set here; +# if you ever need a different set for this worker, edit that script's +# default directly (or run-worker.sh's equivalent for non-docker dev runs) — +# there's no env-var lever for it today. + +# Agent-KV API (backend/backend/settings/base.py) — see docs/agent-kv-api.md +# --------------------------------------------------------------------------- +# Django app settings for the Agent-KV extraction API (unstract/backend's +# `agent_kv` app). Consumed by the Django backend process, not by +# docker-compose itself — mirror these into backend/.env (backend/sample.env) +# for a real deployment; listed here per the Agent-KV env-sample convention. +AGENT_KV_PATH_PREFIX=agent-kv # Public URL prefix (whitelisted past tenant middleware) +AGENT_KV_MAX_FILE_SIZE_MB=50 # Submit-time file size cap (MB) +AGENT_KV_MAX_PAGES=100 # Pre-OCR page cap for PDFs/images +AGENT_KV_MAX_CALCULATIONS_BYTES=20000 # Byte cap on the optional `calculations` field +AGENT_KV_MAX_SCHEMA_BYTES=262144 # Byte cap on the raw `keys` JSON schema (256 KiB) +AGENT_KV_RESULT_TTL_DAYS=7 # Result/input retention window, stamped at submit time +AGENT_KV_MAX_TIMEOUT_SECONDS=300 # Upper bound on the submit `timeout` (sync-wait) field +AGENT_KV_CONCURRENT_LIMIT=5 # Per-org concurrent in-flight job cap +AGENT_KV_KEY_RATE_LIMIT_PER_MINUTE=60 # Per-key request rate limit (submit + validate) +AGENT_KV_LIMITER_FAIL_OPEN=false # true = both limiters above allow the request when Redis is down (default: refuse with 429) +AGENT_KV_SWEEP_GRACE_SECONDS=3600 # Age before a never-dispatched PENDING job is swept +AGENT_KV_STUCK_JOB_GRACE_SECONDS=21600 # Age (from dispatch) before a stuck DISPATCHED/RUNNING job is force-failed (6h) +AGENT_KV_CALCULATIONS_ENABLED=false # Gates the submit `calculations` field; engine can't run it yet +AGENT_KV_STRUCTURED_OUTPUT_ENABLED=false # Gates the submit `structured_output` field; engine can't run it yet + +# Bucket-rooted object-store root for staged inputs/results +# ({AGENT_KV_STORAGE_DIR_PREFIX}/{org_id}/{job_id}/…), like WORKFLOW_EXECUTION_DIR_PREFIX. +# The FIRST segment is the BUCKET (s3fs/gcsfs treat it so) and must already exist — +# the MinIO dev bootstrap creates `unstract`. MUST match the cloud executor's +# AGENT_KV_STORAGE_DIR_PREFIX: the executor keys its OCR cache under the same root. +AGENT_KV_STORAGE_DIR_PREFIX="unstract/agent_kv" + +# Object storage credentials for those files — same shape as +# WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS; required, no default. +AGENT_KV_FILE_STORAGE_CREDENTIALS='{"provider": "minio", "credentials": {"endpoint_url": "http://unstract-minio:9000", "key": "minio", "secret": "minio123"}}' + # Worker Circuit Breaker Settings CIRCUIT_BREAKER_FAILURE_THRESHOLD=5 CIRCUIT_BREAKER_RECOVERY_TIMEOUT=60 diff --git a/docs/agent-kv-api.md b/docs/agent-kv-api.md new file mode 100644 index 0000000000..316cd85a1b --- /dev/null +++ b/docs/agent-kv-api.md @@ -0,0 +1,1143 @@ +# Agent-KV API + +Agent-KV is Unstract's agentic key-value extraction product: upload a document plus a +`keys.json`-style extraction schema, get back a job id, poll a verbose agent-centric +status document, and fetch a rich result (record + per-field QA/challenger audit +trail). + +This document describes the API **as built in this repo** — the OSS scaffold (routes, +validation, dispatch, job lifecycle, key management). The document, `keys.json` schema, +and OCR/LLM extraction pipeline itself are described only insofar as they shape the +contract; the extraction engine ships as a cloud executor plugin (see +[§6, 501 behavior](#10-when-the-engine-is-unavailable-501-behavior) below). + +Design background: `docs/superpowers/specs/2026-08-28-agent-kv-api-design.md` (referred +to as "the spec" throughout this document; section numbers below are spec section +numbers). This document records what shipped, including a handful of places where the +implementation is more specific than — or diverges from — the spec's sketch. Those are +called out explicitly rather than silently reconciled. + +## 1. Authentication + +There are two independent auth schemes on two different URL prefixes: + +| Surface | Prefix | Auth | +|---|---|---| +| Public execution API (submit/status/result/cancel) | `/agent-kv/…` (top-level, not under the tenant subfolder) | `Authorization: Bearer ` — a dedicated per-organization key, validated by the view itself | +| Key management (create/list/rotate/revoke) | `{PATH_PREFIX}/unstract/…/agent-kv/keys/…` (tenant-scoped, alongside Platform API key management) | Session auth, `IsOrganizationAdmin` only | + +The execution prefix (`AGENT_KV_PATH_PREFIX`, default `agent-kv`) sits **outside** the +tenant-subfolder/organization-middleware chain — it is added to `WHITELISTED_PATHS` in +`backend/backend/settings/base.py` — because a caller authenticates with an API key, not +a session/org header. Every view under it therefore owns its own auth via the +`AgentKVKeyValidator.validate_api_key` decorator (`backend/agent_kv/key_validator.py`), +which: + +1. Requires `Authorization: Bearer ` where `` is a UUID. +2. Looks up an active (`is_active=True`) `AgentKVKey` by that UUID. +3. Injects the resolved key object as `agent_kv_key` into the view kwargs, which every + subsequent lookup filters by (`organization_id=agent_kv_key.organization_id`) — this + is what makes a job in another org indistinguishable from an unknown job id (both + 404, see [Status](#4-status--get-agent-kvjob_id)). + +**Drift from the spec:** spec §6.1 describes this as "every endpoint 401s without a +valid key." The implementation returns **403** (`api_v2.exceptions.Forbidden`, reused +as-is rather than adding a new exception class) for both a missing/malformed +`Authorization` header and an unknown/inactive key. This is exercised directly by +`backend/agent_kv/tests/test_auth.py` and `test_job_views.py::test_all_job_views_401_without_key` +(the test name predates the code's actual status; both assert `403`). Document the code: +**403, not 401**, for any auth failure on the public prefix. + +Key management (`backend/agent_kv/views.py`, `AgentKVKeyViewSet`) is a standard DRF +`ModelViewSet` scoped by the global `OrganizationFilterBackend` and gated by +`IsOrganizationAdmin` (`backend/agent_kv/permissions.py`). Fields returned: +`id, name, description, key, is_active, created_at`; write-side fields are +`name, description, is_active` only (`key` and `id` are server-generated). Routes +(`backend/agent_kv/urls.py`, mounted at `agent-kv/` under +`backend/backend/urls_v2.py`): + +``` +GET/POST .../agent-kv/keys/ +GET/PATCH/DELETE .../agent-kv/keys// +POST .../agent-kv/keys//rotate/ # replaces `key` with a new UUID +``` + +Key storage follows the same conventions as Platform API key management (D5) — +plaintext-at-rest UUID keys, not hashed-at-rest; that hardening is explicitly deferred, +not a v1 gap unique to this feature. + +## 2. The extraction schema (`keys` field) + +> **This section describes the `kv` extractor's schema language, which is not +> routable on this deployment** (see [§3](#3-submit--post-agent-kv-multipart)). +> It is documented here because the compiler ships and the language is frozen, +> so an integration can be written against it ahead of the extractor's release. +> The `table` extractor's `keys` is just `{"target_table": ""}` and needs +> none of this. + +The `keys` submit field — spec-called "keys.json" — is a JSON object describing what to +extract. It compiles through `unstract.agent_kv_schema.compile_schema` +(`unstract/agent-kv-schema/src/unstract/agent_kv_schema/`), the **single schema compiler +in the system**: the OSS API validates with it at submit time, and +the cloud engine imports the same package to execute — the schema language cannot drift +between validation and execution by construction (spec §5.1). + +### Node types + +A schema node is one of three shapes, decided structurally: + +- **Interior node** — every value is itself an object → its keys are child nodes, and + nesting continues (dotted path, e.g. `vendor.address.city`). +- **Leaf node** — no value is an object → defined by scalar attributes: + + | Attribute | Required | Meaning | + |---|---|---| + | `description` | **yes** | Prompted to the extractor; also fed to the breadcrumb-prefixed `effective_description`. | + | `format` | no (default `string`) | `string` \| `number` \| `currency` \| `date` \| `enum:v1,v2,…` \| `regex:` \| any other free-text hint. `number`/`currency` parse via a strip-then-`float()` check; `date` accepts a fixed set of common formats; `enum` matches case-insensitively; `regex` is `re.fullmatch`. | + | `required` | no (default `false`) | Completeness is checked separately from format. | + | `aliases` | no (default `[]`) | Alternate names/labels to help the extractor recognize the field. | + | `multivalued` | no (default `false`) | Value is a comma-separated string; each element is format-checked independently. | + + A node mixing object values and scalar attributes is a **compile-time error** + ("mixes object children and scalar attributes"), as is a leaf with unknown attribute + keys or a missing `description`. + +- **Array node** — has an `_array` key → a flat (non-nested) row schema: + + ```json + "line_items": { + "description": "One row per invoice line", + "_key": "sku", + "_array": { + "sku": {"description": "SKU"}, + "total": {"description": "Line total", "format": "currency"} + } + } + ``` + + `_array`'s value is itself a flat object of leaf columns (row-local paths, not + dotted). `description` is optional context; `_key` (optional) names a column used for + row identity in downstream scoring, otherwise rows are matched positionally. A + nested array inside an array's row schema is rejected at compile time ("not supported + in P8a") — arrays are single-level only in v1. + + `_dedup` (optional, default `true`) controls whether rows identical in **every** + extracted column are collapsed to one. Leave it on for documents where + layout-preserving OCR repeats a row once per replicate column — that is what it + exists for. Set it to `false` when your documents can genuinely contain two + identical rows: + + ```json + "line_items": { + "_dedup": false, + "_array": { + "description": {"description": "Line description"}, + "total": {"description": "Line total", "format": "currency"} + } + } + ``` + + With `_dedup: true` an invoice listing the same item twice comes back with one + row, so the row count — and any `_constraints` aggregate over those rows — is + wrong. The default is `true` only because flipping it would break the OCR corpus + the extractor was built against; it is not the safer value for every document. + Must be a JSON boolean: `"false"` as a string is rejected rather than silently + read as true. + +### Constraints (`_constraints`) + +An optional top-level `_constraints` key: a list of string expressions checked against +the extracted record post-hoc. Each expression is parsed and statically validated +against an AST allowlist before it is ever accepted — no arbitrary code execution: +boolean/comparison/arithmetic operators, plus exactly one string-literal argument to one +of `sum`, `count`, `min`, `max`, `avg` (an array path). Example: +`"count('line_items') >= 1"`. + +### Structural caps (spec §6.1) + +Enforced by `compile_schema` regardless of the byte-size cap below (`SchemaCaps` in +`compile.py`): + +| Cap | Default | +|---|---| +| Max leaves | 200 | +| Max arrays | 20 | +| Max columns per array | 40 | +| Max nesting depth | 6 | +| Max regex pattern length | 200 | +| Max aliases per field | 10 | +| Max description length | 500 | +| Max constraints | 30 | + +These caps are hardcoded in the compiler (not env-configurable); what **is** +env-configurable is the raw byte size of the whole `keys` document +(`AGENT_KV_MAX_SCHEMA_BYTES`) and of `calculations` +(`AGENT_KV_MAX_CALCULATIONS_BYTES`) — see [§13 env reference](#12-environment-reference). + +A schema that fails to compile is rejected with a `400` at submit time — **nothing is +billed**. + +## 3. Submit — `POST /agent-kv/` (multipart) + +Copied verbatim from spec §7.1 and cross-checked field-by-field against +`backend/agent_kv/execution_serializers.py::SubmitSerializer` — every row below matches +the shipped serializer exactly; no drift found in the submit contract itself (see the +notes after the table for two spec-vs-code nuances worth flagging). + +The request is **extractor-scoped**: a per-extractor schema and its knobs live inside +`extractors`, and only fields describing the request itself stay top level. + +**Job-level fields:** + +| Field | Required | Notes | +|---|---|---| +| `file` | yes | One document per job (D12). | +| `extractors` | yes | JSON array (file part or inline JSON string), one entry per extractor. Compiled synchronously; invalid ⇒ 400, nothing billed. | +| `page_start`, `page_end` | no | 1-based inclusive range. **Job-level, not per-extractor**: the range drives the shared OCR pass and the page cap, so it cannot differ between extractors reading the same document. | +| `timeout` | no | Omitted or `0` = pure async (immediate 202). `1–300`: the view polls the job row up to the deadline and returns the result inline if the job completes, else 202 with the job id. | +| `tags`, `custom_data` | no | Echoed through. | +| `webhook_url` | no | Terminal-state POST `{job_id, status}` only — no result payload. | + +**Each `extractors` entry:** + +| Field | Required | Notes | +|---|---|---| +| `name` | yes | Which extractor. **`table` is the only supported value on this deployment**; anything else, `kv` included, ⇒ 400 (see the note below). More than one entry ⇒ 400 — the format is in place for multi-extractor jobs, the execution is not built yet. | +| `keys` | yes | This extractor's own request. For `table`, `{"target_table": ""}` — the engine's one required extraction parameter. | +| `adapters` | yes (for `table`) | The platform adapter instances this extraction runs on: `{"llm": "", "lite_llm": "", "x2text": ""}`. All three required. See **§3a** below. | +| `options` | no | This extractor's own knobs (below). An unrecognised option ⇒ 400, so a knob aimed at the wrong extractor fails loudly instead of being silently dropped. | + +### 3a. `adapters` — you supply the models and the OCR credentials + +The `table` extractor runs on **your own platform adapters**, named by id. Create +them once in Unstract (Settings → Adapters, or the adapter API) and pass their ids +on every submit: + +| Role | Adapter type | What it does | +|---|---|---| +| `llm` | `LLM` | Structure detection and extraction — the quality-critical model. | +| `lite_llm` | `LLM` | Per-page table-presence detection. Runs once per page, so this is the cost-critical one; a cheap fast model belongs here. | +| `x2text` | `X2TEXT` | OCR. An LLMWhisperer adapter, which also carries the endpoint — there is no separate base-URL field. | + +Each id is validated at submit, against **your organization** and against the +expected adapter type. A ⇒ 400, before anything is staged or billed: + +- an id that is not your organization's — reported identically to one that does not + exist, deliberately, so the endpoint cannot be used to discover which adapter ids + are real elsewhere +- an id of the wrong type, e.g. an `X2TEXT` adapter in the `llm` slot +- a missing role, or an unknown one + +Page usage and token usage are metered against the adapters you named, so **model +choice and LLM spend sit on your account**, per job. + +> **This is a deliberate reversal of the original design (spec D6).** D6 specified +> that this API would have "no platform adapters by design", with models and OCR +> credentials configured by the operator in environment variables and "no end-user +> model control". That was the right call while the API's consumer was hypothetical; +> it is the wrong one now that the consumers are existing Unstract customers. They +> already own adapters, there was no operator-level LLM credential anywhere in the +> platform to reuse — LLM credentials live in per-organization adapter rows, +> encrypted — and env configuration would have put every customer's model choice and +> every customer's LLM spend on a single shared operator key. +> +> Consequences worth knowing: the API is no longer usable by a caller who has never +> signed in, since adapters must exist first; and the six `AGENT_KV_LLM_*` / +> `AGENT_KV_LLMWHISPERER_*` environment variables are **not read on the table path** +> (they remain in place for the `kv` extractor, which is still env-configured). +> +> One thing this does not honour: adapters carry per-user visibility within an +> organization, and an Agent-KV key resolves to an organization rather than a user. +> So a key may use **any** adapter in its own organization. A key is already an +> organization-level credential, so this is a deliberate decision rather than an +> oversight — but treat key issuance accordingly. + +> **Why `kv` returns 400 here.** This deployment ships the table extractor only. +> The `kv` extractor's engine runs on its own release track and nothing on this +> deployment consumes its queue, so the API refuses a `kv` submit outright +> rather than accepting it with a `202` for a job that would never run. The +> wire format is unchanged: when `kv` ships, it becomes a second accepted value +> of `name` and nothing else about this document changes. + +**`kv` options:** + +| Option | Notes | +|---|---| +| `qa` | Default **on**. | +| `challenge` | Default **on** (~doubles LLM spend; meter records what ran). | +| `extraction_mode` | `whole-doc` (default) \| `per-page`. | +| `structured_output` | Disabled by default; enabled per deployment via `AGENT_KV_STRUCTURED_OUTPUT_ENABLED`. | +| `calculations` | Post-processing instructions; opt-in codegen. Disabled by default; enabled per deployment via `AGENT_KV_CALCULATIONS_ENABLED`. | +| `document_class` | Free-text hint (as CLI). | +| `key_notes` | Free-text notes appended to the prompt (as CLI). | + +**`table` options:** + +| Option | Notes | +|---|---| +| `instructions` | Free-text extraction guidance handed to the engine. | +| `json_structure` | Free-text description of the row structure to return. | +| `enable_header_mapping` | Default **off**. **Changes the result shape** — see [§5](#5-result--get-agent-kvjob_idresult). | +| `correct_number_separators` | Default **off**. Re-reads thousands/decimal separators in numeric cells. | +| `number_format` | `US` (default) \| `EU`. Which convention `correct_number_separators` assumes. | + +Not exposed for `table` (they are meaningless on a blind API): `output_path`, +`enable_highlight`, and the IDE callback hints `prompt_key`/`doc_name`. + +> **No backward compatibility with the pre-`extractors` flat shape.** `keys`, `qa`, +> `challenge`, `extraction_mode`, `structured_output`, `calculations`, +> `document_class` and `key_notes` are no longer accepted at top level — a request +> using the old form gets a 400 for the missing `extractors` field. The API had no +> consumers when this changed, so no alias was kept. + +Not exposed (D6): model choice, challenger model, `parallel_pages`, thinking budgets. + +**Two implementation-level notes not visible in the table:** + +- File type is allowlisted (`.pdf .xlsx .xls`) and size-capped + (`AGENT_KV_MAX_FILE_SIZE_MB`); PDFs are page-counted **locally, pre-OCR**, via + `pdfplumber` and rejected over `AGENT_KV_MAX_PAGES` before any paid work runs. Excel + has no pre-OCR page concept — it is capped by file size only at submit time (the + engine enforces a post-OCR virtual-page cap). + - **The cap bounds the pages this request will PROCESS, not the document's + length.** It is checked against `pages_selected` — the count after + `page_start` / `page_end` are applied — not against `pages_total`. So pages + 1–50 of a 400-page PDF is accepted under a 100-page cap, while the same + document submitted without a page range is refused. That is the intended + reading: the cap exists to bound paid work per job, and a page range is the + caller declaring how much work they are asking for. `pages_total` is still + recorded on the job and returned in the status document, so the document's + real length is never lost. + - Images (`.png .jpg .jpeg .tiff`) are **not** accepted, and no longer count as + 1 page. The engine treats only `.pdf/.xlsx/.xls` as a document and + `ImageLoader.load_pages` has no call site, so an accepted image returned + `success: true` with every key not-found and a page billed. They are refused at + submit until the engine side is wired. +- `timeout`'s upper bound is not the literal `300` in the spec prose — it's + `AGENT_KV_MAX_TIMEOUT_SECONDS` (default `300`, but deployment-configurable). + +### Response + +`202` (async, the default): + +```json +{ + "job_id": "5b6e9b0a-...", + "status": "dispatched", + "status_url": "/agent-kv/5b6e9b0a-...", + "created_at": "2026-08-28T10:15:00.123456+00:00" +} +``` + +`200` (when `timeout` was set and the job finished before the deadline): the full +[result payload](#5-result--get-agent-kvjob_idresult), same shape as the result +endpoint. + +`501`: [engine unavailable](#10-when-the-engine-is-unavailable-501-behavior). +`429`: rate-limited (per-key request rate, or per-org concurrent-job limit — see +[§6](#12-environment-reference)). +`402`: the organization's subscription does not permit the request — its trial has +expired, or the subscription is inactive. Nothing is dispatched and nothing is +billed. This is the same gate, the same policy and the same response bodies an +API deployment execute gets from the cloud subscription middleware; Agent-KV +applies it in the submit view instead, because its URL carries no org segment for +that middleware to resolve an organization from. Cloud-only: an OSS-only +deployment returns `501` before ever reaching it. +`400`: serializer/schema validation failure — nothing billed. + +### Example — `table` (the only extractor this deployment accepts) + +```bash +curl -X POST https://api.unstract.example/agent-kv/ \ + -H "Authorization: Bearer 5c9e2c9e-1234-4a5b-9c6d-abcdef012345" \ + -F "file=@invoice.pdf" \ + -F 'extractors=[{ + "name": "table", + "keys": {"target_table": "Line items"}, + "adapters": { + "llm": "7f3c1a90-0000-4000-8000-000000000001", + "lite_llm": "7f3c1a90-0000-4000-8000-000000000002", + "x2text": "7f3c1a90-0000-4000-8000-000000000003" + }, + "options": { + "instructions": "Ignore the summary block at the foot of the page", + "enable_header_mapping": true + } + }]' \ + -F "page_start=1" -F "page_end=5" \ + -F "webhook_url=https://example.com/hooks/agent-kv" +``` + +The three adapter ids are yours, from your own organization — see §3a. Omit any of +them and the submit is a 400. + +### Example — `kv` + +Shown for the wire format only; **`kv` returns 400 on this deployment** (see the +note above). Note it takes no `adapters` block: the `kv` extractor is +env-configured on the operator side. + +```bash +curl -X POST https://api.unstract.example/agent-kv/ \ + -H "Authorization: Bearer 5c9e2c9e-1234-4a5b-9c6d-abcdef012345" \ + -F "file=@invoice.pdf" \ + -F 'extractors=[{ + "name": "kv", + "keys": { + "quotation_number": {"description": "The quote number", "required": true}, + "customer": {"name": {"description": "Bill-to name"}}, + "line_items": { + "description": "One row per line", + "_key": "sku", + "_array": { + "sku": {"description": "SKU"}, + "total": {"description": "Line total", "format": "currency"} + } + }, + "_constraints": ["count(\"line_items\") >= 1"] + }, + "options": {"challenge": false} + }]' \ + -F "webhook_url=https://example.com/hooks/agent-kv" +``` + +## 4. Status — `GET /agent-kv/{job_id}` + +Verbose, agent-centric, stage-level (`backend/agent_kv/execution_views.py::_status_document`). +**Stage names are per-extractor.** They are wire format, and `qa`/`challenge`/`codegen` +mean nothing to the table extractor, so a job's recorded stages are filtered through the +list for the extractor that actually ran +(`backend/agent_kv/constants.py::STAGE_NAMES_BY_EXTRACTOR`). + +- `kv` (superset; only stages that actually ran for this job appear, in a fixed + order): `document_processing, extraction, qa, challenge, normalize, constraints, + codegen, code_execution`. +- `table`: `table_extraction` — a single stage. The table engine exposes no node-level + progress hooks, so finer stage names would describe progress the executor cannot + actually report. It re-reports `table_extraction` as `running` while the engine works, + carrying a `steps` counter that increases as the engine passes its own progress points. + +```json +{ + "job_id": "5b6e9b0a-...", + "status": "running", + "extractors": { + "kv": { + "stage": "challenge", + "stages": [ + {"name": "document_processing", "status": "done", "seconds": 6.2}, + {"name": "extraction", "status": "done", "seconds": 11.4}, + {"name": "qa", "status": "done", "seconds": 4.1}, + {"name": "challenge", "status": "running"} + ] + } + }, + "adapters": { + "llm": "3f2b1c90-...", + "lite_llm": "a71e4d55-...", + "x2text": "c90aa812-..." + }, + "created_at": "2026-08-28T10:15:00.123456+00:00", + "started_at": "2026-08-28T10:15:01.500000+00:00", + "completed_at": null, + "pages_total": 3 +} +``` + +Each stage entry always carries `status` (`running`|`done`) and, if reported, +`seconds`; any additional flat scalar counters the executor reported for that stage +(e.g. `keys_checked`, `flagged`, `fields_repulled`) ride alongside — nested/list +counters, and any counter using the reserved keys `status`/`seconds`, are dropped by the +internal stage-report endpoint rather than persisted (`backend/agent_kv/internal_views.py::_sanitize_counters`). + +`adapters` echoes the platform adapter ids this job was dispatched with, by role +— the ones the submit named and the gate validated (§3a). It is reported because +the caller chose them and is paying for them: naming the wrong `llm` id yields a +correct-looking result at a different price, and the only other record of what +was spent is in `usage_v2`, which the caller cannot read. **Ids only, never +adapter metadata** — that is where provider credentials live. `{}` for an +env-configured extractor (`kv`) and for any job predating the field. + +`status` is lowercased (`pending`, `dispatched`, `running`, `completed`, `failed`, +`cancelled`); `error` is included (only) when `status == "failed"`. `adapters`, +`pages_total` and `completed_at` are always present (as `null`/`{}` until known) — an addition beyond the +spec's illustrative example, not a drift, since spec §7.2 was explicitly a partial +sample ("Verbose, agent-centric, stage-level"). + +```bash +curl https://api.unstract.example/agent-kv/5b6e9b0a-.../ \ + -H "Authorization: Bearer 5c9e2c9e-1234-4a5b-9c6d-abcdef012345" +``` + +## 5. Result — `GET /agent-kv/{job_id}/result` + +Re-readable until `expires_at` (D11), plus explicit `DELETE`. + +- **Non-terminal job** (`PENDING`/`DISPATCHED`/`RUNNING`): `409` with + `{"status": ""}`. +- **Expired, or `COMPLETED` with an already-swept/blank `result_ref`**: `404`. +- **`COMPLETED`**: `200` with the engine's full result object, unchanged (record, + normalized_record, per-field audit trail, `qa_passed`, `challenge_passed`, + `consistency_violations`, `cost_summary`, `timing` — spec §4; this repo does not + shape or re-validate that payload, it stores and replays exactly what the executor's + finalize call sent). +- **`FAILED`**: `200` with + `{"success": false, "status": "failed", "error": ""}`. +- **`CANCELLED`**: `200` with `{"success": false, "status": "cancelled"}`. + +**Calculation result shape** (when `calculations` is supplied and +`AGENT_KV_CALCULATIONS_ENABLED=true`): a `COMPLETED` result adds these fields to the +engine's standard result object: + +- `calculations_applied: true` — indicates calculation codegen ran. +- `execution: {success: bool, rows_written: int, error: string|null}` — sandbox worker + execution status. `success` is `true` if the generated code ran without error; + `rows_written` is the count of rows the code emitted; `error` is `null` on success, or + a user-safe message (e.g., `"Runtime error: division by zero"`) on failure. +- `codegen_validation_passed: bool` — whether the AST gate (layer 1 of a 5-layer + defense-in-depth model) accepted the generated code syntax. +- `calculation_rows: [...]` — an array of computed JSONL rows (size-capped). If the code + produced more rows than the cap allows, this array is empty and `calculation_rows_truncated: true` + is present instead. + +If the sandbox worker fails to execute, the job completes with `success: false` and a +user-safe `error` (e.g., `"Calculation execution timed out"`); the `execution` field +carries the inner error details. + +**Drift from the spec:** §7.3 sketches the failed-job body as +`{success: false, error, timing}`. The shipped body +(`backend/agent_kv/execution_views_result.py::result_payload`) is +`{"success": false, "status": "failed", "error": ...}` — it adds `status` and does +**not** include a `timing` key (there is nothing to time for a job that never produced +an engine result). Document the code: no `timing` on failure/cancellation bodies. + +```bash +curl https://api.unstract.example/agent-kv/5b6e9b0a-.../result \ + -H "Authorization: Bearer 5c9e2c9e-1234-4a5b-9c6d-abcdef012345" +``` + +The payload is **keyed by extractor**, with usage attributed per extractor +alongside the authoritative total: + +```json +{ + "success": true, + "status": "completed", + "extractors": { + "kv": { + "success": true, + "record": {"quotation_number": "Q-10234", "customer": {"name": "Acme Corp"}}, + "normalized_record": {"quotation_number": "Q-10234", "customer": {"name": "Acme Corp"}}, + "keys": [ + {"key_path": "quotation_number", "qa_status": "pass", "challenge_status": "pass"} + ], + "qa_passed": true, + "challenge_passed": true, + "consistency_violations": [], + "cost_summary": { + "total_cost": 0.0, + "input_tokens": 3120, + "output_tokens": 480, + "agents": { + "kv_extractor": {"input_tokens": 2400, "output_tokens": 360, "cost": 0.0}, + "kv_qa": {"input_tokens": 720, "output_tokens": 120, "cost": 0.0} + } + }, + "timing": {"document_processing": 6.2, "extraction": 11.4} + } + }, + "usage_summary": { + "total": {"pages": 3, "input_tokens": 3120, "output_tokens": 480, "total_cost": 0.0}, + "by_extractor": { + "kv": {"pages": 3, "input_tokens": 3120, "output_tokens": 480, "total_cost": 0.0} + } + } +} +``` + +`success` and `status` are at the **top level on every terminal payload** — +completed, failed and cancelled — so a client branches on one key. They describe +the job; each extractor keeps its own `success` inside its own block. +`usage_summary.total` is the billing figure; `by_extractor` exists so a +multi-extractor job can be attributed. With one extractor they are the same numbers. +Note the **failed** and **cancelled** payloads above are NOT extractor-keyed: a +failure belongs to the job, not to one extractor. + +**`cost_summary` dollars are always `0.0` and are not a billing figure**: the executor +zeroes the engine's per-token price fields (they were stale hardcoded list prices), so +`total_cost` and every per-agent `cost` come back `0.0` — metering is what bills, from +the SDK's own litellm-priced usage records. Read `cost_summary` for its token/agent +breakdown only. `agents` is omitted entirely when no agent recorded usage (a full +cache hit). + +**The per-field audit entries in `keys[]` are keyed `key_path`, not `path`.** A live +entry carries more than the example above shows — this one is verbatim from a real run: +`{"key_path": "invoice_number", "value": "INV-…", "found": true, "qa_status": +"unchecked", "qa_attempts": 0, "line_start": 5, "line_end": 5, "normalized_value": +"INV-…"}` (`qa_status` is `"unchecked"` with `qa_attempts: 0` when the job ran with +`qa=false`). + +(The exact result shape beyond `success` is the cloud engine's contract — spec §4 — not +re-specified or validated by this repo.) + +**The `table` extractor's result** is the table engine's extraction output, filed +under `extractors.table`: + +```json +{ + "success": true, + "status": "completed", + "extractors": { + "table": { + "tables": [{"unit": "A1", "rent": 1200}], + "page_count": 3, + "table_pages": [1, 2], + "headers": ["unit", "rent"], + "row_count": 1, + "timing": {"document_processing": 6.2, "total": 41.7} + } + }, + "usage_summary": { + "total": {"pages": 3, "input_tokens": 3120, "output_tokens": 480, "total_cost": 0.0}, + "by_extractor": { + "table": {"pages": 3, "input_tokens": 3120, "output_tokens": 480, "total_cost": 0.0} + } + } +} +``` + +`tables` carries one object per extracted row; `headers` is the canonical column list +the engine settled on; `table_pages` are the 1-based pages the rows came from; and +`row_count` is the number of rows extracted. `timing` rides inside the extractor block, +and `parse_failures` appears only when some page failed to parse. There is **no +`cost_summary`** on this extractor — the table engine has no per-agent cost tracker — so +`usage_summary` is the only billing figure, as it is for `kv`. + +> **`enable_header_mapping` changes the type of `tables`.** With it **off** (the +> default), `tables` is a flat list of row objects, exactly as above. With it **on**, +> the engine wraps them and `tables` becomes +> `{"header_mapping": {...}, "rows": [...]}`. That is a breaking shape change between +> two settings of one option: a client that sets the flag must read `tables.rows` +> instead of iterating `tables`. + +## 6. Cancel — `POST /agent-kv/{job_id}/cancel` + +Defined precisely per spec §7.4: no task-revocation machinery exists anywhere in the +platform, so cancel flips the job row to `CANCELLED` via the same guarded +`mark_terminal` write gate every other terminalization path uses +(`backend/agent_kv/models.py::AgentKVJob.mark_terminal`) — later callbacks/stage +reports become no-ops; a job not yet picked up is dropped at pickup. + +A job **already running is stopped cooperatively**, not revoked. Cancellation is +detected opportunistically: the executor only learns of it on its next stage report, +which the internal endpoint answers as a no-op once the row is terminal. Both extractors +check at their engine's progress points — `kv` at every graph-node boundary, `table` at +the engine's own progress points (rate-limited, so detection can lag by a few seconds) — +and the job then terminates as `cancelled`. Work already done is still billed and +metered; work not yet started is not. A job that finishes before the cancel is noticed +simply has its result discarded on arrival. + +- Won the race (job was non-terminal): `200` `{"status": "cancelled"}`. +- Already terminal: `409` `{"status": ""}`, lowercased (e.g. + `"completed"`) — consistent with every other `status` value in this API. + + *An earlier version of this document claimed the opposite: that this one field + carried the raw uppercase enum, and that it was "not a typo". It was wrong. + `JobCancelView` returns `job.status.lower()`, and the unit test cited as proof + (`test_job_views.py::test_cancel_on_completed_is_409_and_result_untouched`) + asserts `{"status": "completed"}`. The e2e lane carried the same wrong + assertion and had never been run.* + +```bash +curl -X POST https://api.unstract.example/agent-kv/5b6e9b0a-.../cancel \ + -H "Authorization: Bearer 5c9e2c9e-1234-4a5b-9c6d-abcdef012345" +``` + +## 7. Delete — `DELETE /agent-kv/{job_id}` + +Deletes the result and any residual staged input immediately, and blanks both refs +(`204`, no body). The job row itself is retained (audit trail); a later `result` fetch +404s (blank `result_ref`), and status still works. + +```bash +curl -X DELETE https://api.unstract.example/agent-kv/5b6e9b0a-.../ \ + -H "Authorization: Bearer 5c9e2c9e-1234-4a5b-9c6d-abcdef012345" +``` + +## 8. Webhook delivery + +`webhook_url` (if given at submit) receives exactly one terminal-state POST: +`{"job_id": "...", "status": "completed"|"failed"|"cancelled"}` — a fixed payload +shape, no result content (spec §6.7, §7.1). Delivered from the `ide_callback` worker +(`workers/ide_callback/agent_kv_tasks.py::_maybe_webhook`) only on a *fresh* finalize +(a duplicate/late finalize for an already-terminal job never re-fires it). + +**SSRF controls, as actually implemented** (`workers/shared/utils/webhook_notify.py`, +function `send_webhook`) — this corrects an earlier draft of this document, which +claimed connect-time re-checking and bounded retries; the code does neither, and its +own module docstring says so: + +- **Resolution-time check only.** The host is resolved once via `socket.getaddrinfo` + and every returned address is checked against private/loopback/link-local/ + reserved/multicast/unspecified ranges *before* the request is made + (`_host_is_public`). It is **not** re-checked at connect time — the subsequent + `requests.post` call performs its own, independent DNS resolution, and nothing + pins that second resolution to the address(es) just validated. +- **Residual risk, accepted for v1: DNS-rebinding TOCTOU.** An attacker whose DNS + answer changes between the pre-check's resolution and `requests`' own resolution + a moment later (serving a public IP to the first, a private/internal IP to the + second) slips through this gap. The module's own docstring names this as a known, + accepted risk. **Fast-follow, not shipped:** IP-pinning the validated address into + the actual connection, so the two resolutions can never disagree. +- **Scheme: https-only**, in this feature's actual usage. `send_webhook` exposes a + generic `allow_http` escape hatch (for explicitly configured dev environments + elsewhere), but the Agent-KV caller (`_maybe_webhook`) never passes it — every + Agent-KV webhook is https-only in practice. +- **No redirects** (`allow_redirects=False`). +- **Fixed, minimal payload** (`{job_id, status}` only — see above) and the + **response body is never read** — only the status code is inspected. +- **Short timeout** (10 seconds). +- **No retries.** One POST attempt only; any failure (host refused by the + pre-check, timeout, connection error, non-2xx status) is logged and swallowed — + there is no retry loop anywhere in this delivery path. + +Test/dev stacks only: `AGENT_KV_WEBHOOK_INSECURE_ALLOW_HTTP_PRIVATE=1` on the +ide-callback worker waives both guards (http scheme and non-public host) so the +e2e lane can deliver to a receiver on the compose host. Never set it in +production. + +## 9. Retention and TTL + +- **Input deletion is completion-triggered, not TTL-based** (spec D10: "uploaded + document deleted on job completion"). `FinalizeView.post` + (`backend/agent_kv/internal_views.py`) deletes the staged input file and blanks + `input_ref` the instant a finalize call actually wins the terminal-state guard (the + job just became `COMPLETED` or `FAILED`), via `storage.delete_input` — which only + ever touches `input_ref`, never `result_ref`/the result file. A duplicate/late + finalize call (guard already lost, since the job is already terminal) never deletes + anything: either the winning call already did, or the job reached terminal some + other way (see next bullet). +- **Cancelled jobs' inputs intentionally ride TTL instead.** `JobCancelView` marks a + job `CANCELLED` directly via `AgentKVJob.mark_terminal`, not through `FinalizeView` + — so cancellation never triggers the completion-time delete above. A late finalize + call against an already-`CANCELLED` job also no-ops (it loses the terminal guard, + same as any duplicate). A cancelled job's staged input is therefore cleaned up the + ordinary way: by the TTL sweep once `expires_at` passes, or by an explicit + `DELETE /agent-kv/{job_id}`. +- **Result retention**: `AGENT_KV_RESULT_TTL_DAYS` (default **7** days, D10's + engineering default) — results are re-readable until then, then swept by the + internal `ttl-cleanup` endpoint ([§6](#11-deploy-checklist)), which blanks + `result_ref` (and `input_ref` too, covering the cancelled-job case above, or + defensively for any input that somehow outlives completion). The job row itself is + never deleted (audit trail persists past TTL, only the object-store payloads are + dropped). +- **`expires_at` is stamped at submit time**, not at completion — + `SubmitView.post` sets `expires_at = timezone.now() + timedelta(days=AGENT_KV_RESULT_TTL_DAYS)` + before dispatch even runs. **This is a real divergence from spec D10's framing** + ("results retained 7 days" implicitly from completion): for a job that runs close to + its `timeout` ceiling, the effective post-completion retention window is slightly + under the full 7 days. Documented here as shipped behavior, not silently + reconciled — flagged for anyone tightening D10 later. + +## 10. When the engine is unavailable (501 behavior) + +The Django backend cannot see the executor plugin registry (that lives in the workers +process). Dispatching to a queue no worker consumes would hang, not error — so gating +happens **before enqueue, in the backend**, via a capability probe: + +```python +if not get_plugin("agent_kv"): + raise EngineUnavailable() # 501, "agent-kv engine not available on this deployment" +``` + +`get_plugin(...)` (`backend/plugins/__init__.py`) returns an empty dict — falsy — when +no `plugins/agent_kv/` package is installed, which is the OSS-only case by design (the +plugin is a cloud deliverable, spec §5.1/§11: "OSS scaffold merges dark"). Every +`POST /agent-kv/` on an OSS-only deployment therefore returns: + +```json +{"detail": "agent-kv engine not available on this deployment"} +``` + +with HTTP `501`. Key management and status/result/cancel/delete on +*already-existing* jobs are unaffected by this gate — only `SubmitView` probes it, +since only submit needs the engine. + +### 11a. The cancel-at-pickup contract (internal API — frozen; build the cloud engine against this) + +Spec §7.4 states that "a job not yet picked up is dropped at pickup," but leaves the +actual mechanism unspecified. This is it, and it is load-bearing enough that it is +recorded here explicitly rather than left to be inferred from the internal-API code: + +`POST /internal/v1/agent-kv/jobs/{job_id}/stage/` — the stage-report endpoint +(`StageReportView`, `backend/agent_kv/internal_views.py`) — no-ops with +`{"ok": true, "noop": true}` whenever the target job is already terminal (its +candidate query excludes `AgentKVJob.TERMINAL`, so a terminal job's row is simply +never found). **That response body *is* the drop signal.** The FIRST stage report an +executor makes for a job it has just picked up doubles as a liveness check against a +job that was cancelled (or otherwise reached a terminal state) in the window between +dispatch and pickup — since there is no task-revocation machinery anywhere in the +platform (spec §7.4) to stop a picked-up job any other way. + +On `noop: true`, the executor **must**: + +1. Stop doing any further work on that job immediately. +2. **Still call `FinalizeView`** (`POST .../jobs/{job_id}/finalize/`) with whatever + outcome it has — success or failure, it does not matter which, since the write is + going to no-op either way. + +Step 2 is not optional. `FinalizeView`'s terminal-state guard (spec §5.4) correctly +no-ops the status write for an already-terminal job — but its concurrency-slot +release (`AgentKVConcurrencyLimiter.release`) lives in a `finally`, unconditional on +whether the guard actually won ([§6](#9-retention-and-ttl) covers the input-deletion +side of this same guard). An executor that treats `noop: true` as "nothing left to do +here" and skips the finalize call entirely leaks that job's concurrency slot for the +rest of the limiter's TTL (6 hours) instead of releasing it immediately. This is the +entire mechanism behind a pre-pickup cancel actually being observed and cleaned up +promptly — there is no other signal. + +## 11. Deploy checklist + +Everything below is required (or worth checking) to run this feature for real, +beyond `docker compose up`: + +1. **Cloud plugin**: install/enable the `agent_kv` backend capability plugin + (probed via `plugins.get_plugin("agent_kv")`) — without it, submit always 501s + ([§6](#10-when-the-engine-is-unavailable-501-behavior)). This repo ships gated + dark by design; nothing to do for an OSS-only deployment except accept the 501. + The marker the probe keys off is the mere presence of + `backend/plugins/agent_kv/__init__.py` — a cloud image ships it (via the + plugin-copy step), an OSS-only build does not. +2. **Executor worker fleet queue**: the executor dispatch derives its Celery queue name + from the executor name as `celery_executor_{executor_name}` — for Agent-KV that's + **`celery_executor_agentic_kv`** (`unstract/sdk1/src/unstract/sdk1/execution/dispatcher.py`). + `workers/run-worker.sh`'s executor role default queue list + (`celery_executor_legacy,celery_executor_agentic,celery_executor_agentic_table`) + does **not** include it — a self-hosted/OSS-only deployment running the cloud + `agentic_kv` executor plugin must add `celery_executor_agentic_kv` to that worker's + consumed queues (`CELERY_QUEUES_EXECUTOR` on the Docker path, or the equivalent + `run-worker.sh` queue map entry) or dispatched jobs will never be picked up. + **The cloud Helm chart already wires this**: both `workerExecutorV2.args` + (`--queues=...,celery_executor_agentic_kv`) and + `workerPgExecutor.env.WORKER_PG_QUEUE_CONSUMER_QUEUE` include it, and both fleets + carry `additionalConfigs: [..., agentKv]` (`charts/unstract-platform/values.yaml`, + cloud repo) so the executor pods get the `AGENT_KV_*` LLM/LLMWhisperer env alongside + the queue wiring. +3. **Callback queue**: the `ide_callback` worker must consume `agent_kv_callback` + alongside its own `ide_callback` queue, or jobs dispatch but never finalize. This is + already the hardcoded default in both `workers/run-worker-docker.sh` and + `workers/run-worker.sh` (`WORKER_QUEUES["ide_callback"] = "ide_callback,agent_kv_callback"`) + — nothing to configure for a standard deployment; see the note already in + `docker/sample.env` under "IDE Callback Worker." The cloud Helm chart matches: + `workerIdeCallbackV2.args` (`--queues=ide_callback,agent_kv_callback`) and + `workerPgIdeCallback.env.WORKER_PG_QUEUE_CONSUMER_QUEUE` + (`"ide_callback,agent_kv_callback"`) both already include it + (`charts/unstract-platform/values.yaml`, cloud repo). +4. **Internal periodic maintenance — proxy tasks now exist; nothing schedules them + yet.** Two internal endpoints do the real work + (`backend/agent_kv/internal_urls.py`, `internal_views.py::SweepView`/ + `TTLCleanupView`), and thin scheduler-side proxy tasks now call them — + `workers/scheduler/agent_kv_tasks.py::agent_kv_sweep`/`agent_kv_ttl_cleanup`, + registered under the wire names `agent_kv.sweep`/`agent_kv.ttl_cleanup` — mirroring + the dashboard-metrics periodics (`workers/scheduler/dashboard_metrics_tasks.py`) + this feature was designed to follow (spec §5.4), except the call goes through the + shared `InternalAPIClient` facade (`agent_kv_sweep`/`agent_kv_ttl_cleanup` methods on + `workers/shared/api/internal_client.py`) rather than a second bespoke HTTP client. + **Nothing calls these tasks on a schedule yet** — an operator must register both as + periodic tasks via the same PG-scheduler mechanism the dashboard-metrics periodics + use: a `PgPeriodicTask` row per task (`name`, `task_name` = the wire name above, + `queue`, `cron_string`, `enabled=True`, `pg_owned=True` once ready to go live — see + `backend/dashboard_metrics/migrations/0004_pg_periodic_tasks.py` for the exact row + shape this mirrors) before relying on TTL cleanup or stuck-job recovery: + - `agent_kv.sweep` → `POST /internal/v1/agent-kv/sweep/` — terminalizes `PENDING` + jobs never dispatched within `AGENT_KV_SWEEP_GRACE_SECONDS` (default 1 hour) + **and** `DISPATCHED`/`RUNNING` jobs stuck past `AGENT_KV_STUCK_JOB_GRACE_SECONDS` + (default 6 hours) — two independent phases per call, response + `{"swept": N, "timed_out": M}`. Suggested cadence: every few minutes. + - `agent_kv.ttl_cleanup` → `POST /internal/v1/agent-kv/ttl-cleanup/` — deletes + staged input/result files past `expires_at`. Suggested cadence: hourly or daily + (retention is measured in days). + Both endpoints are idempotent, batch-capped at 500 rows per phase per call, and + safe to call more often than needed. + + **Cloud mechanism — Kubernetes CronJobs, not the PG-scheduler.** The sweep and + TTL-cleanup logic itself lives in `backend/agent_kv/maintenance.py::run_sweep`/ + `run_ttl_cleanup` — `SweepView`/`TTLCleanupView` above are now thin wrappers + around it. Two Django management commands wrap the same functions — + `python manage.py agent_kv_sweep` and `python manage.py agent_kv_ttl_cleanup` + (`backend/agent_kv/management/commands/`) — each printing the JSON counts + (`{"swept": N, "timed_out": M}` / `{"cleaned": N}`) and exiting 0. In the cloud + deployment, a Kubernetes CronJob runs each command on the cadences above instead + of registering a `PgPeriodicTask` row. + + **The CronJob is the only owner.** Migration `0004_pg_periodic_tasks` once also + seeded two `PgPeriodicTask` rows for the same two tasks; it now removes them. + They were seeded `pg_owned: False` so nothing double-fired, but the flag exists + to be flipped, and a flip would have run a `*/10` sweep against the chart's + `*/15` — colliding at `:30` on the same rows — with TTL cadences disagreeing 24x. + Only the CronJob side carries `concurrencyPolicy: Forbid` and a deadline. + + > **Self-hosted OSS: you must schedule these two yourself.** There is no CronJob + > outside the cloud chart and no migration registers them any more. Run + > `manage.py agent_kv_sweep` and `manage.py agent_kv_ttl_cleanup` from cron, a + > systemd timer, or your own `PgPeriodicTask` rows with `pg_owned: True` — the + > scheduler proxy tasks `agent_kv.sweep`/`agent_kv.ttl_cleanup` are still + > registered and still work. Without them a stranded job holds its concurrency + > slot until the 6 h Redis TTL, and `AGENT_KV_RESULT_TTL_DAYS` is advisory: + > staged documents and results stay in the bucket indefinitely, which is a + > retention failure, not just a disk-usage one. + > + > Note this deployment routes `table` only, and that extractor needs the cloud + > `agentic_table` plugin — so a pure-OSS install cannot dispatch an Agent-KV job + > at all, and has no stranded rows to reap until it ships an engine plugin. + + **Chart keys (cloud repo)**: `backend.agentKvCronJobs` in + `charts/unstract-platform/values.yaml` — `enabled: false` there (on-prem never runs + it), flipped to `enabled: true` in `charts/cloud-deployment-values/cloud.values.yaml`. + Its `jobs` list is the two commands above with concrete schedules: + `{name: sweep, command: agent_kv_sweep, schedule: "*/15 * * * *"}` and + `{name: ttl-cleanup, command: agent_kv_ttl_cleanup, schedule: "17 * * * *"}` — offset + from sweep's :00/:15/:30/:45 ticks so the two CronJobs never contend on the same + rows. Neither job needs the `agentKv` shared-config group (LLM/LLMWhisperer + creds — only the executor touches those): sweep only needs `database`/`redis`, + ttl-cleanup additionally needs `storage` for `AGENT_KV_FILE_STORAGE_CREDENTIALS` + (see item 6 below) — both already on the `backend` deployment's config list. +5. **Env vars**: every `AGENT_KV_*` setting plus `AGENT_KV_FILE_STORAGE_CREDENTIALS` — + see [§6](#12-environment-reference) and `docker/sample.env`. + + > **For the `table` extractor, the executor-side credential group below must be + > left UNSET — not merely "is not required".** That path resolves the caller's + > own platform adapters (see §3a), so none of the six `AGENT_KV_LLM_*` / + > `AGENT_KV_LLMWHISPERER_*` values is read on it. The group only has to be + > configured where the `kv` extractor is deployed, which is a separate release + > track. A table-only deployment leaves + > `global.sharedConfigs.agentKv.enabled: false` and the API works. + > + > **Deploy order: cloud worker first, then the backend.** The two repos deploy + > independently and nothing on the wire marks the capability. A new backend + > against an *old* worker sends validated `adapters` in `executor_params` to a + > worker that has no `ParamKeys.ADAPTERS`; it ignores the key and calls + > `ExtractionConfig.from_env()`. If those six values are set it runs on + > **operator** credentials, bills the `agent_kv:{provider}/{model}` sentinel, and + > returns a 202 with a correct-looking result — while the adapters the caller + > named were never used. Nothing in either repo's logs reveals it. With the group + > unset the same old worker fails loudly with "Missing required Agent-KV env + > vars", which is what a half-finished rollout should look like. The reverse + > order is safe by construction: the new worker hard-fails "Missing required + > adapters" rather than falling back to env. + > + > Two operator knobs are still read on the table path, independently of that + > group: `AGENT_KV_STORAGE_DIR_PREFIX` (which must match the backend's — the + > executor derives the job's expected directory from it to enforce the tenant + > boundary on the staged input) and `AGENT_KV_PARALLEL_PAGES`. + + The executor-side vars + (read by the cloud `agentic_kv` plugin, not by the backend settings in §13) are a + separate chart group: `global.sharedConfigs.agentKv` + (`charts/unstract-platform/values.yaml`, cloud repo) — + `AGENT_KV_LLM_PROVIDER`, `AGENT_KV_LITE_MODEL`, `AGENT_KV_ADVANCED_MODEL`, + `AGENT_KV_LLM_API_KEY`, `AGENT_KV_LLMWHISPERER_API_KEY`, + `AGENT_KV_LLMWHISPERER_BASE_URL`, `AGENT_KV_MAX_TOKENS`, `AGENT_KV_PARALLEL_PAGES`, + `AGENT_KV_ENGINE_VERSION`. Of those, two are ESO (External Secrets Operator) + secrets rather than plain values — `AGENT_KV_LLM_API_KEY` and + `AGENT_KV_LLMWHISPERER_API_KEY`, mapped in `global.externalSecrets.groups.agentKv` + to the GCP Secret Manager suffixes `agent-kv-llm-api-key` / + `agent-kv-llmwhisperer-api-key`. `AGENT_KV_LLM_PROVIDER` is a v1 hard allowlist of + two values — `anthropic` and `openai` (`SUPPORTED_PROVIDERS` in the executor's LLM + adapter) — anything else raises `ConfigError` at dispatch, before any LLM call. + + **Rollout order is strict, and getting it wrong takes down the whole executor + fleet.** Both executor fleets mount this group with `envFrom: secretRef: + -agent-kv`, so no executor pod of any kind — not just Agent-KV ones — + can start until that Secret exists: enabling the group ahead of its contents + leaves `workerExecutorV2` and `workerPgExecutor` in `CreateContainerConfigError`. + Per environment, in this order: (1) create the two GCP Secret Manager secrets + `-agent-kv-llm-api-key` and + `-agent-kv-llmwhisperer-api-key`; (2) set the three + non-secret keys — `AGENT_KV_LLM_PROVIDER`, `AGENT_KV_LITE_MODEL`, + `AGENT_KV_ADVANCED_MODEL` — in that environment's own values file (under ESO they + are emitted as inline literals in the ExternalSecret's `target.template`, and a nil + value is silently **dropped**, so leaving them unset ships a Secret without them and + every Agent-KV job then fails at dispatch with `ConfigError: Missing required + Agent-KV env vars`); (3) only then roll the fleets out. `cloud.values.yaml` enables + the group fleet-wide, so step 2 is a per-environment prerequisite, not an optional + extra. In inline (non-ESO) mode the same three keys are enforced at render time by a + `fail` guard in `charts/unstract-platform/templates/shared/agent-kv-secret.yaml`, so + a mistake there is a loud deploy-time failure instead of a per-job runtime one. +6. **Object storage**: `AGENT_KV_FILE_STORAGE_CREDENTIALS` must point at a real bucket + in production (defaults to the same shared MinIO instance as workflow execution). + Paths are prefix- and org-rooted + (`{AGENT_KV_STORAGE_DIR_PREFIX}/{org_id}/{job_id}/input{ext}`, `.../result.json`) — + no document bytes ever ride the message broker. Chart key: + `global.sharedConfigs.storage.AGENT_KV_FILE_STORAGE_CREDENTIALS` + (`charts/unstract-platform/values.yaml`, cloud repo) — derived from + `fileStorageCredentials`/`MINIO_CREDS` unless set explicitly, the same fan-out + pattern as its three siblings (`WORKFLOW_EXECUTION_`, `API_`, + `HITL_FILES_FILE_STORAGE_CREDENTIALS`). +7. **`AGENT_KV_STORAGE_DIR_PREFIX` — must be set on the backend *and* the executor + fleet, to the same value** (default `unstract/agent_kv`). The **first segment is the + bucket** — s3fs/gcsfs read it that way — so it must be a bucket that already exists + (the MinIO dev bootstrap creates `unstract`); a bucket-less prefix makes every submit + fail `NoSuchBucket` and return a 500 with nothing billed. The executor reads the + backend-staged `input_ref` and keys its OCR cache under the same root + (`{prefix}/{org_id}/cache/…`), so a mismatch silently splits the cache and breaks + result read-back. Unlike the executor-only group in item 5, this one has to reach + **both** sides, so in the chart (cloud repo) it must be fanned out to the backend + *and* to `workerExecutorV2`/`workerPgExecutor` from a single source of truth — never + set independently per fleet, where the two can drift apart unnoticed. +8. **Feature flags — `calculations` and `structured_output`**: the cloud engine + cannot execute these yet, so submit rejects them with a 400 until the flags are + flipped. Both `AGENT_KV_CALCULATIONS_ENABLED` and + `AGENT_KV_STRUCTURED_OUTPUT_ENABLED` default `false`; set either to `true` only + once the executor side has shipped support for it. +9. **Executor time limit — the real baseline is 3,600s, not 7,200s.** The executor + fleet (`workerExecutorV2`/`workerPgExecutor`) is not a `FILE_PROCESSING` worker — + `WorkerConfig._get_worker_specific_timeout_defaults()` + (`workers/shared/models/worker_models.py`) falls through to the "conservative + defaults for other workers" branch for `WorkerType.EXECUTOR`: **1 hour / 3,600s** + hard limit, 3,300s soft. The chart's bare `TASK_TIME_LIMIT: "3600"` ConfigMap key + is inert — `get_celery_setting`'s 4-tier resolution (cmdline → + `{WORKER_TYPE}_{SETTING}` → `CELERY_{SETTING}` → default) means only + **`EXECUTOR_TASK_TIME_LIMIT`** (worker-specific) or `CELERY_TASK_TIME_LIMIT` + (global) actually override it; see the comment on `workerExecutorV2`'s + `terminationGracePeriodSeconds` in `charts/unstract-platform/values.yaml`, which + already documents this ("Its ceiling is task_time_limit=3600s... note the chart's + bare TASK_TIME_LIMIT is NOT read; only EXECUTOR_TASK_TIME_LIMIT / CELERY_*"). + At the 100-page cap (`AGENT_KV_MAX_PAGES`) with `extraction_mode="per-page"`, + `parallel_pages=4`, and both QA and challenge enabled (the worst-case option + combination), the arithmetic worked out in Task 9's report is: each of + `key_extractor` → `kv_qa` → `kv_challenger` fans 100 pages across 4 workers (25 + sequential rounds), each round up to ~99s if its slowest call needs a *full* retry + backoff (`SDKLLMClient`'s `_DEFAULT_MAX_RETRIES=8`) — 25 × 99s = 2,475s per stage, + 3 stages sequential = **7,425s**, a gap of **3,825s (over 2x)** against the real + 3,600s baseline, not the 225s a 7,200s baseline would suggest. That 7,425s figure + assumes every single call across all 300 (100 pages × 3 stages) hits the full + 8-retry backoff, which is not realistic — a real call is ~10-30s, so a full run at + the page cap is closer to **~2,000-2,500s** in practice; the 7,425s number is a + theoretical ceiling for capacity planning, not an expected duration. + **Recommendation: do not raise the fleet-wide time limit preemptively.** + `EXECUTOR_TASK_TIME_LIMIT` is fleet-wide — every executor task (every plugin, not + just Agent-KV) inherits it, and it's coupled to `workerPgExecutor`'s PG-visibility + lease/`vt` and `terminationGracePeriodSeconds` (raising one without the others + reintroduces the mid-flight-SIGKILL/silent-redelivery risk `UN-3964`'s comment on + `workerExecutorV2` already guards against). Measure real per-page timings from the + Task 13 integration run first, then either raise `EXECUTOR_TASK_TIME_LIMIT` (and + `terminationGracePeriodSeconds`/the PG lease together, not alone) or lower + `AGENT_KV_MAX_PAGES` for deployments that see per-page-heavy, near-cap documents in + practice. + + **What actually happens at each limit.** At the **soft** limit (3,300s) Celery raises + `SoftTimeLimitExceeded` into the task; the executor catches it, sets a run-level stop + signal that makes every in-flight LLM retry loop abort at its next backoff instead of + sleeping through the remaining ladder (up to ~30s per attempt × 8 attempts), flushes + the usage records accrued so far, and finalizes the job `failed` with the user-safe + error `"timed out"`. At the **hard** limit (3,600s) the worker is SIGKILLed: usage + records still held in memory are lost, so that run's LLM spend goes **unmetered** + (page usage, posted per-OCR-call, is unaffected). That gap is the reason the soft + limit now stops retries promptly — it is what keeps a slow run from crossing into the + hard kill. +10. **Testing lanes — Docker-free unit gate, Docker-required integration lane.** The + cloud rig's `unit-agentic-kv` group (`tests/groups.cloud.yaml`, cloud repo) runs the + executor plugin's own test suite in an isolated venv with no `requires_services` — + no Docker needed, and it's what `tox -e unit` (CI's unit-tier matrix leg) runs. The + integration tier is the one that needs Docker: it authenticates to the container + registry and boots a mock Auth0 sidecar (`.github/workflows/ci-test.yaml`, cloud + repo) for the cross-service suites. A local `unit-agentic-kv` run needs neither + Docker nor the platform stack up. + + **A real-key Agent-KV e2e run cannot go through `tests.rig run`.** Whenever a + selected group needs the platform, the rig sets `UNSTRACT_LLM_MOCK_RESPONSE` + to `MOCK_LLM_OK` (`tests/rig/cli.py`) — and an exported *empty* value counts as + unset, so there is no way to opt out from inside the rig; the lane then exercises + the mock, not a provider. To run it against real LLMWhisperer/LLM keys, boot the + stack first (`python -m tests.rig platform up`, or `docker compose` with + `tests/compose/docker-compose.test.yaml`) and then invoke pytest directly with + `UNSTRACT_BACKEND_URL`, `AGENT_KV_E2E=1` and the `AGENT_KV_*` keys exported: + `AGENT_KV_E2E=1 UNSTRACT_BACKEND_URL=http://localhost:8000 pytest tests/e2e/agent_kv`. + + Two scenarios are operator-gated on top of that: the bad-LLM-key scenario + (`AGENT_KV_E2E_BAD_KEY_JOB=1`, see the test module docstring) and the + completion-webhook scenario, which needs + `AGENT_KV_WEBHOOK_INSECURE_ALLOW_HTTP_PRIVATE=1` set BOTH on the + ide-callback worker (it waives the webhook SSRF guards — https scheme + + public host — for test/dev stacks only; never set it in production) and in + the pytest process env; the receiver is reached via + `host.docker.internal` (the compose host-gateway mapping). The lane also + covers sync-wait submits (`timeout` → 200 with the inline result) and an + `.xlsx` document (`fixtures/invoice.xlsx`, the no-pre-OCR-page-count path). +11. **Sandbox worker — calculations codegen fleet** (when + `AGENT_KV_CALCULATIONS_ENABLED=true`). The optional `calculations` field + invokes out-of-process codegen and execution in a dedicated sandbox worker fleet + consuming the `sandbox_codegen` Celery queue. **Deployment order is strict:** + deploy the sandbox worker (Deployment healthy, actively consuming tasks from + `sandbox_codegen`) **before** flipping `AGENT_KV_CALCULATIONS_ENABLED=true`. + With the flag on but the sandbox down, calculation jobs fail user-safely at the + RPC timeout (they do not hang past it); however, **never leave the flag on + without a healthy sandbox fleet**, as all jobs with `calculations` will eventually + exceed their timeouts and return failures to callers. + + **Security posture:** The sandbox pod carries no LLM, OCR, or storage secrets and + runs with default-deny egress — it only executes generated calculation code against + the extracted record in-memory. This no-secrets/deny-all-egress posture is the + **Kubernetes/chart deployment guarantee**. The reference `docker-compose.yaml` + included here is a best-effort dev mirror only, not that guarantee: its + `worker-sandbox` service shares the same full worker env-file (`workers/.env`) as + every other worker for operational simplicity, so on that stack the untrusted-code + boundary is the runner's per-subprocess env scrub (spec §6.3 layer 2), not full + pod-level secret isolation — don't run untrusted `calculations` against secrets you + care about on docker-compose. User-supplied `calculations` expressions are + validated through an AST gate (layer 1 of a 5-layer defense-in-depth model) at + submit time before ever reaching the sandbox, and further constrained (parsing, + type-checking, capability allowlist) at execution time. The AST gate is a + **best-effort** layer-1 control — a fail-closed denylist that closes the cheap, + demonstrated bypasses and raises the cost of the rest, not a provably airtight + sandbox by itself. The *enforced* containment is layers 2–5: the runner's + env-scrub + rlimits (layer 2) and the pod's non-root/read-only-rootfs/no-secrets/ + deny-all-egress/per-job isolation (layers 3–5), with `runtimeClass: gvisor` as the + deferred syscall-sandbox mitigation (see the accepted v1 residual below). Do not + treat an AST-gate pass alone as proof generated code is safe. + + **Accepted v1 residual (in-pod file reads):** because generated code + legitimately needs the `open()` builtin for its input and output files, untrusted + calculation code can read a world-readable file inside its own sandbox pod and + return the contents to the submitter. The import allowlist is exactly `{json, math, + statistics, decimal, datetime, re, collections, itertools, functools, sys}` — + `pathlib`, `os`, `io`, `socket` and `urllib` are all outside it, so this is a single + read of a known path, not directory traversal and not an exfiltration path. This is contained — not by the AST gate, but + by the pod's layers 2–5: no secrets in the pod (not even the platform + `ENCRYPTION_KEY`), non-root on a read-only rootfs, default-deny egress (no + exfiltration), per-job isolation, and results returned only to the submitting + key — so the worst case is disclosure of a non-sensitive container file to the + customer who submitted the job. The deferred hardening for this is a + `runtimeClass: gvisor` sandbox. Consequently the sandbox pod's default-deny egress + and no-secrets posture are load-bearing and must not be relaxed. + +## 12. Environment reference + +All `AGENT_KV_*` settings (`backend/backend/settings/base.py`), each `os.environ.get` +with the default shown: + +| Variable | Default | Meaning | +|---|---|---| +| `AGENT_KV_PATH_PREFIX` | `agent-kv` | Top-level public URL prefix (whitelisted past tenant middleware). | +| `AGENT_KV_MAX_FILE_SIZE_MB` | `50` | Submit-time file size cap. | +| `AGENT_KV_MAX_PAGES` | `100` | Pre-OCR cap on the pages a request will PROCESS for PDFs — checked against `pages_selected` (after `page_start`/`page_end`), not the document's `pages_total` (§6.1). Excel has no pre-OCR page concept, so it takes the engine's post-OCR virtual-page cap instead; images are not accepted. | +| `AGENT_KV_MAX_CALCULATIONS_BYTES` | `20000` | Byte cap on the optional `calculations` field. | +| `AGENT_KV_MAX_SCHEMA_BYTES` | `262144` | Byte cap on the raw `keys` JSON document (256 KiB). | +| `AGENT_KV_RESULT_TTL_DAYS` | `7` | Result retention window (and a cancelled job's input, which rides the same TTL — a completed/failed job's input is deleted immediately at finalize instead), stamped at submit time (see [§6](#9-retention-and-ttl)). | +| `AGENT_KV_MAX_TIMEOUT_SECONDS` | `300` | Upper bound on the submit `timeout` (synchronous-wait) field. | +| `AGENT_KV_CONCURRENT_LIMIT` | `5` | Per-organization concurrent in-flight job cap (own Redis namespace). Fails **closed** (429) when Redis is unreachable — see `AGENT_KV_LIMITER_FAIL_OPEN`. | +| `AGENT_KV_KEY_RATE_LIMIT_PER_MINUTE` | `60` | Per-key request rate limit (submit + validate). Fails **closed** (429) when Redis is unreachable — see `AGENT_KV_LIMITER_FAIL_OPEN`. | +| `AGENT_KV_LIMITER_FAIL_OPEN` | `false` | What both limiters above do when Redis is unreachable. Default `false` = fail **closed**: the request is refused with a 429. Set `true` to restore the old fail-open behaviour, where an unreachable Redis removed the concurrency ceiling **and** the per-key rate ceiling at once while the API went on accepting billable work — a state that cannot be observed from the outside. A 429 is recoverable by the caller; an unbounded fan-out is not. | +| `AGENT_KV_SWEEP_GRACE_SECONDS` | `3600` | Age (from `created_at`) before a never-dispatched `PENDING` job is eligible for the sweep. | +| `AGENT_KV_STUCK_JOB_GRACE_SECONDS` | `21600` | Age (from `dispatched_at`) before a `DISPATCHED`/`RUNNING` job is eligible for the sweep's stuck-job phase — force-failed as `"Job timed out"` (6 hours). | +| `AGENT_KV_CALCULATIONS_ENABLED` | `false` | Gates the submit `calculations` field; the engine cannot execute it yet, so submit 400s while this is off. | +| `AGENT_KV_STRUCTURED_OUTPUT_ENABLED` | `false` | Gates the submit `structured_output` field; the engine cannot execute it yet, so submit 400s while this is off. | +| `AGENT_KV_STORAGE_DIR_PREFIX` | `unstract/agent_kv` | Bucket-rooted object-store root for staged inputs/results (`{prefix}/{org_id}/{job_id}/…`), mirroring `WORKFLOW_EXECUTION_DIR_PREFIX`/`API_EXECUTION_DIR_PREFIX`. The **first segment is the bucket** and must already exist. The cloud executor reads the same variable and keys its OCR cache under `{prefix}/{org_id}/cache/…`, so both fleets must agree. Trailing `/` is stripped. | +| `AGENT_KV_FILE_STORAGE_CREDENTIALS` | *(none — must be set)* | JSON credentials for the `AGENT_KV` `FileStorageType` (`unstract/filesystem`), same shape as `WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS`: `{"provider": "minio", "credentials": {"endpoint_url": "...", "key": "...", "secret": "..."}}`. | + +`AGENT_KV_FILE_STORAGE_CREDENTIALS` is looked up via +`FILE_STORAGE_CREDENTIALS_TO_ENV_NAME_MAPPING[FileStorageType.AGENT_KV]` +(`unstract/filesystem/src/unstract/filesystem/file_storage_config.py`) and is required +for `stage_input`/`write_result`/`read_result`/`delete_job_files`/`delete_input` +(`backend/agent_kv/storage.py`) to have anywhere to write to. diff --git a/docs/superpowers/plans/2026-10-06-table-extractor-api-carveout.md b/docs/superpowers/plans/2026-10-06-table-extractor-api-carveout.md new file mode 100644 index 0000000000..78cff019db --- /dev/null +++ b/docs/superpowers/plans/2026-10-06-table-extractor-api-carveout.md @@ -0,0 +1,358 @@ +# Table Extractor on the Agent-KV API — Standalone Release Carve-Out + +**Goal:** ship `POST /agent-kv/` with `extractors: [{"name": "table", "keys": {"target_table": "…"}}]` +to customers as **its own OSS PR + cloud PR off current `main`**, excluding the +`agentic_kv` extraction engine, the schema codegen path and the hardened sandbox — +all of which stay on the large PRs ([unstract#2309](https://github.com/Zipstack/unstract/pull/2309), +[unstract-cloud#1816](https://github.com/Zipstack/unstract-cloud/pull/1816)) until those stabilise. + +**Non-goal:** anything that makes the `kv` extractor reachable. This plan +deliberately leaves the KV code in the tree, dormant and tested, so the large PRs +re-enable it by reverting four one-line subtractions rather than by merging content. + +**Ticket:** [UN-4232](https://zipstack.atlassian.net/browse/UN-4232), under epic [UN-4044](https://zipstack.atlassian.net/browse/UN-4044). +**Source branches to copy from:** OSS `Feat/agent-kv-api`, cloud `UN-4044-agent-kv-cloud-executor`. +**Target branches:** OSS `feat/agent-kv-api-table` off `main`; cloud `UN-4232-agent-kv-api-table` off `main`. +**Strategy: copy-forward, not cherry-pick.** The source branches carry 100+ interleaved +commits across 338 files; reconstruct the subset by copying files, then subtracting. + +--- + +## The constraint that defines the carve-out + +`workers/plugins/agentic_table/src/api_binding.py:29-32`: + +```python +from agentic_kv.constants import AgentKVConfig +from agentic_kv.exceptions import PageCapExceeded +from agentic_kv.llm_adapter import build_llm, flush +from agentic_kv.progress import JobCancelled, StageReporter +``` + +`agentic_table/pyproject.toml` declares `unstract-agentic-table-engine` and +`unstract-extraction-seams` — **not** `agentic_kv`. The import resolves only +because both plugins are installed into the same worker image. Drop the +`agentic_kv` plugin and the table API path fails at import. + +The module's own docstring already names this ("**Known debt:** … Spec step 4 +moves all three behind the seams package"). So **Task 1 is spec §5 step 4, +pulled forward** — not throwaway work, just work done in the order the carve-out forces. + +--- + +## Task 1 — Break `agentic_table`'s dependency on `agentic_kv` (cloud) + +Move into `workers/plugins/extraction_seams/`, keeping **every env var name byte-identical** +so no chart, secret or external-secret wiring changes: + +| New seams module | Moved from | Notes | +|---|---|---| +| `src/config.py` | `agentic_kv/src/constants.py` (`AgentKVConfig`, 248 L) | Keep the `AGENT_KV_*` env names. The table path reads `advanced_model`, `lite_model`, `llmwhisperer_api_key`, `llmwhisperer_base_url`, `parallel_pages` | +| `src/llm.py` | `agentic_kv/src/llm_adapter.py` (`build_llm`, `flush`, 502 L) + `engine/llm_client_types.py` (41 L) | `llm_client_types` holds 4 small dataclasses/exceptions; it comes along | +| `src/progress.py` | `agentic_kv/src/progress.py` (`StageReporter`, `JobCancelled`, 187 L) | **Drop `ENGINE_NODE_TO_STAGE`** — it maps KV engine nodes to KV stages and means nothing to the table path | +| `src/exceptions.py` | `agentic_kv/src/exceptions.py` (`PageCapExceeded`, `ConfigError`) | Only the two the table path raises | + +Then: +- `agentic_table/src/api_binding.py` imports from `extraction_seams` — the module-scope + seam pattern is preserved, so `tests/test_executor_api_operation.py` still patches + `api_binding`'s own attributes; re-point the patch targets only. +- `agentic_table/pyproject.toml` is already correct — it never declared `agentic_kv`. +- `extraction_seams` keeps its one-way rule: nothing in it may import `agentic_table`, + `agentic_table_engine` or `agentic_kv`. + +**Cost to the large cloud PR on rebase:** 12 files inside `agentic_kv` import these +four modules; they re-point to `extraction_seams`. Mechanical, and it is step 4's +work regardless of this plan. + +--- + +## Task 2 — Stand up the API app without the KV extractor (OSS) + +**Copy from `Feat/agent-kv-api`:** + +``` +backend/agent_kv/ (the API app; 6.5k L incl. tests) +backend/backend/{base_urls,internal_base_urls,urls_v2,settings/base}.py +backend/{pyproject.toml,sample.env,uv.lock} +unstract/filesystem/ (the AGENT_KV storage type; 4 files) +unstract/agent-kv-schema/ (keep — see §Decision 3) +workers/ide_callback/ (agent_kv_callback queue; 5 files) +workers/{run-worker.sh,run-worker-docker.sh,sample.env,pyproject.toml} +docker/{docker-compose.yaml,sample.env,dockerfiles/worker-unified.Dockerfile} +tests/{compose/docker-compose.test.yaml,groups.yaml}, tox.ini, pyproject.toml, uv.lock +docs/agent-kv-api.md +``` + +**Then subtract — four lines, each independently revertible by the large PR:** + +1. **`backend/agent_kv/constants.py` — drop the `kv` entry from `EXTRACTOR_ROUTES`.** + This is the single most important line in the carve-out. `SUPPORTED_EXTRACTORS = + tuple(EXTRACTOR_ROUTES)` (`execution_serializers.py:33`), so one deletion makes a + `kv` submit return **400 `unknown extractor 'kv'; supported: ['table']`**. + + Without it, a `kv` submit is accepted with a 202 and dispatched to + `celery_executor_agentic_kv`, which has no consumer in this deployment — the job + sits forever, durably and silently. That is the exact failure mode that already + cost the team ~30 hours of firings with zero executions. + + **Leave `V1_EXTRACTOR_NAME`, `STAGE_NAMES`, `KVOptionsSerializer` and + `STAGE_NAMES_BY_EXTRACTOR["kv"]` in place.** Dormant and still under test; + `execution_views.py:45-56` already tolerates a job row naming an extractor this + build cannot describe. + +2. **`workers/run-worker.sh:76`** — remove `celery_executor_agentic_kv` from the + executor role's queue list; same in `docker/docker-compose.yaml`. + +3. **`backend/agent_kv/execution_urls.py`** — unregister the `/validate` route. + `ValidateView` compiles a `kv` keys schema via `compile_schema` + (`execution_views.py:255`) and nothing else; the table extractor's `keys` is + `{"target_table": …}`, validated by `TableKeysSerializer`. Shipping an endpoint + that validates schemas for an extractor the deployment refuses is an incoherent + public contract. Keep the view class and the `unstract/agent-kv-schema` package — + deleting them would force a content merge in files the large PR rewrites. + +4. **Drop `workers/sandbox/`** (3 files), the `worker-sandbox` compose service and + `PG_ROLE_SANDBOX` in `run-worker.sh`. See §Decision 2. + +--- + +## Task 3 — Test coverage: API and end to end (OSS) + +Explicit scope, not a trailing chore. The carve-out's whole claim is "this is the +stable subset", and the tests are what substantiate it. Two tiers, because they +fail for different reasons and run in different places. + +### 3a · API testing — contract level, no model calls + +`backend/agent_kv/tests/`, in CI on every PR, no LLM or OCR credentials. Much of +this already exists (6.5k lines, largely extractor-agnostic: auth, rate limiter, +storage, sweeps, internal views, job views) and `test_table_extractor_routing.py` +(252 L) covers the table route. Audit the submit-serializer and dispatch suites for +`kv` assumptions, then close these gaps: + +| Area | Assertions | +|---|---| +| **Routing** | `SUPPORTED_EXTRACTORS == ("table",)` · a `kv` submit returns **400, not 202** | +| **Submit validation** | missing `target_table` → 400 · a `kv` option on a table entry → 400 · unknown extractor → 400 · unknown keys **or** unknown options → 400 · more than one extractor entry → 400 · oversized `extractors` payload → 400 | +| **Auth and tenancy** | no key → 403 · revoked or wrong-org key → 403 · another org's job → **404, not 403** (no existence disclosure) | +| **Limits** | per-key rate limit → 429 · concurrency cap → 429 · a cancelled job releases its slot | +| **Lifecycle** | status shape for queued / running / completed / failed / cancelled · `extractors.table.stages` non-empty and containing only `table_extraction` · result byte-for-byte stable on re-read · DELETE then GET result → 404 · swept or TTL-expired job → 404 | +| **Dispatch** | `table` routes to `("agentic_table", "table_extract_api")` · every supported extractor has a route, a stage list **and** an options serializer — the three tables stay in step | +| **Internal endpoints** | stage report · finalize · sweep · TTL cleanup, including the cleanup-failure path | +| **Webhook** | delivered once on completion · a delivery failure does not change job state | +| **Subscription gate** | expired trial and inactive subscription both 402, byte-identical to the API-deployment path · a database error surfaces as 500, never as unmetered access | + +The routing row matters most. It is the guard on Task 2's first subtraction, and the +only thing standing between a future rebase and a `kv` submit that returns 202 and +then sits in an unconsumed queue forever. + +The unknown-keys/unknown-options row is not padding either: DRF silently discards +unrecognised keys, so `"Options"` for `"options"` would be dropped whole, `options` +would default to `{}`, and the job would run with roughly double the LLM spend the +caller asked for — with a 202 and no indication anything was ignored. + +### 3b · End-to-end testing — deployed stack, real LLM and OCR + +`tests/e2e/agent_kv/`, run first against compose, then against the dev namespace +(Task 5b). These are the tests that catch the wiring failures nobody has had a +chance to hit yet. + +| Scenario | What it proves | +|---|---| +| **Table happy path** — `rent_roll.pdf` | submit → poll → rows match the known table | +| **Excel path** — `invoice.xlsx` | the engine's own Excel branch extracts, and UN-4219's post-OCR cap applies | +| **Cancel mid-run** | the run stops, billing stops, the concurrency slot is released | +| **Page cap** | an oversized document is rejected before spend | +| **Sync-wait submit** | the result comes back inline | +| **Webhook on completion** | delivered with the right payload | +| **Unreadable PDF** | 400 at submit | +| **Metering** | page usage on the standard `Audit()` path, token and cost sums in `usage_summary` — the same trail as the IDE path | +| **Stages non-empty** | `GET /agent-kv/{job}` returns `stages: [{"name": "table_extraction", …}]` | + +That last row is load-bearing: it is the only check that catches Task 5a's cross-repo +constant drift **from the outside** — from a client's point of view, where the +symptom is a job that completes normally while reporting no progress at all. + +### 3c · Re-pointing the existing scenarios + +| Keep, re-pointed to `table_keys()` | Drop (KV-engine specific) | +|---|---| +| `test_submit_without_key_is_403` | `test_validate_good_and_bad_schema` (route unregistered) | +| `test_cancel_mid_run` | `test_happy_path_extraction` | +| `test_cancelled_job_does_not_leak_its_concurrency_slot` | `test_calculation_happy_path` | +| `test_delete_completed_job_then_result_404` | `test_hostile_calculation_fails_user_safely` | +| `test_page_cap_rejects_oversized_document` | `test_resubmit_same_document_hits_document_cache` (KV's cached `DocumentProcessor`; the table path is uncached until spec step 5) | +| `test_sync_wait_submit_returns_result_inline` | `test_bad_llm_key_ends_failed` (re-point if cheap) | +| `test_webhook_delivered_on_completion` | | +| `test_concurrency_limit_returns_429` | | +| `test_submit_unreadable_pdf_is_400` | | +| `test_excel_submit_extracts` — **keep, re-pointed.** The table engine has its own Excel branch and UN-4219's post-OCR page cap applies there | | + +Already table-native, keep as-is: `test_table_extractor_happy_path`, +`test_table_entry_without_a_target_table_is_400`, `test_a_kv_option_on_a_table_entry_is_400`, +`test_an_unknown_extractor_is_400`. + +Fixtures: keep `rent_roll.pdf` and `invoice.xlsx`; `invoice.pdf` only if a re-pointed +scenario still uses it. + +### 3d · Where each tier runs + +- **3a** — the unit and integration groups in `tests/groups.yaml`, every PR, no credentials. +- **3b** — the e2e group against compose on the PR, and against the dev namespace once + deployed (Task 5b). Gate the model-calling scenarios on the existing `require_llm` + fixture so a credential-less run skips rather than fails. +- **Cloud** — `tests/groups.cloud.yaml` carries `agentic_table`, `agentic_table_engine` + and `extraction_seams`. These go **red by design** until the OSS PR lands, exactly as + the `agentic_kv` groups do today. Say so in the PR description so a reviewer does not + read it as breakage. + +--- + +## Task 4 — The cloud half, minus the engine and the sandbox + +**Include:** + +``` +workers/plugins/agentic_table/ (9 files) +workers/plugins/agentic_table_engine/ (77 files) +workers/plugins/extraction_seams/ (5 files + Task 1's four new modules) +backend/plugins/agent_kv/ (5 files — the subscription gate) +charts/unstract-platform/templates/shared/agent-kv-secret.yaml +charts/unstract-platform/templates/shared/storage-secret.yaml +charts/unstract-platform/templates/backend/agent-kv-cronjobs.yaml +charts/unstract-platform/templates/external-secrets/external-secrets.yaml +charts/unstract-platform/templates/_helpers.tpl +charts/unstract-platform/templates/worker-v2/configmaps-specific.yaml +charts/unstract-platform/{values.yaml,unittests/agent_kv_cronjobs_test.yaml, + unittests/agent_kv_wiring_test.yaml,unittests/pg_worker_fleet_test.yaml, + unittests/render_guards_test.yaml} +charts/cloud-deployment-values/cloud.values.yaml +tests/groups.cloud.yaml, copy_cloud_deps.py (+ its test), .github/workflows/… +``` + +**Exclude:** all 97 `workers/plugins/agentic_kv/` files · `charts/unstract-platform/templates/worker-sandbox/` (5 files) · `unittests/sandbox_wiring_test.yaml` · `unittests/agent_kv_calculations_ga_test.yaml`. + +**Chart edits, not deletions:** +- `values.yaml` keeps the whole `AGENT_KV_*` block — the table path reads the same + env through Task 1's seams config — but drops the `agentic_kv` executor fleet entry + and the sandbox fleet. +- **`pg_worker_fleet_test.yaml` must be edited to assert the fleet actually shipped.** + Left untouched it either fails the build or, worse, passes while guarding a fleet + that no longer exists. The guard is the only thing standing between this deployment + and an unconsumed queue; a guard that guards nothing is worse than no guard. +- `agent_kv_wiring_test.yaml`: the `AKV-Q5` assertion on + `celery_executor_agentic_kv` comes out with the fleet entry. + +No new worker fleet is needed: `celery_executor_agentic_table` is already wired in +production for the IDE table path, which is precisely why the table work chose a +second *operation* over a new executor name (plan ruling R1). + +--- + +## Task 5 — The cross-repo constant, and a real deploy + +**5a. Pin the silent-drift constant.** `STAGE_TABLE_EXTRACTION = "table_extraction"` +(cloud `api_binding.py`) must equal `TABLE_STAGE_NAMES[0]` (OSS +`backend/agent_kv/constants.py:41`). Today **only a pair of comments ties them**. +On drift: `StageReportView` persists whatever the executor sends, `_status_document` +filters it through the OSS list, and every status response returns an empty `stages` +array while jobs complete normally. No test on either side can catch it alone. + +For a first customer-facing release this is the bug that reads as "your API is +broken". Cheapest fix: assert the literal on both sides, each test naming the other +repo's file and value, so a change to one fails the other's suite on the next sync. + +**5b. Deploy and exercise it.** Push to the dev namespace and run the table e2e +against the deployed stack. The PG consumer wiring and the agent-kv CronJobs have +never run outside compose ([UN-4213](https://zipstack.atlassian.net/browse/UN-4213)). +Dropping the sandbox removes the hardest half of that validation — no hardened pod, +no NetworkPolicy, no least-privilege DB role to prove out. + +**Exit criteria for the whole carve-out:** +- a `table` submit against the deployed stack returns rows for `rent_roll.pdf` +- `GET /agent-kv/{job}` reports a non-empty `stages: [{"name": "table_extraction", …}]` +- a `kv` submit returns 400, not 202 +- cancel mid-run stops the run and releases the concurrency slot +- the xlsx path extracts and is capped +- page usage and token/cost land on the usual `Audit()` path +- the full API suite (§3a) is green in CI without credentials +- `helm unittest` green with the reduced fleet, and `pg_worker_fleet_test` + demonstrably **fails** when a fleet entry is removed +- OSS PR merged before the cloud PR — the cloud table groups depend on the OSS routing + +--- + +## Task 6 — Documentation and the public contract + +- `docs/agent-kv-api.md:199` currently reads "Which extractor: `kv` or `table`". + It must say `table` is the only supported value on this deployment, and the + `/validate` section comes out with the route. +- Customer-facing docs in `unstract-docs` (Docusaurus). +- The `202` → poll `GET /{job}` → `GET /{job}/result` flow, the `target_table` + parameter, the `extractors.table` result shape, and the webhook. + +--- + +## Decisions + +**1 · Quota admission — deferred. Decided 2026-10-06.** +The entitlement gate ships: `backend/plugins/agent_kv/service.py` calls the same +`get_subscription` / `verify_subscription` the API-deployment middleware calls, so an +expired trial or inactive subscription gets a byte-identical 402. What is **not** +wired is quota — an org inside a valid subscription but over its page or token +allowance runs and bills anyway ([UN-4216](https://zipstack.atlassian.net/browse/UN-4216), +tagged a GA blocker). This release ships without it; revisit before open availability. + +**2 · The sandbox is excluded — decided 2026-10-06. Ship what exists.** +`agentic_table` generates Python with an LLM and executes it on **every run** +(`runner.py:1158`, `1340`) through its own `code_executor.py`, where an AST gate +blocking four builtins is the entire boundary: no rlimits, no isolated pod, no +NetworkPolicy. That pod holds LLM keys, OCR adapters and storage credentials and has +full egress. **This is exactly what the IDE table path does in production today**, so +the carve-out introduces no new class of risk — but it does widen reach from +"Prompt Studio users in the org" to "any holder of an API key". +Shipping as-is, with [UN-4215](https://zipstack.atlassian.net/browse/UN-4215) as the +immediate fast-follow rather than a backlog item. Pulling the sandbox in now +would also force the module-allowlist reconciliation (the table prompt permits 15 +modules, the gate allows 10; `copy`, `pathlib`, `string`, `textwrap`, `time`, +`unicodedata` are rejected) — which is the thing most likely to destabilise the +table path, i.e. the opposite of what this carve-out is for. + +**3 · Keep `unstract/agent-kv-schema` and `ValidateView`, unregister the route.** +The package is zero-dependency, self-contained and fully tested, and +`execution_serializers.py` imports `compile_schema` at module scope. Deleting it +forces a content merge in the file the large PR rewrites most; keeping it costs +twelve dormant files and leaves the large PR a one-line route restoration. + +**4 · Keep the URL `/agent-kv/`.** It reads oddly for a table-only API, but the first +customers' integrations are the thing a later rename breaks. Name it the agentic +extraction API in the docs and leave the path alone. + +**5 · One extractor per job stays.** `execution_serializers.py:283` rejects more than +one entry. That is exactly right for table and needs no work here. + +--- + +## Risks + +| Risk | Why it bites | Mitigation | +|---|---|---| +| Rebase cost on the large PRs | The carve-out touches `constants.py`, `execution_serializers.py`, `execution_urls.py`, `run-worker.sh`, `values.yaml` — files both large PRs own | Every subtraction is the deletion of a line whose code still exists. The large PRs re-add four lines and re-point 12 imports; no content merge | +| Temporary duplication in `extraction_seams` | Task 1 moves four modules out of `agentic_kv`, which the large cloud PR still carries copies of | Expected and bounded: the large PR's rebase re-points its 12 importers, completing spec §5 step 4 | +| A helm guard edited into uselessness | `pg_worker_fleet_test` is the only thing catching an unconsumed queue | Edit the assertions to the shipped fleet and verify the test **fails** when a fleet entry is removed | +| Uncached OCR on the table API path | Ruling R3: one LLMWhisperer pass per run until spec step 5 / [UN-4096](https://zipstack.atlassian.net/browse/UN-4096) | Cost, not correctness. Size it against expected customer volume before launch | +| Table reports one coarse stage | The engine has no node-level hooks, so `stages` is `["table_extraction"]` | Known and documented; do not invent stage names the executor cannot substantiate | + +--- + +## Sequencing + +Task 1 gates Task 4 (the cloud PR cannot build without it). Tasks 2 and 3 are OSS and +run in parallel with Task 1. Task 4 needs Task 1; Task 5 needs both PRs stood up; +Task 6 runs alongside from the start. OSS merges before cloud — the cloud table test +groups depend on the OSS routing. + +No week numbers: size this on a walkthrough. The honest shape is that Task 1 is the +only genuinely new engineering, Tasks 2–4 are subtraction and re-wiring of code that +already exists and already passes, and Task 5b is the validation nobody has done yet +for any of it. diff --git a/pyproject.toml b/pyproject.toml index d46632ef40..6019f21a17 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ test-rig = [ ] [tool.uv.sources] +unstract-agent-kv-schema = { path = "./unstract/agent-kv-schema", editable = true } unstract-filesystem = { path = "./unstract/filesystem", editable = true } unstract-workflow-execution = { path = "./unstract/workflow-execution", editable = true } unstract-tool-sandbox = { path = "./unstract/tool-sandbox", editable = true } diff --git a/tests/compose/docker-compose.test.yaml b/tests/compose/docker-compose.test.yaml index 742d6b2b5e..7dcce8dd00 100644 --- a/tests/compose/docker-compose.test.yaml +++ b/tests/compose/docker-compose.test.yaml @@ -10,6 +10,18 @@ services: # Workers ask the backend for this, so this is the value that normally # takes effect for fan-out. - MAX_PARALLEL_FILE_BATCHES=${MAX_PARALLEL_FILE_BATCHES:-3} + # Lets the agent-kv e2e lane's page-cap scenario set a cap smaller than + # its 2-page fixture (default matches AGENT_KV_MAX_PAGES's own default + # in docker/sample.env; the test skips unless this is overridden <=2). + - AGENT_KV_MAX_PAGES=${AGENT_KV_MAX_PAGES:-100} + # Bucket-rooted object-store root for staged inputs/results; the first + # segment is the bucket and must exist. Must be identical to the value + # given to worker-executor-v2 below -- the executor reads back what the + # backend staged and caches OCR under the same root. + - AGENT_KV_STORAGE_DIR_PREFIX=${AGENT_KV_STORAGE_DIR_PREFIX:-unstract/agent_kv} + # Flips on the Agent-KV calc e2e lane (Task 12); unset/empty (the + # default) keeps it off so other lanes are unaffected. + - AGENT_KV_CALCULATIONS_ENABLED=${AGENT_KV_CALCULATIONS_ENABLED:-} platform-service: image: unstract/platform-service:${UNSTRACT_TEST_VERSION:-latest} @@ -28,6 +40,32 @@ services: worker-pg-executor: environment: - UNSTRACT_LLM_MOCK_RESPONSE=${UNSTRACT_LLM_MOCK_RESPONSE:-} + # Add the agent-kv executor's queue to the default set (docs/agent-kv-api.md + # §12 item 2) without dropping the queues other e2e lanes' tools dispatch + # to on this same worker -- this is a superset of docker-compose.yaml's + # own default, not a replacement for it. + # + # This MUST be WORKER_PG_QUEUE_CONSUMER_QUEUE, not CELERY_QUEUES_EXECUTOR: + # since UN-4046 made the PG transport unconditional this service runs + # `pg-queue-consumer`, which reads WORKER_PG_QUEUE_CONSUMER_QUEUE. + # CELERY_QUEUES_EXECUTOR is read only by workers/executor/worker.py -- the + # now-disabled Celery executor -- so setting it here drained nothing and + # the agent-kv lane hung in DISPATCHED with no error at the producer. + - WORKER_PG_QUEUE_CONSUMER_QUEUE=${WORKER_PG_QUEUE_CONSUMER_QUEUE:-celery_executor_legacy,celery_executor_agentic,celery_executor_table,celery_executor_smart_table,celery_executor_simple_prompt_studio,celery_executor_agentic_table,celery_executor_lookup_test} + # Agent-KV engine config, read by the cloud `agentic_table` executor plugin + # (not by backend settings) -- docs/agent-kv-api.md §12 item 5 / §13. + # `AGENT_KV_LLM_PROVIDER` is a v1 hard allowlist of `anthropic`|`openai`. + # Model names have no safe default -- Task 13b must supply real ones. + - AGENT_KV_LLM_PROVIDER=${AGENT_KV_LLM_PROVIDER:-anthropic} + - AGENT_KV_LITE_MODEL=${AGENT_KV_LITE_MODEL:-} + - AGENT_KV_ADVANCED_MODEL=${AGENT_KV_ADVANCED_MODEL:-} + - AGENT_KV_LLM_API_KEY=${AGENT_KV_LLM_API_KEY:-} + - AGENT_KV_LLMWHISPERER_API_KEY=${AGENT_KV_LLMWHISPERER_API_KEY:-} + - AGENT_KV_LLMWHISPERER_BASE_URL=${AGENT_KV_LLMWHISPERER_BASE_URL:-https://llmwhisperer-api.us-central.unstract.com/api/v2} + # Same root the backend stages into (see `backend` above) -- the executor + # reads `input_ref` from it and keys its OCR cache under + # `{prefix}/{org_id}/cache/`. Both sides MUST agree. + - AGENT_KV_STORAGE_DIR_PREFIX=${AGENT_KV_STORAGE_DIR_PREFIX:-unstract/agent_kv} worker-pg-fileproc: environment: @@ -39,6 +77,15 @@ services: environment: - MAX_PARALLEL_FILE_BATCHES=${MAX_PARALLEL_FILE_BATCHES:-3} + # Re-anchored onto worker-pg-ide-callback: UN-4046 replaced the Celery + # worker-ide-callback service with its PG-transport twin. + # Completion-webhook e2e: waive the worker's SSRF guards (https + public + # host) so a receiver on the compose host is reachable. Empty/unset (the + # default) keeps the guards mandatory; the gated webhook scenario skips. + worker-pg-ide-callback: + environment: + - AGENT_KV_WEBHOOK_INSECURE_ALLOW_HTTP_PRIVATE=${AGENT_KV_WEBHOOK_INSECURE_ALLOW_HTTP_PRIVATE:-} + # The `ui` Playwright group drives a browser at the frontend's OWN origin, and # the app derives every API URL from window.location.origin. The production # nginx config serves only static files plus an SPA fallback, so `/api/v1/*` diff --git a/tests/e2e/agent_kv/__init__.py b/tests/e2e/agent_kv/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/e2e/agent_kv/conftest.py b/tests/e2e/agent_kv/conftest.py new file mode 100644 index 0000000000..392ccba914 --- /dev/null +++ b/tests/e2e/agent_kv/conftest.py @@ -0,0 +1,341 @@ +"""Fixtures and HTTP helpers shared by the Agent-KV e2e tests. + +This lane exercises the Agent-KV *product* API (``/agent-kv/...``, top-level, +API-key authed -- not the tenant-scoped platform surface most other e2e +lanes hit). It needs a deployment where the cloud ``agentic_kv`` executor +plugin is installed (``plugins.get_plugin("agent_kv")`` truthy) -- an +OSS-only build 501s every submit (see ``docs/agent-kv-api.md`` §11). The +module-level skip in ``test_agent_kv_e2e.py`` keeps this lane from running +by accident under a plain OSS e2e sweep; this conftest assumes that guard +already passed by the time any fixture here actually executes. +""" + +from __future__ import annotations + +import json +import os +import time +import uuid +from dataclasses import dataclass +from pathlib import Path + +import pytest +import requests + +from tests.e2e.conftest import _org_id +from tests.rig.runtime import PlatformEndpoints + +FIXTURES_DIR = Path(__file__).parent / "fixtures" +# A real, OCR-able 2-page invoice, generated for this lane with reportlab (the +# checked-in backend/agent_kv/tests/fixtures/two_page.pdf has no extractable +# text -- pdfplumber reads it as two blank pages -- so it can't stand in for +# an extraction fixture). It's a static binary checked into the repo; the e2e +# venv itself does not depend on reportlab to regenerate it at test time. +INVOICE_PDF = FIXTURES_DIR / "invoice.pdf" +# A small single-sheet invoice workbook (generated with openpyxl, checked in +# as a static binary like invoice.pdf): Excel exercises the no-pre-OCR-page- +# count path -- ``pages_total`` is None at submit and the engine enforces the +# post-OCR virtual-page cap instead. +INVOICE_XLSX = FIXTURES_DIR / "invoice.xlsx" + +# A real, OCR-able 2-page rent roll, generated for this lane with reportlab the +# same way invoice.pdf was, and checked in as a static binary (the e2e venv +# does not depend on reportlab to regenerate it at test time). The invoice +# fixture cannot stand in: it has no tabular region, so a table extraction +# against it would legitimately return no tables and the scenario would assert +# nothing. +RENT_ROLL_PDF = FIXTURES_DIR / "rent_roll.pdf" + +#: The table the rent-roll fixture contains, named the way the engine expects +#: (`target_table` is its one required extraction parameter). +RENT_ROLL_TABLE = "Rent roll" + + +def table_keys(target_table: str = RENT_ROLL_TABLE) -> dict: + """The `table` extractor's `keys` member. + + Every extractor entry carries `keys`; for the table extractor the thing + being asked for is a table, named by `target_table`. + """ + return {"target_table": target_table} + + +def table_result_body(resp: requests.Response) -> dict: + """The `table` extractor's own result out of an extractor-keyed payload.""" + return resp.json()["extractors"]["table"] + + +def table_stages(status_doc: dict) -> list[dict]: + """The `table` extractor's stage list out of an extractor-keyed status doc.""" + return status_doc["extractors"]["table"]["stages"] + + +# The 3 leaves the happy-path schema asks for; also used to build the schema. +INVOICE_FIELDS = ("invoice_number", "vendor_name", "total_amount") + +_DEFAULT_TIMEOUT = 30 + + +@dataclass(frozen=True) +class AgentKVAuth: + """A usable Agent-KV API key plus the backend root it's valid against.""" + + base: str # backend root, e.g. http://localhost:8000 (no trailing slash) + key: str # raw AgentKVKey UUID + org_id: str + + @property + def headers(self) -> dict[str, str]: + return {"Authorization": f"Bearer {self.key}"} + + @property + def exec_url(self) -> str: + return f"{self.base}/agent-kv/" + + def job_url(self, job_id: str, suffix: str = "") -> str: + return f"{self.base}/agent-kv/{job_id}{suffix}" + + +@pytest.fixture(scope="session") +def agent_kv_key( + platform: PlatformEndpoints, authed_session: requests.Session +) -> AgentKVAuth: + """Create a fresh AgentKVKey for this session via the tenant API. + + Session-scoped: one key, reused by every test in the lane (mirrors + ``api_deployment``'s ``api_deployment`` fixture). Key management sits + under the tenant-scoped platform surface + (``{base}/api/v1/unstract/{org_id}/agent-kv/keys/``), unlike the + execution endpoints this key then authenticates against. + """ + base = platform.backend_url.rstrip("/") + org_id = _org_id(authed_session, base) + prefix = f"{base}/api/v1/unstract/{org_id}" + name = f"e2e-agent-kv-{uuid.uuid4().hex[:8]}" + resp = authed_session.post( + f"{prefix}/agent-kv/keys/", + json={"name": name, "description": "agent-kv e2e lane"}, + timeout=_DEFAULT_TIMEOUT, + ) + assert resp.status_code == 201, f"create agent-kv key: {resp.text}" + body = resp.json() + return AgentKVAuth(base=base, key=body["key"], org_id=org_id) + + +#: Submit fields that describe the REQUEST rather than one extractor (§7.1). +#: The page range drives the shared OCR pass and the page cap, so it cannot be +#: per-extractor. +JOB_LEVEL_SUBMIT_FIELDS = frozenset( + {"page_start", "page_end", "timeout", "tags", "custom_data", "webhook_url"} +) + + +def kv_result_body(resp: requests.Response) -> dict: + """The `kv` extractor's own result out of an extractor-keyed payload (§7.3). + + The result endpoint returns `{"extractors": {...}, "usage_summary": {...}}`, + so a test asserting on the engine's own fields has to reach through one + level. Centralised here so the reach is spelled out once. + """ + return resp.json()["extractors"]["kv"] + + +def kv_stages(status_doc: dict) -> list[dict]: + """The `kv` extractor's stage list out of an extractor-keyed status doc. + + Stage reporting moved under `extractors.` (spec §7.2) because stage + names are extractor-specific -- `qa`/`challenge`/`codegen` mean nothing to a + Table Extractor. Centralised here so the reach is spelled out once. + """ + return status_doc["extractors"]["kv"]["stages"] + + +#: The platform adapters an Agent-KV `table` submit must name, read from the +#: environment because they are rows in the deployed stack's database -- an +#: e2e client cannot mint them, and they carry real provider credentials. +#: +#: Unset by default, which makes every submit a 400. The `require_adapters` +#: fixture is what turns that into an explicit skip instead of a confusing +#: failure; scenarios that expect a rejection BEFORE the adapter gate (bad +#: key, rate limit, absent `extractors`) do not need it. +#: +#: Create them once in the target org -- an LLM adapter, a second cheaper LLM +#: adapter, and an LLMWhisperer X2TEXT adapter -- and export their ids. +E2E_ADAPTER_ENV = { + "llm": "AGENT_KV_E2E_LLM_ADAPTER", + "lite_llm": "AGENT_KV_E2E_LITE_LLM_ADAPTER", + "x2text": "AGENT_KV_E2E_X2TEXT_ADAPTER", +} + + +def e2e_adapters() -> dict[str, str]: + """The configured adapter ids, by role; roles with no env var are omitted.""" + return { + role: os.environ[var] + for role, var in E2E_ADAPTER_ENV.items() + if os.environ.get(var) + } + + +def missing_adapter_env() -> list[str]: + """Which adapter env vars are unset, for a skip message that names them.""" + return [var for var in E2E_ADAPTER_ENV.values() if not os.environ.get(var)] + + +def submit_raw( + auth: AgentKVAuth, + file_bytes: bytes, + filename: str, + keys: dict | None, + *, + extractor: str = "table", + adapters: dict[str, str] | None = None, + **fields: object, +) -> requests.Response: + """POST a submit request and return the raw response -- no assertions. + + For tests that need to inspect a non-202 outcome (403, 400, 429, ...). + ``keys=None`` omits ``extractors`` entirely (covers submits that intend to + fail before schema validation is even reached). + + Builds the extractor-scoped wire format (spec §7.0): per-extractor schema + and knobs go inside ``extractors``; only fields describing the request stay + top level. Callers keep passing knobs as plain kwargs and this routes them, + so a test still reads as "submit with this one thing changed". + """ + data: dict[str, object] = {} + options: dict[str, object] = {} + for k, v in fields.items(): + (data if str(k) in JOB_LEVEL_SUBMIT_FIELDS else options)[str(k)] = v + if keys is None: + # `extractors` is omitted entirely, so there is nowhere for per-extractor + # options to go. Silently dropping them would make a test look like it + # exercised a knob it never sent -- and it would still pass, because + # these submits are expected to fail on the missing field anyway. + assert not options, ( + f"submit_raw(keys=None) cannot carry extractor options {sorted(options)}; " + f"pass a schema, or move the field to JOB_LEVEL_SUBMIT_FIELDS if it " + f"belongs there" + ) + else: + entry: dict[str, object] = { + "name": extractor, + "keys": keys, + "options": options, + } + # `adapters` is REQUIRED for `table` -- the caller names the LLM and + # OCR adapters the extraction runs on. Defaults to whatever the + # environment configured; `{}` is sent as-is so a scenario can assert + # the 400 for an absent block. + # + # When the env vars are UNSET this omits the block, and the backend + # then refuses the submit at the adapter field validator -- before any + # file or page-cap check. That is correct behaviour but it is a trap + # for tests asserting a DIFFERENT 400: they still see a 400, then fail + # on the blamed attribute (`extractors`, not `file`). Red, not skipped. + # Every scenario that asserts a specific rejection attribute therefore + # takes `require_adapters`, which turns the unconfigured lane into an + # explicit skip naming the missing vars. + resolved = e2e_adapters() if adapters is None else adapters + if resolved: + entry["adapters"] = resolved + data["extractors"] = json.dumps([entry]) + return requests.post( + auth.exec_url, + headers=auth.headers, + files={"file": (filename, file_bytes, "application/octet-stream")}, + data=data, + timeout=60, + ) + + +def submit( + auth: AgentKVAuth, + file_bytes: bytes, + filename: str, + keys: dict, + *, + extractor: str = "table", + adapters: dict[str, str] | None = None, + **fields: object, +) -> tuple[str, str]: + """POST a submit request expected to succeed; return (job_id, status_url). + + Asserts the 202 handshake (spec §7.1 / docs §3) so every caller only ever + polls a genuinely dispatched job. A 501 here almost always means the + cloud ``agentic_table`` executor plugin isn't installed on this deployment + -- this whole lane requires it (docs §11). A 400 naming `adapters` + means the `AGENT_KV_E2E_*_ADAPTER` ids are unset or do not belong to this + key's organization -- gate such a scenario on `require_adapters`. + """ + resp = submit_raw( + auth, file_bytes, filename, keys, extractor=extractor, adapters=adapters, **fields + ) + assert resp.status_code == 202, ( + f"submit: HTTP {resp.status_code} (expected 202; a 501 means the " + f"agentic_table executor plugin isn't installed on this deployment): {resp.text}" + ) + body = resp.json() + job_id = body["job_id"] + status_url = body["status_url"] + assert job_id, body + return job_id, status_url + + +TERMINAL_JOB_STATUSES = {"completed", "failed", "cancelled"} + + +def poll(auth: AgentKVAuth, job_id: str, timeout_s: float = 600) -> dict: + """Poll the status document until the job reaches a terminal state. + + Returns the last status document. Fails loudly (rather than looping + forever) once ``timeout_s`` elapses. + """ + deadline = time.monotonic() + timeout_s + last: dict = {} + while time.monotonic() < deadline: + resp = requests.get( + auth.job_url(job_id), headers=auth.headers, timeout=_DEFAULT_TIMEOUT + ) + assert resp.status_code == 200, f"status poll: {resp.status_code}: {resp.text}" + last = resp.json() + if last.get("status") in TERMINAL_JOB_STATUSES: + return last + time.sleep(2) + pytest.fail(f"job {job_id} not terminal within {timeout_s}s; last status: {last}") + + +def result(auth: AgentKVAuth, job_id: str) -> requests.Response: + """GET the result endpoint and return the raw response (any status).""" + return requests.get( + auth.job_url(job_id, "/result"), headers=auth.headers, timeout=_DEFAULT_TIMEOUT + ) + + +def cancel(auth: AgentKVAuth, job_id: str) -> requests.Response: + return requests.post( + auth.job_url(job_id, "/cancel"), headers=auth.headers, timeout=_DEFAULT_TIMEOUT + ) + + +def delete(auth: AgentKVAuth, job_id: str) -> requests.Response: + return requests.delete( + auth.job_url(job_id), headers=auth.headers, timeout=_DEFAULT_TIMEOUT + ) + + +def invoice_schema() -> dict: + """A 3-leaf schema matching ``fixtures/invoice.pdf``'s obvious fields.""" + return { + "invoice_number": {"description": "The invoice number", "required": True}, + "vendor_name": {"description": "The vendor or supplier name issuing the invoice"}, + "total_amount": { + "description": "The total amount due on the invoice", + "format": "currency", + }, + } + + +def invalid_schema() -> dict: + """A schema that fails to compile: a leaf missing the required 'description'.""" + return {"total": {"format": "currency"}} diff --git a/tests/e2e/agent_kv/fixtures/invoice.pdf b/tests/e2e/agent_kv/fixtures/invoice.pdf new file mode 100644 index 0000000000..a79fdf55c4 Binary files /dev/null and b/tests/e2e/agent_kv/fixtures/invoice.pdf differ diff --git a/tests/e2e/agent_kv/fixtures/invoice.xlsx b/tests/e2e/agent_kv/fixtures/invoice.xlsx new file mode 100644 index 0000000000..0e90596d5c Binary files /dev/null and b/tests/e2e/agent_kv/fixtures/invoice.xlsx differ diff --git a/tests/e2e/agent_kv/fixtures/rent_roll.pdf b/tests/e2e/agent_kv/fixtures/rent_roll.pdf new file mode 100644 index 0000000000..0e471c9ea4 Binary files /dev/null and b/tests/e2e/agent_kv/fixtures/rent_roll.pdf differ diff --git a/tests/e2e/agent_kv/test_agent_kv_e2e.py b/tests/e2e/agent_kv/test_agent_kv_e2e.py new file mode 100644 index 0000000000..8790f99843 --- /dev/null +++ b/tests/e2e/agent_kv/test_agent_kv_e2e.py @@ -0,0 +1,793 @@ +"""E2E: the Agent-KV product API against a real running platform. + +This lane needs the cloud ``agentic_kv`` executor plugin (docs/agent-kv-api.md +§11) -- an OSS-only deployment 501s every submit, so it is NOT part of a plain +OSS e2e sweep. The whole module skips unless both: + +- ``UNSTRACT_BACKEND_URL`` is set (the usual e2e platform-URL gate every + other lane in this tree relies on), AND +- ``AGENT_KV_E2E=1`` is set explicitly -- so this lane never runs by + accident just because a platform happens to be up; it must be opted into + for a run that actually ships the plugin. + +A second gate (the ``require_adapters`` fixture) applies to every scenario +whose submit must be ACCEPTED. `table` submits now name the caller's own +``llm`` / ``lite_llm`` / ``x2text`` platform adapters, and the API validates +each id against the Bearer key's organization and its adapter type -- so +without ``AGENT_KV_E2E_LLM_ADAPTER`` / ``_LITE_LLM_ADAPTER`` / +``_X2TEXT_ADAPTER`` every submit is a 400 before dispatch. Scenarios that +expect a rejection BEFORE that gate (a bad key, the rate limiter, an absent +``extractors``) do not take it. + +There is no third gate for "can a job actually reach COMPLETED". There used to +be (``require_llm``, on ``AGENT_KV_LLM_API_KEY`` or +``UNSTRACT_LLM_MOCK_RESPONSE``), and on the table path it was wrong: the LLM +credential lives on the ADAPTER now, not in the executor's environment, so a +correctly configured lane with adapters and no exported key SKIPPED the only +two scenarios that poll a job through to a genuine COMPLETED result -- the most +valuable assertions in the file. ``require_adapters`` is the whole gate: +configuring adapter ids is what declares this lane can run real extraction, and +(unless the adapters point at a mock) that it will spend real provider money. + +One scenario needs the opposite -- a job that dispatches and then FAILS inside +the engine -- and takes ``bad_llm_adapter``, the id of an adapter deliberately +carrying an invalid credential. That also moved: it was gated on +``AGENT_KV_E2E_BAD_KEY_JOB=1`` plus an invalid ``AGENT_KV_LLM_API_KEY``, a +variable this path does not read, so an operator who set the stack up exactly +as documented got a job that ran on the valid adapter and reached +``completed``. + +Values extracted under ``UNSTRACT_LLM_MOCK_RESPONSE`` are whatever the mock +config returns, not real answers -- assertions here deliberately check +*structure* (the requested keys are present) rather than field values, for +both the mock and a real-key run, to keep this lane's outcome independent of +model behavior. + +Test order matters: the concurrency scenario is deliberately defined LAST +(see its own docstring) and ``tests/groups.yaml``'s ``e2e-agent-kv`` entry +pins ``parallel: false`` for the same reason -- see that scenario for why. + +**A third, operator-driven gate: ``AGENT_KV_E2E_BAD_KEY_JOB=1``.** One +scenario here (``test_bad_llm_key_ends_failed``) needs a *dispatched* job to +end FAILED, which no request-time check can produce -- the only deterministic +way to get there is to run the platform with a deliberately invalid +``AGENT_KV_LLM_API_KEY`` so the executor's LLM calls all fail. That is a +whole-stack configuration, not something a test can arrange, so the scenario +skips unless an operator sets ``AGENT_KV_E2E_BAD_KEY_JOB=1`` to declare "this +stack is running with a bad Agent-KV LLM key on purpose". Note that every +other scenario in this module that needs a job to COMPLETE will (correctly) +fail on such a stack -- run this one on its own, e.g.:: + + AGENT_KV_E2E=1 AGENT_KV_E2E_BAD_KEY_JOB=1 pytest tests/e2e/agent_kv -k bad_llm_key + +What it asserts is the *shape* of the failure, not its text: the job reaches +``failed`` and the ``error`` the customer sees is user-safe -- no filesystem +paths, no ``Permission denied``, no provider response body (spec Sec 8's +``" failed: "`` contract). +""" + +from __future__ import annotations + +import concurrent.futures +import json +import os +import time + +import pytest +import requests + +from tests.e2e.agent_kv.conftest import ( + INVOICE_PDF, + INVOICE_XLSX, + RENT_ROLL_PDF, + AgentKVAuth, + cancel, + delete, + e2e_adapters, + missing_adapter_env, + poll, + result, + submit, + submit_raw, + table_keys, + table_result_body, + table_stages, +) + +if not (os.environ.get("UNSTRACT_BACKEND_URL") and os.environ.get("AGENT_KV_E2E") == "1"): + pytest.skip( + "agent-kv e2e lane needs UNSTRACT_BACKEND_URL and AGENT_KV_E2E=1 " + "(this lane requires the cloud agentic_kv executor plugin -- see " + "docs/agent-kv-api.md §11)", + allow_module_level=True, + ) + +pytestmark = [pytest.mark.e2e] + +_INVOICE_BYTES = INVOICE_PDF.read_bytes() +_RENT_ROLL_BYTES = RENT_ROLL_PDF.read_bytes() + +# Bounded wait for the org's concurrency slots to free up again. Used by the +# cancel scenario (a cancelled job must not leak its slot) and by the +# concurrency scenario's own cleanup. +_SLOT_DRAIN_TIMEOUT_S = 120 +_SLOT_DRAIN_POLL_INTERVAL_S = 3 + + +def _validation_error_attrs(body: object) -> set[str]: + """The field names a 400 body blames, across BOTH error envelopes. + + This lane runs against whichever deployment is up, and the two differ: + + - OSS DRF default -- ``{"file": ["Unreadable PDF"]}``: the field names are + the top-level keys. + - Cloud (``drf-standardized-errors``) -- + ``{"type": "validation_error", "errors": [{"code": "invalid", + "detail": "Unreadable PDF", "attr": "file"}]}``: the field names are the + ``attr`` of each entry under ``errors``. + + Returns an empty set for anything that is neither shape, so a caller's + ``assert "file" in _validation_error_attrs(body), body`` fails loudly with + the real body rather than raising a ``KeyError``/``TypeError``. + """ + if not isinstance(body, dict): + return set() + errors = body.get("errors") + if isinstance(errors, list): + return { + entry["attr"] + for entry in errors + if isinstance(entry, dict) and isinstance(entry.get("attr"), str) + } + return {key for key in body if isinstance(key, str)} + + +def _cheap_submit(auth: AgentKVAuth) -> requests.Response: + """One raw submit with QA and challenge off -- the cheapest job this API + will accept, used purely to occupy/probe a concurrency slot. + """ + return submit_raw( + auth, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + ) + + +def _drain_until_slot_free( + auth: AgentKVAuth, timeout_s: float = _SLOT_DRAIN_TIMEOUT_S +) -> bool: + """Probe with cheap submits until one is ACCEPTED (proof a slot is free) + or the bounded wait runs out. Any accepted probe is cancelled again so it + doesn't hold the slot it just proved was available. + + Returns True if a probe was accepted, False if the window expired or the + probes stopped being answerable. Never raises -- callers decide whether a + False is an assertion failure (the cancel scenario) or merely a + best-effort cleanup giving up (the concurrency scenario). + """ + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + try: + probe = _cheap_submit(auth) + except Exception: # noqa: BLE001 - probing only + return False + if probe.status_code == 202: + try: + cancel(auth, probe.json()["job_id"]) + except Exception: # noqa: BLE001 - cleanup only + pass + return True + time.sleep(_SLOT_DRAIN_POLL_INTERVAL_S) + return False + + +@pytest.fixture +def require_adapters() -> None: + """Skip a scenario that needs a submit to be ACCEPTED (202). + + `table` submits must name the caller's own `llm` / `lite_llm` / `x2text` + platform adapters, and the API validates each id against the Bearer key's + organization and against its adapter type. Those are rows in the deployed + stack's database carrying real provider credentials -- an e2e client cannot + create them, so the ids come from the environment. + + Without them every submit is a 400 before dispatch, so this gate turns an + unconfigured lane into an explicit skip rather than a wall of confusing + failures. Scenarios that expect a rejection BEFORE the adapter gate -- a + bad key (403), the rate limiter (429), an absent `extractors` field -- + deliberately do not take this fixture. + + This replaces what `AGENT_KV_LLM_API_KEY` used to provide for the table + path: the credentials now live on the adapters, not in the executor's + environment. + """ + missing = missing_adapter_env() + if missing: + pytest.skip( + "needs the caller's platform adapters so a submit can be accepted; " + f"unset: {', '.join(missing)}. Create an LLM adapter, a cheaper " + "LLM adapter and an LLMWhisperer X2TEXT adapter in the target org " + "and export their ids." + ) + + +@pytest.fixture +def bad_llm_adapter() -> str: + """An adapter id whose provider credential is deliberately INVALID. + + Replaces the old `AGENT_KV_E2E_BAD_KEY_JOB=1` + invalid + `AGENT_KV_LLM_API_KEY` arrangement, which stopped working when the table + path moved onto caller-supplied adapters: that env var is not read on this + path at all, so the scenario it gated could no longer be produced. A test + that declared a precondition the stack could not satisfy -- and then ran + against a VALID credential, reaching `completed` and failing its own + `status == "failed"` assertion. + + The credential now lives on an adapter row, so the trigger has to as well. + Create one LLM adapter with a junk API key in the target org and export its + id; nothing else about the lane changes. + """ + adapter_id = os.environ.get("AGENT_KV_E2E_BAD_LLM_ADAPTER") + if not adapter_id: + pytest.skip( + "needs AGENT_KV_E2E_BAD_LLM_ADAPTER: the id of an LLM adapter in " + "this key's organization carrying a deliberately INVALID provider " + "credential, so a job dispatches and then fails inside the engine" + ) + return adapter_id + + +# --------------------------------------------------------------------------- +# 1. Auth +# --------------------------------------------------------------------------- + + +def test_submit_without_key_is_403(agent_kv_key: AgentKVAuth) -> None: + resp = requests.post( + agent_kv_key.exec_url, + files={"file": ("invoice.pdf", _INVOICE_BYTES, "application/pdf")}, + data={"keys": '{"total": {"description": "Grand total"}}'}, + timeout=30, + ) + assert resp.status_code == 403, resp.text + + +# --------------------------------------------------------------------------- +# 2. Validate +# --------------------------------------------------------------------------- + + +def test_table_extractor_happy_path( + require_adapters: None, agent_kv_key: AgentKVAuth +) -> None: + """A `table` entry runs the table engine and files its result under `table`. + + Before the job row recorded its extractor, every response keyed by the + hardcoded `kv` name -- a table job's output would have been filed under the + wrong extractor, which no unit test on either side would have caught. + """ + job_id, status_url = submit( + agent_kv_key, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + extractor="table", + ) + assert status_url.endswith(job_id), status_url + + status_doc = poll(agent_kv_key, job_id, timeout_s=600) + assert status_doc["status"] == "completed", status_doc + # Keyed by the extractor that ran, and by that one only. + assert set(status_doc["extractors"]) == {"table"}, status_doc + + # R7: the stage is recorded AND visible. StageReportView persists whatever + # name the executor sends, so a KV-only stage-name filter would leave this + # list empty while the job still completed. + stages = {s["name"]: s for s in table_stages(status_doc)} + assert "table_extraction" in stages, table_stages(status_doc) + assert stages["table_extraction"]["status"] == "done", stages + + resp = result(agent_kv_key, job_id) + assert resp.status_code == 200, resp.text + payload = resp.json() + assert payload["success"] is True, payload + assert payload["status"] == "completed", payload + + body = table_result_body(resp) + # Structure, not values: this lane runs under a mock LLM as well as a real + # key, so field values are not assertable (see the module docstring). + assert isinstance(body["tables"], list), body + assert "page_count" in body, body + assert "row_count" in body, body + + assert set(payload["usage_summary"]["by_extractor"]) == {"table"}, payload + + # Re-readable until TTL (D11), same contract as the kv path. + again = result(agent_kv_key, job_id) + assert again.status_code == 200, again.text + assert again.json() == payload, "result must be stable on re-read" + + +def test_table_entry_without_a_target_table_is_400(agent_kv_key: AgentKVAuth) -> None: + """Rejected at submit, before any paid work -- `target_table` is the + engine's one required extraction parameter. + """ + resp = submit_raw( + agent_kv_key, _RENT_ROLL_BYTES, "rent_roll.pdf", {}, extractor="table" + ) + assert resp.status_code == 400, resp.text + assert "target_table" in resp.text, resp.text + + +def test_a_kv_option_on_a_table_entry_is_400(agent_kv_key: AgentKVAuth) -> None: + """Per-extractor options exist so an option aimed at the wrong extractor + cannot be silently dropped and change what the caller is billed for. + """ + resp = submit_raw( + agent_kv_key, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + extractor="table", + qa=True, + ) + assert resp.status_code == 400, resp.text + assert "qa" in resp.text, resp.text + + +def test_an_unknown_extractor_is_400(agent_kv_key: AgentKVAuth) -> None: + resp = submit_raw( + agent_kv_key, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + extractor="nonexistent", + ) + assert resp.status_code == 400, resp.text + + +# --------------------------------------------------------------------------- +# 4. Failure body (submit-time rejection) +# --------------------------------------------------------------------------- + + +def test_submit_unreadable_pdf_is_400( + require_adapters: None, agent_kv_key: AgentKVAuth +) -> None: + """A ``.pdf``-named file that isn't a real PDF is rejected before dispatch. + + Takes ``require_adapters`` even though it asserts a REJECTION. The claim + that "scenarios expecting a rejection before the adapter gate keep working + unchanged" holds only for the three raised inside ``validate_extractors``; + this rejection comes from ``SubmitSerializer.validate()``, and DRF raises + field-level errors from ``to_internal_value`` BEFORE ``validate()`` runs -- + with ``_validated_adapter_shape`` called from the field validator. So with + the adapter vars unset (the documented default) the status assert still + passed on a 400 and the attr assert then FAILED, blaming ``extractors`` + instead of ``file``. Red, not skipped. + + This is the submit-serializer's own page-count check + (``pdfplumber.open()`` raising -> ``{"file": "Unreadable PDF"}``), which + is deterministic and needs no LLM/engine work to reach. The executor's + own FAILED-terminal path (a job that *dispatches* successfully but the + engine later fails) has no deterministic trigger reachable without a + real key/schema combination that's guaranteed to fail -- that path is + covered by the unit suites instead (e.g. + ``backend/agent_kv/tests/test_job_views.py``, + ``backend/agent_kv/tests/test_dispatch.py``). + """ + garbage = b"this is not a pdf file, just garbage bytes" * 20 + resp = submit_raw(agent_kv_key, garbage, "garbage.pdf", table_keys()) + assert resp.status_code == 400, ( + f"expected 400 (a 501 here means the agent-kv engine plugin isn't " + f"installed on this deployment): {resp.text}" + ) + body = resp.json() + assert "file" in _validation_error_attrs(body), body + + +# --------------------------------------------------------------------------- +# 5. Cancel +# --------------------------------------------------------------------------- + + +def test_cancel_mid_run(require_adapters: None, agent_kv_key: AgentKVAuth) -> None: + job_id, _ = submit( + agent_kv_key, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + ) + cancel_resp = cancel(agent_kv_key, job_id) + # Either cancel won the race (200) or the job was already terminal by the + # time it landed (409) -- both are correct per docs §7; either way the + # status endpoint must agree with what cancel just reported. + assert cancel_resp.status_code in (200, 409), cancel_resp.text + cancel_body = cancel_resp.json() + + status_resp = requests.get( + agent_kv_key.job_url(job_id), headers=agent_kv_key.headers, timeout=30 + ) + assert status_resp.status_code == 200, status_resp.text + final_status = status_resp.json()["status"] # lowercased, per docs §4 + + if cancel_resp.status_code == 200: + assert cancel_body == {"status": "cancelled"}, cancel_body + assert final_status == "cancelled", status_resp.json() + else: + # Lowercased, like every other status field in this API. The docs and + # this assertion both used to claim the 409 body carried the RAW + # uppercase enum -- it does not: `JobCancelView` returns + # `job.status.lower()`, and the unit test the docs cited as proof + # (`test_cancel_on_completed_is_409_and_result_untouched`) asserts + # `{"status": "completed"}`. This lane had never run, so the wrong + # assertion was never executed. (2.6 in the branch review.) + assert cancel_body["status"] == final_status, (cancel_body, final_status) + + +def test_cancelled_job_does_not_leak_its_concurrency_slot( + require_adapters: None, + agent_kv_key: AgentKVAuth, +) -> None: + """A cancelled job must give its concurrency slot back. + + ``JobCancelView`` marks the job ``CANCELLED`` directly and never touches + ``AgentKVConcurrencyLimiter``; the slot is released by the job's own + *finalize* call, in a ``finally`` (docs/agent-kv-api.md §11a). So the + release is asynchronous -- it lands whenever the executor that picked the + job up notices the cancellation and calls back -- which is exactly the + kind of path where a leak hides: nothing in the cancel request itself + would fail if the release were dropped, and the only other backstop is + the limiter's 6-hour slot TTL. + + This is a bounded probe, not a saturation test: it proves the limiter is + still handing out slots after a cancel (i.e. a cancel cannot wedge it), + and it reuses the same drain helper the concurrency scenario ends with. + It deliberately does NOT saturate the org first -- doing so here would + make every scenario defined after it flaky, which is why the one + saturating scenario in this module is defined last. The saturated form of + this check is that last scenario's own drain. + """ + job_id, _ = submit( + agent_kv_key, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + ) + cancel_resp = cancel(agent_kv_key, job_id) + assert cancel_resp.status_code in (200, 409), cancel_resp.text + + assert _drain_until_slot_free(agent_kv_key), ( + f"no submit was accepted within {_SLOT_DRAIN_TIMEOUT_S}s of cancelling " + f"job {job_id}; the limiter is refusing new work, which means a slot " + f"was leaked rather than released by finalize" + ) + + +# --------------------------------------------------------------------------- +# 6. Delete +# --------------------------------------------------------------------------- + + +def test_delete_completed_job_then_result_404( + require_adapters: None, agent_kv_key: AgentKVAuth +) -> None: + job_id, _ = submit( + agent_kv_key, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + ) + status_doc = poll(agent_kv_key, job_id, timeout_s=600) + assert status_doc["status"] == "completed", status_doc + + del_resp = delete(agent_kv_key, job_id) + assert del_resp.status_code == 204, del_resp.text + + after = result(agent_kv_key, job_id) + assert after.status_code == 404, after.text + + +# --------------------------------------------------------------------------- +# 7. Page cap +# --------------------------------------------------------------------------- + + +def test_page_cap_rejects_oversized_document( + require_adapters: None, agent_kv_key: AgentKVAuth +) -> None: + """The page cap refuses an oversized document at submit. + + Takes ``require_adapters`` even though it asserts a REJECTION. The claim + that "scenarios expecting a rejection before the adapter gate keep working + unchanged" holds only for the three raised inside ``validate_extractors``; + this rejection comes from ``SubmitSerializer.validate()``, and DRF raises + field-level errors from ``to_internal_value`` BEFORE ``validate()`` runs -- + with ``_validated_adapter_shape`` called from the field validator. So with + the adapter vars unset (the documented default) the status assert still + passed on a 400 and the attr assert then FAILED, blaming ``extractors`` + instead of ``file``. Red, not skipped. + + Same defect as ``test_submit_unreadable_pdf_is_400``, but masked: the + ``AGENT_KV_MAX_PAGES`` skip below usually fires first, so it would only + have gone red on a lane that configured the cap. + """ + fixture_pages = 2 # fixtures/invoice.pdf + cap_raw = os.environ.get("AGENT_KV_MAX_PAGES") + if cap_raw is None or int(cap_raw) > fixture_pages: + pytest.skip( + f"AGENT_KV_MAX_PAGES={cap_raw!r} exceeds the fixture's " + f"{fixture_pages} pages; set AGENT_KV_MAX_PAGES<={fixture_pages} " + "on the backend service to exercise the cap" + ) + resp = submit_raw(agent_kv_key, _RENT_ROLL_BYTES, "rent_roll.pdf", table_keys()) + assert resp.status_code == 400, resp.text + body = resp.json() + assert "file" in _validation_error_attrs(body), body + + +# --------------------------------------------------------------------------- +# 8. Executor-side failure (operator-configured bad LLM key) +# --------------------------------------------------------------------------- + + +def _assert_user_safe_error(error: str) -> None: + """No internals in a customer-visible error string (spec §8).""" + lowered = error.lower() + for leak in ("permission denied", "traceback", "errno", "sk-", "api_key"): + assert leak not in lowered, f"internal detail {leak!r} leaked into: {error!r}" + # No filesystem paths: the engine's own text used to carry the executor's + # tempfile/page-image directories verbatim. + for leak in ("/app/", "/tmp/", "/var/", "storage/processing", "\\"): + assert leak not in error, f"a filesystem path leaked into: {error!r}" + + +def test_a_bad_adapter_credential_ends_failed( + require_adapters: None, agent_kv_key: AgentKVAuth, bad_llm_adapter: str +) -> None: + """A job that DISPATCHES and then fails in the executor ends ``failed`` + with a user-safe error. + + Every other failure scenario in this module is a request-time rejection + (400/403/429) that never reaches the engine. This one covers the other + half of the contract -- the terminal FAILED path -- and it needs the whole + stack to be running with an adapter whose provider credential is + deliberately invalid, which no test can mint for itself. Hence the explicit + operator gate (see ``bad_llm_adapter``). + + It used to gate on ``AGENT_KV_E2E_BAD_KEY_JOB=1`` plus an invalid + ``AGENT_KV_LLM_API_KEY``. That env var is not read on the table path any + more, so an operator who set the stack up exactly as documented got a job + that ran on the VALID adapter credential, reached ``completed``, and failed + this test's own assertion. The trigger moved to where the credential now + lives. + + The assertion is about SHAPE, not text: whatever went wrong inside the + engine, the customer-visible ``error`` must not carry a filesystem path, a + ``Permission denied``, or a provider response body -- spec §8 mandates the + user-safe ``" failed: "`` form (or one of the fixed strings + ``cancelled`` / ``timed out``), which is what the executor now builds from + the node listener's exception rather than echoing the engine's own text. + """ + # The advanced-LLM role carries the junk credential; the other two stay + # valid, so the run gets past OCR and fails where this test says it does. + adapters = {**e2e_adapters(), "llm": bad_llm_adapter} + + job_id, _ = submit( + agent_kv_key, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + adapters=adapters, + ) + + status_doc = poll(agent_kv_key, job_id, timeout_s=600) + assert status_doc["status"] == "failed", status_doc + assert "error" in status_doc, status_doc + error = status_doc["error"] + assert isinstance(error, str), status_doc + assert error, status_doc + + _assert_user_safe_error(error) + + # The result endpoint reports the same failure, in the docs §5 shape. + resp = result(agent_kv_key, job_id) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["success"] is False, body + assert body["status"] == "failed", body + _assert_user_safe_error(body["error"]) + + +# --------------------------------------------------------------------------- +# 9. Document cache (spec §6.5 document tier) +# --------------------------------------------------------------------------- + + +def test_sync_wait_submit_returns_result_inline( + require_adapters: None, agent_kv_key: AgentKVAuth +) -> None: + """A submit with ``timeout`` seconds sync-waits: if the job reaches a + terminal state inside the window, the response is 200 with the result + payload itself (docs §3) -- no polling leg. Uses the same invoice as the + happy path, so a warm document cache keeps the wait well inside the + window; even a cold cache (~15 s pipeline) fits comfortably. + """ + resp = submit_raw( + agent_kv_key, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + timeout=120, + ) + assert resp.status_code == 200, ( + f"sync-wait submit: HTTP {resp.status_code} (expected 200 with the " + f"inline result): {resp.text}" + ) + body = resp.json() + # `success`/`status` are top level on every terminal payload; the engine's + # own record is namespaced per extractor (spec §7.3). + assert body.get("success") is True, body + body = table_result_body(resp) + assert isinstance(body["tables"], list), body + assert "row_count" in body, body + + +# --------------------------------------------------------------------------- +# 9c. Excel input +# --------------------------------------------------------------------------- + + +def test_excel_submit_extracts( + require_adapters: None, agent_kv_key: AgentKVAuth +) -> None: + """An .xlsx submit runs the whole pipeline. Excel has no pre-OCR page + count (``pages_total`` stays None at submit) and takes the post-OCR + virtual-page cap path instead. Structure-only assertions, like the + happy path. + """ + job_id, _ = submit( + agent_kv_key, INVOICE_XLSX.read_bytes(), "invoice.xlsx", table_keys("Line items") + ) + doc = poll(agent_kv_key, job_id) + assert doc["status"] == "completed", doc + body = table_result_body(result(agent_kv_key, job_id)) + assert body.get("success") is True, body + assert isinstance(body["tables"], list), body + assert "page_count" in body, body + + +# --------------------------------------------------------------------------- +# 9d. Webhook delivery (operator-gated) +# --------------------------------------------------------------------------- + + +def test_webhook_delivered_on_completion( + require_adapters: None, agent_kv_key: AgentKVAuth +) -> None: + """Completion-webhook delivery, end to end: submit with a ``webhook_url`` + pointing at a receiver on the compose host and wait for the POST. + + Operator-gated: the worker's SSRF guards (https + public host, spec + §6.7) refuse a host-local receiver unless the stack runs with + ``AGENT_KV_WEBHOOK_INSECURE_ALLOW_HTTP_PRIVATE=1`` on the ide-callback + worker (test/dev only), reached via ``host.docker.internal`` (compose + host-gateway mapping). Skips unless the operator set the same env for + the test run to declare that setup. + """ + if os.environ.get("AGENT_KV_WEBHOOK_INSECURE_ALLOW_HTTP_PRIVATE", "").lower() not in ( + "1", + "true", + "yes", + ): + pytest.skip( + "webhook e2e needs AGENT_KV_WEBHOOK_INSECURE_ALLOW_HTTP_PRIVATE=1 " + "on both the ide-callback worker and this test run" + ) + + import http.server + import threading + + hits: list[dict] = [] + + class _Receiver(http.server.BaseHTTPRequestHandler): + def do_POST(self): # noqa: N802 - stdlib API name + length = int(self.headers.get("Content-Length", "0")) + hits.append(json.loads(self.rfile.read(length) or b"{}")) + self.send_response(200) + self.end_headers() + + def log_message(self, *args): # keep pytest output pristine + pass + + server = http.server.HTTPServer(("0.0.0.0", 0), _Receiver) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + job_id, _ = submit( + agent_kv_key, + _RENT_ROLL_BYTES, + "rent_roll.pdf", + table_keys(), + webhook_url=f"http://host.docker.internal:{port}/agent-kv-hook", + ) + doc = poll(agent_kv_key, job_id) + assert doc["status"] == "completed", doc + deadline = time.time() + 30 + while time.time() < deadline and not hits: + time.sleep(0.5) + assert hits, "no webhook POST arrived within 30s of completion" + assert hits[0] == {"job_id": job_id, "status": "completed"}, hits + finally: + server.shutdown() + thread.join(timeout=5) + + +def test_concurrency_limit_returns_429( + require_adapters: None, agent_kv_key: AgentKVAuth +) -> None: + """Saturate the org's concurrency limiter; every other submit-based + scenario in this module runs before this one, and ``tests/groups.yaml`` + pins this whole group ``parallel: false`` -- both for the same reason: + + ``JobCancelView`` marks a job ``CANCELLED`` directly via + ``AgentKVJob.mark_terminal`` and never touches + ``AgentKVConcurrencyLimiter`` -- only a job's own *finalize* call + releases its slot (in a ``finally``, unconditional on whether the + terminal-state guard actually won -- docs/agent-kv-api.md §11a), and + that only happens once the executor that picked the job up actually + runs it to completion/failure and calls back. So immediately after this + test oversaturates the limiter, the org's concurrency budget stays + genuinely exhausted for however long those accepted jobs take to + finalize -- anywhere from seconds to the executor's task time limit. + Any sibling submit-based test running concurrently with, or soon after, + this one would risk being 429'd by *this* test's saturation rather than + by anything it did itself. Running this scenario last (and the group + serially) removes that risk entirely for this module; the cleanup below + is a courtesy for whatever runs against this org after it, not a + correctness requirement for this test itself -- the limiter's own 6h + slot TTL (``AgentKVConcurrencyLimiter._SLOT_TTL_SECONDS``) is the + ultimate backstop either way. + """ + limit = int(os.environ.get("AGENT_KV_CONCURRENT_LIMIT", "5")) + n = limit + 1 + + def _submit(_i: int): + return _cheap_submit(agent_kv_key) + + accepted_job_ids: list[str] = [] + try: + with concurrent.futures.ThreadPoolExecutor(max_workers=n) as pool: + responses = list(pool.map(_submit, range(n))) + + statuses = [r.status_code for r in responses] + assert all(s in (202, 429) for s in statuses), ( + f"unexpected status among {n} concurrent submits " + f"(only 202/429 are valid outcomes here): {statuses}" + ) + assert 429 in statuses, ( + f"submitting {n} concurrent jobs against a limit of {limit} " + f"produced no 429: {statuses}" + ) + for r in responses: + if r.status_code == 202: + accepted_job_ids.append(r.json()["job_id"]) + assert len(accepted_job_ids) <= limit, accepted_job_ids + finally: + # Cancelling stops each job's eventual result from ever being + # readable, but -- per the docstring above -- does NOT itself free + # its concurrency slot; only that job's own finalize call does. + for job_id in accepted_job_ids: + try: + cancel(agent_kv_key, job_id) + except Exception: # noqa: BLE001 - cleanup only, never fail the test on it + pass + + # Best-effort drain: keep probing with a cheap submission until one + # is actually accepted again (proof a slot freed up) or the bounded + # wait runs out. The RESULT is deliberately ignored here -- see the + # 6h-TTL backstop note in the docstring above; this is a courtesy for + # whatever runs next, not a correctness requirement for this test. + _drain_until_slot_free(agent_kv_key) diff --git a/tests/groups.yaml b/tests/groups.yaml index cc2e5beb73..652c4885c9 100644 --- a/tests/groups.yaml +++ b/tests/groups.yaml @@ -89,7 +89,8 @@ groups: unit-workers: tier: unit workdir: workers - paths: [tests, shared/tests] + paths: [tests, shared/tests, shared/infrastructure/config/tests, + ide_callback/tests] # Negative filter: worker unit tests aren't tagged `unit` (~1100 of them), # so select by exclusion. `conftest.py` marks real-Postgres tests # `integration` (routed to the `integration-workers` group); `slow` tests are @@ -128,6 +129,25 @@ groups: install_editable: true coverage_source: src + # Both of these had tests but NO group, so nothing collected them. For + # agent-kv-schema that mattered specifically: test_compile.py is the + # regression net for the schema depth/ReDoS caps, so the caps were asserted + # only on a developer's machine. A PR reporting "165 passed" for this package + # was reporting a local figure, not a CI one. + unit-agent-kv-schema: + tier: unit + workdir: unstract/agent-kv-schema + paths: [tests] + install_editable: true + coverage_source: src + + unit-filesystem: + tier: unit + workdir: unstract/filesystem + paths: [tests] + install_editable: true + coverage_source: src + # runner/tests was registered in NO group, so it was collected by nothing. Its # only suite guards that LOG_TRANSPORT and LOG_STREAM_QUEUE_NAME reach the tool # sidecar's hand-picked env allowlist — a list that drops anything not named in @@ -254,6 +274,29 @@ groups: depends_on: [e2e-smoke] optional: true + # Agent-KV product API (/agent-kv/...). Needs the cloud `agentic_kv` + # executor plugin installed (docs/agent-kv-api.md §11) -- an OSS-only image + # 501s every submit, so this is optional: an OSS-only CI run must not go + # red just because the plugin isn't there. The module itself additionally + # self-gates on AGENT_KV_E2E=1 so it never runs by accident even when the + # platform happens to be up (see tests/e2e/agent_kv/test_agent_kv_e2e.py). + # Longer budget than the other platform groups: the happy-path and delete + # scenarios poll a real extraction job through to COMPLETED. + e2e-agent-kv: + tier: e2e + paths: [tests/e2e/agent_kv] + requires_platform: true + depends_on: [e2e-smoke] + optional: true + timeout_seconds: 1200 + # Serial, not xdist-parallel: the concurrency scenario deliberately + # saturates the org's AgentKVConcurrencyLimiter slots, which + # JobCancelView never releases (only a finalize callback does, by + # design -- see the scenario's own docstring). Running sibling + # submit-based tests on other xdist workers at the same moment would + # let that saturation 429 them too. + parallel: false + ui: # Playwright browser e2e over the real UI. Scaffolding only for now — the # spec directory holds a single smoke test; extend it here. diff --git a/tox.ini b/tox.ini index b1d66781a9..cd4bc9c8fd 100644 --- a/tox.ini +++ b/tox.ini @@ -52,6 +52,11 @@ passenv = GITHUB_* DOCKER_HOST UNSTRACT_* + # The Agent-KV e2e lane's stack config (engine keys/models, storage prefix, + # the caps its scenarios override) is read by + # tests/compose/docker-compose.test.yaml from the caller's environment; tox + # would otherwise strip it and the lane would silently run misconfigured. + AGENT_KV_* # ComposeRuntime's `docker compose up` resolves ${VERSION} image tags in # docker/docker-compose.yaml from the caller's environment. VERSION diff --git a/unstract/agent-kv-schema/pyproject.toml b/unstract/agent-kv-schema/pyproject.toml new file mode 100644 index 0000000000..6723d4478f --- /dev/null +++ b/unstract/agent-kv-schema/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "unstract-agent-kv-schema" +version = "0.1.0" +description = "Agent-KV keys.json schema compiler and submit-time validation gate (the engine calls kv_schema.compile directly and does not re-apply the caps — see compile.py)" +requires-python = ">=3.12" +dependencies = [] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/unstract"] diff --git a/unstract/agent-kv-schema/src/unstract/agent_kv_schema/__init__.py b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/__init__.py new file mode 100644 index 0000000000..07087e3f47 --- /dev/null +++ b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/__init__.py @@ -0,0 +1,15 @@ +from .compile import CompiledSchema, SchemaCaps, SchemaError, compile_schema +from .constraints import evaluate_constraints +from .dataclasses import ArraySpec, KeySpec +from .validators import validate_format + +__all__ = [ + "ArraySpec", + "CompiledSchema", + "KeySpec", + "SchemaCaps", + "SchemaError", + "compile_schema", + "evaluate_constraints", + "validate_format", +] diff --git a/unstract/agent-kv-schema/src/unstract/agent_kv_schema/compile.py b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/compile.py new file mode 100644 index 0000000000..b3aa72e075 --- /dev/null +++ b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/compile.py @@ -0,0 +1,387 @@ +"""Cap-enforcing, syntax-validating wrapper over the ported compiler. + +This is the entry point the API uses for submit-time validation. Anything +``compile_schema`` accepts, the engine must execute; anything it rejects never +reaches OCR or an LLM. + +It is NOT, despite what this docstring said before, "the single entry point +both the API and the cloud engine use". The engine calls the raw +``kv_schema.compile``/``compile_arrays`` directly (see the import note in the +cloud plugin's ``engine/kv_extractor.py``), so the caps below are enforced on +the submit path only. That is sound while the backend is the sole producer of a +compiled schema -- which it is today -- but it means the caps are a gate, not +an invariant the engine itself re-checks. Anything that ever hands the engine a +schema from another source has to apply them, or this module has to become what +it claimed to be. +""" + +import ast +import re +from dataclasses import dataclass, field + +from . import kv_schema +from .dataclasses import ArraySpec, KeySpec + + +class SchemaError(ValueError): + """User-facing schema rejection; message is safe to return in a 400.""" + + +@dataclass(frozen=True) +class SchemaCaps: + max_leaves: int = 200 + max_arrays: int = 20 + max_columns_per_array: int = 40 + max_depth: int = 6 + max_regex_len: int = 200 + max_aliases: int = 10 + max_description_len: int = 500 + max_constraints: int = 30 + + +@dataclass(frozen=True) +class CompiledSchema: + key_specs: list[KeySpec] = field(default_factory=list) + array_specs: list[ArraySpec] = field(default_factory=list) + constraints: list[str] = field(default_factory=list) + + +_ALLOWED_CALLS = {"sum", "count", "min", "max", "avg"} +_ALLOWED_NODES = ( + ast.Expression, + ast.BoolOp, + ast.And, + ast.Or, + ast.UnaryOp, + ast.Not, + ast.USub, + ast.Compare, + ast.Eq, + ast.NotEq, + ast.Lt, + ast.LtE, + ast.Gt, + ast.GtE, + ast.BinOp, + ast.Add, + ast.Sub, + ast.Mult, + ast.Div, + ast.Call, + ast.Name, + ast.Attribute, + ast.Constant, + ast.Load, +) + + +def _check_constraint_syntax(expr: str) -> None: + """Static allowlist mirroring constraints._evaluate_one's grammar.""" + try: + tree = ast.parse(expr, mode="eval") + except SyntaxError as e: + raise SchemaError(f"constraint does not parse: {expr!r} ({e.msg})") from e + for node in ast.walk(tree): + if not isinstance(node, _ALLOWED_NODES): + raise SchemaError( + f"constraint uses disallowed syntax ({type(node).__name__}): {expr!r}" + ) + if isinstance(node, ast.Call): + if not isinstance(node.func, ast.Name) or node.func.id not in _ALLOWED_CALLS: + raise SchemaError(f"constraint calls a disallowed function: {expr!r}") + if ( + node.keywords + or len(node.args) != 1 + or not isinstance(node.args[0], ast.Constant) + or not isinstance(node.args[0].value, str) + ): + raise SchemaError( + f"constraint aggregate needs one string literal arg: {expr!r}" + ) + + +def _max_depth(node: object, max_depth: int, depth: int = 0) -> int: + # HARD ceiling FIRST -- before the `_array` short-circuit and before any + # further recursion. Two bugs this closes: (a) a pathologically deep plain + # schema (256 KiB of JSON nests ~40k deep, far past Python's ~1000 + # recursion limit) used to raise an uncaught RecursionError (500) here; + # now it fails fast as a clean SchemaError (400). (b) The `_array` + # short-circuit below cannot be reached to bypass the cap once `depth` + # has already blown past it. + if depth > max_depth: + raise SchemaError(f"schema exceeds max_depth={max_depth}") + if not isinstance(node, dict): + return depth + if "_array" in node: + return depth + 1 # array columns are row-local, not nesting + child = [v for v in node.values() if isinstance(v, dict)] + if not child: + return depth + 1 + return max(_max_depth(v, max_depth, depth + 1) for v in child) + + +# Quantified groups that themselves contain a quantifier, e.g. `(a+)+`, `(a*)*`, +# `(.+)*`, `(?:x+)+`. This is the shape behind catastrophic backtracking. +_NESTED_QUANTIFIER = re.compile( + r"\((?:\?[:=!]|\?<[=!]|\?P<[^>]+>)?" # group open, incl. non-capturing/named + r"[^()]*[+*}]\??" # ...containing a quantifier + r"[^()]*\)" # ...group close + r"\s*[+*]|\)\s*\{\d+,\d*\}" # ...itself quantified +) + + +# A quantified group whose alternatives OVERLAP, e.g. `(a|aa)+`. Found in +# review after the nested-quantifier check shipped: `^(a|aa)+$` passes that +# check and still backtracks catastrophically, because at each position the +# engine can consume one `a` or two and must try both on failure. +# +# Only LITERAL branches are compared, and only by the prefix relation: if one +# branch is a prefix of another (`a` of `aa`), the group is ambiguous and +# refused. `(foo|bar)+` is left alone -- distinct first characters mean no +# position admits two parses, so it is linear. Branches containing +# metacharacters are not analysed (`.`/classes/nested groups need real regex +# analysis, which is what UN-4225 is for); this closes the demonstrated family +# without pretending to be a decision procedure. +_QUANTIFIED_GROUP = re.compile(r"\((\?:)?([^()]*)\)\s*(?:[+*]|\{\d+,\d*\})") +_LITERAL_BRANCH = re.compile(r"^[\w\-/ ]*$") + + +def _overlapping_alternation(pattern: str) -> str | None: + """Return a human-readable overlap if a quantified group is ambiguous.""" + for match in _QUANTIFIED_GROUP.finditer(pattern): + body = match.group(2) + if "|" not in body: + continue + branches = body.split("|") + if not all(_LITERAL_BRANCH.match(b) for b in branches): + continue + for i, a in enumerate(branches): + for j, b in enumerate(branches): + if i != j and a and b.startswith(a): + return f"{a!r} is a prefix of {b!r}" + return None + + +def _reject_unsafe_regex(path: str, pattern: str) -> None: + """Refuse an author-supplied pattern at SUBMIT rather than at match time. + + Two separate problems, both found in review: + + 1. The pattern was never compiled here, so a syntactically invalid one was + accepted and only discovered per-value in the engine's QA pass -- where + ``_check_one`` swallows ``re.error`` and returns True, silently passing + validation the author thought they had configured. + 2. ``validate_format`` runs the pattern against extracted values with no + time budget, so a catastrophically-backtracking pattern is a DoS. The + length cap is NOT a mitigation: ``^(a+)+$`` is 7 characters and takes + ~1.9s on 26 ``a``s, ~4x per character added (measured), so a 40-char + value runs for hours. With ``AGENT_KV_CONCURRENT_LIMIT=5`` one org can + pin five shared worker slots from a single submit. + + The nested-quantifier check is a CONSERVATIVE HEURISTIC, not a proof. It + rejects the shape responsible for the realistic cases (a quantified group + whose body is itself quantified) and will reject some safe patterns that + happen to look like it -- an explicit trade, since the author gets an + immediate, actionable error instead of a job that hangs. It does not catch + every pathological pattern; the complete fix is a linear-time engine (RE2), + which cannot be added here because this package deliberately has zero + dependencies and is installed by both repos. Tracked as UN-4225. + """ + try: + re.compile(pattern) + except re.error as e: + raise SchemaError(f"'{path}' has an invalid regex: {e}") from None + if _NESTED_QUANTIFIER.search(pattern): + raise SchemaError( + f"'{path}' regex has a nested quantifier (e.g. '(a+)+'), which can " + "backtrack catastrophically and stall extraction. Rewrite it " + "without a quantifier inside a quantified group." + ) + overlap = _overlapping_alternation(pattern) + if overlap: + raise SchemaError( + f"'{path}' regex quantifies a group whose alternatives overlap " + f"({overlap}), e.g. '(a|aa)+', which can backtrack " + "catastrophically and stall extraction. Make the alternatives " + "mutually exclusive, or use a character class." + ) + + +#: The formats ``validators._check_one`` actually implements. Anything else is +#: passed through as a free-text LLM hint, which is deliberate (see the +#: ``format`` comment on ``KeySpec``) -- and is also why a typo in one of these +#: six is invisible: ``format: "Number"`` is not a known format, so it becomes a +#: hint, so ``validate_format`` returns True for every value forever and +#: ``/validate`` reports ``{"valid": true}``. The author configured validation +#: and got none, with no error anywhere to find it by. +_KNOWN_FORMATS = frozenset({"string", "number", "date", "currency", "enum", "regex"}) + + +def _reject_unusable_formats(key_specs, array_specs) -> None: + """Refuse declared formats that can never validate anything, or always fail. + + Four shapes, all accepted before review and all verified by execution: + + * ``format: "Number"`` -- a case variant of a known format. Silently + degrades to a free-text hint, disabling validation for that key. Rejected + with the intended spelling, rather than guessed at: a schema author who + meant the hint can lower-case it or reword it, and one who meant the + format gets told. + * ``format: "Enum:paid,unpaid"`` / ``"REGEX:^[0-9]+$"`` -- the same defect + class for the two formats that take an ARGUMENT, and the one the first + round's fix missed. ``_parse_format`` matches the ``enum:`` / ``regex:`` + prefix case-sensitively, so these never reach kind ``"enum"``/``"regex"``; + they arrive as the whole raw string, which no bare-name comparison + matches. Validation was silently off for the key -- the precise outcome + the first bullet exists to prevent, reached by a route it did not cover. + * ``format: "enum:"`` -- no values, so ``_check_one`` tests membership of + the empty set and EVERY non-empty value fails QA for the life of the job. + * ``format: "regex:"`` -- empty pattern, so ``re.fullmatch("", v)`` matches + nothing but the empty string, which ``validate_format`` already passes + before reaching the pattern. Same outcome: nothing can ever conform. + + The last two are worse than the first: they do not disable validation, they + invert it, and the result is a document that fails QA no matter what is on + the page. + """ + for kspec in _every_leaf(key_specs, array_specs): + fmt = kspec.format + lowered = fmt.casefold() + # `lowered` catches a bare case variant (`"Number"`). `kind_only` + # catches the two formats that take an ARGUMENT, which the bare check + # misses entirely: `_parse_format` matches the `enum:` / `regex:` + # prefix case-SENSITIVELY, so `"Enum:paid,unpaid"` never becomes kind + # `"enum"` -- it arrives here as the whole raw string, and + # `"enum:paid,unpaid".casefold()` is not a member of `_KNOWN_FORMATS` + # (which holds the bare `"enum"`). So the mis-cased form slipped this + # guard, parsed as a free-text kind, and validation was silently off + # for that key -- exactly the failure this function closes for bare + # formats. + # + # Splitting on the colon is safe for genuine free text: a hint like + # `"customer:id"` yields `"customer"`, which is not a known format. + # + # Rejected rather than auto-corrected, to match the bare case above: a + # schema author who meant the format gets told the spelling, and one + # who meant a free-text hint can reword it. Silently honouring `Enum:` + # would make case significant in one direction only. + kind_only = lowered.split(":", 1)[0] + if fmt not in _KNOWN_FORMATS and ( + lowered in _KNOWN_FORMATS or kind_only in _KNOWN_FORMATS + ): + raise SchemaError( + f"'{kspec.path}' declares format {fmt!r}, which is not a known " + f"format and is therefore treated as free text -- no validation " + f"would run. Did you mean {lowered!r}?" + ) + if fmt == "enum" and not kspec.enum_values: + raise SchemaError( + f"'{kspec.path}' declares an enum with no values " + f"(e.g. 'enum:paid,unpaid'); as written no value can ever conform." + ) + if fmt == "regex" and not kspec.regex_pattern: + raise SchemaError( + f"'{kspec.path}' declares an empty regex; as written no value " + f"can ever conform." + ) + + +def _reject_unknown_key_columns(array_specs) -> None: + """Refuse an array's ``_key`` that names a column the array does not declare. + + ``key_column`` selects row identity for scoring. A name with no matching + column is not an error anywhere downstream -- every row simply misses, and + the array silently falls back to positional identity. So a typo costs + accuracy on exactly the arrays the author cared enough about to key, and + reports nothing. + """ + for aspec in array_specs: + if not aspec.key_column: + continue # '' is the documented "positional identity" default + columns = {s.path for s in aspec.item_specs} + if aspec.key_column not in columns: + raise SchemaError( + f"array '{aspec.path}' declares _key " + f"{aspec.key_column!r}, which is not one of its columns " + f"({sorted(columns)})." + ) + + +def _every_leaf(key_specs, array_specs): + """Scalar leaves and array columns, which carry the same per-leaf rules.""" + return list(key_specs) + [s for a in array_specs for s in a.item_specs] + + +def compile_schema(spec: dict, caps: SchemaCaps | None = None) -> CompiledSchema: + caps = caps or SchemaCaps() + if not isinstance(spec, dict): + raise SchemaError("Top-level key schema must be a JSON object") + cleaned = {k: v for k, v in spec.items() if k != "_constraints"} + if _max_depth(cleaned, caps.max_depth) > caps.max_depth: + raise SchemaError(f"schema exceeds max_depth={caps.max_depth}") + # The compile.py `_max_depth` pre-check does not count array-column + # nesting (arrays are row-local there, by design) and cannot see a decoy + # top-level `_array` field's real nesting -- so the actual recursive walk + # (`kv_schema._walk`) carries its own max_depth ceiling too, both to close + # that bypass and to guarantee a clean SchemaError instead of an uncaught + # RecursionError on a deeply-nested input. + try: + key_specs = kv_schema.compile(spec, max_depth=caps.max_depth) + array_specs = kv_schema.compile_arrays(spec, max_depth=caps.max_depth) + except ValueError as e: + raise SchemaError(str(e)) from e + + _enforce_shape_caps(key_specs, array_specs, caps) + _enforce_leaf_caps(key_specs, array_specs, caps) + _reject_unusable_formats(key_specs, array_specs) + _reject_unknown_key_columns(array_specs) + constraints = _validated_constraints(spec, caps) + + return CompiledSchema( + key_specs=key_specs, array_specs=array_specs, constraints=list(constraints) + ) + + +def _enforce_shape_caps(key_specs, array_specs, caps: SchemaCaps) -> None: + """Bound how much structure the schema declares.""" + if len(key_specs) > caps.max_leaves: + raise SchemaError(f"schema exceeds max_leaves={caps.max_leaves}") + if len(array_specs) > caps.max_arrays: + raise SchemaError(f"schema exceeds max_arrays={caps.max_arrays}") + for aspec in array_specs: + if len(aspec.item_specs) > caps.max_columns_per_array: + raise SchemaError( + f"array '{aspec.path}' exceeds " + f"max_columns_per_array={caps.max_columns_per_array}" + ) + + +def _enforce_leaf_caps(key_specs, array_specs, caps: SchemaCaps) -> None: + """Bound every leaf's author-supplied text, across scalars and array columns.""" + for kspec in _every_leaf(key_specs, array_specs): + _reject_unsafe_regex(kspec.path, kspec.regex_pattern) + if len(kspec.regex_pattern) > caps.max_regex_len: + raise SchemaError( + f"'{kspec.path}' regex exceeds max_regex_len={caps.max_regex_len}" + ) + if len(kspec.aliases) > caps.max_aliases: + raise SchemaError(f"'{kspec.path}' exceeds max_aliases={caps.max_aliases}") + if len(kspec.effective_description) > caps.max_description_len: + raise SchemaError( + f"'{kspec.path}' description exceeds " + f"max_description_len={caps.max_description_len}" + ) + + +def _validated_constraints(spec: dict, caps: SchemaCaps) -> list: + """Return the schema's `_constraints`, rejecting a malformed or oversized list.""" + constraints = spec.get("_constraints", []) + if not isinstance(constraints, list) or not all( + isinstance(c, str) for c in constraints + ): + raise SchemaError("_constraints must be a list of strings") + if len(constraints) > caps.max_constraints: + raise SchemaError(f"schema exceeds max_constraints={caps.max_constraints}") + for expr in constraints: + _check_constraint_syntax(expr) + return constraints diff --git a/unstract/agent-kv-schema/src/unstract/agent_kv_schema/constraints.py b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/constraints.py new file mode 100644 index 0000000000..a64b8bc9c9 --- /dev/null +++ b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/constraints.py @@ -0,0 +1,343 @@ +"""Cross-field consistency (spec §18.3): an ADVISORY evaluator for schema-author-declared +constraints over the NORMALIZED key values. No eval/exec — a static AST allowlist (Compare/BoolOp/ +BinOp/UnaryOp/Name/Attribute-path/Constant only). Operands resolve to normalized values; a missing/ +empty/un-coercible operand SKIPS the constraint, never crashes. Returns violated exprs. + +Two senses of "fail-closed" were conflated here, and this docstring used to +claim the wrong one. What IS fail-closed is the *grammar*: anything outside the +allowlist is refused rather than executed, so no constraint can run code. What +is NOT fail-closed is the *outcome*: a constraint that cannot be evaluated -- +missing operand, mixed types, or an unexpected exception -- is dropped, and +``evaluate_constraints`` records only ``is False`` as a violation. So a dropped +constraint is indistinguishable, in the return value, from one that passed. + +That matters because dropping is REACHABLE, not theoretical: `_BIN` includes +`truediv`, so `unit_price == line_total / quantity` is an accepted constraint, +and a row where `quantity` normalizes to `0.0` raises `ZeroDivisionError`. The +QA result then reports ``violations: []`` -- a positive assurance that nothing +was checked. + +Until the result shape can carry a third state (skipped-with-reason), every +drop is LOGGED, so the condition is at least diagnosable from the worker logs +rather than silent. Surfacing it in the QA document is a public contract change +on the `kv` extractor, which this deployment does not ship; tracked separately. +""" + +import ast +import logging +import math +import operator + +from .validators import coerce_number + +logger = logging.getLogger(__name__) + +# Equality on these values is TOLERANT, not exact. Operands are normalized +# currency/number values that arrived as floats, so binary floating point makes +# exact `==` wrong for the thing constraints exist to express. The feature's own +# headline example reports a false violation on a CORRECT invoice: +# three line items of 8230.4 sum to 24691.199999999997, which `operator.eq` +# says is not 24691.2 (reproduced). +# +# The tolerance has to sit BELOW a cent at every realistic magnitude while +# staying ABOVE the float noise floor. Those two bounds are what set the +# numbers, and the first attempt at this got it wrong: `rel_tol=1e-9` is 0.1 at +# a total of $100,000,000, so it silently absorbed a one-cent reconciliation +# error on a large invoice -- the exact failure this check exists to catch +# (found in review; the original test only covered a ~$25,000 total). +# +# Noise floor: summing N values of magnitude M accumulates roughly +# N * 2.2e-16 * M of float error. M=$1e8 over 1,000 rows is ~2.2e-5, i.e. about +# a five-hundredth of a cent. +# +# rel_tol=1e-12 -> 1e-4 at $1e8 (a hundredth of a cent): absorbs that noise, +# and a one-cent error is 100x larger, so it is still +# reported. Relative rather than absolute-only so the bound +# tracks magnitude -- an absolute epsilon suited to invoice +# totals is meaningless against unit prices in thousandths. +# abs_tol=1e-6 -> keeps comparisons against exact zero working, where +# relative tolerance is useless (every non-zero value is +# infinitely far from 0 in relative terms), and is itself +# well under a cent. +# +# Known ceiling: past ~$1e12 with ~10,000 rows the noise floor (~2.2) exceeds a +# cent and no float tolerance can separate the two. That is a float problem, +# not a tolerance problem, and the fix is Decimal end to end -- +# `normalizers.coerce_number`, UN-4226. This makes the COMPARISON correct for +# the float values that exist today, and stays correct afterwards. +_REL_TOL = 1e-12 +_ABS_TOL = 1e-6 + + +def _num_eq(a, b) -> bool: + if isinstance(a, (int, float)) and isinstance(b, (int, float)): + return math.isclose(a, b, rel_tol=_REL_TOL, abs_tol=_ABS_TOL) + return operator.eq(a, b) + + +def _num_ne(a, b) -> bool: + return not _num_eq(a, b) + + +_CMP = { + ast.Eq: _num_eq, + ast.NotEq: _num_ne, + # Ordering comparisons are left exact on purpose: a tolerant `<` would make + # `a < b` and `a == b` both true at the boundary, and the failure mode these + # express (a total that is too large/small) is not a rounding artefact. + ast.Lt: operator.lt, + ast.LtE: operator.le, + ast.Gt: operator.gt, + ast.GtE: operator.ge, +} +_BIN = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, +} + +# P8c: the ONLY callable names permitted in a constraint, and only over a single +# string-literal `'array_path'` / `'array_path.column'` argument. This is NOT a general +# function-call capability — it is a closed allowlist of pure numeric reductions evaluated +# in Python (never via eval/exec) so a scalar key can be reconciled against an array column +# (e.g. grand_total == sum('line_items.line_total')). +_AGG = {"sum", "count", "min", "max", "avg"} + + +class _Skip(Exception): + """Operand missing/empty/incomparable — skip this constraint (advisory).""" + + +def _path(node) -> str: + """Reconstruct a dotted path from a Name / Attribute-chain (pure attribute access only).""" + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return f"{_path(node.value)}.{node.attr}" + raise _Skip() + + +def _coerce(raw: str): + """Normalized value -> finite float when numeric, else the string (ISO dates compare + lexicographically). Uses the canonical coercer so thousands-separated values ('9,000') + parse numerically (a bare float() would leave them as strings and silently drop a real + violation via the like-type guard); 'nan'/'inf' tokens stay strings (not comparable numbers). + """ + n = coerce_number( + raw + ) # strips commas/$/%; returns None for non-numeric AND non-finite + return n if n is not None else raw + + +def _parse_agg_call(node: ast.Call) -> tuple[str, str, str]: + """Validate `NAME('literal')` and split it into (fn, array_path, column). + + FAIL-CLOSED by construction: the only thing accepted is `NAME('literal')` + where NAME is in `_AGG`, the func is a bare `ast.Name` (no attribute access + -- blocks `os.system` / `x.__class__`), there is exactly one positional arg, + no keywords, and that arg is a string Constant. Any deviation raises `_Skip`, + so the constraint is skipped and never run. + """ + if not isinstance(node.func, ast.Name) or node.func.id not in _AGG: + raise _Skip() # not a whitelisted aggregate name + if node.keywords or len(node.args) != 1: + raise _Skip() # exactly one positional arg, no keywords + arg = node.args[0] + if not (isinstance(arg, ast.Constant) and isinstance(arg.value, str)): + raise _Skip() # arg must be a literal string path + ref = arg.value + # Split on the LAST dot: an ArraySpec.path may itself be dotted (e.g. + # 'invoice.lines'), while the column is always a row-LOCAL bare name. A bare + # ref (no dot) is a whole-array row count -> array_path=ref, column=''. + if "." in ref: + array_path, _, column = ref.rpartition(".") + else: + array_path, column = ref, "" + return node.func.id, array_path, column + + +def _aggregate(node: ast.Call, arrays: dict[str, list[dict[str, str]]]): + """Evaluate one of the five allowlisted aggregates over an array column. + + FAIL-CLOSED by construction: the only thing this accepts is `NAME('literal')` where + NAME is in `_AGG`, the func is a bare ast.Name (no attribute access — blocks + `os.system`/`x.__class__`), there is exactly one positional arg, no keywords, and that + arg is a string Constant. Any deviation raises `_Skip` (-> constraint skipped, never run). + + Argument is `'array_path'` (only valid for count = row count) or `'array_path.column'`. + Numeric cells (sum/min/max/avg) are pulled with the SAME `coerce_number` used elsewhere + (strips commas/$/%, drops non-finite/empty -> None); such cells are skipped (not zeroed). + `count('a.col')` counts rows whose column value is NON-EMPTY (text columns like + sku/description count too); `count('a')` is the row count. A sum/min/max/avg over zero + usable cells raises `_Skip` (advisory) rather than guessing 0, and sum/avg additionally + require EVERY row to contribute -- see the comment at that check. + + "FAIL-CLOSED" above is about the grammar, not the outcome: a deviation is + refused rather than executed. A `_Skip` is then reported by + `evaluate_constraints` as "no violation", which is the module docstring's + subject. + """ + fn, array_path, column = _parse_agg_call(node) + rows = arrays.get(array_path) + if rows is None: # no such array available + raise _Skip() + + if fn == "count": + if not column: # count('array') -> number of rows + return float(len(rows)) + # count('array.col') -> rows with a non-empty value for that column (text columns + # like sku/description/name are valid to count; numeric coercion would zero them out) + return float(sum(1 for r in rows if (r.get(column) or "").strip() != "")) + + if not column: # sum/min/max/avg need a column + raise _Skip() + nums = [n for r in rows if (n := coerce_number(r.get(column))) is not None] + if not nums: # nothing usable -> advisory skip + raise _Skip() + # A PARTIAL sum or average is not the column's sum or average -- it is a + # smaller number that will then be compared against a scalar key covering + # the whole column, and report a violation on a CORRECT document. One + # blank or un-parseable cell in a 40-row invoice is enough. That is the + # same false-positive class the comparison tolerance above exists to + # eliminate, arriving by a different route, so sum/avg require every row + # to contribute and otherwise skip (advisory) rather than guess. + # + # min/max/count are deliberately NOT subject to this: they are meaningful + # over a subset by construction (the smallest value present, the largest + # value present, the number of rows carrying one), and `count` has already + # returned above. + if fn in ("sum", "avg") and len(nums) != len(rows): + raise _Skip() + # `math.fsum`, not the builtin `sum`: exactly-rounded, so accumulation error + # does not grow with row count. The builtin accumulates left to right, and + # on a long array that drift is the one thing that could push a CORRECT + # total outside the comparison tolerance and report a false violation. + # + # Review raised exactly that risk for a 100,000-row array. It did not + # reproduce -- 60 randomised value/row combinations up to 100,000 rows and + # ~$100k per line produced zero false positives, and the specific figure + # quoted (`sum([8230.40] * 100_000) == 823039999.9985306`) is not what + # CPython computes; it is exactly 823040000.0, because for identical + # addends the partial sums stay representable. But "does not reproduce + # across 60 samples" is a weaker guarantee than "cannot happen", the fix + # costs one function call, and this is the third review pass over these few + # lines -- so the error source is removed rather than argued about. + if fn == "sum": + return float(math.fsum(nums)) + if fn == "min": + return float(min(nums)) + if fn == "max": + return float(max(nums)) + return float(math.fsum(nums) / len(nums)) # avg + + +def _operand(node, values: dict[str, str], arrays: dict[str, list[dict[str, str]]]): + if isinstance(node, ast.Call): + return _aggregate(node, arrays) # ONLY the _AGG allowlist; else _Skip + if isinstance(node, ast.Constant): + return node.value + if isinstance(node, ast.BinOp) and type(node.op) in _BIN: + a, b = _operand(node.left, values, arrays), _operand(node.right, values, arrays) + if not (isinstance(a, (int, float)) and isinstance(b, (int, float))): + raise _Skip() + return _BIN[type(node.op)](a, b) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + v = _operand(node.operand, values, arrays) + if not isinstance(v, (int, float)): + raise _Skip() + return -v + # else: a key reference + path = _path(node) + raw = values.get(path) + if raw is None or raw == "": + raise _Skip() + return _coerce(raw) + + +def _truth(node, values: dict[str, str], arrays: dict[str, list[dict[str, str]]]) -> bool: + if isinstance(node, ast.BoolOp): + sub = [_truth(v, values, arrays) for v in node.values] + return all(sub) if isinstance(node.op, ast.And) else any(sub) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not): + return not _truth(node.operand, values, arrays) + if isinstance(node, ast.Compare) and len(node.ops) == len(node.comparators): + return _compare_chain(node, values, arrays) + raise _Skip() + + +def _compare_chain( + node: ast.Compare, values: dict[str, str], arrays: dict[str, list[dict[str, str]]] +) -> bool: + """Evaluate a (possibly chained) comparison, Python-style: `a < b < c`. + + Each link is checked against the previous operand, so the chain short- + circuits on the first False exactly as the language does. + """ + left = _operand(node.left, values, arrays) + for op, comp in zip(node.ops, node.comparators, strict=False): + if type(op) not in _CMP: + raise _Skip() + right = _operand(comp, values, arrays) + # only compare like-typed operands (number<->number, str<->str); else skip + if isinstance(left, (int, float)) != isinstance(right, (int, float)): + raise _Skip() + if not _CMP[type(op)](left, right): + return False + left = right + return True + + +def _evaluate_one( + expr: str, values: dict[str, str], arrays: dict[str, list[dict[str, str]]] +): + """Return True/False, or None to skip (missing operand / unsupported / unsafe).""" + try: + tree = ast.parse(expr, mode="eval") + except SyntaxError: + # Unreachable from the API: `compile._check_constraint_syntax` parses + # every constraint at submit. Still logged, because reaching it means + # a constraint got in past that gate. + logger.warning("agent-kv constraint does not parse, skipping: %r", expr) + return None + try: + return _truth(tree.body, values, arrays) + except _Skip: + # The ordinary, expected drop: an operand is missing, empty, or not + # like-typed. Debug rather than warning -- on a document with optional + # keys this is normal and would otherwise be per-constraint noise. + logger.debug("agent-kv constraint skipped (operand unusable): %r", expr) + return None + except Exception: + # NOT fail-closed, whatever the previous comment on this line said: the + # caller turns None into "no violation", so this returns a pass for a + # check that never ran. `ZeroDivisionError` on a quantity of 0 reaches + # here. Keeping the catch is still right -- one bad constraint must not + # fail the whole extraction -- but it is logged at exception level so + # the drop is findable, which is the part that was missing. + logger.exception( + "agent-kv constraint raised and was DROPPED (reported as no " + "violation, not as a failure): %r", + expr, + ) + return None + + +def evaluate_constraints( + constraints: list[str], + values: dict[str, str], + arrays: dict[str, list[dict[str, str]]] | None = None, +) -> list[str]: + """Return the list of constraint expressions that evaluated to False (violations). + Skipped (missing operand / unsupported / unsafe) constraints are NOT violations. + + `arrays` (optional) maps an ArraySpec path to its rendered rows (list of {column: value}) + so the five allowlisted aggregates (sum/count/min/max/avg) can reconcile a scalar key + against an array column. Defaults to {} -> aggregates become no-op skips (back-compat). + """ + arrays = arrays or {} + violations = [] + for expr in constraints or []: + if _evaluate_one(expr, values, arrays) is False: + violations.append(expr) + return violations diff --git a/unstract/agent-kv-schema/src/unstract/agent_kv_schema/dataclasses.py b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/dataclasses.py new file mode 100644 index 0000000000..544c9c5efb --- /dev/null +++ b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/dataclasses.py @@ -0,0 +1,61 @@ +"""Dataclasses for the key-value extractor. + +ExecutionResult is intentionally NOT defined here — code execution reuses the +imported src/core code_executor's result type (see spec §4, §10). +""" + +from dataclasses import dataclass, field + +# Both specs are FROZEN. They are compile output: produced once by +# `kv_schema._walk` and then read by the prompt generator, the QA pass and the +# constraint evaluator. Nothing assigns to a field (the one derived variant, +# `normalizers`' scalar view of a multivalued leaf, already goes through +# `dataclasses.replace`), so freezing costs nothing and makes a stray mutation +# in a later stage a loud AttributeError instead of a spec that disagrees with +# the schema the caller submitted. + + +@dataclass(frozen=True) +class KeySpec: + """One compiled leaf key from the user's nested key schema.""" + + path: str # dotted path, e.g. "vendor.address.city" + effective_description: str # breadcrumb + leaf description + format: str = "string" # kind: string|number|date|currency|enum|regex| + enum_values: list[str] = field(default_factory=list) # set when format == "enum" + regex_pattern: str = "" # set when format == "regex" + required: bool = False + aliases: list[str] = field(default_factory=list) + multivalued: bool = False # value is a comma-separated string + + +@dataclass(frozen=True) +class ArraySpec: + """One compiled FLAT-array node (P8a). `path` = dotted array location (e.g. 'line_items', + 'invoice.lines'). `item_specs` = the declared columns as row-LOCAL scalar KeySpecs (their + `path` is the bare column name). `key_column` (optional) is a column used for row identity in + scoring; '' = positional. Nested arrays inside an item are P8b and rejected at compile. + """ + + path: str + description: str = "" + item_specs: list[KeySpec] = field(default_factory=list) + key_column: str = "" + required: bool = False + #: Collapse rows identical in EVERY extracted cell? Declared per array as + #: `"_dedup": false`. + #: + #: Defaults True, which is the pre-existing behaviour and is wanted for the + #: corpus the extractor was built against: layout-preserving OCR repeats a + #: label once per replicate column, producing rows identical in every cell. + #: + #: But it is NOT lossless, and the docstring that claimed it was has been + #: corrected. A document that genuinely contains two identical line items + #: comes back with one, so the row count and any calculation summing or + #: counting the rows are wrong -- and the codegen path consumes exactly + #: these rows. Being unable to tell two real records from one OCR artefact + #: is a reason not to guess, so a schema author who knows their documents + #: can say so. Default left at True rather than flipped: changing it would + #: trade a known-wrong case for an untested one on the corpus this was + #: built for. + dedup_rows: bool = True diff --git a/unstract/agent-kv-schema/src/unstract/agent_kv_schema/kv_schema.py b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/kv_schema.py new file mode 100644 index 0000000000..04235763d2 --- /dev/null +++ b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/kv_schema.py @@ -0,0 +1,251 @@ +"""Compile a user-authored nested key schema into a flat, ordered KeySpec list, +and reassemble flat extracted values back into the nested shape. + +Leaf/interior rule (spec §7): a node is INTERIOR iff ALL of its values are objects +(its keys are child nodes); otherwise it is a LEAF defined by its scalar attributes. +A node mixing object and scalar values is a compile-time error. Reserved attribute +words cannot be used as child-node key names. +""" + +from typing import Any + +from .dataclasses import ArraySpec, KeySpec + +RESERVED = {"description", "format", "required", "aliases", "multivalued"} +RESERVED_NODE = { + "_array", + "description", + "_key", + "_dedup", +} # node-level (distinct from leaf-attr RESERVED) + +# Absolute structural-depth ceiling for callers that don't pass an explicit +# max_depth (the public compile/compile_arrays entry points). It only has to +# sit well below Python's ~1000 recursion limit so a hostile keys.json fails +# fast as a clean ValueError (-> SchemaError) instead of a RecursionError; +# compile_schema always passes the real, much smaller cap from SchemaCaps. +_DEFAULT_MAX_DEPTH = 100 + + +def _parse_dedup(raw: Any, path: str) -> bool: + """Validate the per-array `_dedup` flag. + + Strict about the type rather than truthy: `"_dedup": "false"` is a string + and would be TRUE under a truthiness check, silently leaving dedup on for + an author who explicitly asked for it off -- and silently dropping their + duplicate rows. A wrong type is a schema error, not a default. + """ + if isinstance(raw, bool): + return raw + raise ValueError(f"'{path}' _dedup must be true or false (got {type(raw).__name__})") + + +def _parse_format(raw: str): + """Return (kind, enum_values, regex_pattern) for a declared format string.""" + if raw.startswith("enum:"): + values = [v.strip() for v in raw[len("enum:") :].split(",") if v.strip()] + return "enum", values, "" + if raw.startswith("regex:"): + return "regex", [], raw[len("regex:") :] + return raw, [], "" + + +def _walk( + node: dict[str, Any], + path_parts: list[str], + out: list[KeySpec], + arrays: list[ArraySpec], + *, + max_depth: int = _DEFAULT_MAX_DEPTH, + depth: int = 0, +) -> None: + # HARD structural-depth ceiling at the TOP of the recursive walk. This is + # the real backstop the compile.py `_max_depth` pre-check can't provide: + # that check treats array columns as row-local (not nesting) and is blind + # to a decoy top-level `_array` field's real nesting, so a schema can pass + # it yet still recurse arbitrarily deep here. `depth` increments on EVERY + # recursion (interior children AND array columns), so exceeding the cap + # raises a clean ValueError -- surfaced as SchemaError by compile_schema -- + # instead of an uncaught RecursionError. + if depth > max_depth: + raise ValueError(f"schema exceeds max_depth={max_depth}") + if not isinstance(node, dict) or not node: + raise ValueError( + f"Schema node at {'.'.join(path_parts) or ''} must be a non-empty object" + ) + + if "_array" in node: + _walk_array(node, path_parts, arrays, max_depth=max_depth, depth=depth) + return + + object_values = [v for v in node.values() if isinstance(v, dict)] + scalar_keys = [k for k, v in node.items() if not isinstance(v, dict)] + + is_interior = len(object_values) == len(node) # every value is an object + is_leaf = len(object_values) == 0 # no value is an object + + if not is_interior and not is_leaf: + raise ValueError( + f"Schema node at {'.'.join(path_parts) or ''} mixes object children " + f"and scalar attributes; an interior node's values must all be objects" + ) + + if is_interior: + for child_name, child_node in node.items(): + _walk( + child_node, + path_parts + [child_name], + out, + arrays, + max_depth=max_depth, + depth=depth + 1, + ) + return + + _append_leaf(node, path_parts, scalar_keys, out) + + +def _walk_array( + node: dict[str, Any], + path_parts: list[str], + arrays: list[ArraySpec], + *, + max_depth: int, + depth: int, +) -> None: + """Compile an `_array` node into one ArraySpec. + + Column paths are row-LOCAL: each column walks with its own name as the + whole path and an empty `arrays` sink, so a nested array would have nowhere + to land -- which is why it is rejected explicitly first rather than left to + disappear into that sink. + """ + extra = set(node.keys()) - RESERVED_NODE + if extra: + raise ValueError( + f"Array node at {'.'.join(path_parts) or ''} has unexpected keys {sorted(extra)}; " + f"allowed: {sorted(RESERVED_NODE)}" + ) + item_schema = node["_array"] + if not isinstance(item_schema, dict) or not item_schema: + raise ValueError( + f"Array item schema at {'.'.join(path_parts)} must be a non-empty object" + ) + item_specs: list[KeySpec] = [] + for col_name, col_node in item_schema.items(): + if isinstance(col_node, dict) and "_array" in col_node: + raise ValueError( + f"Nested array at {'.'.join(path_parts)}.{col_name} is P8b (not supported in P8a)" + ) + _walk( + col_node, + [col_name], + item_specs, + [], + max_depth=max_depth, + depth=depth + 1, + ) + arrays.append( + ArraySpec( + path=".".join(path_parts), + description=str(node.get("description", "")), + item_specs=item_specs, + key_column=str(node.get("_key", "")), + dedup_rows=_parse_dedup(node.get("_dedup", True), ".".join(path_parts)), + ) + ) + + +def _append_leaf( + node: dict[str, Any], + path_parts: list[str], + scalar_keys: list[str], + out: list[KeySpec], +) -> None: + """Compile a leaf node (scalar attributes only) into one KeySpec.""" + unknown = [k for k in scalar_keys if k not in RESERVED] + if unknown: + raise ValueError( + f"Leaf at {'.'.join(path_parts)} has unknown attribute(s) {unknown}; " + f"allowed: {sorted(RESERVED)}" + ) + if "description" not in node: + raise ValueError(f"Leaf at {'.'.join(path_parts)} is missing a 'description'") + + kind, enum_values, regex_pattern = _parse_format(str(node.get("format", "string"))) + breadcrumb = " > ".join(path_parts) + out.append( + KeySpec( + path=".".join(path_parts), + effective_description=f"{breadcrumb}: {node['description']}", + format=kind, + enum_values=enum_values, + regex_pattern=regex_pattern, + required=bool(node.get("required", False)), + aliases=list(node.get("aliases", [])), + multivalued=bool(node.get("multivalued", False)), + ) + ) + + +def _compile_both(spec_json: dict[str, Any], max_depth: int = _DEFAULT_MAX_DEPTH): + if not isinstance(spec_json, dict): + raise ValueError("Top-level key schema must be a JSON object") + out: list[KeySpec] = [] + arrays: list[ArraySpec] = [] + for top_name, top_node in spec_json.items(): + if top_name == "_constraints": + continue + _walk(top_node, [top_name], out, arrays, max_depth=max_depth) + return out, arrays + + +def compile( + spec_json: dict[str, Any], max_depth: int = _DEFAULT_MAX_DEPTH +) -> list[KeySpec]: + """Compile to an ordered flat list of SCALAR leaf KeySpecs (array nodes excluded).""" + return _compile_both(spec_json, max_depth)[0] + + +def compile_arrays( + spec_json: dict[str, Any], max_depth: int = _DEFAULT_MAX_DEPTH +) -> list[ArraySpec]: + """Compile the top-level/interior-nested array nodes to an ordered list of ArraySpecs.""" + return _compile_both(spec_json, max_depth)[1] + + +def reassemble( + values: dict[str, Any], + specs: list[KeySpec], + arrays: list[ArraySpec] = None, + array_values: dict[str, Any] = None, +) -> dict[str, Any]: + """Rebuild the nested dict from flat {dotted_path: value}. + + Iterates `specs` (so output key order follows the schema) and places each + present value at its dotted path. Multi-valued leaves are comma-separated + strings and pass through verbatim. Paths absent from `values` are skipped. + + Arrays (P8a): for each ArraySpec in `arrays`, place its pre-rendered + list-of-dicts (from `array_values[path]`) at its dotted path. Scalar-only + callers (arrays/array_values=None) behave identically to before. + """ + out: dict[str, Any] = {} + for spec in specs: + if spec.path not in values: + continue + parts = spec.path.split(".") + node = out + for part in parts[:-1]: + node = node.setdefault(part, {}) + node[parts[-1]] = values[spec.path] + for aspec in arrays or []: + rows = (array_values or {}).get(aspec.path) + if rows is None: + continue + parts = aspec.path.split(".") + node = out + for part in parts[:-1]: + node = node.setdefault(part, {}) + node[parts[-1]] = rows # list-of-dicts + return out diff --git a/unstract/agent-kv-schema/src/unstract/agent_kv_schema/validators.py b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/validators.py new file mode 100644 index 0000000000..e05fabdd3b --- /dev/null +++ b/unstract/agent-kv-schema/src/unstract/agent_kv_schema/validators.py @@ -0,0 +1,115 @@ +"""Deterministic format/type validation for extracted values (QA, spec §9 check 2). + +An empty value passes here (presence is the separate completeness check). For a +multivalued leaf each comma-separated element is validated against the format. +""" + +import math +import re +from datetime import datetime + +from .dataclasses import KeySpec + +_NUMERIC_STRIP = re.compile(r"[,\s$%]") +_DATE_FORMATS = ( + "%m/%d/%Y", + "%m/%d/%y", + "%Y-%m-%d", + "%d/%m/%Y", + "%m-%d-%Y", + "%B %d, %Y", + "%b %d, %Y", + "%m/%Y", + "%Y", +) + + +def _is_number(value: str) -> bool: + cleaned = _NUMERIC_STRIP.sub("", value) + if cleaned in ("", "-", "."): + return False + try: + float(cleaned) + return True + except ValueError: + return False + + +def _is_date(value: str) -> bool: + v = value.strip() + for fmt in _DATE_FORMATS: + try: + datetime.strptime(v, fmt) + return True + except ValueError: + continue + return False + + +def _check_one(value: str, spec: KeySpec) -> bool: + fmt = spec.format + if fmt in ("number", "currency"): + return _is_number(value) + if fmt == "date": + return _is_date(value) + if fmt == "enum": + return value.strip().casefold() in {e.casefold() for e in spec.enum_values} + if fmt == "regex": + try: + return re.fullmatch(spec.regex_pattern, value.strip()) is not None + except re.error: + return True # a broken author-supplied pattern should not fail extraction + # "string" or any free-text hint: nothing to validate. + return True + + +def validate_format(value: str, spec: KeySpec) -> bool: + """True if `value` conforms to `spec.format`. Empty value -> True (completeness is + checked separately). Multivalued -> every element must conform. + """ + if value is None or value.strip() == "": + return True + if spec.multivalued: + elements: list[str] = [e.strip() for e in value.split(",") if e.strip()] + return all(_check_one(e, spec) for e in elements) if elements else True + return _check_one(value, spec) + + +def coerce_number(value: str) -> float | None: + """Parse a numeric/currency string to a float, or None. Strips commas/spaces/$/%.""" + if value is None: + return None + cleaned = _NUMERIC_STRIP.sub("", value) + if cleaned in ("", "-", "."): + return None + try: + n = float(cleaned) + except ValueError: + return None + # float() accepts 'inf'/'nan'/'1e400' as literals — reject non-finite so callers + # (output normalization, constraint comparison) never see inf/nan. + return n if math.isfinite(n) else None + + +def coerce_date(value: str) -> str | None: + """Parse a date string to ISO 'YYYY-MM-DD', or None.""" + if value is None: + return None + v = value.strip() + for fmt in _DATE_FORMATS: + try: + return datetime.strptime(v, fmt).strftime("%Y-%m-%d") + except ValueError: + continue + return None + + +def coerce_enum(value: str, spec: KeySpec) -> "str | None": + """Return the canonical `spec.enum_values` element matching `value` case-insensitively, or None.""" + if value is None: + return None + cf = value.strip().casefold() + for e in spec.enum_values: + if e.casefold() == cf: + return e + return None diff --git a/unstract/agent-kv-schema/tests/__init__.py b/unstract/agent-kv-schema/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/unstract/agent-kv-schema/tests/test_compile.py b/unstract/agent-kv-schema/tests/test_compile.py new file mode 100644 index 0000000000..471df824c7 --- /dev/null +++ b/unstract/agent-kv-schema/tests/test_compile.py @@ -0,0 +1,463 @@ +"""compile_schema: structural validation + caps on top of the ported compiler.""" + +from dataclasses import FrozenInstanceError + +import pytest + +from unstract.agent_kv_schema import ( + CompiledSchema, + SchemaCaps, + SchemaError, + compile_schema, +) + +VALID = { + "quotation_number": {"description": "The quote number", "required": True}, + "customer": {"name": {"description": "Bill-to name"}}, + "line_items": { + "description": "One row per line", + "_key": "sku", + "_array": { + "sku": {"description": "SKU"}, + "total": {"description": "Line total", "format": "currency"}, + }, + }, + "_constraints": ["count('line_items') >= 1"], +} + + +def test_valid_schema_compiles(): + out = compile_schema(VALID) + assert isinstance(out, CompiledSchema) + assert [s.path for s in out.key_specs] == ["quotation_number", "customer.name"] + assert out.array_specs[0].path == "line_items" + assert out.constraints == ["count('line_items') >= 1"] + + +def test_missing_description_is_schema_error(): + with pytest.raises(SchemaError, match="missing a 'description'"): + compile_schema({"a": {"format": "string"}}) + + +def test_mixed_node_is_schema_error(): + with pytest.raises(SchemaError, match="mixes object children"): + compile_schema({"a": {"description": "x", "b": {"description": "y"}}}) + + +def test_leaf_cap_enforced(): + spec = {f"k{i}": {"description": "d"} for i in range(5)} + caps = SchemaCaps(max_leaves=4) + with pytest.raises(SchemaError, match="max_leaves"): + compile_schema(spec, caps=caps) + + +def test_depth_cap_enforced(): + spec = {"a": {"b": {"c": {"d": {"description": "deep"}}}}} + caps = SchemaCaps(max_depth=3) + with pytest.raises(SchemaError, match="max_depth"): + compile_schema(spec, caps=caps) + + +def test_regex_length_cap(): + spec = {"a": {"description": "d", "format": "regex:" + "x" * 300}} + with pytest.raises(SchemaError, match="max_regex_len"): + compile_schema(spec) + + +def test_bad_constraint_syntax_rejected(): + spec = { + "a": {"description": "d"}, + "_constraints": ["__import__('os').system('true')"], + } + with pytest.raises(SchemaError, match="constraint"): + compile_schema(spec) + + +def test_constraint_call_allowlist(): + spec = {"a": {"description": "d"}, "_constraints": ["foo('a.b') > 1"]} + with pytest.raises(SchemaError, match="constraint"): + compile_schema(spec) + + +def test_constraints_cap(): + spec = {"a": {"description": "d"}, "_constraints": ["a > 0"] * 31} + with pytest.raises(SchemaError, match="max_constraints"): + compile_schema(spec) + + +def test_non_dict_top_level_rejected(): + with pytest.raises(SchemaError): + compile_schema(["not", "a", "dict"]) + + +# --------------------------------------------------------------------------- +# Depth-cap bypass + unbounded-recursion DoS (pre-Greptile critical #2). +# --------------------------------------------------------------------------- + + +def _nest(levels: int) -> dict: + """Build a plain interior schema `levels` deep ending in a leaf.""" + node = {"description": "deep"} + for i in range(levels): + node = {f"l{i}": node} + return node + + +def test_deeply_nested_raises_schema_error_not_recursion_error(): + # Well past max_depth but nowhere near Python's recursion limit: must be a + # clean SchemaError (400), never accepted. + deep, caps = _nest(30), SchemaCaps(max_depth=6) + with pytest.raises(SchemaError, match="max_depth"): + compile_schema(deep, caps=caps) + + +def test_decoy_array_key_cannot_bypass_depth_cap(): + # A top-level field literally named "_array" short-circuited the old + # depth walk (`if "_array" in node: return depth + 1`), so its real + # nesting was never counted. The compiler treats "_array" as a plain + # field name here (an array NODE is one CONTAINING an `_array` key), so + # the nesting under it is real -- and must still be capped. + decoy = {"_array": _nest(30)} + caps = SchemaCaps(max_depth=6) + with pytest.raises(SchemaError, match="max_depth"): + compile_schema(decoy, caps=caps) + + +def test_valid_array_within_limits_still_compiles(): + # Genuine array-column spec (node CONTAINING `_array`) with shallow + # columns must still compile -- the ceiling bounds total structural + # depth, it does not break real arrays. + spec = { + "rows": { + "description": "line items", + "_array": { + "sku": {"description": "SKU"}, + "qty": {"description": "Quantity", "format": "number"}, + }, + } + } + out = compile_schema(spec, caps=SchemaCaps(max_depth=6)) + assert isinstance(out, CompiledSchema) + assert out.array_specs[0].path == "rows" + assert [s.path for s in out.array_specs[0].item_specs] == ["sku", "qty"] + + +def test_pathologically_deep_raises_fast_no_recursion_error(): + # ~5000 levels: 256 KiB of JSON nests far past Python's ~1000 recursion + # limit, so the OLD code raised an uncaught RecursionError (500). Must be + # a clean, fast SchemaError instead. + deep, caps = _nest(5000), SchemaCaps(max_depth=6) + with pytest.raises(SchemaError, match="max_depth"): + compile_schema(deep, caps=caps) + + +def test_pathologically_deep_under_decoy_array_raises_fast(): + # Same DoS but hidden under a decoy `_array` key so the compile.py depth + # pre-check short-circuits -- the guard inside the real recursive walk + # (`kv_schema._walk`) must still catch it as SchemaError, not + # RecursionError. + decoy, caps = {"_array": _nest(5000)}, SchemaCaps(max_depth=6) + with pytest.raises(SchemaError, match="max_depth"): + compile_schema(decoy, caps=caps) + + +# --------------------------------------------------------------------------- +# Author-supplied regex is refused at SUBMIT, not discovered at match time. +# +# Two defects, both found in review: +# 1. The pattern was never compiled here, so an invalid one was accepted and +# only hit per-value in the engine's QA pass -- where `_check_one` swallows +# `re.error` and returns True, silently passing validation the author +# thought they had configured. +# 2. `validate_format` runs the pattern with no time budget, so catastrophic +# backtracking is a DoS. The length cap is NOT a mitigation: `^(a+)+$` is 7 +# characters and takes ~1.9s against 26 `a`s, ~4x per character added +# (measured on this code), so a 40-char value runs for hours -- and +# AGENT_KV_CONCURRENT_LIMIT=5 lets one org pin five shared worker slots +# from a single submit. +# +# The detector is a conservative heuristic, not a proof: it rejects the shape +# behind the realistic cases (a quantifier inside a quantified group). The +# complete fix is a linear-time engine (RE2), which cannot live here because +# this package deliberately has zero dependencies -- UN-4225. +# --------------------------------------------------------------------------- + + +def _schema(pattern: str) -> dict: + return {"amount": {"description": "x", "format": f"regex:{pattern}"}} + + +@pytest.mark.parametrize( + "pattern", + [ + "^(a+)+$", # the classic + "^(a*)*$", + "(?:x+)+", # non-capturing group is not an escape hatch + "^(a+){2,}$", # counted repetition of a quantified group + ], +) +def test_nested_quantifier_is_refused_at_compile(pattern): + spec = _schema(pattern) + with pytest.raises(SchemaError, match="nested quantifier"): + compile_schema(spec) + + +def test_invalid_regex_is_refused_at_compile_not_swallowed_at_match_time(): + spec = _schema("[") + with pytest.raises(SchemaError, match="invalid regex"): + compile_schema(spec) + + +@pytest.mark.parametrize( + "pattern", + [ + r"^\d{3}-\d{4}$", + r"^[A-Z]{2}\d+$", + r"^INV-\d+$", + r"^\$?[\d,]+\.\d{2}$", + ], +) +def test_ordinary_patterns_still_compile(pattern): + """The other half of the heuristic's contract. A detector that rejected + real-world patterns would just push authors off the regex format entirely. + """ + compiled = compile_schema(_schema(pattern)) + assert compiled.key_specs[0].regex_pattern == pattern + + +# The overlapping-alternation family, found in review AFTER the +# nested-quantifier check shipped: `^(a|aa)+$` passes that check and still +# backtracks catastrophically, because at each position the engine can consume +# one `a` or two and must try both on failure. +@pytest.mark.parametrize("pattern", ["^(a|aa)+$", "^(a|ab)+$", "(?:x|xx)+"]) +def test_overlapping_alternation_is_refused(pattern): + spec = _schema(pattern) + with pytest.raises(SchemaError, match="alternatives overlap"): + compile_schema(spec) + + +@pytest.mark.parametrize( + "pattern", + [ + "^(foo|bar)+$", # distinct first chars -> no position admits two parses + r"^(\d|[A-Z])+$", # metacharacter branches are not analysed + "^INV-(A|B)+$", + ], +) +def test_non_overlapping_alternation_still_compiles(pattern): + """The detector compares only LITERAL branches, by the prefix relation. + Refusing every quantified alternation would reject a lot of safe, ordinary + patterns -- `(?:a|b)+` is linear. + """ + compile_schema(_schema(pattern)) + + +# -------------------------------------------------------------------------- +# 2.18: formats that can never validate anything, or always fail. +# +# Free-text `format` stays an LLM hint -- that is deliberate, see the `format` +# comment on KeySpec. The three shapes below are not hints, they are mistakes +# with no observable symptom, all three verified by execution before the fix. +# -------------------------------------------------------------------------- + + +def _one_key(fmt: str) -> dict: + return {"amount": {"description": "An amount", "format": fmt}} + + +@pytest.mark.parametrize( + ("fmt", "intended"), + [ + ("Number", "number"), + ("NUMBER", "number"), + ("Date", "date"), + ("Currency", "currency"), + ("String", "string"), + ("Enum", "enum"), + ("Regex", "regex"), + ], +) +def test_a_case_variant_of_a_known_format_is_refused(fmt, intended): + """`format: "Number"` silently became a free-text hint. + + `validate_format("not a number at all", spec)` then returned True for + every value for the life of the job, and `/validate` reported + `{"valid": true}` -- the author configured validation and got none. + """ + spec = _one_key(fmt) + with pytest.raises(SchemaError, match=f"Did you mean '{intended}'"): + compile_schema(spec) + + +@pytest.mark.parametrize( + ("fmt", "intended"), + [ + ("Enum:paid,unpaid", "enum:paid,unpaid"), + ("ENUM:paid,unpaid", "enum:paid,unpaid"), + ("Regex:^[0-9]+$", "regex:^[0-9]+$"), + ("REGEX:^[0-9]+$", "regex:^[0-9]+$"), + ], +) +def test_a_case_variant_of_an_ARGUMENT_taking_format_is_refused(fmt, intended): + """The residual the first round's fix left behind. + + `test_a_case_variant_of_a_known_format_is_refused` above covers the BARE + names, including bare `"Enum"` and `"Regex"` -- but not the forms that + actually carry an argument, which is how anyone real writes them. + `_parse_format` matches the `enum:` / `regex:` prefix case-sensitively, so + `"Enum:paid,unpaid"` never becomes kind `"enum"`; it reached the typo guard + as the whole raw string, and `"enum:paid,unpaid"` is not a member of + `_KNOWN_FORMATS` (which holds the bare `"enum"`). + + Net effect before this: the key compiled as a free-text hint, `/validate` + reported `{"valid": true}`, and the enum was never enforced -- identical to + the `"Number"` failure, through the one gap the fix for it did not close. + """ + with pytest.raises(SchemaError, match="Did you mean"): + compile_schema(_one_key(fmt)) + + # And the suggestion has to be the usable spelling, argument included -- + # `"enum"` alone would be rejected by the no-values check one line later, + # sending the author in a circle. + with pytest.raises(SchemaError) as caught: + compile_schema(_one_key(fmt)) + assert intended in str(caught.value) + + +@pytest.mark.parametrize( + "fmt", + ["enum:paid,unpaid", "regex:^[0-9]+$", "enum:a,b,c"], +) +def test_correctly_cased_argument_formats_still_compile(fmt): + """The other half: the fix must not reject the valid spellings. + + These reach the guard with kind already resolved to `"enum"`/`"regex"` by + `_parse_format`, so the colon-splitting branch must not fire on them. + """ + compile_schema(_one_key(fmt)) + + +@pytest.mark.parametrize( + "fmt", + [ + "customer:id", # a colon in genuine free text + "total:in USD", + "ref:12345", + ], +) +def test_free_text_containing_a_colon_is_not_mistaken_for_a_format(fmt): + """Splitting on `:` must not capture unrelated hints. + + The prefix before the colon is not a known format in any of these, so they + stay free text -- which is what a schema author writing a descriptive hint + expects. + """ + compile_schema(_one_key(fmt)) + + +@pytest.mark.parametrize("fmt", ["number", "date", "currency", "string"]) +def test_the_known_formats_still_compile(fmt): + compile_schema(_one_key(fmt)) + + +@pytest.mark.parametrize( + "fmt", + [ + "the invoice total in USD", # a real hint, which is the point of free text + "amount", + "numeric-ish", # near a known format but not a case variant of one + "dates", + ], +) +def test_free_text_format_is_still_accepted_as_an_llm_hint(fmt): + compile_schema(_one_key(fmt)) + + +@pytest.mark.parametrize("fmt", ["enum:", "enum", "enum: , ,"]) +def test_an_enum_with_no_values_is_refused(fmt): + """`enum_values=[]` means `_check_one` tests membership of the empty set, + so every non-empty value fails QA forever. This does not disable + validation, it inverts it. + """ + spec = _one_key(fmt) + with pytest.raises(SchemaError, match="enum with no values"): + compile_schema(spec) + + +def test_an_enum_with_values_still_compiles(): + compiled = compile_schema(_one_key("enum:paid,unpaid")) + assert compiled.key_specs[0].enum_values == ["paid", "unpaid"] + + +@pytest.mark.parametrize("fmt", ["regex:", "regex"]) +def test_an_empty_regex_is_refused(fmt): + spec = _one_key(fmt) + with pytest.raises(SchemaError, match="empty regex"): + compile_schema(spec) + + +def test_the_same_rules_apply_to_array_columns(): + """Array columns are leaves too, and were checked by nothing.""" + spec = { + "lines": { + "description": "One row per line", + "_array": {"total": {"description": "Line total", "format": "Currency"}}, + } + } + with pytest.raises(SchemaError, match="Did you mean 'currency'"): + compile_schema(spec) + + +# -------------------------------------------------------------------------- +# 2.18 (second half): `_key` naming a column that does not exist. +# -------------------------------------------------------------------------- + + +def test_a_key_column_that_does_not_exist_is_refused(): + """A `_key` typo is not an error anywhere downstream. + + Every row misses, the array silently falls back to positional identity, + and the cost is accuracy on exactly the arrays the author cared enough + about to key. Nothing is reported. + """ + spec = { + "lines": { + "description": "One row per line", + "_key": "skew", # the column is 'sku' + "_array": { + "sku": {"description": "SKU"}, + "total": {"description": "Line total"}, + }, + } + } + with pytest.raises(SchemaError, match="_key 'skew'"): + compile_schema(spec) + + +def test_a_key_column_that_exists_still_compiles(): + compiled = compile_schema(VALID) + assert [a.key_column for a in compiled.array_specs] == ["sku"] + + +def test_an_array_with_no_key_column_is_positional_and_accepted(): + spec = { + "lines": { + "description": "One row per line", + "_array": {"sku": {"description": "SKU"}}, + } + } + compiled = compile_schema(spec) + assert compiled.array_specs[0].key_column == "" + + +# -------------------------------------------------------------------------- +# 2.18 (third): the specs are compile OUTPUT and must not be mutable. +# -------------------------------------------------------------------------- + + +def test_compiled_specs_are_frozen(): + compiled = compile_schema(VALID) + with pytest.raises(FrozenInstanceError): + compiled.key_specs[0].format = "number" + with pytest.raises(FrozenInstanceError): + compiled.array_specs[0].key_column = "total" diff --git a/unstract/agent-kv-schema/tests/test_constraints.py b/unstract/agent-kv-schema/tests/test_constraints.py new file mode 100644 index 0000000000..43fa014923 --- /dev/null +++ b/unstract/agent-kv-schema/tests/test_constraints.py @@ -0,0 +1,258 @@ +"""The constraint evaluator's operator/aggregate matrix, and its drop behaviour. + +This file was 0 bytes against a 258-line module -- an empty file with that name +reads as coverage, which is worse than no file (reported as 2.20 in the branch +review). What it covers now is the part the review found wrong (2.19): the +module called itself fail-closed while a constraint that raises is reported as +"no violation", and a partial sum over a column with one blank cell produced a +false violation on a correct document. + +`test_constraints_tolerance.py` covers the float-tolerance derivation; this +covers everything else. +""" + +import ast +import logging + +import pytest + +from unstract.agent_kv_schema import evaluate_constraints +from unstract.agent_kv_schema.compile import _ALLOWED_CALLS, _ALLOWED_NODES +from unstract.agent_kv_schema.constraints import _AGG, _BIN, _CMP + +_ROWS = { + "lines": [ + {"amount": "100.00", "qty": "2", "sku": "A"}, + {"amount": "50.00", "qty": "1", "sku": "B"}, + ] +} + + +# -------------------------------------------------------------------------- +# The comparison and arithmetic matrix +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("expr", "values", "violated"), + [ + ("a == b", {"a": "1", "b": "1"}, False), + ("a == b", {"a": "1", "b": "2"}, True), + ("a != b", {"a": "1", "b": "2"}, False), + ("a != b", {"a": "1", "b": "1"}, True), + ("a < b", {"a": "1", "b": "2"}, False), + ("a < b", {"a": "2", "b": "2"}, True), + ("a <= b", {"a": "2", "b": "2"}, False), + ("a > b", {"a": "3", "b": "2"}, False), + ("a >= b", {"a": "2", "b": "2"}, False), + # chained, Python-style + ("a < b < c", {"a": "1", "b": "2", "c": "3"}, False), + ("a < b < c", {"a": "1", "b": "3", "c": "2"}, True), + # arithmetic + ("total == net + tax", {"total": "110", "net": "100", "tax": "10"}, False), + ("total == net + tax", {"total": "111", "net": "100", "tax": "10"}, True), + ("net == total - tax", {"net": "100", "total": "110", "tax": "10"}, False), + ("total == qty * price", {"total": "20", "qty": "4", "price": "5"}, False), + ("price == total / qty", {"price": "5", "total": "20", "qty": "4"}, False), + ("a == -b", {"a": "-5", "b": "5"}, False), + # boolean composition + ("a > 0 and b > 0", {"a": "1", "b": "1"}, False), + ("a > 0 and b > 0", {"a": "1", "b": "-1"}, True), + ("a > 0 or b > 0", {"a": "-1", "b": "1"}, False), + ("not (a > b)", {"a": "1", "b": "2"}, False), + # thousands separators and currency symbols coerce numerically + ("a == b", {"a": "9,000", "b": "9000"}, False), + ("a == b", {"a": "$1,234.50", "b": "1234.5"}, False), + # ISO dates compare lexicographically as strings + ("start < end", {"start": "2026-01-01", "end": "2026-02-01"}, False), + ("start < end", {"start": "2026-03-01", "end": "2026-02-01"}, True), + ], +) +def test_operator_matrix(expr, values, violated): + assert evaluate_constraints([expr], values, {}) == ([expr] if violated else []) + + +@pytest.mark.parametrize( + "expr", + [ + "a == b", # 'a' missing entirely + "missing == 1", # operand absent + "a == 1", # operand empty string + "a < b", # number vs string: not like-typed + "a == 'nan'", # non-finite token stays a string + ], +) +def test_unusable_operands_are_skipped_not_violations(expr): + values = {"a": "", "b": "text"} + assert evaluate_constraints([expr], values, {}) == [] + + +# -------------------------------------------------------------------------- +# Aggregates +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("expr", "values", "violated"), + [ + ("total == sum('lines.amount')", {"total": "150.00"}, False), + ("total == sum('lines.amount')", {"total": "150.01"}, True), + ("n == count('lines')", {"n": "2"}, False), + ("n == count('lines')", {"n": "3"}, True), + ("n == count('lines.sku')", {"n": "2"}, False), + ("lo == min('lines.amount')", {"lo": "50"}, False), + ("hi == max('lines.amount')", {"hi": "100"}, False), + ("mean == avg('lines.amount')", {"mean": "75"}, False), + ("mean == avg('lines.amount')", {"mean": "76"}, True), + ], +) +def test_aggregate_matrix(expr, values, violated): + assert evaluate_constraints([expr], values, _ROWS) == ([expr] if violated else []) + + +@pytest.mark.parametrize( + "expr", + [ + "x == sum('nosucharray.amount')", # array not supplied + "x == sum('lines')", # sum needs a column + "x == avg('lines')", + "x == min('lines')", + "x == sum('lines.sku')", # no numeric cell at all + "x == median('lines.amount')", # not in the allowlist + "x == sum('lines.amount', 2)", # two args + "x == sum(column)", # arg is not a string literal + ], +) +def test_unusable_aggregates_are_skipped_not_violations(expr): + assert evaluate_constraints([expr], {"x": "150"}, _ROWS) == [] + + +def test_count_of_a_column_counts_non_empty_cells_only(): + rows = {"lines": [{"sku": "A"}, {"sku": ""}, {"sku": " "}, {"sku": "C"}]} + assert evaluate_constraints(["n == count('lines.sku')"], {"n": "2"}, rows) == [] + assert evaluate_constraints(["n == count('lines')"], {"n": "4"}, rows) == [] + + +# -------------------------------------------------------------------------- +# 2.19: a PARTIAL sum/avg must skip, not report a false violation +# -------------------------------------------------------------------------- + +_PARTIAL = { + "lines": [ + {"amount": "100.00"}, + {"amount": ""}, # blank cell: OCR missed it, or the row has no amount + {"amount": "50.00"}, + ] +} + + +def test_a_partial_sum_skips_rather_than_reporting_a_false_violation(): + """The document is CORRECT: 100 + + 50 against a stated 200. + + Summing the two readable cells gives 150, which is not 200 -- so before + the fix this reported a violation on a correct invoice. One blank cell in + a long array was enough. The sum is not the column's sum, so there is + nothing to compare and the constraint is advisory-skipped. + """ + expr = "total == sum('lines.amount')" + assert evaluate_constraints([expr], {"total": "200.00"}, _PARTIAL) == [] + + +def test_a_partial_avg_also_skips(): + expr = "mean == avg('lines.amount')" + assert evaluate_constraints([expr], {"mean": "66.67"}, _PARTIAL) == [] + + +def test_an_un_parseable_cell_counts_as_missing_for_sum(): + """`total` is the document's real total, including the cell OCR mangled. + + Summing only the readable cell gives 100 against a stated 250, so without + the guard this is a violation on a correct document. + """ + rows = {"lines": [{"amount": "100.00"}, {"amount": "n/a"}]} + expr = "total == sum('lines.amount')" + assert evaluate_constraints([expr], {"total": "250.00"}, rows) == [] + + +@pytest.mark.parametrize("fn", ["min", "max", "count"]) +def test_min_max_and_count_stay_total_over_a_subset(fn): + """These are meaningful over the rows that HAVE a value, so they must not skip.""" + expected = {"min": "50.00", "max": "100.00", "count": "2"}[fn] + expr = f"x == {fn}('lines.amount')" + assert evaluate_constraints([expr], {"x": expected}, _PARTIAL) == [] + + +# -------------------------------------------------------------------------- +# 2.19: a constraint that RAISES is dropped and reported as no violation. +# That stays true (one bad constraint must not fail an extraction) but it is +# now logged, which is the part that was missing. +# -------------------------------------------------------------------------- + + +def test_a_dropped_constraint_is_logged_at_exception_level(caplog): + """`quantity` normalizes to 0 -> ZeroDivisionError -> dropped. + + `_BIN` includes truediv and `compile._ALLOWED_NODES` includes ast.Div, so + this expression is accepted at submit: the path is reachable, not + theoretical. + """ + expr = "unit_price == line_total / quantity" + values = {"unit_price": "5", "line_total": "20", "quantity": "0"} + with caplog.at_level(logging.ERROR, logger="unstract.agent_kv_schema.constraints"): + assert evaluate_constraints([expr], values, {}) == [] + assert any( + "DROPPED" in r.getMessage() and expr in r.getMessage() for r in caplog.records + ), ( + "a constraint that raised was dropped silently; the only signal that " + "nothing was checked is this log line" + ) + + +def test_an_ordinary_skip_is_not_logged_as_an_error(caplog): + """A missing operand is normal on a document with optional keys.""" + with caplog.at_level(logging.WARNING, logger="unstract.agent_kv_schema.constraints"): + assert evaluate_constraints(["a == b"], {"a": "1"}, {}) == [] + assert caplog.records == [] + + +def test_one_dropped_constraint_does_not_stop_the_others(): + exprs = ["unit_price == line_total / quantity", "total == net"] + values = { + "unit_price": "5", + "line_total": "20", + "quantity": "0", + "total": "1", + "net": "2", + } + assert evaluate_constraints(exprs, values, {}) == ["total == net"] + + +# -------------------------------------------------------------------------- +# 2.20: the two hand-maintained allowlists live in two files with nothing +# enforcing agreement. compile.py decides what a submit ACCEPTS; constraints.py +# decides what the evaluator can RUN. A node accepted but not runnable is a +# constraint that is silently skipped for every document; a node runnable but +# not accepted is dead code. +# -------------------------------------------------------------------------- + + +def test_the_submit_allowlist_and_the_evaluator_agree_on_comparisons(): + accepted = {n for n in _ALLOWED_NODES if issubclass(n, ast.cmpop)} + assert accepted == set(_CMP), ( + "compile._ALLOWED_NODES and constraints._CMP disagree: a comparison " + "accepted at submit but absent from _CMP is skipped for every document" + ) + + +def test_the_submit_allowlist_and_the_evaluator_agree_on_arithmetic(): + accepted = {n for n in _ALLOWED_NODES if issubclass(n, ast.operator)} + assert accepted == set( + _BIN + ), "compile._ALLOWED_NODES and constraints._BIN disagree on binary operators" + + +def test_the_submit_allowlist_and_the_evaluator_agree_on_aggregates(): + assert _ALLOWED_CALLS == _AGG, ( + "compile._ALLOWED_CALLS and constraints._AGG disagree: a function name " + "accepted at submit but absent from _AGG is skipped for every document" + ) diff --git a/unstract/agent-kv-schema/tests/test_constraints_tolerance.py b/unstract/agent-kv-schema/tests/test_constraints_tolerance.py new file mode 100644 index 0000000000..6bfdbf4d77 --- /dev/null +++ b/unstract/agent-kv-schema/tests/test_constraints_tolerance.py @@ -0,0 +1,131 @@ +"""Equality in a constraint is tolerant, not exact. + +Operands are normalized currency/number values that arrive as floats, so +binary floating point made exact `operator.eq` wrong for the one thing +constraints exist to express. The feature's own headline example reported a +false violation on a CORRECT invoice -- which is worse than no check, because +it trains reviewers to ignore the output. +""" + +import pytest + +from unstract.agent_kv_schema import evaluate_constraints + +# Three identical line items. 8230.4 * 3 is 24691.199999999997 in binary +# floating point, so `sum(...) == 24691.2` was False and the constraint was +# reported violated. +_ROWS = {"line_items": [{"line_total": "8230.40"} for _ in range(3)]} + + +def test_the_headline_example_is_not_a_false_violation(): + violations = evaluate_constraints( + ["grand_total == sum('line_items.line_total')"], + {"grand_total": "24691.20"}, + _ROWS, + ) + assert violations == [] + + +def test_a_real_mismatch_is_still_reported(): + """The tolerance must not swallow genuine errors -- an invoice off by a + cent is exactly what this check is for. + """ + violations = evaluate_constraints( + ["grand_total == sum('line_items.line_total')"], + {"grand_total": "24691.21"}, + _ROWS, + ) + assert violations == ["grand_total == sum('line_items.line_total')"] + + +def test_not_equal_is_the_negation_of_the_tolerant_equality(): + """`!=` must agree with `==`, or a schema author gets both reported.""" + expr = "grand_total != sum('line_items.line_total')" + assert evaluate_constraints([expr], {"grand_total": "24691.20"}, _ROWS) == [expr] + assert evaluate_constraints([expr], {"grand_total": "24691.21"}, _ROWS) == [] + + +def test_comparison_against_exact_zero_still_works(): + """Relative tolerance alone is useless at zero -- every non-zero value is + infinitely far from 0 in relative terms -- which is why abs_tol is also set. + """ + assert evaluate_constraints(["balance == 0"], {"balance": "0.00"}, {}) == [] + assert evaluate_constraints(["balance == 0"], {"balance": "1.00"}, {}) == [ + "balance == 0" + ] + + +def test_ordering_comparisons_stay_exact(): + """Left deliberately exact: a tolerant `<` would make `a < b` and `a == b` + both true at the boundary, and "a total that is too large" is not a + rounding artefact. + """ + assert evaluate_constraints(["total > 100"], {"total": "100.00"}, {}) == [ + "total > 100" + ] + assert evaluate_constraints(["total >= 100"], {"total": "100.00"}, {}) == [] + + +# The tolerance must sit BELOW a cent at every realistic magnitude, not just at +# invoice scale. `rel_tol=1e-9` (the first attempt) is 0.1 at a $100,000,000 +# total, so it silently absorbed a one-cent reconciliation error -- the exact +# failure the check exists to catch. The original cent-mismatch test only +# covered a ~$25,000 total and so could not see it. + + +def test_a_one_cent_error_is_caught_at_one_hundred_million(): + rows = {"line_items": [{"line_total": "50000000.00"} for _ in range(2)]} + expr = "grand_total == sum('line_items.line_total')" + assert evaluate_constraints([expr], {"grand_total": "100000000.01"}, rows) == [expr] + + +def test_a_correct_one_hundred_million_total_is_not_a_false_violation(): + """The other side: tightening must not start reporting correct documents.""" + rows = {"line_items": [{"line_total": "50000000.00"} for _ in range(2)]} + assert ( + evaluate_constraints( + ["grand_total == sum('line_items.line_total')"], + {"grand_total": "100000000.00"}, + rows, + ) + == [] + ) + + +def test_float_noise_is_still_absorbed_over_many_rows(): + """The lower bound on the tolerance. 1,000 rows of 8230.40 do not sum to + exactly 8230400.0 in binary floating point, and that must not be a + violation. + """ + rows = {"line_items": [{"line_total": "8230.40"} for _ in range(1000)]} + assert ( + evaluate_constraints( + ["grand_total == sum('line_items.line_total')"], + {"grand_total": "8230400.00"}, + rows, + ) + == [] + ) + + +# Long arrays. The aggregate uses `math.fsum`, so accumulation error does not +# grow with row count and a correct total stays correct at any realistic array +# length -- while a real one-cent error is still reported at the same length. +@pytest.mark.parametrize("rows", [1_000, 10_000, 100_000]) +def test_a_correct_total_is_not_a_violation_however_many_rows(rows): + items = {"line_items": [{"line_total": "8230.40"}] * rows} + total = f"{8230.40 * rows:.2f}" + assert ( + evaluate_constraints( + ["grand_total == sum('line_items.line_total')"], + {"grand_total": total}, + items, + ) + == [] + ) + + +def test_a_one_cent_error_is_still_caught_at_one_hundred_thousand_rows(): + items = {"line_items": [{"line_total": "8230.40"}] * 100_000} + expr = "grand_total == sum('line_items.line_total')" + assert evaluate_constraints([expr], {"grand_total": "823040000.01"}, items) == [expr] diff --git a/unstract/agent-kv-schema/uv.lock b/unstract/agent-kv-schema/uv.lock new file mode 100644 index 0000000000..cef30e8cf2 --- /dev/null +++ b/unstract/agent-kv-schema/uv.lock @@ -0,0 +1,8 @@ +version = 1 +revision = 1 +requires-python = ">=3.12" + +[[package]] +name = "unstract-agent-kv-schema" +version = "0.1.0" +source = { editable = "." } diff --git a/unstract/filesystem/src/unstract/filesystem/file_storage_config.py b/unstract/filesystem/src/unstract/filesystem/file_storage_config.py index 870ed30a02..06ba625154 100644 --- a/unstract/filesystem/src/unstract/filesystem/file_storage_config.py +++ b/unstract/filesystem/src/unstract/filesystem/file_storage_config.py @@ -29,10 +29,12 @@ def get_provider(var_name: str, default: str = "minio") -> FileStorageProvider: FileStorageType.WORKFLOW_EXECUTION: StorageType.SHARED_TEMPORARY, FileStorageType.API_EXECUTION: StorageType.SHARED_TEMPORARY, FileStorageType.HITL_FILES: StorageType.SHARED_TEMPORARY, # Use shared temporary for HITL files + FileStorageType.AGENT_KV: StorageType.SHARED_TEMPORARY, } FILE_STORAGE_CREDENTIALS_TO_ENV_NAME_MAPPING = { FileStorageType.WORKFLOW_EXECUTION: "WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS", FileStorageType.API_EXECUTION: "API_FILE_STORAGE_CREDENTIALS", FileStorageType.HITL_FILES: "HITL_FILES_FILE_STORAGE_CREDENTIALS", + FileStorageType.AGENT_KV: "AGENT_KV_FILE_STORAGE_CREDENTIALS", } diff --git a/unstract/filesystem/src/unstract/filesystem/file_storage_types.py b/unstract/filesystem/src/unstract/filesystem/file_storage_types.py index 2b6956fe69..b27241071c 100644 --- a/unstract/filesystem/src/unstract/filesystem/file_storage_types.py +++ b/unstract/filesystem/src/unstract/filesystem/file_storage_types.py @@ -5,3 +5,4 @@ class FileStorageType(Enum): WORKFLOW_EXECUTION = "WORKFLOW_EXECUTION" API_EXECUTION = "API_EXECUTION" HITL_FILES = "HITL_FILES" + AGENT_KV = "AGENT_KV" diff --git a/unstract/filesystem/tests/__init__.py b/unstract/filesystem/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/unstract/filesystem/tests/test_agent_kv_storage_type.py b/unstract/filesystem/tests/test_agent_kv_storage_type.py new file mode 100644 index 0000000000..63f44ae4da --- /dev/null +++ b/unstract/filesystem/tests/test_agent_kv_storage_type.py @@ -0,0 +1,15 @@ +from unstract.filesystem.file_storage_config import ( + FILE_STORAGE_CREDENTIALS_TO_ENV_NAME_MAPPING, + STORAGE_MAPPING, +) +from unstract.filesystem.file_storage_types import FileStorageType + + +def test_agent_kv_type_exists_and_is_mapped(): + t = FileStorageType.AGENT_KV + assert t.value == "AGENT_KV" + assert t in STORAGE_MAPPING + assert ( + FILE_STORAGE_CREDENTIALS_TO_ENV_NAME_MAPPING[t] + == "AGENT_KV_FILE_STORAGE_CREDENTIALS" + ) diff --git a/unstract/sdk1/tests/test_page_usage_metering.py b/unstract/sdk1/tests/test_page_usage_metering.py new file mode 100644 index 0000000000..6d242e4764 --- /dev/null +++ b/unstract/sdk1/tests/test_page_usage_metering.py @@ -0,0 +1,172 @@ +"""`push_usage_details` is the single point where a page-usage row is billed. + +Pinned because two independent paths depend on the counts it produces and one +of them disagrees with it today (UN-4232): + +* PDFs are counted by `pdfplumber`, one row per real page. +* **Everything else is counted as ONE page**, with a standing + `TODO: Calculate page usage for other file types` at the call site. + +The second is a real under-count for spreadsheets. The agentic table engine +splits an Excel workbook into *virtual* pages and extracts each one, so a +12-sheet workbook does twelve pages of work, reports twelve in the job result, +and bills one. That affects both the IDE table path +(`agentic_table/src/executor.py`) and the Agent-KV blind API +(`agentic_table/src/api_binding.py`) identically -- they share this function -- +so it is a product-wide billing question, not a per-path bug, and correcting it +here moves both at once. + +These tests do not assert that the current numbers are RIGHT. They assert what +they currently are, so that: + +1. A change to the non-PDF count surfaces as a failing test with this + explanation attached, instead of silently re-pricing every spreadsheet + extraction in the product. +2. `file_name` and `run_id` keep flowing from `usage_kwargs` onto the row -- + the two fields a billing dispute is reconstructed from. The blind API + shipped without `file_name` once already. + +Excel page metering has moved twice on this branch (an over-count from blank +chunks, then this under-count), which is why it is pinned rather than +described. +""" + +from collections.abc import Callable +from pathlib import Path +from typing import Any +from unittest import mock + +import pytest +from unstract.sdk1.constants import MimeType +from unstract.sdk1.x2txt import X2Text + + +def _x2text(usage_kwargs: dict[str, str]) -> X2Text: + """An `X2Text` with its `__init__` bypassed. + + Constructing one for real needs a platform round-trip to resolve an adapter; + `push_usage_details` touches only `self._tool` and `self._usage_kwargs`. + """ + inst = object.__new__(X2Text) + inst._usage_kwargs = usage_kwargs + inst._tool = mock.Mock(get_env_or_die=mock.Mock(return_value="pk")) + return inst + + +def _minimal_pdf(pages: int) -> bytes: + """A structurally valid PDF with `pages` empty pages. + + Built by hand rather than with a fixture file so the page count is visible + in the test that asserts on it. The PDF branch really does open the file + with `pdfplumber`, so a placeholder byte string is not enough. + """ + objs = [ + "<< /Type /Catalog /Pages 2 0 R >>", + "<< /Type /Pages /Kids [{}] /Count {} >>".format( + " ".join(f"{3 + i} 0 R" for i in range(pages)), pages + ), + *[ + "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>" + for _ in range(pages) + ], + ] + out = bytearray(b"%PDF-1.4\n") + offsets = [] + for n, body in enumerate(objs, start=1): + offsets.append(len(out)) + out += f"{n} 0 obj\n{body}\nendobj\n".encode() + xref_at = len(out) + out += f"xref\n0 {len(objs) + 1}\n0000000000 65535 f \n".encode() + for off in offsets: + out += f"{off:010d} 00000 n \n".encode() + out += ( + f"trailer\n<< /Size {len(objs) + 1} /Root 1 0 R >>\n" + f"startxref\n{xref_at}\n%%EOF\n" + ).encode() + return bytes(out) + + +XLSX_MIME = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + + +@pytest.fixture +def pushed(tmp_path: Path) -> Callable[..., dict[str, Any]]: + """Capture the single `push_page_usage_data` call as a kwargs dict.""" + + def _run( + mime_type: str, + usage_kwargs: dict[str, str], + *, + pdf_pages: int = 1, + name: str = "book.xlsx", + ) -> dict[str, Any]: + f = tmp_path / name + if mime_type == MimeType.PDF: + f.write_bytes(_minimal_pdf(pdf_pages)) + else: + f.write_bytes(b"a real file with a real size, but not a workbook") + with mock.patch("unstract.sdk1.x2txt.Audit") as m_audit: + _x2text(usage_kwargs).push_usage_details(str(f), mime_type) + assert m_audit.return_value.push_page_usage_data.call_count == 1 + return m_audit.return_value.push_page_usage_data.call_args.kwargs + + return _run + + +def test_a_pdf_is_billed_per_real_page(pushed: Callable[..., dict[str, Any]]) -> None: + """The contrast that makes the spreadsheet case a gap rather than a policy. + + PDFs are counted by `pdfplumber`, so pages of work and pages billed agree. + """ + call = pushed(MimeType.PDF, {"run_id": "job-1"}, pdf_pages=3, name="rent-roll.pdf") + + assert call["page_count"] == 3 + + +def test_a_spreadsheet_is_billed_as_a_single_page( + pushed: Callable[..., dict[str, Any]], +) -> None: + """The under-count, pinned with its consequence spelled out. + + If this starts failing because the count became sheet- or virtual-page + aware, that is an INTENTIONAL billing change: update this test, and check + that `agentic_table`'s reported `pages` now agrees with it (the job result + and the usage row currently disagree for Excel). + """ + call = pushed(XLSX_MIME, {"run_id": "job-1", "file_name": "book.xlsx"}) + + assert call["page_count"] == 1, ( + "non-PDF inputs are metered as one page regardless of how many pages " + "of work they cause; the table engine splits a workbook into virtual " + "pages and bills this single row for all of them" + ) + + +def test_the_billing_identifiers_reach_the_row( + pushed: Callable[..., dict[str, Any]], +) -> None: + """`run_id` and `file_name` are carried only by `usage_kwargs`. + + `Audit.push_page_usage_data` reads both off that dict and defaults each to + `""`, so a caller that omits one gets a billed row with no way back to the + job or the input. The Agent-KV blind API omitted `file_name` on its first + cut (UN-4232); nothing failed, which is why this exists. + """ + call = pushed(XLSX_MIME, {"run_id": "job-7", "file_name": "rent-roll.xlsx"}) + + assert call["kwargs"]["run_id"] == "job-7" + assert call["kwargs"]["file_name"] == "rent-roll.xlsx" + + +def test_a_missing_file_name_is_not_rejected_here( + pushed: Callable[..., dict[str, Any]], +) -> None: + """Documents WHY the omission was silent: this layer tolerates it. + + Not an endorsement -- the point is that no exception, no warning and no log + marks the gap, so the only defence is the caller passing the field. Both + table paths now do. + """ + call = pushed(XLSX_MIME, {"run_id": "job-7"}) + + assert call["kwargs"].get("file_name") is None diff --git a/uv.lock b/uv.lock index cc248a07dc..bd30279e26 100644 --- a/uv.lock +++ b/uv.lock @@ -3719,6 +3719,11 @@ test-rig = [ ] workers = [{ name = "unstract-workers", editable = "workers" }] +[[package]] +name = "unstract-agent-kv-schema" +version = "0.1.0" +source = { editable = "unstract/agent-kv-schema" } + [[package]] name = "unstract-connectors" version = "0.0.3" @@ -3968,6 +3973,7 @@ dependencies = [ { name = "python-socketio" }, { name = "redis" }, { name = "requests" }, + { name = "unstract-agent-kv-schema" }, { name = "unstract-connectors" }, { name = "unstract-core" }, { name = "unstract-filesystem" }, @@ -3993,6 +3999,7 @@ requires-dist = [ { name = "python-socketio", specifier = ">=5.9.0" }, { name = "redis", specifier = ">=4.5.0,<6.0.0" }, { name = "requests", specifier = ">=2.33.0,<3.0.0" }, + { name = "unstract-agent-kv-schema", editable = "unstract/agent-kv-schema" }, { name = "unstract-connectors", editable = "unstract/connectors" }, { name = "unstract-core", editable = "unstract/core" }, { name = "unstract-filesystem", editable = "unstract/filesystem" }, diff --git a/workers/executor/tasks.py b/workers/executor/tasks.py index 16b97c2c33..f6545c2300 100644 --- a/workers/executor/tasks.py +++ b/workers/executor/tasks.py @@ -42,12 +42,26 @@ # ``executor/worker.py`` reaches it by importing this module. register_all() +#: Operations that are expected to emit usage records. An op in this set that +#: finishes successfully with NO records logs below -- it is the only signal +#: that billing silently produced nothing. +#: +#: Keeping a paid op out of this set is not a missing log line, it is a missing +#: alarm: the cloud `flush()` returns `[]` on a skipped LLM without raising, so +#: total loss of a run's billing rows is otherwise indistinguishable from a run +#: that legitimately made no LLM calls. Add every new LLM-bearing operation here +#: at the same time as the operation itself. _LLM_BEARING_OPS = frozenset( { "answer_prompt", "single_pass_extraction", "summarize", "structure_pipeline", + # The Agent-KV API's table extractor (cloud `agentic_table`'s blind-API + # operation). Every run drives two LLMs; a run of this op with no usage + # records means the billing chain broke, never that there was nothing + # to bill. + "table_extract_api", } ) diff --git a/workers/ide_callback/agent_kv_tasks.py b/workers/ide_callback/agent_kv_tasks.py new file mode 100644 index 0000000000..65255ce61c --- /dev/null +++ b/workers/ide_callback/agent_kv_tasks.py @@ -0,0 +1,326 @@ +"""Agent-KV terminal callbacks (spec §5.3). Thin: parse, finalize, webhook. + +These are the Celery ``link``/``link_error`` callbacks that +``agent_kv.dispatch.dispatch_job`` (backend, Task 8) attaches to the +executor dispatch: ``agent_kv_complete`` on success, ``agent_kv_error`` on +an unhandled executor exception. Both run on the dedicated +``agent_kv_callback`` queue (see +``shared/infrastructure/config/registry.py``) served by the IDE_CALLBACK +worker, and both terminalize the job via +``InternalAPIClient.agent_kv_finalize`` (Task 11's internal endpoint) +before firing the completion webhook (Task 13). +""" + +import logging +import os +from typing import Any + +from queue_backend import worker_task +from shared.utils.webhook_notify import send_webhook + +logger = logging.getLogger(__name__) + +_UNKNOWN = "Executor failed without an error message" + + +def _get_api_client(): + """Lazily build an InternalAPIClient. + + Lazy import + plain instantiation (no shared config/session): mirrors + ``queue_backend/pg_queue/consumer.py`` and ``.../reaper.py``'s + ``_get_api_client`` helpers, which use this exact pattern to keep the + module import-cycle-free at load time. Unlike ``ide_callback.tasks``'s + ``_get_api_client`` (which returns a ``PromptStudioAPIClient`` scoped to + a handful of prompt-studio endpoints), this callback talks to a general + internal endpoint, so it goes through the general-purpose + ``InternalAPIClient`` facade that owns ``agent_kv_finalize``. + """ + from shared.api import InternalAPIClient + + return InternalAPIClient() + + +def _resolve_error(failed_task_id: str, explicit: str | None = None) -> str: + """Resolve the real error text for the ``agent_kv_error`` link_error callback. + + Mirrors ``ide_callback.tasks._get_task_error``'s precedence exactly: + prefer ``explicit`` when the caller already has it. The PG-queue + transport's self-chained error path (``queue_backend/pg_queue/consumer.py``'s + ``_chain_continuation``) hands the real exception through + ``callback_kwargs["error"]`` -- on that path the executor ran eagerly and + never wrote a Celery result backend entry under ``failed_task_id``, so + the ``AsyncResult`` lookup below would come back empty. The Celery + ``link_error`` path (``agent_kv.dispatch.dispatch_job``) passes no + explicit error and relies on the result backend, then ``_UNKNOWN``. + + Falsy-aware (``explicit or ...``), matching ``_get_task_error``'s own + ``if explicit is not None`` intent in practice: an empty-string + ``explicit`` (e.g. an executor that raised with no message) must not be + persisted verbatim as a blank error -- it falls through to the result + backend lookup, then ``_UNKNOWN``, same as an absent explicit error. + + Kept as a sibling implementation rather than importing + ``ide_callback.tasks._get_task_error`` directly: ``ide_callback/tasks.py`` + already imports *this* module at its own bottom (so both of + ``workers/worker.py``'s task-loading mechanisms register + ``agent_kv_complete``/``agent_kv_error`` -- see that file's comment). + Importing back from ``tasks.py`` here would still resolve (Python + tolerates the resulting redundant re-import under the two different + ``sys.modules`` names that mechanism ends up using), but it adds a + two-way coupling between the modules for ~10 lines of small, stable + logic that isn't worth the added complexity to reason about. + """ + if explicit: + return explicit + try: + from celery import current_app as app + from celery.result import AsyncResult + + res = AsyncResult(failed_task_id, app=app) + if res.result: + return str(res.result) + except Exception: + # Was a bare `pass`. The common case here is a kombu deserialization + # failure: the executor raised a cloud-plugin exception class that is + # not importable in the OSS ide_callback image, so the result backend + # has the error and this process cannot read it. The job is then + # finalized with "Executor failed without an error message" -- the exact + # useless error this lookup exists to avoid -- and silently, so the real + # cause was unrecoverable even from logs. + logger.warning( + "Could not read the executor error from the result backend for " + "task %s; the job will be finalized with a generic message", + failed_task_id, + exc_info=True, + ) + return _UNKNOWN + + +@worker_task(name="agent_kv_complete") +def agent_kv_complete( + result_dict: dict[str, Any], + callback_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Celery link callback after the agent-kv executor task returns. + + ``result_dict["success"]`` reflects the *executor's* own outcome: it can + still be ``False`` here even though this is the success ``link`` (the + executor caught its error and returned a structured failure instead of + raising, which would have routed to ``agent_kv_error`` instead). Either + way the job is finalized via the internal API, then the webhook fires + if the response says to. + """ + cb = callback_kwargs or {} + job_id = cb.get("job_id", "") + org_id = cb.get("org_id", "") + api = _get_api_client() + + try: + if not result_dict.get("success", False): + error = result_dict.get("error") or _UNKNOWN + logger.error( + "agent_kv executor reported failure: job_id=%s error=%s", job_id, error + ) + out = api.agent_kv_finalize(job_id, org_id, success=False, error=error) + else: + data = result_dict.get("data") or {} + out = api.agent_kv_finalize( + job_id, + org_id, + success=True, + result=data.get("output") or {}, + usage_summary=data.get("usage_summary"), + ) + + _maybe_webhook(out, job_id) + return {"job_id": job_id, "finalized": out.get("finalized", False)} + + except Exception: + logger.exception( + "agent_kv_complete callback failed: job_id=%s org_id=%s", job_id, org_id + ) + raise + + +@worker_task( + bind=True, + name="agent_kv_error", + autoretry_for=(Exception,), + max_retries=3, + retry_backoff=True, + retry_backoff_max=300, + retry_jitter=True, +) +def agent_kv_error( + self, + failed_task_id: str, + callback_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any] | None: + """Celery link_error callback when the agent-kv executor task raises. + + This is the SOLE terminalizer of a failed job, so a finalize that does not + happen is a job stuck in RUNNING. It used to catch, log and ``return None`` + -- which the consumer records as SUCCESS, deletes the message, and moves on: + no retry, no dead letter, no failed-task record. A momentary backend blip + during finalize therefore stranded the job until the sweep eventually + terminalized it with "Job timed out", overwriting the real executor error, + which existed only in the log line above. + + Now retries and RAISES on exhaustion, matching its sibling + ``agent_kv_complete`` and the ``process_batch_callback_api`` precedent in + ``workers/callback/tasks.py``. The PG consumer relies on task-level + autoretry for exactly this (see its own comment in + ``queue_backend/pg_queue/consumer.py``). + + **What ``retry_backoff`` actually does here: nothing.** The PG consumer runs + tasks via ``task.apply(..., throw=True)`` + (``queue_backend/pg_queue/consumer.py``) -- i.e. EAGER. Celery's eager + retry recurses synchronously and ignores ``countdown``, so + ``retry_backoff=True`` / ``retry_backoff_max=300`` / ``retry_jitter=True`` + produce three back-to-back attempts in microseconds, which rides out no + "momentary blip" at all. They are kept because they are correct for the + Celery transport and cost nothing under the PG one -- not because they + space the attempts out here. + + The delay that genuinely helps comes AFTER exhaustion, from vt-redelivery + once this raises and the message's visibility timeout lapses. So the value + of the raise is not just the failed-task record: it is the only thing that + buys any real elapsed time before the next attempt. An earlier version of + this docstring credited the decorator with that, which would mislead anyone + tuning the backoff to fix a slow finalize. + + ``autoretry_for=(Exception,)`` rather than a narrow transport type: the + whole point is that this task must not quietly give up, and a programming + error that burns three retries and then surfaces loudly is a far better + outcome than one that returns None and reports success. + """ + cb = callback_kwargs or {} + job_id = cb.get("job_id", "") + org_id = cb.get("org_id", "") + api = _get_api_client() + + try: + error = _resolve_error(failed_task_id, explicit=cb.get("error")) + logger.error("agent_kv executor task failed: job_id=%s error=%s", job_id, error) + out = api.agent_kv_finalize(job_id, org_id, success=False, error=error) + + _maybe_webhook(out, job_id) + return {"job_id": job_id, "finalized": out.get("finalized", False)} + + except Exception: + logger.exception( + "agent_kv_error callback failed: job_id=%s org_id=%s (retry %s of %s)", + job_id, + org_id, + self.request.retries, + self.max_retries, + ) + raise + + +@worker_task( + bind=True, + name="agent_kv_cancelled", + autoretry_for=(Exception,), + max_retries=3, + retry_backoff=True, + retry_backoff_max=300, + retry_jitter=True, +) +def agent_kv_cancelled( + self, + callback_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any] | None: + """Fire the terminal webhook for a job cancelled by the API. + + Cancellation does not go through finalize: ``JobCancelView`` (and DELETE on + a non-terminal job) terminalizes the row directly, releases the slot and + returns. So nothing in the finalize path ever fired a webhook for it, and a + late executor callback cannot either -- it loses the terminal guard and + ``_maybe_webhook`` correctly declines a non-fresh finalize. The caller who + supplied ``webhook_url`` was simply never told, despite docs §8 promising + delivery on terminal states. + + Enqueued by the backend ONLY when its guarded cancel actually won, which is + what keeps this from double-notifying: if the cancel lost to a finalize, the + backend does not enqueue and the finalize callback sends instead. Exactly + one of the two paths owns the notification for any given job. + """ + cb = callback_kwargs or {} + job_id = cb.get("job_id", "") + url = cb.get("webhook_url", "") + if not url: + return None + try: + delivered = _send_webhook(url, job_id, "cancelled") + except Exception: + logger.exception("agent_kv_cancelled webhook failed: job_id=%s", job_id) + raise + if not delivered: + # `send_webhook` returns False for a non-2xx response or a connection + # failure. Returning success here acknowledged the queue message, so + # the caller simply never received the terminal notification this task + # exists to deliver -- and nothing retried or recorded it. Raise so the + # retry budget applies, exactly as a transport exception would. + logger.error( + "agent_kv_cancelled: webhook not delivered for job %s (retry %s of %s)", + job_id, + self.request.retries, + self.max_retries, + ) + raise RuntimeError(f"cancellation webhook not delivered for job {job_id}") + return {"job_id": job_id, "status": "cancelled"} + + +def _send_webhook(url: str, job_id: str, status: str) -> bool: + """Deliver one terminal notification. + + Single place the SSRF waiver is read, so the cancel path cannot drift from + the finalize path on what it permits. + """ + # Test/dev stacks only (e2e lane): waive the SSRF guards so a receiver on + # the compose host is reachable. Unset/false in production. + allow_insecure = os.environ.get( + "AGENT_KV_WEBHOOK_INSECURE_ALLOW_HTTP_PRIVATE", "" + ).lower() in ("1", "true", "yes") + return send_webhook( + url, + {"job_id": job_id, "status": status}, + allow_insecure=allow_insecure, + ) + + +def _maybe_webhook(finalize_response: dict[str, Any], job_id: str) -> None: + """Fire the completion webhook, but only for a fresh finalize. + + ``finalized`` is ``False`` for a duplicate/late finalize call (the job + was already terminal) as well as for an unknown job -- either way the + webhook already fired (or never should), so firing again here would + double-notify the caller. + + A job cancelled through the API is the "already terminal" case, and its + notification is owned by ``agent_kv_cancelled``, enqueued by the backend at + the moment its guarded cancel won. + """ + if not finalize_response.get("finalized"): + return + url = finalize_response.get("webhook_url") or "" + if not url: + return + status = finalize_response.get("status", "") + if not _send_webhook(url, job_id, status): + # `send_webhook` returns False for a non-2xx or a connection failure. + # Discarding that meant a terminal notification could fail to land with + # no retry, no record on the job, and -- for a non-2xx -- not even a log + # line. Logged at ERROR here rather than raised: unlike the cancellation + # task, this runs AFTER finalize has already terminalized the job, so + # raising would re-run a finalize that is no longer idempotent-free. + # + # Durable delivery (attempt counts, a `webhook_delivered_at` the status + # document can expose) is the open design question in 2.10. + logger.error( + "agent-kv webhook not delivered for job %s (status=%s); the caller " + "was not notified and nothing will retry", + job_id, + status, + ) diff --git a/workers/ide_callback/tasks.py b/workers/ide_callback/tasks.py index 1a298dffe3..7f2ed70e7c 100644 --- a/workers/ide_callback/tasks.py +++ b/workers/ide_callback/tasks.py @@ -115,7 +115,24 @@ def _get_task_error( if res.result: return str(res.result) except Exception: - pass + # Was a bare `pass`, the same defect fixed in the Agent-KV twin + # (`agent_kv_tasks.py::_error_link`) -- carried over here because the + # cause and the cost are identical, not because this PR owns the IDE + # path. + # + # The common case is a kombu deserialization failure: the executor + # raised an exception class that is not importable in this image, so + # the result backend HAS the error and this process cannot read it. The + # execution is then finalized with the generic `default` -- the exact + # useless message this lookup exists to avoid -- and silently, leaving + # the real cause unrecoverable even from logs. + logger.warning( + "Could not read the executor error from the result backend for " + "task %s; finalizing with the generic message %r", + failed_task_id, + default, + exc_info=True, + ) return default @@ -693,3 +710,19 @@ def extraction_error( source, file_id, ) + + +# ------------------------------------------------------------------ +# Agent-KV terminal callbacks (spec §5.3) +# +# ``workers/worker.py``'s ``load_worker_tasks()`` registers this worker +# type's tasks by loading THIS file directly (by path, under the bare +# module name ``"tasks"``) rather than importing the ``ide_callback`` +# package -- so ``agent_kv_complete``/``agent_kv_error`` only bind to the +# Celery app if something imports them as a side effect of importing this +# module. Import at the bottom (not the top) to avoid a circular import: +# ``agent_kv_tasks`` doesn't currently import back from here, but keeping +# task-module imports last matches this file's own load-bearing position +# in the worker boot sequence. +# ------------------------------------------------------------------ +from ide_callback import agent_kv_tasks # noqa: E402, F401 diff --git a/workers/ide_callback/tests/__init__.py b/workers/ide_callback/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/workers/ide_callback/tests/test_agent_kv_callbacks.py b/workers/ide_callback/tests/test_agent_kv_callbacks.py new file mode 100644 index 0000000000..d78f7c670e --- /dev/null +++ b/workers/ide_callback/tests/test_agent_kv_callbacks.py @@ -0,0 +1,574 @@ +"""Unit tests for Agent-KV terminal callback tasks (spec §5.3). + +Mirrors ``tests/test_ide_callback.py``'s style: task functions are called +directly (bypassing Celery task machinery) since we're testing callback +logic, not Celery routing. ``_get_api_client`` and ``send_webhook`` are +mocked at the ``ide_callback.agent_kv_tasks`` import site. +""" + +import logging +from unittest.mock import MagicMock, patch + +import pytest + +_PATCH_GET_CLIENT = "ide_callback.agent_kv_tasks._get_api_client" +_PATCH_SEND_WEBHOOK = "ide_callback.agent_kv_tasks.send_webhook" +_PATCH_ASYNC_RESULT = "celery.result.AsyncResult" + +_UNKNOWN = "Executor failed without an error message" + + +@pytest.fixture +def mock_api(): + """A mocked InternalAPIClient whose agent_kv_finalize succeeds with a webhook.""" + api = MagicMock() + api.agent_kv_finalize.return_value = { + "finalized": True, + "webhook_url": "https://example.com/hook", + "status": "completed", + } + return api + + +@pytest.fixture +def cb_kwargs(): + """Standard callback_kwargs, matching agent_kv.dispatch.dispatch_job's shape.""" + return {"job_id": "job-1", "org_id": "org-1"} + + +# --------------------------------------------------------------------------- +# agent_kv_complete +# --------------------------------------------------------------------------- + + +class TestAgentKvComplete: + def _call(self, result_dict, callback_kwargs=None): + from ide_callback.agent_kv_tasks import agent_kv_complete + + return agent_kv_complete(result_dict, callback_kwargs) + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_success_path_finalizes_true_with_engine_result( + self, mock_get_client, mock_send_webhook, mock_api, cb_kwargs + ): + mock_get_client.return_value = mock_api + result_dict = { + "success": True, + "data": { + "output": {"field_1": "value_1"}, + "usage_summary": {"total_tokens": 42}, + }, + "error": None, + } + + result = self._call(result_dict, cb_kwargs) + + mock_api.agent_kv_finalize.assert_called_once_with( + "job-1", + "org-1", + success=True, + result={"field_1": "value_1"}, + usage_summary={"total_tokens": 42}, + ) + assert result == {"job_id": "job-1", "finalized": True} + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_missing_output_and_usage_summary_default_to_empty( + self, mock_get_client, mock_send_webhook, mock_api, cb_kwargs + ): + mock_get_client.return_value = mock_api + result_dict = {"success": True, "data": {}, "error": None} + + self._call(result_dict, cb_kwargs) + + mock_api.agent_kv_finalize.assert_called_once_with( + "job-1", "org-1", success=True, result={}, usage_summary=None + ) + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_executor_reported_failure_finalizes_false_with_error( + self, mock_get_client, mock_send_webhook, mock_api, cb_kwargs + ): + mock_get_client.return_value = mock_api + result_dict = {"success": False, "data": None, "error": "boom"} + + self._call(result_dict, cb_kwargs) + + mock_api.agent_kv_finalize.assert_called_once_with( + "job-1", "org-1", success=False, error="boom" + ) + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_executor_failure_with_no_error_message_uses_fallback( + self, mock_get_client, mock_send_webhook, mock_api, cb_kwargs + ): + mock_get_client.return_value = mock_api + result_dict = {"success": False, "data": None, "error": None} + + self._call(result_dict, cb_kwargs) + + mock_api.agent_kv_finalize.assert_called_once_with( + "job-1", "org-1", success=False, error=_UNKNOWN + ) + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_none_callback_kwargs_uses_empty_defaults( + self, mock_get_client, mock_send_webhook, mock_api + ): + mock_get_client.return_value = mock_api + result_dict = {"success": True, "data": {"output": {}}, "error": None} + + result = self._call(result_dict, None) + + mock_api.agent_kv_finalize.assert_called_once_with( + "", "", success=True, result={}, usage_summary=None + ) + assert result["job_id"] == "" + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_returns_finalized_false_when_finalize_says_so( + self, mock_get_client, mock_send_webhook, cb_kwargs + ): + """A duplicate/late callback: the backend reports finalized=False.""" + api = MagicMock() + api.agent_kv_finalize.return_value = { + "finalized": False, + "webhook_url": "", + "status": "completed", + } + mock_get_client.return_value = api + result_dict = {"success": True, "data": {"output": {}}, "error": None} + + result = self._call(result_dict, cb_kwargs) + + assert result == {"job_id": "job-1", "finalized": False} + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_finalize_raising_is_logged_and_reraised( + self, mock_get_client, mock_send_webhook, cb_kwargs, caplog + ): + """agent_kv_complete mirrors ide_index_complete: log, then re-raise.""" + import logging + + api = MagicMock() + api.agent_kv_finalize.side_effect = RuntimeError("backend unreachable") + mock_get_client.return_value = api + result_dict = {"success": True, "data": {"output": {}}, "error": None} + + with caplog.at_level(logging.ERROR, logger="ide_callback.agent_kv_tasks"): + with pytest.raises(RuntimeError, match="backend unreachable"): + self._call(result_dict, cb_kwargs) + + assert "agent_kv_complete callback failed" in caplog.text + mock_send_webhook.assert_not_called() + + +# --------------------------------------------------------------------------- +# agent_kv_error +# --------------------------------------------------------------------------- + + +class TestAgentKvError: + def _call(self, failed_task_id, callback_kwargs=None): + from ide_callback.agent_kv_tasks import agent_kv_error + + return agent_kv_error(failed_task_id, callback_kwargs) + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_error_link_finalizes_false_with_result_backend_error( + self, mock_get_client, mock_send_webhook, mock_api, cb_kwargs + ): + mock_get_client.return_value = mock_api + mock_async_result = MagicMock() + mock_async_result.result = RuntimeError("executor crashed") + + with patch(_PATCH_ASYNC_RESULT, return_value=mock_async_result): + result = self._call("failed-task-1", cb_kwargs) + + mock_api.agent_kv_finalize.assert_called_once_with( + "job-1", "org-1", success=False, error="executor crashed" + ) + assert result == {"job_id": "job-1", "finalized": True} + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_error_link_falls_back_when_no_result_available( + self, mock_get_client, mock_send_webhook, mock_api, cb_kwargs + ): + mock_get_client.return_value = mock_api + mock_async_result = MagicMock() + mock_async_result.result = None + + with patch(_PATCH_ASYNC_RESULT, return_value=mock_async_result): + self._call("failed-task-2", cb_kwargs) + + mock_api.agent_kv_finalize.assert_called_once_with( + "job-1", "org-1", success=False, error=_UNKNOWN + ) + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_error_link_falls_back_when_lookup_raises( + self, mock_get_client, mock_send_webhook, mock_api, cb_kwargs + ): + mock_get_client.return_value = mock_api + + with patch(_PATCH_ASYNC_RESULT, side_effect=RuntimeError("backend down")): + self._call("failed-task-3", cb_kwargs) + + mock_api.agent_kv_finalize.assert_called_once_with( + "job-1", "org-1", success=False, error=_UNKNOWN + ) + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_a_swallowed_lookup_failure_is_logged_with_its_traceback( + self, mock_get_client, mock_send_webhook, mock_api, cb_kwargs, caplog + ): + """The handler was a bare `except Exception: pass`. + + `test_error_link_falls_back_when_lookup_raises` above asserts only the + `_UNKNOWN` fallback -- which the bare `pass` also produced, so that test + is byte-identical before and after the fix and cannot detect it. + + What actually changed is that the cause is now recoverable. The common + case is a kombu deserialization failure: the executor raised a + cloud-plugin exception class that this OSS image cannot import, so the + result backend HAS the error and this process cannot read it. The job + is finalized with "Executor failed without an error message" -- and + before, with nothing anywhere saying why. + + `exc_info` is asserted, not just the message: without the traceback the + log says a lookup failed but not what failed, which for an unimportable + exception class is the entire diagnosis. + """ + with caplog.at_level(logging.WARNING): + with patch(_PATCH_ASYNC_RESULT, side_effect=RuntimeError("backend down")): + self._call("failed-task-3b", cb_kwargs) + + # Matched by CONTENT, not by position. The callback also logs its own + # terminal ERROR ("agent_kv executor task failed: ..."), so taking + # `records[0]` would let this pass against a bare `pass` by picking up + # an unrelated record -- which it did on the first cut of this test. + matching = [ + r + for r in caplog.records + if r.levelno >= logging.WARNING and "result backend" in r.getMessage() + ] + assert len(matching) == 1, ( + "expected exactly one warning about the unreadable result backend; " + f"got {[r.getMessage() for r in caplog.records]}. Restoring the " + "bare `pass` makes this the failing assertion." + ) + record = matching[0] + assert "failed-task-3b" in record.getMessage(), record.getMessage() + assert record.exc_info is not None, ( + "logged without exc_info -- the traceback is the diagnosis when the " + "cause is an exception class this image cannot import" + ) + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_none_callback_kwargs_uses_empty_defaults( + self, mock_get_client, mock_send_webhook, mock_api + ): + mock_get_client.return_value = mock_api + + with patch(_PATCH_ASYNC_RESULT, return_value=MagicMock(result=None)): + result = self._call("failed-task-4", None) + + mock_api.agent_kv_finalize.assert_called_once_with( + "", "", success=False, error=_UNKNOWN + ) + assert result["job_id"] == "" + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_pg_transport_explicit_error_preferred_over_async_result( + self, mock_get_client, mock_send_webhook, mock_api + ): + """PG-queue self-chained path: callback_kwargs carries the real error. + + ``queue_backend/pg_queue/consumer.py``'s ``_chain_continuation`` injects + the executor's real error into ``callback_kwargs["error"]`` because the + PG path runs the executor eagerly and never writes a Celery result + backend entry under ``failed_task_id`` -- so this must be preferred + over (and must skip) the ``AsyncResult`` lookup entirely. + """ + mock_get_client.return_value = mock_api + pg_cb_kwargs = {"job_id": "job-1", "org_id": "org-1", "error": "real cause"} + + with patch(_PATCH_ASYNC_RESULT) as mock_async_result_cls: + result = self._call("failed-task-pg-1", pg_cb_kwargs) + + mock_async_result_cls.assert_not_called() + mock_api.agent_kv_finalize.assert_called_once_with( + "job-1", "org-1", success=False, error="real cause" + ) + assert result == {"job_id": "job-1", "finalized": True} + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_pg_transport_empty_explicit_error_falls_back_to_async_result( + self, mock_get_client, mock_send_webhook, mock_api + ): + """An empty-string explicit error must not be persisted verbatim -- + it's falsy, so it falls through to the result-backend lookup (then + ``_UNKNOWN``), same as an absent explicit error. + """ + mock_get_client.return_value = mock_api + pg_cb_kwargs = {"job_id": "job-1", "org_id": "org-1", "error": ""} + mock_async_result = MagicMock() + mock_async_result.result = RuntimeError("real backend cause") + + with patch(_PATCH_ASYNC_RESULT, return_value=mock_async_result): + self._call("failed-task-empty-error", pg_cb_kwargs) + + mock_api.agent_kv_finalize.assert_called_once_with( + "job-1", "org-1", success=False, error="real backend cause" + ) + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_finalize_raising_propagates_instead_of_reporting_success( + self, mock_get_client, mock_send_webhook, cb_kwargs, caplog + ): + """A finalize that fails must NOT be reported as a completed callback. + + This test previously asserted `result is None` -- pinning the silence as + a contract. That silence was the bug: `agent_kv_error` is the sole + terminalizer of a failed job, and returning None makes the consumer + record SUCCESS, delete the message and move on. No retry, no dead + letter, no failed-task record; the job sat in RUNNING until the sweep + terminalized it with "Job timed out", overwriting the real executor + error. Inverted rather than deleted, because the old assertion is + exactly what a regression would restore. + + Reported as 2.4 in the branch review. + """ + import logging + + api = MagicMock() + api.agent_kv_finalize.side_effect = RuntimeError("backend unreachable") + mock_get_client.return_value = api + + with ( + patch(_PATCH_ASYNC_RESULT, return_value=MagicMock(result=None)), + caplog.at_level(logging.ERROR, logger="ide_callback.agent_kv_tasks"), + pytest.raises(Exception, match="backend unreachable"), + ): + self._call("failed-task-6", cb_kwargs) + + assert "agent_kv_error callback failed" in caplog.text + mock_send_webhook.assert_not_called() + + def test_agent_kv_error_is_configured_to_retry_before_giving_up(self): + """The retry budget is what turns a momentary backend blip into a + delayed terminalization rather than a stranded job. + """ + from ide_callback.agent_kv_tasks import agent_kv_error + + assert agent_kv_error.max_retries == 3 + assert Exception in agent_kv_error.autoretry_for + + @patch(_PATCH_SEND_WEBHOOK) + @patch(_PATCH_GET_CLIENT) + def test_agent_kv_error_matches_its_siblings_failure_posture( + self, mock_get_client, mock_send_webhook, cb_kwargs + ): + """Both callbacks must RAISE when finalize fails -- neither may swallow. + + These two are the success and failure halves of one contract: whichever + one runs is the SOLE terminalizer of the job, so a finalize that fails + silently leaves the row in RUNNING until the sweep overwrites the real + error with "Job timed out". `agent_kv_error` used to catch, log and + `return None`, which the consumer records as SUCCESS -- no retry, no + dead letter, no failed-task record. + + The previous version of this test asserted + `agent_kv_complete is not None` and `agent_kv_error.max_retries is not + None`, which is true of essentially any task object and could not fail + for the reason the test names. It is the behaviour that has to be + pinned, so both callbacks are actually driven here with a finalize that + blows up. + """ + from ide_callback.agent_kv_tasks import agent_kv_complete, agent_kv_error + + api = MagicMock() + api.agent_kv_finalize.side_effect = RuntimeError("backend unreachable") + mock_get_client.return_value = api + + with pytest.raises(RuntimeError): + agent_kv_complete( + {"success": True, "data": {"output": {}}, "error": None}, cb_kwargs + ) + + with pytest.raises(RuntimeError): + with patch(_PATCH_ASYNC_RESULT, return_value=MagicMock(result="boom")): + agent_kv_error.run("failed-task-posture", cb_kwargs) + + # And neither may fire a webhook off a finalize that did not happen -- + # telling the caller a job is terminal when the row never moved is + # worse than telling them nothing. + assert not mock_send_webhook.called + + +# --------------------------------------------------------------------------- +# Webhook firing rules (shared by both callbacks) +# --------------------------------------------------------------------------- + + +class TestWebhookFiring: + """Webhook fires iff finalize returns finalized=True AND a non-empty webhook_url.""" + + def _call_complete(self, callback_kwargs, api): + from ide_callback.agent_kv_tasks import agent_kv_complete + + with patch(_PATCH_GET_CLIENT, return_value=api): + result_dict = {"success": True, "data": {"output": {"x": 1}}, "error": None} + return agent_kv_complete(result_dict, callback_kwargs) + + @patch(_PATCH_SEND_WEBHOOK) + def test_webhook_fires_when_finalized_true_and_url_present( + self, mock_send_webhook, cb_kwargs + ): + api = MagicMock() + api.agent_kv_finalize.return_value = { + "finalized": True, + "webhook_url": "https://example.com/hook", + "status": "completed", + } + + self._call_complete(cb_kwargs, api) + + mock_send_webhook.assert_called_once_with( + "https://example.com/hook", + {"job_id": "job-1", "status": "completed"}, + allow_insecure=False, + ) + + @patch(_PATCH_SEND_WEBHOOK) + def test_webhook_not_fired_on_duplicate_finalize_false( + self, mock_send_webhook, cb_kwargs + ): + """Duplicate/late finalize: finalized=False must not re-fire the webhook, + even though the job (from an earlier finalize) still carries a webhook_url. + """ + api = MagicMock() + api.agent_kv_finalize.return_value = { + "finalized": False, + "webhook_url": "https://example.com/hook", + "status": "completed", + } + + self._call_complete(cb_kwargs, api) + + mock_send_webhook.assert_not_called() + + @patch(_PATCH_SEND_WEBHOOK) + def test_webhook_not_fired_when_url_empty(self, mock_send_webhook, cb_kwargs): + """finalized=True but the job has no webhook configured: nothing to call.""" + api = MagicMock() + api.agent_kv_finalize.return_value = { + "finalized": True, + "webhook_url": "", + "status": "completed", + } + + self._call_complete(cb_kwargs, api) + + mock_send_webhook.assert_not_called() + + @patch(_PATCH_SEND_WEBHOOK) + def test_webhook_fires_from_error_link_too(self, mock_send_webhook, cb_kwargs): + from ide_callback.agent_kv_tasks import agent_kv_error + + api = MagicMock() + api.agent_kv_finalize.return_value = { + "finalized": True, + "webhook_url": "https://example.com/hook", + "status": "failed", + } + + with ( + patch(_PATCH_GET_CLIENT, return_value=api), + patch(_PATCH_ASYNC_RESULT, return_value=MagicMock(result=None)), + ): + agent_kv_error("failed-task-5", cb_kwargs) + + mock_send_webhook.assert_called_once_with( + "https://example.com/hook", + {"job_id": "job-1", "status": "failed"}, + allow_insecure=False, + ) + + +class TestWebhookDeliveryResult: + """A webhook that does not land must not pass silently. + + `send_webhook` returns False for a non-2xx response or a connection + failure, and both call sites used to discard it. The cancellation task now + raises (so its retry budget applies); the finalize path logs, because it + runs after the job is already terminalized and re-running finalize is not + free. Reported as 2.10 in the branch review and by Greptile on #2317. + """ + + @patch(_PATCH_SEND_WEBHOOK, return_value=False) + def test_a_failed_finalize_webhook_is_logged(self, mock_send, caplog): + import logging + + from ide_callback import agent_kv_tasks as akt + + with caplog.at_level(logging.ERROR, logger="ide_callback.agent_kv_tasks"): + akt._maybe_webhook( + { + "finalized": True, + "webhook_url": "https://hook.example/x", + "status": "completed", + }, + "job-1", + ) + + assert "not delivered" in caplog.text + assert "job-1" in caplog.text + + @patch(_PATCH_SEND_WEBHOOK, return_value=True) + def test_a_delivered_finalize_webhook_is_quiet(self, mock_send, caplog): + import logging + + from ide_callback import agent_kv_tasks as akt + + with caplog.at_level(logging.ERROR, logger="ide_callback.agent_kv_tasks"): + akt._maybe_webhook( + { + "finalized": True, + "webhook_url": "https://hook.example/x", + "status": "completed", + }, + "job-1", + ) + + assert caplog.records == [] + + @patch(_PATCH_SEND_WEBHOOK, return_value=False) + def test_a_failed_cancellation_webhook_raises_so_it_retries(self, mock_send): + from ide_callback.agent_kv_tasks import agent_kv_cancelled + + with pytest.raises(Exception, match="not delivered"): + agent_kv_cancelled.apply( + kwargs={ + "callback_kwargs": { + "job_id": "job-1", + "webhook_url": "https://hook.example/x", + } + }, + throw=True, + ) diff --git a/workers/ide_callback/worker.py b/workers/ide_callback/worker.py index cb1bb4200c..36a2c25ca8 100644 --- a/workers/ide_callback/worker.py +++ b/workers/ide_callback/worker.py @@ -56,3 +56,16 @@ def check_ide_callback_health(): WorkerRegistry.register_health_check( WorkerType.IDE_CALLBACK, "ide_callback_health", check_ide_callback_health ) + +# Import task modules so Celery registers them on this app -- mirrors how +# other single-app workers (e.g. executor/worker.py) import their tasks +# module explicitly. NOTE: this module is not the actual Celery entrypoint +# in production today -- both docker-compose and run-worker(.sh|-docker.sh) +# always launch the unified `workers/worker.py` (`celery -A worker worker`) +# for this worker type, which loads `ide_callback/tasks.py` itself via +# `load_worker_tasks()` (see that file's bottom-of-file import of +# `agent_kv_tasks` for the mechanism that actually applies in production). +# These imports are kept here anyway so this module stays correct/self- +# registering if it's ever wired up as a real `-A` target. +import ide_callback.agent_kv_tasks # noqa: E402, F401 +import ide_callback.tasks # noqa: E402, F401 diff --git a/workers/pyproject.toml b/workers/pyproject.toml index 00a682fa04..a5399042d8 100644 --- a/workers/pyproject.toml +++ b/workers/pyproject.toml @@ -32,6 +32,7 @@ dependencies = [ "unstract-tool-sandbox", "unstract-workflow-execution", "unstract-filesystem", + "unstract-agent-kv-schema", # Caching "redis>=4.5.0,<6.0.0", # Redis client for worker cache access # PG Queue — direct Postgres access for the bespoke SKIP LOCKED queue @@ -223,6 +224,7 @@ exclude_lines = [ [tool.uv.sources] unstract-filesystem = { path = "../unstract/filesystem", editable = true } +unstract-agent-kv-schema = { path = "../unstract/agent-kv-schema", editable = true } unstract-workflow-execution = { path = "../unstract/workflow-execution", editable = true } unstract-tool-sandbox = { path = "../unstract/tool-sandbox", editable = true } unstract-tool-registry = { path = "../unstract/tool-registry", editable = true } diff --git a/workers/run-worker-docker.sh b/workers/run-worker-docker.sh index 05d2b9398f..eeafc65dfc 100755 --- a/workers/run-worker-docker.sh +++ b/workers/run-worker-docker.sh @@ -62,7 +62,10 @@ declare -A WORKER_QUEUES=( ["log_consumer"]="celery_log_task_queue" ["scheduler"]="scheduler" ["${EXECUTOR_WORKER_TYPE}"]="celery_executor_legacy" - ["ide_callback"]="ide_callback" + # agent_kv_callback carries the Agent-KV terminal callbacks + # (agent_kv_complete/agent_kv_error, spec §5.3) dispatched by + # backend/agent_kv/dispatch.py; ide_callback owns both queues. + ["ide_callback"]="ide_callback,agent_kv_callback" ) # Worker health ports diff --git a/workers/run-worker.sh b/workers/run-worker.sh index 5bd41dfbeb..e538831c8d 100755 --- a/workers/run-worker.sh +++ b/workers/run-worker.sh @@ -87,7 +87,7 @@ declare -rA PG_CONSUMER_ROLES=( # succeeds, so nothing errors at the producer. ["$PG_ROLE_NOTIFICATION"]="notification;notifications,notifications_webhook,notifications_email,notifications_sms,notifications_priority" # Prompt Studio IDE callbacks (ide_index_*/ide_prompt_*/extraction_*). - ["$PG_ROLE_IDE_CALLBACK"]="ide_callback;ide_callback" + ["$PG_ROLE_IDE_CALLBACK"]="ide_callback;ide_callback,agent_kv_callback" ) declare -rA PG_QUEUE_MEMBERS=( ["$PG_QUEUE_CONSUMER_TYPE"]=1 @@ -168,7 +168,10 @@ declare -A WORKER_QUEUES=( ["notification"]="notifications,notifications_webhook,notifications_email,notifications_sms,notifications_priority" ["scheduler"]="scheduler" ["${EXECUTOR_WORKER_TYPE}"]="celery_executor_legacy" - ["${IDE_CALLBACK_WORKER_TYPE}"]="${IDE_CALLBACK_WORKER_TYPE}" + # agent_kv_callback carries the Agent-KV terminal callbacks + # (agent_kv_complete/agent_kv_error, spec §5.3) dispatched by + # backend/agent_kv/dispatch.py; ide_callback owns both queues. + ["${IDE_CALLBACK_WORKER_TYPE}"]="${IDE_CALLBACK_WORKER_TYPE},agent_kv_callback" # The PG queue (in pg_queue_message) this consumer polls — exported as # WORKER_PG_QUEUE_CONSUMER_QUEUE, not a Celery --queues value. ["$PG_QUEUE_CONSUMER_TYPE"]="notifications" diff --git a/workers/sample.env b/workers/sample.env index b6c1078571..bffe936a26 100644 --- a/workers/sample.env +++ b/workers/sample.env @@ -398,6 +398,18 @@ UNSTRACT_RUNNER_API_BACKOFF_FACTOR=3 # File Storage Credentials (MinIO) WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS='{"provider": "minio", "credentials": {"endpoint_url": "http://unstract-minio:9000", "key": "minio", "secret": "minio123"}}' API_FILE_STORAGE_CREDENTIALS='{"provider": "minio", "credentials": {"endpoint_url": "http://unstract-minio:9000", "key": "minio", "secret": "minio123"}}' +# Bucket-rooted root the backend stages Agent-KV inputs under +# ({AGENT_KV_STORAGE_DIR_PREFIX}/{org_id}/{job_id}/...); the executor plugin reads +# `input_ref` from it and keys its OCR cache under {prefix}/{org_id}/cache/. +# MUST be identical to backend/sample.env's value -- the first segment is the +# bucket (must exist; the MinIO dev bootstrap creates `unstract`). +AGENT_KV_STORAGE_DIR_PREFIX="unstract/agent_kv" +# Read directly by the cloud agentic_table executor plugin's blind-API +# operation (workers/plugins/agentic_table, +# not checked in here) to stage/fetch Agent-KV job input+result files -- same +# shape and bucket as its siblings above. Must match docker/sample.env's copy +# (the backend, which also reads this key, gets its value from there). +AGENT_KV_FILE_STORAGE_CREDENTIALS='{"provider": "minio", "credentials": {"endpoint_url": "http://unstract-minio:9000", "key": "minio", "secret": "minio123"}}' # Remote storage for Prompt Studio / IDE flows. Must match backend/sample.env. # Required by executor and ide-callback workers; missing/empty values raise FileStorageError in EnvHelper.get_storage(). @@ -527,6 +539,7 @@ GOOGLE_OAUTH2_SECRET= # UNSTRACT_RUNNER_HOST=http://localhost # WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS={"provider": "minio", "credentials": {"endpoint_url": "http://localhost:9000", "key": "minio", "secret": "minio123"}} # API_FILE_STORAGE_CREDENTIALS={"provider": "minio", "credentials": {"endpoint_url": "http://localhost:9000", "key": "minio", "secret": "minio123"}} +# AGENT_KV_FILE_STORAGE_CREDENTIALS={"provider": "minio", "credentials": {"endpoint_url": "http://localhost:9000", "key": "minio", "secret": "minio123"}} # Flipt Service FLIPT_SERVICE_AVAILABLE=False @@ -549,3 +562,8 @@ PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python # BACKOFF = base seconds, doubling each attempt, with each sleep capped at 5.0s. # WORKER_PG_QUEUE_CONNECT_RETRIES=3 # WORKER_PG_QUEUE_CONNECT_BACKOFF=0.5 + +# Test/dev ONLY: waive the completion-webhook SSRF guards (https scheme + +# public host) on the ide-callback worker so the e2e lane can deliver to a +# receiver on the compose host. NEVER set in production. +# AGENT_KV_WEBHOOK_INSECURE_ALLOW_HTTP_PRIVATE=false diff --git a/workers/scheduler/agent_kv_tasks.py b/workers/scheduler/agent_kv_tasks.py new file mode 100644 index 0000000000..117723f3cc --- /dev/null +++ b/workers/scheduler/agent_kv_tasks.py @@ -0,0 +1,87 @@ +"""Thin scheduler-side tasks for the Agent-KV maintenance periodics (spec §5.4). + +Mirrors ``dashboard_metrics_tasks.py``'s shape and role in this worker (UN-3796): +these do **no** work themselves, each just calls a backend internal endpoint that +does the real work, because the workers image has no Django and the ORM-heavy +job-row work (terminalizing stuck/never-dispatched jobs, deleting staged files) +can only run there. + +Unlike the dashboard-metrics proxies, the calls here go through the shared +``InternalAPIClient`` facade (``shared.api.InternalAPIClient``) rather than a +bespoke ``httpx`` client -- ``agent_kv_finalize`` already lives on that facade +(Task 11) for the Agent-KV terminal callbacks +(``ide_callback/agent_kv_tasks.py::_get_api_client``), and the sweep/TTL-cleanup +calls added here (``agent_kv_sweep``/``agent_kv_ttl_cleanup``, +``workers/shared/api/internal_client.py``) are its siblings on the same client, +not a second HTTP stack for the same feature. + +Two internal endpoints, both platform-wide (no ``org_id``, spec §5.4): + +- ``POST /internal/v1/agent-kv/sweep/`` -- terminalizes never-dispatched + ``PENDING`` jobs and stuck ``DISPATCHED``/``RUNNING`` jobs. +- ``POST /internal/v1/agent-kv/ttl-cleanup/`` -- deletes staged input/result + files for jobs past ``expires_at`` and blanks their refs. + +Both handlers are idempotent and batch-capped server-side +(``backend/agent_kv/internal_views.py``), so redelivery/backlog from a missed +tick is safe without extra guards here -- exactly like the dashboard-metrics +proxies' own redelivery story. + +**Not registered under a dotted Beat-mirroring name for a Beat row that +exists** (contrast ``dashboard_metrics_tasks.py``, whose names are pinned to +verbatim Beat rows it mirrors): there is no pre-existing Beat/PG-scheduler +entry for these two tasks to match, since Agent-KV shipped with the internal +endpoints but no periodic registration (spec §5.4, this feature's own +deploy-checklist gap). The ``agent_kv.sweep``/``agent_kv.ttl_cleanup`` wire +names below are this task pair's own new contract -- an operator registers a +``PgPeriodicTask`` row against them (see `docs/agent-kv-api.md`'s deploy +checklist for the exact shape), the same PG-scheduler mechanism the +dashboard-metrics periodics use. +""" + +import logging +from typing import Any + +from queue_backend import worker_task + +logger = logging.getLogger(__name__) + + +def _get_api_client(): + """Lazily build an InternalAPIClient. + + Mirrors ``ide_callback.agent_kv_tasks._get_api_client``'s lazy import + + plain instantiation (no shared config/session) for the same reason: keeps + this module import-cycle-free at load time. + """ + from shared.api import InternalAPIClient + + return InternalAPIClient() + + +@worker_task(name="agent_kv.sweep") +def agent_kv_sweep() -> dict[str, Any]: + """Terminalize never-dispatched PENDING and stuck Agent-KV jobs (spec §5.4).""" + result = _get_api_client().agent_kv_sweep() + # Logged here as well as backend-side: this is the only place that records + # the periodic actually FIRED. The counts alone cannot distinguish "ran and + # found nothing" from "never ran", which is precisely the failure this + # task pair shipped with. + # + # Nothing in this repo schedules these. The cloud chart's + # `backend.agentKvCronJobs` runs the equivalent management commands and is + # the single owner (migration `0004_pg_periodic_tasks` used to seed + # PgPeriodicTask rows as a second owner and now removes them). A + # self-hosted OSS operator who wants these tasks driven through the PG + # scheduler registers the rows themselves with `pg_owned: True` -- see + # `docs/agent-kv-api.md` §12. + logger.info("agent_kv.sweep completed: %s", result) + return result + + +@worker_task(name="agent_kv.ttl_cleanup") +def agent_kv_ttl_cleanup() -> dict[str, Any]: + """Delete staged Agent-KV input/result files past their TTL (spec §5.4).""" + result = _get_api_client().agent_kv_ttl_cleanup() + logger.info("agent_kv.ttl_cleanup completed: %s", result) + return result diff --git a/workers/scheduler/tasks.py b/workers/scheduler/tasks.py index d79f7bf758..56bcebcfb5 100644 --- a/workers/scheduler/tasks.py +++ b/workers/scheduler/tasks.py @@ -23,7 +23,10 @@ # worker-general (this happened — the bare form shipped and broke the general # worker at flag-off, where PG is not even involved). # `/app` is on PYTHONPATH (run-worker-docker.sh), so `scheduler.` resolves under both. -from scheduler import dashboard_metrics_tasks # noqa: F401, E402 (side-effect import) +# +# The same reasoning covers agent_kv_tasks, which registers the Agent-KV +# sweep/TTL-cleanup periodics (spec §5.4) on this worker type. +from scheduler import agent_kv_tasks, dashboard_metrics_tasks from shared.enums.status_enums import PipelineStatus from shared.enums.worker_enums import QueueName from shared.infrastructure.config import WorkerConfig @@ -47,6 +50,28 @@ # Import the exact backend logic to ensure consistency +# Both modules above are imported ONLY for the side effect of registering their +# @worker_task entries on this worker; nothing here calls into them. Naming them +# in a module-level tuple is what keeps them imported. +# +# A bare `noqa` directive is NOT enough, and that is not a style preference — it +# reproduced cause of this exact line being deleted once already. Two hooks +# interact: +# 1. ruff's isort (I001) merges two `from scheduler import X` statements into +# one parenthesised statement, which relocates each trailing noqa directive onto +# a MEMBER line. +# 2. F401 is reported against the statement, whose first line now carries no +# directive, so pycln (`[tool.pycln] all = true` — it removes side-effect +# imports by design) deletes both members. +# That is what happened in a79e9d69: pre-commit.ci silently dropped BOTH this +# import and the pre-existing dashboard-metrics one, un-registering the +# dashboard-metrics periodics that had been live since UN-3796 as collateral. +# A genuine reference cannot be fixed away by either hook. +# +# Deleting this tuple re-arms that failure, and it fails SILENTLY: the worker +# starts fine and only the scheduled messages are rejected as unknown tasks. +_PERIODIC_TASK_MODULES = (dashboard_metrics_tasks, agent_kv_tasks) + logger = WorkerLogger.get_logger(__name__) # Initialize worker configuration diff --git a/workers/shared/api/internal_client.py b/workers/shared/api/internal_client.py index ee9fbca7ac..329659dd30 100644 --- a/workers/shared/api/internal_client.py +++ b/workers/shared/api/internal_client.py @@ -1608,6 +1608,124 @@ def get_configuration( error=str(e), ) + # Agent-KV client methods (spec §5.3/§5.4) + def agent_kv_finalize( + self, + job_id: str, + org_id: str, + success: bool, + result: dict[str, Any] | None = None, + error: str = "", + usage_summary: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Terminalize an Agent-KV job via the internal finalize endpoint (Task 11). + + Idempotent server-side: a job already in a terminal state short-circuits + to ``{"finalized": False, ...}`` instead of rewriting its result. + + Args: + job_id: Agent-KV job ID. + org_id: Organization ID (required by the endpoint's body, not the + URL — mirrors ``StageReportView``/``FinalizeView``). + success: Whether the executor run succeeded. Must be a real bool — + the endpoint 400s on a non-bool value. + result: Engine result to persist as the job's output (success only). + error: Failure reason to persist (failure only). + usage_summary: Optional usage/cost summary to persist (success only). + + Returns: + Backend response: ``{"finalized": bool, "webhook_url": str, "status": str}``. + """ + payload: dict[str, Any] = {"org_id": org_id, "success": success} + if success: + payload["result"] = result or {} + if usage_summary is not None: + payload["usage_summary"] = usage_summary + else: + payload["error"] = error + return self.post( + f"v1/agent-kv/jobs/{job_id}/finalize/", + data=payload, + organization_id=org_id, + ) + + def agent_kv_stage_report( + self, + job_id: str, + org_id: str, + stage: str, + status: str, + seconds: float | None = None, + counters: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Report a pipeline stage for an Agent-KV job via the internal + stage-report endpoint (spec §5.3/§5.4). + + Idempotent server-side, mirroring ``agent_kv_finalize``: a job already + in a terminal state short-circuits to ``{"ok": True, "noop": True}`` + instead of recording the stage. + + Args: + job_id: Agent-KV job ID. + org_id: Organization ID (required by the endpoint's body, not the + URL — mirrors ``StageReportView``/``FinalizeView``). + stage: Pipeline stage name. Must be one of the stage names the + job's extractor declares (``agent_kv.constants + .STAGE_NAMES_BY_EXTRACTOR``) -- ``"table_extraction"`` for + ``table``, one of ``STAGE_NAMES`` for ``kv``. The endpoint + only checks the name is non-empty, but ``_status_document`` + filters a job's recorded stages through that list, so an + off-list name is stored and then dropped from every status + response. + status: Stage status -- ``"running"`` or ``"done"``, the only two + values ``StageReportView`` accepts. (These examples read + ``"started"``/``"completed"`` before review; both are 400s.) + seconds: Optional stage duration in seconds. + counters: Optional stage counters to record. + + Returns: + Backend response: ``{"ok": True}`` or ``{"ok": True, "noop": True}``. + """ + payload: dict[str, Any] = {"org_id": org_id, "stage": stage, "status": status} + if seconds is not None: + payload["seconds"] = seconds + if counters: + payload["counters"] = counters + return self.post( + f"v1/agent-kv/jobs/{job_id}/stage/", + data=payload, + organization_id=org_id, + ) + + def agent_kv_sweep(self) -> dict[str, Any]: + """Trigger the Agent-KV maintenance sweep via the internal endpoint + (spec §5.4, Fix 5/Fix 8). + + Platform-wide, like the endpoint itself (``SweepView``): no + ``org_id`` -- it terminalizes never-dispatched ``PENDING`` jobs and + stuck ``DISPATCHED``/``RUNNING`` jobs across every org in one call. + Idempotent and batch-capped server-side, so safe to call on a tight + schedule. + + Returns: + Backend response: ``{"swept": int, "timed_out": int}``. + """ + return self.post("v1/agent-kv/sweep/", data={}) + + def agent_kv_ttl_cleanup(self) -> dict[str, Any]: + """Trigger the Agent-KV TTL cleanup via the internal endpoint + (spec §5.4, Fix 5). + + Platform-wide, like the endpoint itself (``TTLCleanupView``): no + ``org_id`` -- it deletes staged input/result files for jobs past + ``expires_at`` and blanks their refs, across every org in one call. + Idempotent and batch-capped server-side. + + Returns: + Backend response: ``{"cleaned": int}``. + """ + return self.post("v1/agent-kv/ttl-cleanup/", data={}) + # Usage client methods (delegate to UsageAPIClient) def get_aggregated_token_count( self, file_execution_id: str, organization_id: str | None = None diff --git a/workers/shared/enums/worker_enums_base.py b/workers/shared/enums/worker_enums_base.py index 65465659c0..2620b8f5fc 100644 --- a/workers/shared/enums/worker_enums_base.py +++ b/workers/shared/enums/worker_enums_base.py @@ -185,6 +185,10 @@ class QueueName(str, Enum): # IDE callback queue (prompt studio post-execution callbacks) IDE_CALLBACK = "ide_callback" + # Agent-KV terminal-callback queue (agent_kv_complete / agent_kv_error; + # spec §5.3). Consumed by the IDE_CALLBACK worker alongside IDE_CALLBACK. + AGENT_KV_CALLBACK = "agent_kv_callback" + def to_env_var_name(self) -> str: """Convert queue name to environment variable name. diff --git a/workers/shared/infrastructure/config/registry.py b/workers/shared/infrastructure/config/registry.py index bfdfcf2e50..56fc84ff80 100644 --- a/workers/shared/infrastructure/config/registry.py +++ b/workers/shared/infrastructure/config/registry.py @@ -70,6 +70,9 @@ class WorkerRegistry: ), WorkerType.IDE_CALLBACK: WorkerQueueConfig( primary_queue=QueueName.IDE_CALLBACK, + # Agent-KV terminal callbacks (agent_kv_complete/agent_kv_error) + # ride the same worker on their own dedicated queue (spec §5.3). + additional_queues=[QueueName.AGENT_KV_CALLBACK], ), } @@ -156,6 +159,12 @@ class WorkerRegistry: TaskRoute("ide_prompt_complete", QueueName.IDE_CALLBACK), TaskRoute("ide_prompt_error", QueueName.IDE_CALLBACK), TaskRoute("ide_callback.tasks.*", QueueName.IDE_CALLBACK), + TaskRoute("agent_kv_complete", QueueName.AGENT_KV_CALLBACK), + TaskRoute("agent_kv_error", QueueName.AGENT_KV_CALLBACK), + # Cancellation does not go through finalize, so its terminal + # webhook has no callback to ride on; the backend enqueues this + # directly when its guarded cancel wins. + TaskRoute("agent_kv_cancelled", QueueName.AGENT_KV_CALLBACK), ], ), } diff --git a/workers/shared/infrastructure/config/tests/__init__.py b/workers/shared/infrastructure/config/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/workers/shared/infrastructure/config/tests/test_registry.py b/workers/shared/infrastructure/config/tests/test_registry.py new file mode 100644 index 0000000000..80a7c63117 --- /dev/null +++ b/workers/shared/infrastructure/config/tests/test_registry.py @@ -0,0 +1,44 @@ +"""Tests for worker registry configuration. + +The registry is where a queue stops being a string in a shell script and +becomes something a worker actually subscribes to. For Agent-KV that matters +more than usual: `agent_kv_callback` carries the terminal callbacks, and a job +whose callback is never consumed runs to SUCCESS in the executor and then sits +in RUNNING forever -- the result is never persisted, the staged input is never +deleted, the concurrency slot is never released and the webhook never fires. +Nothing errors, because nothing failed. +""" + +from shared.enums.worker_enums import WorkerType +from shared.enums.worker_enums_base import QueueName +from shared.infrastructure.config.registry import WorkerRegistry + + +def test_ide_callback_worker_also_subscribes_to_the_agent_kv_callback_queue(): + cfg = WorkerRegistry._QUEUE_CONFIGS[WorkerType.IDE_CALLBACK] + assert QueueName.AGENT_KV_CALLBACK in cfg.additional_queues, ( + "the ide_callback worker does not subscribe to agent_kv_callback; " + "Agent-KV jobs would execute and then never terminalize" + ) + + +def test_both_agent_kv_terminal_callbacks_route_to_that_queue(): + """`dispatch_job` attaches both as Celery link/link_error. A route missing + for either leaves one half of the outcome space unhandled -- most visibly, + a failed job that never moves out of RUNNING. + """ + routing = WorkerRegistry._TASK_ROUTES[WorkerType.IDE_CALLBACK] + routed = {r.pattern: r.queue for r in routing.routes} + + assert routed["agent_kv_complete"] == QueueName.AGENT_KV_CALLBACK + assert routed["agent_kv_error"] == QueueName.AGENT_KV_CALLBACK + + +def test_no_worker_type_claims_the_unshipped_sandbox_queue(): + """This deployment ships no codegen sandbox (UN-4215 is the fast-follow). + + A registry entry for a worker that is never deployed is how a queue ends up + advertised and undrained, so the absence is asserted rather than assumed. + """ + assert not hasattr(QueueName, "SANDBOX_CODEGEN") + assert not hasattr(WorkerType, "SANDBOX") diff --git a/workers/shared/tests/test_agent_kv_client.py b/workers/shared/tests/test_agent_kv_client.py new file mode 100644 index 0000000000..86f02b8f35 --- /dev/null +++ b/workers/shared/tests/test_agent_kv_client.py @@ -0,0 +1,90 @@ +"""Tests for ``InternalAPIClient.agent_kv_stage_report`` (Agent-KV cloud +executor Task 2) and the worker package's dependency on +``unstract-agent-kv-schema``. + +Mirrors ``workers/tests/test_agent_kv_scheduler_tasks.py``'s +``TestInternalClientAgentKvMethods`` pattern for patching the client's HTTP +layer: ``patch.object(InternalAPIClient, "post", ...)`` plus +``InternalAPIClient.__new__(InternalAPIClient)`` to avoid running the real +``__init__`` (session/config setup unrelated to this method). +""" + +from unittest.mock import patch + +from shared.api import InternalAPIClient + + +class TestAgentKvStageReport: + def test_posts_minimal_body_to_the_stage_endpoint(self): + """seconds/counters are omitted from the body when not provided.""" + with patch.object(InternalAPIClient, "post", return_value={"ok": True}) as m_post: + client = InternalAPIClient.__new__(InternalAPIClient) + result = client.agent_kv_stage_report("job-1", "org-1", "extract", "started") + + m_post.assert_called_once_with( + "v1/agent-kv/jobs/job-1/stage/", + data={"org_id": "org-1", "stage": "extract", "status": "started"}, + organization_id="org-1", + ) + assert result == {"ok": True} + + def test_posts_full_body_with_seconds_and_counters(self): + with patch.object(InternalAPIClient, "post", return_value={"ok": True}) as m_post: + client = InternalAPIClient.__new__(InternalAPIClient) + result = client.agent_kv_stage_report( + "job-2", + "org-2", + "extract", + "completed", + seconds=1.5, + counters={"rows": 3}, + ) + + m_post.assert_called_once_with( + "v1/agent-kv/jobs/job-2/stage/", + data={ + "org_id": "org-2", + "stage": "extract", + "status": "completed", + "seconds": 1.5, + "counters": {"rows": 3}, + }, + organization_id="org-2", + ) + assert result == {"ok": True} + + def test_returns_the_parsed_noop_response_for_a_terminal_job(self): + """The backend answers {"ok": true, "noop": true} for a terminal job; + the client returns it unchanged (idempotent, no special-casing). + """ + with patch.object( + InternalAPIClient, "post", return_value={"ok": True, "noop": True} + ): + client = InternalAPIClient.__new__(InternalAPIClient) + result = client.agent_kv_stage_report("job-3", "org-3", "extract", "started") + + assert result == {"ok": True, "noop": True} + + def test_empty_counters_dict_is_omitted_like_none(self): + """``if counters:`` (not ``is not None``) mirrors the brief's + implementation exactly -- an empty dict is falsy and left out. + """ + with patch.object(InternalAPIClient, "post", return_value={"ok": True}) as m_post: + client = InternalAPIClient.__new__(InternalAPIClient) + client.agent_kv_stage_report( + "job-4", "org-4", "extract", "started", counters={} + ) + + m_post.assert_called_once_with( + "v1/agent-kv/jobs/job-4/stage/", + data={"org_id": "org-4", "stage": "extract", "status": "started"}, + organization_id="org-4", + ) + + +class TestAgentKvSchemaImport: + def test_agent_kv_schema_package_importable(self): + """The worker venv depends on unstract-agent-kv-schema (path dependency, + editable) so the schema compiler is importable inside worker code. + """ + import unstract.agent_kv_schema # noqa: F401 diff --git a/workers/shared/tests/test_webhook_notify.py b/workers/shared/tests/test_webhook_notify.py new file mode 100644 index 0000000000..1c0fddfb9b --- /dev/null +++ b/workers/shared/tests/test_webhook_notify.py @@ -0,0 +1,221 @@ +"""Tests for the SSRF-guarded webhook sender (spec §6.7). + +``requests`` is patched at the ``webhook_notify`` import site; DNS is patched +inside ``unstract.core.network.ssrf``, which is where the egress decision now +lives (this module no longer has a resolver of its own -- Greptile review #4 +replaced its hand-rolled ``_host_is_public`` with the shared guard). So no real +network/DNS activity occurs. + +The behaviours asserted here are therefore the SINK's contract -- scheme +policy, the ``allow_insecure`` waiver, status-code handling, and the fact that +an unsafe destination never reaches ``requests.post``. The guard's own +address-classification rules are covered by +``unstract/core/tests/test_ssrf_guard.py``; the cases kept below are the ones +this sink must keep refusing regardless of which implementation is behind it. +""" + +import socket +from unittest import mock + +from shared.utils import webhook_notify as wn + +from unstract.core.network import ssrf + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_private_ip_refused(m_gai, m_requests): + m_gai.return_value = [(2, 1, 6, "", ("10.0.0.5", 443))] + assert wn.send_webhook("https://internal.example/x", {"a": 1}) is False + assert not m_requests.post.called + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_metadata_ip_refused(m_gai, m_requests): + m_gai.return_value = [(2, 1, 6, "", ("169.254.169.254", 80))] + assert wn.send_webhook("https://md.example/x", {}) is False + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_cgnat_shared_address_space_refused(m_gai, m_requests): + """RFC 6598 100.64.0.0/10. `ipaddress.is_private` does not cover this + range, which is why the old local guard needed a hand-maintained + `_CGNAT_RANGE` constant for it (and why SonarCloud flagged that hardcoded + address). The shared guard gets it from `is_global`, an allowlist kept + against the IANA registries, so the range is covered without naming it. + """ + m_gai.return_value = [(2, 1, 6, "", ("100.64.0.1", 443))] + assert wn.send_webhook("https://cgnat.example/x", {}) is False + assert not m_requests.post.called + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_nat64_embedded_loopback_refused(m_gai, m_requests): + """``64:ff9b::7f00:1`` is the RFC 6052 NAT64 prefix carrying 127.0.0.1 in + its low 32 bits; on a deployment with a NAT64 route, that is what gets + dialled. + + A REGRESSION GUARD, not a fixed bug: the previous local guard refused this + too, via ``ipaddress.is_reserved`` (checked on 3.12.9 -- the claim in + Greptile review #4 that it was delivered to does not hold). It is pinned + here because that refusal was incidental to a version-dependent flag, + whereas the shared guard refuses it deliberately by re-checking the + embedded IPv4. Both implementations pass; only one of them means to. + """ + m_gai.return_value = [(socket.AF_INET6, 1, 6, "", ("64:ff9b::7f00:1", 443, 0, 0))] + assert wn.send_webhook("https://nat64.example/x", {}) is False + assert not m_requests.post.called + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_nat64_embedded_public_ip_still_allowed(m_gai, m_requests): + """The companion to the test above, and the case that actually CHANGES + behaviour. + + ``64:ff9b::5db8:d822`` embeds the public 93.184.216.34 -- a legitimate + destination through a NAT64 translator. The old guard refused it, by the + same blanket ``is_reserved`` that refused the loopback one: it could not + tell the two apart, so it rejected the whole prefix. This test fails + against that implementation and passes against the shared guard, which + judges the embedded address. + """ + m_gai.return_value = [(socket.AF_INET6, 1, 6, "", ("64:ff9b::5db8:d822", 443, 0, 0))] + m_requests.post.return_value.status_code = 200 + assert wn.send_webhook("https://nat64-ok.example/hook", {}) is True + + +@mock.patch.object(wn, "requests") +def test_localhost_subdomain_refused_without_dns(m_requests): + """RFC 6761 reserves ``*.localhost`` to loopback, and the shared guard + decides that WITHOUT a resolver. + + No resolver is patched in this test, which is the point: the old guard had + no such rule and could only reach the same answer through DNS, so its + verdict depended on the image's libc (glibc's nss-myhostname answers for + the whole subtree, musl only for the bare apex, a minimal container for + neither). Same outcome, reached by luck rather than by rule. + """ + assert wn.send_webhook("https://api.localhost/hook", {}) is False + assert not m_requests.post.called + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_normal_public_ip_still_allowed(m_gai, m_requests): + m_gai.return_value = [(2, 1, 6, "", ("93.184.216.34", 443))] + m_requests.post.return_value.status_code = 200 + assert wn.send_webhook("https://cgnat-sibling.example/x", {}) is True + + +def test_http_scheme_refused_by_default(): + assert wn.send_webhook("http://example.com/x", {}) is False + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_public_host_posted_no_redirects(m_gai, m_requests): + m_gai.return_value = [(2, 1, 6, "", ("93.184.216.34", 443))] + m_requests.post.return_value.status_code = 200 + assert wn.send_webhook("https://example.com/hook", {"job_id": "j"}) is True + kw = m_requests.post.call_args.kwargs + assert kw["allow_redirects"] is False + assert kw["timeout"] == 10 + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_delivery_error_returns_false_never_raises(m_gai, m_requests): + m_gai.return_value = [(2, 1, 6, "", ("93.184.216.34", 443))] + m_requests.post.side_effect = Exception("boom") + assert wn.send_webhook("https://example.com/hook", {}) is False + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_ipv6_loopback_refused(m_gai, m_requests): + """Pins the 4-tuple IPv6 sockaddr shape: (addr, port, flowinfo, scope_id).""" + m_gai.return_value = [(socket.AF_INET6, 1, 6, "", ("::1", 443, 0, 0))] + assert wn.send_webhook("https://v6.example/x", {}) is False + assert not m_requests.post.called + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_ipv6_public_posted(m_gai, m_requests): + m_gai.return_value = [(socket.AF_INET6, 1, 6, "", ("2606:4700::1111", 443, 0, 0))] + m_requests.post.return_value.status_code = 200 + assert wn.send_webhook("https://v6.example/hook", {}) is True + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_empty_resolution_refused(m_gai, m_requests): + m_gai.return_value = [] + assert wn.send_webhook("https://nowhere.example/x", {}) is False + assert not m_requests.post.called + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_multi_record_any_unsafe_refuses(m_gai, m_requests): + """One public + one private record for the same host: must refuse (ALL + resolved addresses must be public, not just the first). + """ + m_gai.return_value = [ + (2, 1, 6, "", ("93.184.216.34", 443)), + (2, 1, 6, "", ("10.0.0.5", 443)), + ] + assert wn.send_webhook("https://mixed.example/x", {}) is False + assert not m_requests.post.called + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_server_error_status_returns_false(m_gai, m_requests): + m_gai.return_value = [(2, 1, 6, "", ("93.184.216.34", 443))] + m_requests.post.return_value.status_code = 500 + assert wn.send_webhook("https://example.com/hook", {}) is False + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_not_found_status_returns_false(m_gai, m_requests): + m_gai.return_value = [(2, 1, 6, "", ("93.184.216.34", 443))] + m_requests.post.return_value.status_code = 404 + assert wn.send_webhook("https://example.com/hook", {}) is False + + +@mock.patch.object(wn, "requests") +@mock.patch.object(ssrf.socket, "getaddrinfo") +def test_allow_http_true_posts_public_host(m_gai, m_requests): + m_gai.return_value = [(2, 1, 6, "", ("93.184.216.34", 80))] + m_requests.post.return_value.status_code = 200 + assert wn.send_webhook("http://example.com/hook", {}, allow_http=True) is True + assert m_requests.post.called + + +class TestAllowInsecure: + """``allow_insecure`` waives scheme + public-host guards (test stacks only).""" + + @mock.patch.object(wn, "requests") + def test_http_private_host_delivered_when_insecure(self, m_requests): + m_requests.post.return_value = mock.Mock(status_code=200) + ok = wn.send_webhook( + "http://host.docker.internal:18099/hook", + {"job_id": "j", "status": "completed"}, + allow_insecure=True, + ) + assert ok is True + m_requests.post.assert_called_once() + + @mock.patch.object(wn, "requests") + def test_default_still_refuses_http_and_private(self, m_requests): + assert ( + wn.send_webhook("http://host.docker.internal:18099/hook", {"job_id": "j"}) + is False + ) + assert not m_requests.post.called diff --git a/workers/shared/utils/webhook_notify.py b/workers/shared/utils/webhook_notify.py new file mode 100644 index 0000000000..da41233119 --- /dev/null +++ b/workers/shared/utils/webhook_notify.py @@ -0,0 +1,85 @@ +"""Terminal-state webhook delivery for Agent-KV (spec §6.7). + +The egress decision is NOT made here. It belongs to +``unstract.core.network.ssrf``, the single place in the codebase that decides +whether a tenant-supplied URL may be dialled -- the same guard behind the +prompt-postprocessing and pipeline-notification sinks, which exists (per its own +docstring) "so a new sink does not carry its own copy of the rules". + +This module shipped with its own ``_host_is_public``, built from negative +``ipaddress`` flags. Greptile review #4 asked for the shared guard on the +grounds that a NAT64 literal such as ``64:ff9b::7f00:1`` -- the RFC 6052 +prefix carrying 127.0.0.1 in its low 32 bits -- slipped past it. **That +specific claim does not hold**: on the runtime this worker uses (3.12), +``ipaddress`` reports ``is_reserved`` for ``64:ff9b::/96``, so the local guard +already refused it. Verified against 3.12.9 rather than assumed. + +The move is still right, for the reasons the local copy's existence was wrong: + +- It refused in the WRONG DIRECTION too. ``64:ff9b::5db8:d822`` embeds the + public 93.184.216.34 -- a legitimate destination for a deployment with a + NAT64 route -- and the same ``is_reserved`` flag refused that as well. The + shared guard re-checks the embedded IPv4 (``_EMBEDS_IPV4``) and so separates + the two cases instead of rejecting the prefix wholesale. +- Its NAT64 refusal was ACCIDENTAL. ``is_reserved`` is not an SSRF control; it + is a version-dependent table (its IPv6 entries have moved between CPython + releases), and nothing recorded that this guard leaned on it. The shared + guard refuses on ``is_global`` -- an allowlist maintained against the IANA + registries -- plus an explicit embedded-IPv4 re-check. +- Enumerating what to refuse misses whatever belongs to none of the flags, and + this copy had already been patched once for RFC 6598 shared address space for + exactly that reason. The shared guard gets that range from ``is_global`` for + free, which also retires the hardcoded ``100.64.0.0`` SonarCloud flagged. +- It also brings rules this sink never had: ``*.localhost`` decided without a + resolver, credentials-in-URL, and the urlparse/urllib3 host disagreement that + determines which host the socket actually connects to. + +Residual accepted risk, unchanged and shared with every other sink: DNS is +resolved for the check and again by ``requests``, so a rebinding window exists. +The control for that is pod egress policy, not application code. The payload +carries only {job_id, status} and the response body is never read. +""" + +import json +import logging + +import requests + +from unstract.core.network.ssrf import is_safe_webhook_url, safe_host + +logger = logging.getLogger(__name__) + +_TIMEOUT = 10 + + +def send_webhook( + url: str, payload: dict, *, allow_http: bool = False, allow_insecure: bool = False +) -> bool: + """``allow_insecure`` waives BOTH guards (http scheme and non-public host). + + Test/dev stacks only -- it exists so the e2e lane can deliver to a + receiver on the compose host (host.docker.internal is a private + address). Production never sets it; the SSRF guards stay mandatory. + """ + try: + if allow_insecure: + # Scheme and destination both unchecked: this is the e2e escape + # hatch, and the only path on which `requests.post` is reached + # without the shared guard having approved the URL. + pass + elif not is_safe_webhook_url( + url, allowed_schemes=("http", "https") if allow_http else ("https",) + ): + # is_safe_webhook_url already logged the refusal reason and host. + return False + resp = requests.post( + url, + data=json.dumps(payload), + headers={"Content-Type": "application/json"}, + timeout=_TIMEOUT, + allow_redirects=False, + ) + return 200 <= resp.status_code < 300 + except Exception: + logger.warning("webhook delivery failed (host=%s)", safe_host(url), exc_info=True) + return False diff --git a/workers/tests/test_agent_kv_scheduler_tasks.py b/workers/tests/test_agent_kv_scheduler_tasks.py new file mode 100644 index 0000000000..1cd3f939c7 --- /dev/null +++ b/workers/tests/test_agent_kv_scheduler_tasks.py @@ -0,0 +1,98 @@ +"""Tests for the thin Agent-KV maintenance-periodic proxy tasks (spec §5.4). + +Mirrors ``test_dashboard_metrics_tasks.py``'s style: these tasks do nothing but +call the shared ``InternalAPIClient`` facade, so what's worth pinning is the +registered wire name (a mismatch means the PG consumer drops the message as an +unknown task -- the failure mode with no error at the enqueue site) and which +client method each task calls. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +# Import via the PACKAGE, matching `scheduler/tasks.py`'s +# `from scheduler import agent_kv_tasks` -- see that file's comment and +# `test_dashboard_metrics_tasks.py`'s identical note on why this form (not a +# bare `import agent_kv_tasks`) is the one both runtime import mechanisms +# converge on. +_WORKERS_ROOT = Path(__file__).resolve().parent.parent +if str(_WORKERS_ROOT) not in sys.path: + sys.path.insert(0, str(_WORKERS_ROOT)) + +from scheduler import agent_kv_tasks as akt # noqa: E402 +from shared.api import InternalAPIClient # noqa: E402 + + +class TestRegistration: + """The wire name is this task pair's own new contract (no pre-existing + Beat/PG-scheduler row to mirror) -- a rename here silently breaks whatever + PgPeriodicTask row an operator registered against the old name. + """ + + @pytest.mark.parametrize( + "name,func", + [ + ("agent_kv.sweep", "agent_kv_sweep"), + ("agent_kv.ttl_cleanup", "agent_kv_ttl_cleanup"), + ], + ) + def test_task_is_registered_under_its_wire_name(self, name, func): + assert getattr(akt, func).name == name + + +class TestApiClientAcquisition: + def test_get_api_client_returns_an_internal_api_client(self): + with patch("shared.api.InternalAPIClient") as mock_cls: + akt._get_api_client() + mock_cls.assert_called_once_with() + + +class TestCallContract: + def test_sweep_calls_the_sweep_client_method(self): + mock_api = MagicMock() + mock_api.agent_kv_sweep.return_value = {"swept": 3, "timed_out": 1} + with patch.object(akt, "_get_api_client", return_value=mock_api): + result = akt.agent_kv_sweep() + mock_api.agent_kv_sweep.assert_called_once_with() + assert result == {"swept": 3, "timed_out": 1} + + def test_ttl_cleanup_calls_the_ttl_cleanup_client_method(self): + mock_api = MagicMock() + mock_api.agent_kv_ttl_cleanup.return_value = {"cleaned": 2, "retained": 1} + with patch.object(akt, "_get_api_client", return_value=mock_api): + result = akt.agent_kv_ttl_cleanup() + mock_api.agent_kv_ttl_cleanup.assert_called_once_with() + # `retained` is part of the contract: it counts jobs whose files could not + # be deleted and whose refs are kept for the next pass. Pinned so it cannot + # be dropped silently. + assert result == {"cleaned": 2, "retained": 1} + + +class TestInternalClientAgentKvMethods: + """The client methods (``workers/shared/api/internal_client.py``) the tasks + above call -- pinned separately so a change to either the endpoint path or + the call contract fails at the right layer. + """ + + def test_agent_kv_sweep_posts_to_the_sweep_endpoint(self): + with patch.object( + InternalAPIClient, "post", return_value={"swept": 1, "timed_out": 0} + ) as m_post: + client = InternalAPIClient.__new__(InternalAPIClient) + result = client.agent_kv_sweep() + m_post.assert_called_once_with("v1/agent-kv/sweep/", data={}) + assert result == {"swept": 1, "timed_out": 0} + + def test_agent_kv_ttl_cleanup_posts_to_the_ttl_cleanup_endpoint(self): + with patch.object( + InternalAPIClient, "post", return_value={"cleaned": 1, "retained": 0} + ) as m_post: + client = InternalAPIClient.__new__(InternalAPIClient) + result = client.agent_kv_ttl_cleanup() + m_post.assert_called_once_with("v1/agent-kv/ttl-cleanup/", data={}) + assert result == {"cleaned": 1, "retained": 0} diff --git a/workers/tests/test_llm_bearing_ops.py b/workers/tests/test_llm_bearing_ops.py new file mode 100644 index 0000000000..7b7e4c8a33 --- /dev/null +++ b/workers/tests/test_llm_bearing_ops.py @@ -0,0 +1,69 @@ +"""Every paid operation must be in the no-usage-records backstop. + +`execute_extraction` logs when an op in `_LLM_BEARING_OPS` finishes +successfully having emitted no usage records. That log line is the ONLY signal +that a run's billing produced nothing: the cloud `flush()` returns an empty list +rather than raising, so total loss of a job's usage rows is otherwise +indistinguishable from a job that legitimately made no LLM calls. + +So leaving a paid op out of the set is not a missing log line, it is a missing +alarm — and the gap is invisible, because the symptom is silence. + +`table_extract_api` (the Agent-KV API's table extractor) was absent, which is +how the whole Agent-KV table path shipped with no billing backstop at all. +Reported as 1.2 in the branch review. +""" + +from executor.tasks import _LLM_BEARING_OPS + +#: Operations known to drive at least one LLM call on every successful run. +#: Add to this list and the set together. +#: +#: **Known limitation of this file.** These are two literals, so the +#: assertions below can only fail when the copies DISAGREE -- never when both +#: are wrong together. A new paid op declared in neither still ships with no +#: billing alarm. The derived version, which asserts every operation in the +#: backend's `EXTRACTOR_ROUTES` is in `_LLM_BEARING_OPS` and so cannot be +#: satisfied by forgetting both, is +#: `backend/agent_kv/tests/test_queue_wiring_is_derived.py` +#: (`test_every_routed_operation_has_a_billing_backstop`). It lives there +#: because `EXTRACTOR_ROUTES` is not importable from this test venv, and +#: `executor.tasks` is not importable from the backend's -- neither side can +#: import both, which is why this pair exists at all. +#: +#: Kept rather than replaced: it covers the four non-Agent-KV paid ops, which +#: `EXTRACTOR_ROUTES` says nothing about. +PAID_OPERATIONS = { + "answer_prompt", + "single_pass_extraction", + "summarize", + "structure_pipeline", + # Drives two LLMs (advanced + lite) on every run. + "table_extract_api", +} + + +def test_the_table_api_operation_is_covered(): + """The regression this file was added for.""" + assert "table_extract_api" in _LLM_BEARING_OPS + + +def test_every_known_paid_operation_is_covered(): + missing = PAID_OPERATIONS - _LLM_BEARING_OPS + assert not missing, ( + f"{sorted(missing)} drive LLM calls but are absent from " + f"_LLM_BEARING_OPS, so a run of one that emits zero usage records " + f"logs nothing and the lost billing is silent" + ) + + +def test_the_set_has_not_quietly_grown_past_what_is_declared_here(): + """The reverse direction, so the two lists stay a real pair rather than one + drifting into a superset of the other. + """ + unexpected = _LLM_BEARING_OPS - PAID_OPERATIONS + assert not unexpected, ( + f"{sorted(unexpected)} were added to _LLM_BEARING_OPS without being " + f"declared here; add them to PAID_OPERATIONS with a note on what they " + f"spend" + ) diff --git a/workers/tests/test_queue_consumer_wiring.py b/workers/tests/test_queue_consumer_wiring.py new file mode 100644 index 0000000000..14fc562cc7 --- /dev/null +++ b/workers/tests/test_queue_consumer_wiring.py @@ -0,0 +1,325 @@ +"""Guard: queues and routes must match the fleet this deployment actually runs. + +The OSS analogue of the cloud chart's ``validate-pg-worker-fleet.yaml``, and it +exists because the failure it catches is **silent**: since UN-4046 made the PG +transport unconditional, a queue whose consumer is not configured still +*accepts* work -- ``enqueue_task`` succeeds, rows land in ``pg_queue_message``, +and nothing errors at the producer. The job simply sits in DISPATCHED forever. + +The assertions run in BOTH directions, which is the part to read carefully: + +* ``celery_executor_agentic_table`` and ``agent_kv_callback`` must HAVE a + consumer -- they are what this deployment dispatches onto. +* ``celery_executor_agentic_kv`` must NOT be advertised. This deployment ships + no ``agentic_kv`` plugin, so a fleet listing that queue would accept work + nothing can drain -- the same silent failure, arrived at from the other side. + +**These names are literals, and that is a known limitation.** Re-enabling the +``kv`` extractor is one ``EXTRACTOR_ROUTES`` entry, after which +``celery_executor_agentic_kv`` must be SERVED -- and the inverse assertion +above would have to be deleted to go green, which is the one moment a guard +should not be getting deleted. The derived version of both directions lives in +``backend/agent_kv/tests/test_queue_wiring_is_derived.py``, where +``EXTRACTOR_ROUTES`` is actually importable (the backend package is not +importable from this test venv -- that asymmetry is why the duplication +exists). That suite also covers ``run-worker-docker.sh``, which this one does +not reach. + +Keeping both is deliberate: two independent statements of the same fact, one of +which does not depend on the routing table being correct. + +An earlier version of this docstring described only the first direction and +narrated how ``celery_executor_agentic_kv`` ought to be wired in. That is the +one thing a guard with an inverted assertion must not say: a reader checking +the file against its own description would have concluded the test was wrong. + +Each test asserts against the variable the *running* consumer actually reads, +at every site that configures one -- plus the traefik rules, where the same +"configured-looking but unreachable" failure shows up as a request served by +the SPA instead of Django. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +yaml = pytest.importorskip("yaml") + +REPO_ROOT = Path(__file__).resolve().parents[2] +DEV_COMPOSE = REPO_ROOT / "docker" / "docker-compose.yaml" +TEST_COMPOSE = REPO_ROOT / "tests" / "compose" / "docker-compose.test.yaml" +RUN_WORKER = REPO_ROOT / "workers" / "run-worker.sh" + +# The queue the `kv` extractor WOULD dispatch onto, and which this deployment +# deliberately does not serve: the `agentic_kv` plugin ships on a separate PR. +# `kv` is absent from EXTRACTOR_ROUTES so a submit is refused at the serializer +# with a 400 -- these assertions are the second half of that, pinning that no +# fleet advertises a consumer for it. A queue that is listed but unserved +# accepts work and drains nothing, with no error at the producer: the job sits +# in DISPATCHED forever. That failure has already cost this team ~30 hours of +# firings with zero executions, which is why it is guarded from both ends. +UNSERVED_KV_QUEUE = "celery_executor_agentic_kv" + +# The terminal-callback queue. Its consumer is what moves a job out of +# RUNNING: finalize persists the result, deletes the staged input, releases +# the concurrency slot and fires the webhook. +AGENT_KV_CALLBACK_QUEUE = "agent_kv_callback" + +# The queue the backend dispatches Agent-KV TABLE extraction onto. Same queue +# the IDE table path already uses -- ExecutionContext(executor_name= +# "agentic_table") -> celery_executor_agentic_table -- which is exactly why the +# API path reuses that executor name rather than introducing one of its own. +# It was already wired for the IDE; these assertions keep it wired now that a +# second, paid entry point depends on it. +AGENT_TABLE_QUEUE = "celery_executor_agentic_table" + +# The env var the live PG consumer reads. Named here so a future rename has to +# touch this constant rather than silently bypassing every assertion below. +PG_QUEUE_VAR = "WORKER_PG_QUEUE_CONSUMER_QUEUE" + +# Read only by the disabled Celery executor. Setting a queue here does NOT give +# it a consumer -- asserting its absence is what makes this suite catch the +# original bug rather than a cosmetic rename of it. +DEAD_CELERY_VAR = "CELERY_QUEUES_EXECUTOR" + + +def _service_env(compose_path: Path, service: str) -> dict[str, str]: + """Return ``service``'s environment as a dict, from a compose file.""" + doc = yaml.safe_load(compose_path.read_text()) + env = (doc.get("services", {}).get(service) or {}).get("environment") or [] + if isinstance(env, dict): # mapping form + return {k: str(v) for k, v in env.items()} + out: dict[str, str] = {} + for item in env: # list form: "KEY=value" + key, _, value = str(item).partition("=") + out[key] = value + return out + + +def _queues(raw: str) -> set[str]: + """Queue names from a consumer queue list, ignoring any ``${VAR:-default}`` + wrapper the compose file uses to keep the value overridable. + """ + inner = re.sub(r"^\$\{[^:}]+:-(.*)\}$", r"\1", raw.strip()) + return {q.strip() for q in inner.split(",") if q.strip()} + + +@pytest.mark.parametrize( + ("compose_path", "label"), + [(DEV_COMPOSE, "dev compose"), (TEST_COMPOSE, "e2e test compose")], +) +def test_no_fleet_advertises_the_unserved_kv_queue(compose_path, label): + """The executor fleet must not list a queue with no plugin behind it.""" + env = _service_env(compose_path, "worker-pg-executor") + raw = env.get(PG_QUEUE_VAR) + assert raw is not None, ( + f"{label}: worker-pg-executor sets no {PG_QUEUE_VAR}; the running " + f"pg-queue-consumer would fall back to a default of unknown content" + ) + assert UNSERVED_KV_QUEUE not in _queues(raw), ( + f"{label}: {UNSERVED_KV_QUEUE} is advertised but this deployment " + f"carries no agentic_kv plugin to drain it. Work routed there is " + f"accepted and never runs, silently. Remove it from {PG_QUEUE_VAR}, " + f"or ship the plugin." + ) + + +def test_test_compose_does_not_wire_the_queue_onto_the_dead_celery_var(): + """Regression: the e2e override once set the Celery-only variable. + + That drained nothing on the PG transport. If someone re-adds it, the queue + looks configured while the lane still hangs -- so fail loudly. + + Scoped to TEST_COMPOSE deliberately, and NOT parametrized over both files + -- review asked for the parametrization and this is why it is wrong here. + `docker/docker-compose.yaml` legitimately sets `CELERY_QUEUES_EXECUTOR` on + `worker-pg-executor` (it is consumed by `run-worker-docker.sh:281` and by + `workers/executor/worker.py` for anyone still launching the Celery + executor), so parametrizing would turn this into an immediate red on a + config that is vestigial rather than wrong. + + What generalizes is not "the variable must be absent" but "the variable + must never advertise a queue the PG consumer does not serve" -- which is + the next test, and which DOES run against both files. + """ + env = _service_env(TEST_COMPOSE, "worker-pg-executor") + assert DEAD_CELERY_VAR not in env, ( + f"{DEAD_CELERY_VAR} is read only by the disabled Celery executor " + f"(workers/executor/worker.py). Queues belong in {PG_QUEUE_VAR}." + ) + + +@pytest.mark.parametrize( + ("compose_path", "label"), + [(DEV_COMPOSE, "dev compose"), (TEST_COMPOSE, "e2e test compose")], +) +def test_the_dead_celery_var_never_advertises_an_unserved_queue(compose_path, label): + """The generalized form of the test above, over both compose files. + + Wherever `CELERY_QUEUES_EXECUTOR` survives, it must stay a SUBSET of what + the PG consumer actually drains. The hazard is not the variable existing -- + it is the variable drifting into a superset, because then a reader + comparing the two lists concludes a queue is served when nothing drains it, + and the Celery path that would have drained it is disabled. + + Concretely this fails if someone adds `celery_executor_agentic_kv` to the + Celery list while the PG list correctly omits it. + """ + env = _service_env(compose_path, "worker-pg-executor") + dead = env.get(DEAD_CELERY_VAR) + if dead is None: + pytest.skip(f"{label} does not set {DEAD_CELERY_VAR}") + + pg_raw = env.get(PG_QUEUE_VAR) + assert pg_raw is not None, f"{label}: no {PG_QUEUE_VAR} to compare against" + + extra = _queues(dead) - _queues(pg_raw) + assert not extra, ( + f"{label}: {DEAD_CELERY_VAR} advertises {sorted(extra)}, which " + f"{PG_QUEUE_VAR} does not serve. {DEAD_CELERY_VAR} is read only by the " + f"disabled Celery executor, so nothing drains those -- work routed " + f"there is accepted and never runs. Add them to {PG_QUEUE_VAR} or drop " + f"them here." + ) + + +def test_run_worker_pg_executor_role_omits_the_unserved_kv_queue(): + """The host-run fleet's hardcoded role default must match the compose stacks. + + ``run-worker.sh`` carries its own PG_CONSUMER_ROLES map; a queue listed + here but not in compose is just as unconsumed for anyone running the fleet + directly. + """ + text = RUN_WORKER.read_text() + match = re.search(r'\["\$PG_ROLE_EXECUTOR"\]="executor;([^"]+)"', text) + assert match, "could not find the PG_ROLE_EXECUTOR entry in run-worker.sh" + assert UNSERVED_KV_QUEUE not in _queues(match.group(1)), ( + f"run-worker.sh's pg-executor role advertises {UNSERVED_KV_QUEUE}; a " + f"host-run fleet would accept Agent-KV work and never drain it." + ) + + +@pytest.mark.parametrize( + ("compose_path", "label"), + [(DEV_COMPOSE, "dev compose"), (TEST_COMPOSE, "e2e test compose")], +) +def test_pg_executor_consumes_the_agent_table_queue(compose_path, label): + """Agent-KV table jobs ride the IDE table executor's queue (spec R1).""" + env = _service_env(compose_path, "worker-pg-executor") + raw = env.get(PG_QUEUE_VAR) + assert raw is not None, ( + f"{label}: worker-pg-executor sets no {PG_QUEUE_VAR}; the running " + f"pg-queue-consumer would fall back to a default that omits " + f"{AGENT_TABLE_QUEUE}" + ) + assert AGENT_TABLE_QUEUE in _queues(raw), ( + f"{label}: {AGENT_TABLE_QUEUE} has no consumer -- both IDE table " + f"prompts and Agent-KV table jobs will sit in DISPATCHED forever with " + f"no error at the producer. Add it to {PG_QUEUE_VAR}." + ) + + +def test_run_worker_pg_executor_role_lists_the_agent_table_queue(): + text = RUN_WORKER.read_text() + match = re.search(r'\["\$PG_ROLE_EXECUTOR"\]="executor;([^"]+)"', text) + assert match, "could not find the PG_ROLE_EXECUTOR entry in run-worker.sh" + assert AGENT_TABLE_QUEUE in _queues(match.group(1)), ( + f"run-worker.sh's pg-executor role omits {AGENT_TABLE_QUEUE}; a " + f"host-run fleet would accept table work and never drain it." + ) + + +def test_pg_ide_callback_drains_the_agent_kv_callback_queue(): + """Without this the job runs to SUCCESS and then never completes. + + The executor enqueues its terminal continuation onto `agent_kv_callback`; + if no consumer drains it the row sits in pg_queue_message and the job stays + RUNNING forever, with nothing logged at the producer. The Celery + ide_callback worker has carried both queues since the callbacks landed -- + its PG twin did not, and the cloud chart wires it correctly, so only the + OSS stacks were affected. Found by running a real job end to end. + """ + env = _service_env(DEV_COMPOSE, "worker-pg-ide-callback") + raw = env.get(PG_QUEUE_VAR) + assert raw is not None, "worker-pg-ide-callback sets no consumer queue" + assert AGENT_KV_CALLBACK_QUEUE in _queues(raw), ( + f"{AGENT_KV_CALLBACK_QUEUE} has no consumer -- Agent-KV jobs will " + f"finish executing and never leave RUNNING." + ) + + +def test_run_worker_pg_ide_callback_role_drains_the_callback_queue(): + text = RUN_WORKER.read_text() + match = re.search(r'\["\$PG_ROLE_IDE_CALLBACK"\]="ide_callback;([^"]+)"', text) + assert match, "could not find the PG_ROLE_IDE_CALLBACK entry in run-worker.sh" + assert AGENT_KV_CALLBACK_QUEUE in _queues(match.group(1)) + + +# -------------------------------------------------------------------------- +# Traefik routing. Same family of failure as an unconsumed queue: the service +# is up, the route looks configured, and the request lands somewhere that +# cannot serve it. +# +# `/agent-kv/` is mounted by `backend/backend/base_urls.py`, but traefik's +# backend rule matched only `/api/v1`, `/deployment` and `/public`, and the +# frontend rule is the negation of exactly those -- so every Agent-KV request +# through the compose stack was served by the SPA's nginx and never reached +# Django. The e2e lane did not catch it because it talks to +# `UNSTRACT_BACKEND_URL` (port 8000) directly, bypassing traefik entirely. +# +# Reported as 2.2 in the branch review. +# -------------------------------------------------------------------------- + +AGENT_KV_PREFIX = "/agent-kv" + + +def _router_rule(compose_path: Path, router: str) -> str: + svc_labels = [] + for svc in yaml.safe_load(compose_path.read_text())["services"].values(): + svc_labels.extend(svc.get("labels") or []) + prefix = f"traefik.http.routers.{router}.rule=" + for label in svc_labels: + if isinstance(label, str) and label.startswith(prefix): + return label[len(prefix) :] + pytest.fail(f"no traefik rule for router {router!r} in {compose_path.name}") + + +def test_traefik_routes_agent_kv_to_the_backend(): + rule = _router_rule(DEV_COMPOSE, "backend") + assert f"PathPrefix(`{AGENT_KV_PREFIX}`)" in rule, ( + "traefik does not route /agent-kv to the backend; the request is served " + "by the frontend SPA and never reaches Django" + ) + + +def test_traefik_excludes_agent_kv_from_the_frontend(): + """The frontend rule is a negation list, so a prefix missing from it is + claimed by the SPA even once the backend also matches -- whichever router + wins, one of them is wrong. + """ + rule = _router_rule(DEV_COMPOSE, "frontend") + assert f"!PathPrefix(`{AGENT_KV_PREFIX}`)" in rule, ( + "the frontend router still claims /agent-kv; it must be excluded " + "explicitly, exactly as /api/v1, /deployment and /public are" + ) + + +def test_every_backend_prefix_is_excluded_from_the_frontend(): + """The two rules are each other's complement by construction. Asserting the + relationship rather than one hardcoded prefix means the next mount point + cannot be added to one side only. + """ + backend_rule = _router_rule(DEV_COMPOSE, "backend") + frontend_rule = _router_rule(DEV_COMPOSE, "frontend") + + backend_prefixes = set(re.findall(r"PathPrefix\(`([^`]+)`\)", backend_rule)) + excluded = set(re.findall(r"!PathPrefix\(`([^`]+)`\)", frontend_rule)) + + missing = backend_prefixes - excluded + assert not missing, ( + f"{sorted(missing)} route to the backend but are not excluded from the " + f"frontend router; the SPA will claim them" + ) diff --git a/workers/tests/test_worker_task_module_imports.py b/workers/tests/test_worker_task_module_imports.py index df54943500..021fd8e0e1 100644 --- a/workers/tests/test_worker_task_module_imports.py +++ b/workers/tests/test_worker_task_module_imports.py @@ -27,6 +27,7 @@ from __future__ import annotations import json +import os import subprocess import sys from pathlib import Path @@ -81,8 +82,11 @@ def _run_probe(source: str, workers: list[str] | None = None) -> dict[str, str]: proc = subprocess.run( - [sys.executable, "-c", - source.format(root=str(_WORKERS_ROOT), workers=workers or _TASK_MODULES)], + [ + sys.executable, + "-c", + source.format(root=str(_WORKERS_ROOT), workers=workers or _TASK_MODULES), + ], capture_output=True, text=True, cwd=str(_WORKERS_ROOT), @@ -90,7 +94,9 @@ def _run_probe(source: str, workers: list[str] | None = None) -> dict[str, str]: ) marker = [ln for ln in proc.stdout.splitlines() if ln.startswith("RESULT")] if not marker: - pytest.fail(f"probe did not complete.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}") + pytest.fail( + f"probe did not complete.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" + ) return json.loads(marker[-1][len("RESULT") :]) @@ -122,35 +128,55 @@ def test_every_tasks_module_imports_by_file_path() -> None: ) -def test_scheduler_registers_the_metrics_proxies_under_their_wire_names() -> None: - """The broken import is a SIDE-EFFECT import — it must actually register. - - Asserting only that scheduler.tasks imports would still pass if someone 'fixed' it - by deleting the import, silently unregistering the three dashboard-metrics proxies - that the PG metrics consumer resolves BY NAME. +def test_scheduler_tasks_registers_every_periodic_under_its_wire_name() -> None: + """Loading ``scheduler/tasks.py`` must REGISTER the periodic proxies, not + merely import without error. + + This replaces a test that stated this exact intent -- "asserting only that + scheduler.tasks imports would still pass if someone 'fixed' it by deleting + the import" -- and then imported ``scheduler.dashboard_metrics_tasks`` + DIRECTLY, which registers the tasks by itself. So it passed no matter what + ``scheduler/tasks.py`` contained, and it passed through the very regression + it was written to catch: in a79e9d69 pre-commit.ci's auto-fix deleted BOTH + side-effect imports from scheduler/tasks.py (ruff's isort merged the two + `from scheduler import ...` statements, which relocated each `# noqa` onto a + member line where it no longer suppresses the statement-level F401, and + pycln -- `all = true` -- then removed both). Nothing went red. + + The failure mode is silent: the worker boots fine and only the scheduled + messages are rejected as unknown tasks, so dashboard metrics stop updating + and Agent-KV jobs stop being swept with no error anywhere near the cause. + Hence going through ``scheduler.tasks`` -- what a booting worker actually + loads -- and asserting against the Celery registry rather than the modules. """ source = """ import importlib, json, sys sys.path.insert(0, {root!r}) -m = importlib.import_module("scheduler.dashboard_metrics_tasks") -print("RESULT" + json.dumps({{ - "aggregate": m.dashboard_metrics_aggregate.name, - "hourly": m.dashboard_metrics_cleanup_hourly.name, - "daily": m.dashboard_metrics_cleanup_daily.name, -}})) +importlib.import_module("scheduler.tasks") +from celery import current_app +print("RESULT" + json.dumps(sorted( + n for n in current_app.tasks + if n.startswith("dashboard_metrics.") or n.startswith("agent_kv.") +))) """ proc = subprocess.run( [sys.executable, "-c", source.format(root=str(_WORKERS_ROOT))], capture_output=True, text=True, cwd=str(_WORKERS_ROOT), + # Resolved at import with no default; unset, scheduler.tasks raises + # before any registration and this would fail for the wrong reason. + env={**os.environ, "INTERNAL_API_BASE_URL": "http://localhost:8000"}, timeout=300, ) marker = [ln for ln in proc.stdout.splitlines() if ln.startswith("RESULT")] assert marker, f"probe failed:\n{proc.stdout}\n{proc.stderr}" - names = json.loads(marker[-1][len("RESULT") :]) - assert names == { - "aggregate": "dashboard_metrics.aggregate_from_sources", - "hourly": "dashboard_metrics.cleanup_hourly_data", - "daily": "dashboard_metrics.cleanup_daily_data", - } + assert json.loads(marker[-1][len("RESULT") :]) == [ + # Pinned to verbatim Beat/PG-scheduler rows -- the consumer resolves + # these BY NAME, so a rename here is a silent un-registration too. + "agent_kv.sweep", + "agent_kv.ttl_cleanup", + "dashboard_metrics.aggregate_from_sources", + "dashboard_metrics.cleanup_daily_data", + "dashboard_metrics.cleanup_hourly_data", + ] diff --git a/workers/uv.lock b/workers/uv.lock index 76fda50140..c22c4ba892 100644 --- a/workers/uv.lock +++ b/workers/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = "==3.12.*" [[package]] @@ -3544,6 +3544,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8f/13/b688a905653871b10b4ff0403c2ff562c17a0bd50be0d44324f3c85ca48f/ujson-5.12.1-graalpy312-graalpy250_312_native-win_amd64.whl", hash = "sha256:4ea0c490c702c20495e97345acfcf0c2f3153e658ef537ff111929c48b89e10a", size = 45988, upload-time = "2026-05-05T22:04:43.36Z" }, ] +[[package]] +name = "unstract-agent-kv-schema" +version = "0.1.0" +source = { editable = "../unstract/agent-kv-schema" } + [[package]] name = "unstract-connectors" version = "0.0.3" @@ -3793,6 +3798,7 @@ dependencies = [ { name = "python-socketio" }, { name = "redis" }, { name = "requests" }, + { name = "unstract-agent-kv-schema" }, { name = "unstract-connectors" }, { name = "unstract-core" }, { name = "unstract-filesystem" }, @@ -3844,6 +3850,7 @@ requires-dist = [ { name = "python-socketio", specifier = ">=5.9.0" }, { name = "redis", specifier = ">=4.5.0,<6.0.0" }, { name = "requests", specifier = ">=2.33.0,<3.0.0" }, + { name = "unstract-agent-kv-schema", editable = "../unstract/agent-kv-schema" }, { name = "unstract-connectors", editable = "../unstract/connectors" }, { name = "unstract-core", editable = "../unstract/core" }, { name = "unstract-filesystem", editable = "../unstract/filesystem" },