kite/.planning/phases/07-redo-and-optimize-llm-integration/07-04-PLAN.md

phase, plan, type, wave, depends_on, files_modified, autonomous, requirements, must_haves

phase	plan	type	wave	depends_on	files_modified	autonomous	requirements	must_haves
07-redo-and-optimize-llm-integration	04	execute	4	07-03	backend/tasks/document_tasks.py backend/api/documents.py backend/tests/test_document_tasks.py backend/tests/test_documents.py	true	D-09 D-10 D-11	truths artifacts key_links extract_and_classify task decorator carries bind=True and max_retries=3 Classification failures raise _ClassificationError from _run; the outer sync task catches it and calls self.retry(countdown=[30,90,270][min(retries,2)]) MaxRetriesExceededError handler writes doc.status = 'classification_failed' via _mark_classification_failed POST /api/documents/{id}/classify sets doc.status='processing', commits, calls extract_and_classify.delay(doc_id), returns {'document_id': str, 'status': 'processing'} Non-classification failures (extract_failed, invalid_id) still return a dict and are not retried path provides contains backend/tasks/document_tasks.py Celery retry harness + _ClassificationError + _mark_classification_failed class _ClassificationError path provides contains backend/api/documents.py Re-queue behaviour on POST /{id}/classify extract_and_classify.delay from to via pattern backend/tasks/document_tasks.py::extract_and_classify Celery self.retry exponential backoff countdowns self.retry(.*countdown from to via pattern backend/api/documents.py POST /{id}/classify extract_and_classify.delay Celery enqueue extract_and_classify.delay from to via pattern extract_and_classify MaxRetriesExceededError handler _mark_classification_failed final status writeback classification_failed

Make the classification pipeline self-healing by adding Celery exponential-backoff retry around `extract_and_classify`, and convert the synchronous `POST /api/documents/{id}/classify` endpoint into a re-queue trigger so the new "Re-analyze" UI button (Plan 05) drives the same retry loop.

Purpose: Deliver D-09 (30s/90s/270s exponential backoff with max_retries=3), D-10 (final state remains classification_failed), and D-11 backend half (re-classify endpoint re-queues Celery instead of running synchronously). Apply Pitfall 3 from RESEARCH.md — self.retry() must escape the sync task layer, not the inner asyncio.run() — by introducing a _ClassificationError sentinel. Output: Refactored extract_and_classify task with bind=True, _ClassificationError sentinel, _mark_classification_failed helper, retry-with-countdown logic, and a refactored POST /classify endpoint that calls .delay(); matching tests promoted from xfail.

<execution_context> @$HOME/.claude/get-shit-done/workflows/execute-plan.md @$HOME/.claude/get-shit-done/templates/summary.md </execution_context>

@.planning/PROJECT.md @.planning/phases/07-redo-and-optimize-llm-integration/07-CONTEXT.md @.planning/phases/07-redo-and-optimize-llm-integration/07-RESEARCH.md @.planning/phases/07-redo-and-optimize-llm-integration/07-PATTERNS.md @.planning/phases/07-redo-and-optimize-llm-integration/07-03-SUMMARY.md @backend/tasks/document_tasks.py @backend/api/documents.py @backend/services/classifier.py @backend/tests/test_documents.py @backend/tests/test_document_tasks.py

From backend/tasks/document_tasks.py (current):

• @celery_app.task(name="tasks.document_tasks.extract_and_classify") def extract_and_classify(document_id: str) -> dict: calls return asyncio.run(_run(document_id))
• async def _run(document_id: str) -> dict performs UUID parse, DB load, MinIO fetch, extract text, then await classify_document(...)
• Current behaviour: on classification exception, sets doc.status = "classification_failed" and returns {"document_id": ..., "status": "classification_failed"} directly (no retry)
• A second task cleanup_abandoned_uploads already follows the asyncio.run bridge pattern

From backend/api/documents.py (current POST /{doc_id}/classify):

• Authenticated route via Depends(get_regular_user); ownership check returns 404 on mismatch (per STATE.md "Cross-user doc access returns 404 not 403")
• Currently await classifier.classify_document(doc_id=..., session=..., ai_provider=..., ai_model=...) synchronously and returns a dict containing assigned topics
• Receives only doc_id and current_user from auth context

Celery retry API (RESEARCH.md Pattern 4):

• @celery_app.task(bind=True, max_retries=3, name=...) signature requires self as first positional arg
• self.request.retries reports the current retry attempt (0 on first try)
• raise self.retry(exc=exc, countdown=N) schedules a retry; raises Retry exception
• self.MaxRetriesExceededError is the exception type raised by Celery when retries are exhausted

Important — asyncio.run + AsyncMock interaction (test contract):

• Production code calls asyncio.run(_mark_classification_failed(document_id)). asyncio.run() invokes the coroutine object returned by calling _mark_classification_failed(...); this is a __call__() on the patched function, NOT a __await__().
• An AsyncMock substitute for _mark_classification_failed therefore records the invocation via mock.call_args / mock.call_count, NOT via await_count. Tests must assert via assert_called_once_with(document_id), not assert_awaited_once_with(...).

Task 1: Celery retry harness + _ClassificationError + classification_failed writeback backend/tasks/document_tasks.py backend/services/classifier.py backend/db/models.py backend/db/session.py backend/tests/test_document_tasks.py - Module-level `class _ClassificationError(Exception)` exists - `async def _mark_classification_failed(document_id: str) -> None` exists, opens an AsyncSession, loads the Document, sets status="classification_failed", commits - `_run()` raises `_ClassificationError(str(exc))` from the existing classification try/except (lines ~109-124 per 07-PATTERNS.md) — for classification exceptions only - Extract/storage/validation failures still `return {... "status": "extract_failed" ...}` or `return {... "status": "invalid_id" ...}` and are not retried - The task decorator becomes `@celery_app.task(name="tasks.document_tasks.extract_and_classify", bind=True, max_retries=3)` - Task signature is `def extract_and_classify(self, document_id: str) -> dict` - try/except outside asyncio.run handles `_ClassificationError`: countdown = [30, 90, 270][min(self.request.retries, 2)]; raise self.retry(exc=exc, countdown=countdown) - Separate handler catches `self.MaxRetriesExceededError` (or `MaxRetriesExceededError` imported from celery.exceptions) and calls `asyncio.run(_mark_classification_failed(document_id))`; returns {"document_id": document_id, "status": "classification_failed"} - test_document_tasks.py::test_retry_backoff and ::test_exhaustion_sets_failed_status promoted to passing Refactor backend/tasks/document_tasks.py: 1. Add `from celery.exceptions import MaxRetriesExceededError` near the existing celery imports. 2. Add `class _ClassificationError(Exception): pass` at module level above the task. 3. Add `async def _mark_classification_failed(document_id: str) -> None:` that opens an AsyncSession via the existing session factory, calls `await session.get(Document, _uuid.UUID(document_id))`, sets doc.status = "classification_failed", commits. Use the same imports the existing _run uses for session + Document. 4. Modify _run(): in the existing try/except around the classification call (the block that currently sets `doc.status = "classification_failed"` and returns a dict on classification exception), replace the except body with `raise _ClassificationError(str(exc)) from exc`. The non-classification failure returns (invalid_id, extract_failed, storage failures) are preserved unchanged. 5. Change the task decorator from `@celery_app.task(name="tasks.document_tasks.extract_and_classify")` to `@celery_app.task(name="tasks.document_tasks.extract_and_classify", bind=True, max_retries=3)`. 6. Change the task signature from `def extract_and_classify(document_id: str)` to `def extract_and_classify(self, document_id: str)`. 7. Wrap the body as: try: return asyncio.run(_run(document_id)) except _ClassificationError as exc: countdowns = [30, 90, 270]; countdown = countdowns[min(self.request.retries, 2)]; raise self.retry(exc=exc, countdown=countdown) except MaxRetriesExceededError: asyncio.run(_mark_classification_failed(document_id)); return {"document_id": document_id, "status": "classification_failed"}.

Promote backend/tests/test_document_tasks.py::test_retry_backoff: use unittest.mock.patch on `tasks.document_tasks._run` so it raises `_ClassificationError("fake")`; construct a MagicMock `self` with `request.retries` settable; assert that calling `extract_and_classify.run(self, "<uuid>")` raises celery.exceptions.Retry; iterate retries=0,1,2 and assert the Retry's `countdown` attribute equals 30, 90, 270 respectively. (Use celery's `extract_and_classify.run` to bypass the Celery messaging layer and invoke the wrapped function directly.)

