UN-4044 [FEAT] Agent-KV extraction engine — public API, schema compiler and hardened code sandbox

[vishnuszipstack]

What

The Agent-KV extraction engine's OSS half: the agent_kv Django app (public API, auth, dispatch, storage, rate limiting, job lifecycle), the unstract-agent-kv-schema workspace package, the hardened code sandbox worker, and the compose/test wiring for both.

Built by @Arun; merged up to current main and raised as part of taking the work forward. Pairs with the cloud half — see Related.

Why

Agent-KV is a multi-stage key/value extraction engine. Agentic Prompt Studio v2 already declares it as the Multi-Stage Extractor — MULTI_STAGE → executor="agentic_kv", operation="kv_extract" is on main today with available=False. Merging this is what lets that flag flip; without it, M2 would have to rebuild the same pipeline as a fourth copy of an engine that already exists.

Reference: Agent-KV — Architecture Review (UN-4044).

How

Size: 116 files, ~16.1k lines, including the agent_kv app, its 16 test modules, and the schema-compiler package.

Security posture — generated Python runs behind five independent layers, and the AST gate is deliberately not the boundary (layers 2–5 are). With no transport configured the engine errors rather than executing locally; there is no fallback path:

1. AST gate — 9-module allowlist; closes aliasing, attribute-reference binding, from sys import, dunder subscripts
2. Scrubbed subprocess — empty env, python -I -S -E, start_new_session, five rlimits, CPU backstop above the wall clock
3. Hardened pod — non-root, read-only rootfs, all caps dropped, RuntimeDefault seccomp, no service-account token
4. NetworkPolicy — ingress []; egress only broker, db-proxy, DNS
5. No secrets — only its own transport credential; never the platform ENCRYPTION_KEY; no tenant id reaches the pod

The accepted v1 residual is narrow: generated code can open() a world-readable file in its own pod and return it to the caller who submitted the job. pathlib, os, socket and urllib are outside the allowlist, so there is no path to a secret, a cross-tenant read, or exfiltration. Layers 4–5 are what make that true and must not be relaxed — weakening either turns a local file read into an exfiltration path.

Review findings — state at this head

The architecture review (6 Sep) raised four. Re-verified against this branch, post-merge:

#	Finding	Severity	State
1	Executor queue has no consumer — jobs hang in DISPATCHED silently	Blocking	Fixed — celery_executor_agentic_kv is in the executor role's queue list; tests/test_queue_consumer_wiring.py (13 tests) guards it
2	Test compose overrides three services main deleted	Blocking	Fixed — all 9 overridden services resolve; worker-sandbox is defined by this branch
3	Public wire format is single-extractor	Freeze risk	Open — deliberately deferred, see below
4	Sandbox consumes over the retired RabbitMQ broker	Operational	Fixed — now pg-sandbox on WORKER_PG_QUEUE_CONSUMER_QUEUE: sandbox_codegen, with fleet-guard wiring

Finding 3 is knowingly out of scope here

The review recommends moving the KV knobs under extractors: [{name, options}], keying the result document by extractor, and promoting usage_summary to per-extractor, with today's flat fields kept as a deprecated alias. An extractors field exists (execution_serializers.py:209) but qa / challenge / extraction_mode / calculations are still top-level.

Shipping first is deliberate: nothing consumes this API yet, APS v2 keeps Multi-Stage dark behind available=False, and a wire-format redesign tangled into a 16k-line PR is harder to review than either change alone. The review's own step 4 is "fix findings 1–2 during the rebase; open both PRs" — 1, 2 and 4 are done.

Note the review's ordering constraint ("must precede M0/M0b merging") is already moot: M0/M0b merged on 30 Sep, and did so having already adopted the review's §5.3 recommendation — the registry points at agentic_kv, not at a rebuilt pipeline.

Can this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)

Low. The work is additive — a new Django app, a new workspace package, a new worker and its compose/test entries. No existing endpoint changes shape and no existing model is altered.

Two things a reviewer should check rather than take on trust:

• Compose bring-up. Finding 2 was a platform-wide break: an override for a non-existent service makes compose reject the whole project, not just Agent-KV. Verified here — every overridden service is defined — but it is the failure worth re-confirming in CI.
• Queue consumption. Finding 1's failure mode is silent: the producer's enqueue succeeds and rows sit unconsumed forever, with no error. test_queue_consumer_wiring.py is the guard.

Database Migrations