Promote test_exhaustion_sets_failed_status: patch `tasks.document_tasks._run` to raise `_ClassificationError`; patch `tasks.document_tasks._mark_classification_failed` with an AsyncMock; configure `self.retry` to raise `MaxRetriesExceededError`; call extract_and_classify.run(self, "<uuid>"); assert `mock_mark_failed.assert_called_once_with(document_id)` — NOT `assert_awaited_once_with`. Rationale: production code is `asyncio.run(_mark_classification_failed(document_id))`. The coroutine factory call (`_mark_classification_failed(document_id)`) is recorded as `call`, not `await`, on the AsyncMock; asyncio.run() then runs the resulting coroutine to completion (which on an AsyncMock returns immediately). Therefore the assertion must inspect `call_args`/`call_count`, not `await_count`. Also assert the return dict has status="classification_failed". Remove the xfail decorators from both tests.

cd backend && pytest tests/test_document_tasks.py::test_retry_backoff tests/test_document_tasks.py::test_exhaustion_sets_failed_status -x -v - Source assertion: backend/tasks/document_tasks.py contains `class _ClassificationError(Exception)`, `bind=True`, `max_retries=3`, `def extract_and_classify(self, document_id`, `countdowns = [30, 90, 270]`, `MaxRetriesExceededError`, `async def _mark_classification_failed` - Source assertion: backend/tasks/document_tasks.py contains `raise self.retry(exc=` once - Source assertion (test contract): backend/tests/test_document_tasks.py test_exhaustion_sets_failed_status uses `assert_called_once_with` (not `assert_awaited_once_with`) on the _mark_classification_failed mock — `grep "assert_called_once_with" backend/tests/test_document_tasks.py` returns a match and `grep "assert_awaited_once_with" backend/tests/test_document_tasks.py` returns no match for the _mark_classification_failed assertion line - Behavior: pytest backend/tests/test_document_tasks.py::test_retry_backoff exits 0 and asserts countdown values [30, 90, 270] - Behavior: pytest backend/tests/test_document_tasks.py::test_exhaustion_sets_failed_status exits 0 - Behavior: pytest backend/tests/ -v shows no new failures Celery retry harness in place; classification failures retry with 30/90/270 backoff; exhaustion writes classification_failed; two test stubs promoted (test_exhaustion_sets_failed_status uses assert_called_once_with to match the asyncio.run(coro) calling convention); rest of suite green. Task 2: POST /api/documents/{id}/classify → re-queue Celery backend/api/documents.py backend/tasks/document_tasks.py backend/services/classifier.py backend/tests/test_documents.py backend/db/models.py - POST /api/documents/{doc_id}/classify is `async def`, depends on get_db + get_regular_user (unchanged) - Loads the document via existing _get_owned_doc helper (raises 404 on wrong owner per STATE.md "Cross-user doc access returns 404 not 403") — keep the helper, do not duplicate ownership logic - Sets doc.status = "processing", awaits session.commit() - Calls extract_and_classify.delay(str(doc.id)) - Returns {"document_id": str(doc.id), "status": "processing"} with HTTP 200 - Existing synchronous behaviour (await classifier.classify_document(...) + return assigned topics) is removed - test_documents.py::test_reclassify_requeues_celery promoted to passing with extract_and_classify.delay monkeypatched Edit backend/api/documents.py: locate the existing POST /{doc_id}/classify route. Import `from tasks.document_tasks import extract_and_classify` (lazy import inside the function body to avoid circular import at module load — match the existing deferred-import pattern used elsewhere in the codebase per STATE.md decision "Deferred Celery import in /password-reset"). Replace the function body with: doc = await _get_owned_doc(doc_id, current_user, session) (use whatever the existing helper name is in documents.py — verify via grep before editing; if no helper exists, replicate the same select+ownership check the existing POST /{doc_id}/classify already performs). Then `doc.status = "processing"`. Then `await session.commit()`. Then `from tasks.document_tasks import extract_and_classify` followed by `extract_and_classify.delay(str(doc.id))`. Then `return {"document_id": str(doc.id), "status": "processing"}`. Remove any call to `await classifier.classify_document(...)` from this endpoint body. Remove any code that returns assigned topics from this endpoint — the topics will be written by the Celery task and clients should re-fetch the document after polling.

Promote backend/tests/test_documents.py::test_reclassify_requeues_celery: as an authenticated regular user, POST to /api/documents/{doc_id}/classify with monkeypatch on `tasks.document_tasks.extract_and_classify.delay` (a MagicMock). Confirm: response status is 200; response JSON contains status="processing" and document_id matching the doc; mock_delay was called exactly once with str(doc.id); after the request the Document row in DB has status="processing". Use the existing auth_user/admin_user fixtures from conftest.py per STATE.md "Celery mock required in /confirm tests". Remove the xfail decorator.

cd backend && pytest tests/test_documents.py::test_reclassify_requeues_celery -x -v - Source assertion: backend/api/documents.py POST /{doc_id}/classify body contains `extract_and_classify.delay(` - Source assertion: backend/api/documents.py POST /{doc_id}/classify body contains `doc.status = "processing"` and `await session.commit()` - Source assertion: backend/api/documents.py POST /{doc_id}/classify body does NOT contain `await classifier.classify_document(` - Behavior: pytest backend/tests/test_documents.py::test_reclassify_requeues_celery exits 0 - Behavior: full suite pytest backend/tests/ -v has no new failures - Behavior: cross-user POST /api/documents/{other_users_doc_id}/classify still returns 404 (regression of existing IDOR test) POST /classify re-queues instead of running synchronously; doc moves to processing; Celery task takes over; test_reclassify_requeues_celery passes; existing IDOR/auth coverage intact.