• New agent_kv app migrations (new tables only; no alterations to existing tables).

Env Config

• AGENT_KV_STORAGE_DIR_PREFIX (defaults to unstract/agent_kv)
• AGENT_KV_CALCULATIONS_ENABLED — defaults false on absence, so OSS/on-prem installs stay flag-off

Notes on Testing

• cd backend && pytest agent_kv/tests --no-migrations — 165 passed
• cd workers && pytest tests/test_queue_consumer_wiring.py — 13 passed
• Compose bring-up with tests/compose/docker-compose.test.yaml

Related

• Cloud half: UN-4044-agent-kv-cloud-executor (raised alongside this)
• UN-4044 · UN-3833 (APS v2 M2)

Checklist

I have read and understood the Contribution Guidelines.

🤖 Generated with Claude Code

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ vishnuszipstack
❌ arunvenkataswamy
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[greptile-apps]

via Greptile

Confidence Score: 4/5

[Critical risk] Adds public API, database schema, and executor dispatch for a new extraction engine.

The PR does not appear safe to merge until failed file deletions can make progress under sustained TTL cleanup load.

Findings

1.  Failed deletions can starve ▶
2.  Index cannot serve cleanup order ▶

Fix with agent prompt

### Issue 1
backend/agent_kv/maintenance.py:167
If at least 500 never-attempted jobs expire between cleanup runs, they fill every batch because rows with no failure timestamp sort first. Jobs whose file deletion failed are then never retried, so their retained files can remain in storage indefinitely. Reserve some cleanup capacity for retries even when new expirations keep arriving.

### Issue 2
backend/agent_kv/models.py:101
The new index uses PostgreSQL’s default ascending null order, but cleanup requests `cleanup_failed_at ASC NULLS FIRST`. PostgreSQL therefore cannot use this index to return rows in cleanup order and must sort the matching expired jobs before applying the 500-row limit. That adds growing work to each cleanup run when there is a large backlog.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

The PR adds the Agent-KV API, schema compiler, sandbox worker, and deployment wiring. Since the previous review, it adds a timestamp-based ordering intended to stop failed deletions from blocking newer TTL cleanup. That ordering makes progress for fresh jobs but can starve retries under sustained load, and its new index does not match the requested null ordering.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  E[Expired jobs with file references] --> O[Order by cleanup failure time, NULLS FIRST]
  O --> L[Select up to 500]
  L --> D[Attempt file deletion]
  D -->|Failure| R[Keep reference and stamp failure time]
  R --> O
  D -->|Success| C[Clear references]

Loading

_{Reviews (4) · Last reviewed commit: "UN-4044 [FIX] Stop failed TTL cleanups f..."}

[sonarqubecloud]

Quality Gate passed

Issues
36 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[greptile-apps]

Failed deletions can starve

If at least 500 never-attempted jobs expire between cleanup runs, they fill every batch because rows with no failure timestamp sort first. Jobs whose file deletion failed are then never retried, so their retained files can remain in storage indefinitely. Reserve some cleanup capacity for retries even when new expirations keep arriving.

Prompt To Fix With AI

This is a comment left during a code review.
Path: backend/agent_kv/maintenance.py
Line: 167

Comment:
**Failed deletions can starve**

If at least 500 never-attempted jobs expire between cleanup runs, they fill every batch because rows with no failure timestamp sort first. Jobs whose file deletion failed are then never retried, so their retained files can remain in storage indefinitely. Reserve some cleanup capacity for retries even when new expirations keep arriving.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[greptile-apps]

Index cannot serve cleanup order

The new index uses PostgreSQL’s default ascending null order, but cleanup requests cleanup_failed_at ASC NULLS FIRST. PostgreSQL therefore cannot use this index to return rows in cleanup order and must sort the matching expired jobs before applying the 500-row limit. That adds growing work to each cleanup run when there is a large backlog.

Prompt To Fix With AI

This is a comment left during a code review.
Path: backend/agent_kv/models.py
Line: 101

Comment:
**Index cannot serve cleanup order**

The new index uses PostgreSQL’s default ascending null order, but cleanup requests `cleanup_failed_at ASC NULLS FIRST`. PostgreSQL therefore cannot use this index to return rows in cleanup order and must sort the matching expired jobs before applying the 500-row limit. That adds growing work to each cleanup run when there is a large backlog.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[github-actions]

Unstract test results

Per-group results

Status	Group	Tier	Passed	Failed	Errors	Skipped	Duration (s)
⚪	e2e-agent-kv	e2e	0	0	0	1	0.1
✅	e2e-api-deployment	e2e	3	0	0	0	16.3
✅	e2e-coowners	e2e	1	0	0	0	1.2
✅	e2e-etl	e2e	1	0	0	0	14.4
✅	e2e-login	e2e	2	0	0	0	1.2
✅	e2e-prompt-studio	e2e	1	0	0	0	7.4
✅	e2e-smoke	e2e	2	0	0	0	1.1
✅	e2e-workflow	e2e	1	0	0	0	12.2
✅	frontend	unit	620	0	0	0	18.8
✅	integration-backend	integration	602	0	0	27	49.7
✅	integration-connectors	integration	1	0	0	7	6.6
✅	integration-workers	integration	164	0	0	1	33.7
❌	ui	e2e	0	1	0	0	0.0
✅	unit-backend	unit	1595	0	0	1	52.6
✅	unit-connectors	unit	72	0	0	0	10.4
✅	unit-core	unit	276	0	0	0	3.3
✅	unit-platform-service	unit	15	0	0	0	2.8
✅	unit-rig	unit	120	0	0	0	4.7
✅	unit-runner	unit	10	0	0	0	3.1
✅	unit-sdk1	unit	718	0	0	0	34.1
✅	unit-workers	unit	1506	0	0	1	131.7
	TOTAL		5710	1	0	38	405.4

Critical paths

⚠️ Critical paths not yet covered

• workflow-execution-fan-out — Multi-file workflow execution fans out to file-processing workers and rejoins. (declared coverage: no groups declared)

✅ Covered critical paths

• auth-login — covered by e2e-login
• adapter-register-llm — covered by integration-backend
• workflow-author — covered by integration-backend
• co-owner-manage — covered by integration-backend, e2e-coowners
• workflow-create-execute — covered by e2e-workflow
• api-deployment-provision — covered by integration-backend
• api-deployment-auth — covered by integration-backend
• api-deployment-run — covered by e2e-api-deployment
• mcp-server-auth — covered by integration-backend
• mcp-platform-auth — covered by integration-backend
• platform-key-whoami — covered by integration-backend
• prompt-studio-author — covered by integration-backend
• prompt-studio-fetch-response — covered by e2e-prompt-studio
• connector-register-test — covered by integration-backend
• pipeline-etl-execute — covered by e2e-etl
• usage-aggregate-read — covered by integration-backend
• usage-token-tracking — covered by e2e-api-deployment
• callback-result-delivery — covered by e2e-api-deployment

[vishnuszipstack]

Automated multi-agent review — findings

This review was produced by a multi-agent code review of #2309 and its cloud counterpart Zipstack/unstract-cloud#1816. Each finding below was re-verified against the code (several by execution). Grouped as must-fix before merge, should-fix, and suggestions. Tenant isolation, mark_terminal CAS, the subscription gate, the TTL-starvation fix (cb6f7f2), and the no-local-execution-fallback property were all checked and confirmed correct — not repeated here.

Merge order: this PR must merge before the cloud half; the cloud plugin imports unstract-agent-kv-schema from this branch and its CI is red until this lands.

Inline comments follow.

[vishnuszipstack]

[Must-fix] Stranded-job race no sweep can recover. This guarded UPDATE persists dispatched_at only for a still-PENDING row, but StageReportView (internal_views.py:112) promotes PENDING→RUNNING on the executor's first stage report. If that report wins the race, this UPDATE matches 0 rows and dispatched_at stays NULL. Neither sweep phase can then see the row: phase 1 requires status=PENDING; phase 2 requires 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. Fix: stamp dispatched_at for any non-terminal row, or add a dispatched_at__isnull=True arm to sweep phase 2.

[vishnuszipstack]

[Must-fix] Both limiters fail OPEN. check_and_acquire here and check_key_rate (:79) both return True on any Redis error. A Sentinel failover or pool exhaustion removes the concurrency ceiling and the per-key rate ceiling simultaneously, and the API keeps returning 202s for billable LLM work — with only a logger.warning per request. The Lua script itself is correct; it's the catch that's the defect. Fail closed (429), or gate the choice behind a setting so the waiver is deliberate and observable.

[vishnuszipstack]