<threat_model>

Trust Boundaries

Boundary	Description
Authenticated user → POST /classify	doc_id crosses untrusted boundary; ownership re-checked via _get_owned_doc → 404 on mismatch
Celery worker process → DB	retry exponential backoff bounded (max_retries=3) — prevents runaway loops

STRIDE Threat Register

Threat ID	Category	Component	Disposition	Mitigation Plan
T-07-09	Denial of Service	Re-classify endpoint enqueue without rate limit	accept	per-user rate-limit infrastructure lives on auth endpoints (Phase 2) and is being expanded by Phase 6; reclassify endpoint is gated by ownership check and inherits Phase 6 per-account limiter once that ships; documenting accepted risk for v1
T-07-10	Tampering	Cross-user reclassify	mitigate	_get_owned_doc returns 404 on cross-user (matches STATE.md cross-user 404 policy); existing IDOR test in test_documents.py validates this
T-07-11	Information Disclosure	Retry exception payload	accept	_ClassificationError carries only the str(exc); no document content; Celery serializes exc message into the broker but RabbitMQ/Redis is internal-only
T-07-12	Tampering	self.retry mid-asyncio.run	mitigate	Pitfall 3 — retry raised from the outer sync layer, never inside asyncio.run; _ClassificationError sentinel enforces this boundary
T-07-SC	Tampering	No new packages	accept	celery already pinned; pip audit re-runs at phase gate
</threat_model>

- grep -F 'bind=True' backend/tasks/document_tasks.py returns the task decorator line. - grep -F 'extract_and_classify.delay' backend/api/documents.py returns the endpoint enqueue line. - pytest backend/tests/test_document_tasks.py backend/tests/test_documents.py::test_reclassify_requeues_celery exits 0. - pytest backend/tests/ -v shows no new failures (test_documents.py IDOR + auth tests still green). - grep -F 'await classifier.classify_document' backend/api/documents.py POST /{doc_id}/classify body returns 0. - test_exhaustion_sets_failed_status uses assert_called_once_with (asyncio.run(coro_factory(...)) calls the factory, not awaits it — see note).

<success_criteria>

• Celery retry harness operational with 30/90/270 backoff and final classification_failed status.
• POST /{id}/classify is a re-queue endpoint; clients receive status=processing.
• 3 previously-xfailed tests now pass (test_retry_backoff, test_exhaustion_sets_failed_status, test_reclassify_requeues_celery).
• No regressions in existing test_documents.py IDOR/auth coverage. </success_criteria>

Create `.planning/phases/07-redo-and-optimize-llm-integration/07-04-SUMMARY.md` when done.