[Must-fix] ReDoS via author-supplied regex. compile_schema caps the pattern length (200) but never compiles it, and this runs it with no time budget. Measured on this code: format: "regex:^(a+)+$" (7 chars) takes ~7s on 28 as and ~4x per added char — a 40-char value runs for hours. validate_format runs per key per value in the QA pass, with no timeout/cancellation, and AGENT_KV_CONCURRENT_LIMIT=5 lets one org pin five shared worker slots. Fix: re.compile() at submit (raise SchemaError on re.error) and run user patterns under a hard wall-clock budget (or re2). The length cap is not a mitigation — ReDoS needs ~10 chars.

[vishnuszipstack]

[Must-fix] Failed input-delete still blanks the ref. delete_input(job) swallows every exception (storage.py), and this update(input_ref="") runs unconditionally. TTL cleanup selects only rows with a non-blank ref, so a transient object-store error here orphans the customer's uploaded document permanently. This is exactly the bug the PR fixed for delete_job_files + its two callers (Greptile #3) — this third site was missed. Make delete_input return a confirmed-cleared signal and blank only on success; log with exc_info=True.

[vishnuszipstack]

[Must-fix: correct the stated security posture] The AST gate is fully bypassable. Verified by execution under the production python -I -S -E: statistics.random._os.system('id') ran id; collections._sys is the sys module; collections._sys._getframe(0).f_builtins['__imp'+'ort__']('os').system(...) reaches arbitrary import+exec — all GATE-PASS. Allowlisted modules re-export sys/os as non-dunder attributes the gate never inspects, and the dunder-subscript check only catches constant strings (so '__imp'+'ort__' slips through). This does not by itself break the model (layers 2–5 still contain it), but the module/PR text claiming the gate 'closes aliasing / from-sys-import / dunder subscripts' overstates it to ~zero marginal value. Represent the gate as best-effort and lean the design explicitly on layers 2–5.

[vishnuszipstack]

[Should-fix] Unbounded capture can OOM the worker. communicate() buffers the full child stdout+stderr in the parent's memory; the [:_CAPTURE_LIMIT] truncation happens after. RLIMIT_FSIZE doesn't apply to pipes and RLIMIT_AS bounds only the child, so while True: print('A'*100000) (no imports, passes the gate trivially) streams GB into the parent and can OOM the Celery worker before the timeout fires — a gate-agnostic DoS. Cap captured bytes during the read, not after. Also: both broad except blocks here (:164, :229) log nothing, so an OSError/fork-EAGAIN is reported to the engine as a codegen failure.

[vishnuszipstack]

[Must-fix — cross-repo contract] Image uploads are accepted here but extract nothing in the engine. ALLOWED_EXTENSIONS accepts .png/.jpg/.jpeg/.tiff and this sets pages_total=1, so the job dispatches normally. But the cloud engine's _build_agent_graph (agentic_kv kv_extractor.py) only treats .pdf/.xlsx/.xls as a document, and ImageLoader.load_pages has no call site anywhere — so an image returns success:true, every key not-found, 1 page billed. Either drop the image extensions here, or wire images through ImageLoader.load_pages on the engine side. The two repos currently disagree about what file types the product supports, and it fails open (billed empty result).

[vishnuszipstack]

[Must-fix] Three new test dirs are in no CI group. This line adds sandbox/tests and shared/infrastructure/config/tests to unit-workers but omits ide_callback/tests (401-line callback suite). Separately, no group covers unstract/agent-kv-schema/tests or unstract/filesystem/tests. So test_compile.py — the regression net for the schema ReDoS/depth findings — and the callback suite run nowhere in CI, and the PR's '165 passed' is a local figure, not CI-verified. Add unit-agent-kv-schema + unit-filesystem groups and append ide_callback/tests here.

[vishnuszipstack]

[Must-fix before sign-off] The documented risk acceptance is factually wrong. This R12 text describes the residual as 'a single read of a known path … non-sensitive … not an exfiltration path.' Verified false by execution: generated code open()s /proc/<ppid>/environ (child runs as the same UID as the worker) and recovers the worker's full environment incl. DB_PASSWORD; with the sandbox pod mounting the database secret group and the NetworkPolicy allowing egress to db-proxy:5432 (cloud half), that credential is directly usable. So the residual is credential disclosure + DB reach beyond the job, not a benign file read. This is the text an admin signs off against — correcting it is a prerequisite independent of when the hardening (numeric runAsUser / scoped DB role / hidepid) lands. Same paragraph is duplicated in the sandbox design spec.