Files
DocsGPT/tests/api/user/sources/test_upload_idempotency.py
T
Alex e351f45d88 Feat notification system (#2472)
* feat: SSE notification system

Adds a per-user SSE pipe (GET /api/events) plus a per-message
chat-stream reconnect endpoint (GET /api/messages/<id>/events).

Backend substrate:
- application/events/ — durable journal (Redis Streams) + live
  pub/sub for user-scoped events, with publish_user_event() as
  the worker-side entrypoint.
- application/streaming/ — broadcast_channel for pub/sub fanout
  and event_replay for the per-message snapshot+tail path.
- application/storage/db/repositories/message_events.py +
  alembic 0007 — Postgres journal for chat-stream events.
- application/worker.py — ingest/reingest/remote/connector/
  attachment/mcp_oauth tasks publish queued/progress/completed/
  failed envelopes alongside their existing status updates.

Frontend client:
- frontend/src/events/ — connect/reconnect, Last-Event-ID cursor,
  backoff with jitter. Each tab runs its own connection; no
  cross-tab dedup (future work).
- frontend/src/notifications/ — recentEvents ring, cursor
  tracking, tool-approval toast.
- frontend/src/upload/uploadSlice.ts — extraReducers for
  source.ingest.* and attachment.* events.

Coverage: 132 SSE tests across events substrate, replay, journal,
routes, and worker publishes.

* refactor(attachments): remove polling, SSE-only

frontend/src/components/MessageInput.tsx no longer runs a 2s
setInterval against getTaskStatus for every processing
attachment. The attachment.* SSE reducers in uploadSlice.ts are
now the sole driver of attachment state transitions.

* feat(connector): consume source.ingest.* SSE, remove polling

frontend/src/components/ConnectorTree.tsx now mirrors FileTree's
slice-walking pattern: it watches notifications.recentEvents
for source.ingest.{completed,failed} envelopes matching the
sync's source id, and no longer polls /task_status every 2s.

* refactor(source-ingest): remove polling, SSE-only

frontend/src/upload/Upload.tsx and
frontend/src/components/FileTree.tsx no longer run getTaskStatus
polling fallbacks. The source.ingest.* SSE reducers in
uploadSlice.ts and FileTree's slice walk are now the sole
drivers of upload/reingest state transitions.

* refactor(mcp-oauth): carry authorization_url in SSE, remove polling

application/worker.py::mcp_oauth now publishes
authorization_url on the mcp.oauth.awaiting_redirect envelope.
frontend/src/modals/MCPServerModal.tsx consumes it from SSE
instead of polling /oauth_status/<task_id> every 1s.

The URL is generated inside DocsGPTOAuth.redirect_handler when
the FastMCP client triggers OAuth. The worker now plumbs a
publish callback through tool_config -> MCPTool -> DocsGPTOAuth
so the awaiting_redirect publish fires from inside the handler
at the exact point the URL becomes known. The legacy Redis
mcp_oauth_status setex writes and the GET
/api/mcp_server/oauth_status/<task_id> endpoint are kept as
belt-and-suspenders; nothing in the frontend reads them now.

* feat(source-ingest): plumb limited flag through SSE for token-cap UX

application/worker.py::ingest_worker and remote_worker now publish
``limited: bool`` on the source.ingest.completed envelope.
uploadSlice routes ``payload.limited === true`` to a failed status
with a ``tokenLimitReached`` flag, and UploadToast surfaces the
translated tokenLimit i18n string. No worker code path sets
limited=true today; this is a forward-looking contract so when
token-cap detection lands, the UX is already wired.

* refactor(mcp-oauth): read status from SSE journal, drop polling endpoint

MCPOAuthManager.get_oauth_status now walks the per-user SSE Streams
journal (user:{user_id}:stream) for the latest mcp.oauth.* envelope
matching the task id, returning the status string derived from the
event type suffix and the payload fields. The worker is the single
source of truth — its publish_user_event calls write the same
record the SSE client receives live.

Removed:
- /api/mcp_server/oauth_status/<task_id> route in
  application/api/user/tools/mcp.py
- mcp_oauth_status worker function and mcp_oauth_status_task Celery
  wrapper
- All mcp_oauth_status:{task_id} Redis setex writes (4 in mcp_oauth,
  2 in DocsGPTOAuth.redirect_handler / callback_handler)
- The update_status closure in mcp_oauth that wrote the polling
  payload

Tests updated:
- get_oauth_status now takes (task_id, user_id); new coverage walks
  a fake xrevrange response for the completed envelope, the no-match
  case, and a Redis-down case
- Removed TestMCPOAuthStatus route tests and TestMcpOauthStatusTask
  celery-wrapper test
- Removed the two oauth_status methods from the integration runner

mcp_oauth:auth_url/state/code/error Redis keys remain — they are
the OAuth flow's own state (not the dropped polling payload).

* chore(mcp-oauth): delete orphaned getMCPOAuthStatus client

The /api/mcp_server/oauth_status/<task_id> endpoint was removed in
the prior commit; the corresponding userService method and the
MCP_OAUTH_STATUS endpoint constant had no remaining callers in the
frontend, so they're deleted along with it.

* fix(events): drop live publish when journal write fails

application/events/publisher.py returned an envelope to live
pubsub subscribers even when the XADD to the durable journal
failed. The envelope had no ``id`` field, which bypassed the SSE
route's dedup floor and broke ``Last-Event-ID`` semantics for any
reconnecting client.

Best-effort delivery means dropping consistently, not delivering
inconsistent state. Now: if the journal write fails the publisher
returns None and skips the live publish entirely.

* fix(notifications): dedupe sseEventReceived against immediate dupes

Snapshot replay + live tail can both deliver the same id when the
live pubsub frame and the replay XRANGE overlap. The route's own
dedup floor catches the common case, but consumers walking
``recentEvents`` (FileTree, ConnectorTree, MCPServerModal,
ToolApprovalToast) would otherwise act on the same envelope
twice when a duplicate slipped through.

Belt-and-suspenders: short-circuit when the most recent id in
the ring matches the incoming one.

* fix(events): skip replay budget INCR when no snapshot work possible

_allow_replay incremented the per-user counter on every
/api/events GET, including no-op connects from a fresh client
with no cursor against an empty backlog. React StrictMode dev
double-mounts plus a few tabs trivially tripped the default
30-per-60s budget on idle reconnects.

XLEN pre-check: when last_event_id is None and the user stream
is empty, the connect can't do snapshot work — return True
without INCR. Cursor-bearing connects still INCR unconditionally
(probing the cursor's relationship to stream contents would
require a redundant XRANGE).

* fix(streaming): tighten journal contract + recover from seq collisions

Two related fixes to application/streaming/message_journal.py.

1. record_event now rejects non-dict payloads at the gate. The
   live path (base.py::_emit) wrapped non-dicts as
   {"value": payload}; the replay path in event_replay synthesized
   {"type": event_type}. A reconnecting client would receive a
   different envelope than the one originally streamed. Now both
   paths see byte-identical envelopes because non-dicts can't be
   journaled at all. The corresponding event_replay fallback is
   replaced with a warn-and-skip for any legacy rows.

2. record_event handles IntegrityError on (message_id, sequence_no)
   collisions by reading latest_sequence_no and retrying once with
   latest+1. The most likely cause is a stale seq seed on a
   continuation retry where the route read MAX(seq) from a
   separate connection before another writer committed past it.
   Previously the error was swallowed and the event silently
   dropped from the journal; now it lands at the next available
   seq. The live pubsub publish uses the materialised seq so the
   journal row and the live frame agree.

* perf(streaming): batch message_events INSERTs per stream

complete_stream previously opened a fresh db_session() per yielded
event, doing one Postgres INSERT + commit per chunk on the WSGI
thread. Streaming answers emit ~100s of answer chunks per response,
so the route was paying ~100 PG roundtrips per stream serialized on
commit latency.

New BatchedJournalWriter in application/streaming/message_journal.py
accumulates rows per stream and flushes on three triggers:
- size: buffer reaches 16 entries
- time: 100ms elapsed since the last flush
- lifecycle: close() at end-of-stream

Live pubsub publishes still fire synchronously per record(), so
subscribers see events in real time — only the durable journal write
is amortized. On bulk INSERT IntegrityError the writer falls back to
per-row record() with the existing seq+1 retry so a single colliding
seq doesn't drop the rest of the batch.

complete_stream wires journal_writer.close() into every exit path
(happy end, tool-approval-paused end, GeneratorExit, error handler)
so the terminal event is committed before the generator returns —
otherwise a reconnecting client could snapshot up to the last flush
boundary and live-tail waiting for an end that's still in memory.

Repository gets bulk_record() — one SQLAlchemy executemany INSERT
for the bulk path. All-or-nothing on collision (Postgres aborts the
whole batch); the writer's per-row fallback handles recovery.

* chore(upload): drop dead UploadTask.lastEventAt field

The lastEventAt field on UploadTask had no remaining consumers — the
matching Attachment.lastEventAt was cleaned up earlier. Remove the
field declaration and the slice write site.

* chore(frontend): drop orphaned getTaskStatus client

After the polling-removal sweep no caller in frontend/src/ references
userService.getTaskStatus or endpoints.USER.TASK_STATUS. The backend
route /api/task_status itself stays — agents, webhooks, e2e specs,
and the public docs still depend on it.

* docs(repo): remove stale planning docs from repo root

notification-channel-design.md, plan.md, and reminder-tool-design.md
were leftover Claude planning artifacts from the SSE substrate work
that landed accidentally. CLAUDE.md prohibits creating planning docs
unless asked — delete them.

* docs(message-events): clarify repo vs wrapper payload contract

MessageEventsRepository.record accepts any JSONB-compatible value; the
streaming wrapper record_event tightens this to dicts only because the
live and replay paths reconstruct non-dict payloads differently. Spell
the split out so the next reader of the repo method doesn't assume the
wrapper's contract applies here.

* refactor(events): raise on malformed stream id instead of lex fallback

stream_id_compare's lex-fallback branch was a footgun: a malformed id
that sorts lex-greater than a real one would pin live-tail dedup
forever, dropping every subsequent legitimate event silently. Both
current callers in application/api/events/routes.py pre-validate
inputs against _STREAM_ID_RE before calling, so changing the function
to raise ValueError is a no-op on the happy path and turns the future-
caller footgun into a loud failure.

* test(tasks): cover cleanup_message_events task body

Adds skipped-when-no-POSTGRES_URI and happy-path coverage for the
Celery janitor. The skipped path returns the documented short-circuit
shape without touching the repo. The happy path seeds a backdated
row, runs the task against the pg_conn fixture, and asserts the
retention window's row is deleted while in-window rows survive.
Mirrors the TestCleanupPendingToolState pattern.

* fix(notifications): treat /c/new as no current conversation

useMatch('/c/:conversationId') treats the literal URL /c/new as a
real conversation id, so the toast suppression check confused
'user is on /c/new' with 'user is on the conversation needing
approval'. Explicit guard: when the matched id is 'new', fall
through to the no-match case so approval toasts still surface.

* docs(events): enumerate publish_user_event None-return paths

The function returns Optional[str] today, with None conflating five
distinct outcomes (missing args / push disabled / unserialisable /
Redis down / XADD failed). Every current call site is fire-and-
forget and ignores the return, so the right move is to document the
five cases rather than promote to an enum return — keeps the API
small while making the diagnostic surface (logs) obvious. If a
future caller needs to react differently per reason, promote then.

* refactor(sources): move source-id derivation out of worker module

application/api/user/sources/upload.py imported _derive_source_id
from application.worker — pulling the entire Celery worker module
into the API process at import time just for a two-line helper.

Move DOCSGPT_INGEST_NAMESPACE and the derivation function to a
new application/storage/db/source_ids.py module that both layers
can import without that dependency edge. worker.py re-exports the
old names (_derive_source_id, DOCSGPT_INGEST_NAMESPACE) for
backward-compatible imports from tests and any other in-tree
callers; new code should import from the new module directly.

* fix(cache): enable Redis health_check_interval to surface half-open TCP

Without health_check_interval, a half-open TCP socket (NAT silently
dropped state, ELB idle-close) can leave pubsub.get_message hanging
past the SSE generator's keepalive cadence — the kernel never
surfaces the dead socket because no payload is in flight. Setting
health_check_interval=10 makes redis-py ping every 10s when
otherwise idle, so the next get_message after the dead window
raises and the SSE loop falls into its reconnect path instead of
silently freezing on the user.

* chore(events): rename attachment.processing.progress to attachment.progress

The event-type taxonomy was inconsistent: source ingest emits
source.ingest.progress (three segments) while attachments emitted
attachment.processing.progress (four segments). Drops the
.processing. infix for parity. Worker publish sites, the slice
reducer's match, and the worker tests all flip together.

No external consumers — the event type is purely internal between
the publisher and the in-tab slice; safe to rename in one commit.

* feat: events cleanup

* fix: better docs

* fix: e2e tests
2026-05-15 12:23:31 +01:00

1156 lines
44 KiB
Python

"""Idempotency-Key behavior on the /api/upload and /api/remote routes."""
import io
import json
from contextlib import contextmanager
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
@pytest.fixture
def app():
return Flask(__name__)
@contextmanager
def _patch_db(conn):
@contextmanager
def _yield():
yield conn
with patch(
"application.api.user.sources.upload.db_session", _yield
), patch(
"application.api.user.sources.upload.db_readonly", _yield
):
yield
def _apply_async_mock():
"""Mock for ``ingest.apply_async``; ``task.id`` mirrors the predetermined id."""
def _side_effect(*args, **kwargs):
return MagicMock(id=kwargs.get("task_id") or "auto-task-id")
m = MagicMock(side_effect=_side_effect)
return m
class TestUploadIdempotency:
def test_no_header_enqueues_normally(self, app, pg_conn):
from application.api.user.sources.upload import UploadFile
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.sources.upload.ingest.apply_async",
apply_mock,
), app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"content"), "doc.txt"),
},
content_type="multipart/form-data",
):
from flask import request
request.decoded_token = {"sub": "alice"}
response = UploadFile().post()
assert response.status_code == 200
assert apply_mock.call_count == 1
# No key → no predetermined id was passed.
assert "task_id" not in apply_mock.call_args.kwargs
def test_header_first_post_records_row(self, app, pg_conn):
from sqlalchemy import text
from application.api.user.sources.upload import UploadFile
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.sources.upload.ingest.apply_async",
apply_mock,
), app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"content"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "up-key-1"},
):
from flask import request
request.decoded_token = {"sub": "alice"}
response = UploadFile().post()
assert response.status_code == 200
assert apply_mock.call_count == 1
predetermined_id = apply_mock.call_args.kwargs["task_id"]
assert response.json["task_id"] == predetermined_id
# The dedup row is keyed on the *scoped* form ``"{user}:{key}"``
# so two users sending the same raw header don't collapse.
row = pg_conn.execute(
text(
"SELECT task_id, task_name, status FROM task_dedup "
"WHERE idempotency_key = :k"
),
{"k": "alice:up-key-1"},
).fetchone()
assert row is not None
assert row[0] == predetermined_id
assert row[1] == "ingest"
assert row[2] == "pending"
def test_header_forwards_idempotency_key_to_delay(self, app, pg_conn):
"""The Celery task body needs the key so ``with_idempotency`` can
record terminal status and ``_derive_source_id`` can pick it up.
"""
from application.api.user.sources.upload import UploadFile
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.sources.upload.ingest.apply_async",
apply_mock,
), app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"c"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "up-fwd"},
):
from flask import request
request.decoded_token = {"sub": "alice"}
UploadFile().post()
# Worker sees the scoped form so its ``with_idempotency`` row
# and ``_derive_source_id`` are also user-distinct.
assert (
apply_mock.call_args.kwargs["kwargs"]["idempotency_key"]
== "alice:up-fwd"
)
def test_same_header_second_post_returns_cached(self, app, pg_conn):
from application.api.user.sources.upload import UploadFile
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.sources.upload.ingest.apply_async",
apply_mock,
):
with app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"content"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "up-rep"},
):
from flask import request
request.decoded_token = {"sub": "alice"}
first = UploadFile().post()
with app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"content"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "up-rep"},
):
from flask import request
request.decoded_token = {"sub": "alice"}
second = UploadFile().post()
assert first.status_code == 200
assert second.status_code == 200
assert first.json == second.json
assert apply_mock.call_count == 1
def test_same_key_different_users_does_not_collide(self, app, pg_conn):
"""Cross-user collision regression: two users sending the same
raw ``Idempotency-Key`` must each get their own dedup row, both
requests enqueue, and the responses carry distinct ``task_id``s.
(Pre-fix, the second user's request was silently deduplicated
against the first user's row.)
"""
from sqlalchemy import text as sql_text
from application.api.user.sources.upload import UploadFile
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
def _fire(user):
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.sources.upload.ingest.apply_async",
apply_mock,
), app.test_request_context(
"/api/upload", method="POST",
data={
"user": user, "name": "j",
"file": (io.BytesIO(b"c"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "shared-key"},
):
from flask import request
request.decoded_token = {"sub": user}
return UploadFile().post()
first = _fire("alice")
second = _fire("bob")
assert first.status_code == 200
assert second.status_code == 200
assert first.json["task_id"] != second.json["task_id"]
assert apply_mock.call_count == 2
rows = pg_conn.execute(
sql_text(
"SELECT idempotency_key FROM task_dedup "
"WHERE idempotency_key LIKE :pat ORDER BY idempotency_key"
),
{"pat": "%:shared-key"},
).fetchall()
assert {r[0] for r in rows} == {"alice:shared-key", "bob:shared-key"}
def test_concurrent_same_key_only_one_apply_async(self, app, pg_engine):
"""Race test (M3): N parallel POSTs with same key → only ONE apply_async.
Uses ``pg_engine`` (not ``pg_conn``) so each thread can check out
its own DB connection — sharing a single Connection across
threads serializes at the driver level and defeats the race.
"""
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from application.api.user.sources.upload import UploadFile
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
@contextmanager
def _engine_session():
with pg_engine.begin() as conn:
yield conn
@contextmanager
def _engine_readonly():
with pg_engine.connect() as conn:
yield conn
def fire(idx):
# Patches sit OUTSIDE the threads (see below); only the
# per-thread Flask request context is set up inside.
with app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"x"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "up-race"},
):
from flask import request
request.decoded_token = {"sub": "alice"}
return UploadFile().post()
# ``unittest.mock.patch`` is not thread-safe — concurrent
# ``__enter__`` calls race on saving/restoring the module
# attribute and can leave threads pointing at the real
# function instead of the mock. Set up patches once, share
# across threads.
with patch(
"application.api.user.sources.upload.db_session",
_engine_session,
), patch(
"application.api.user.sources.upload.db_readonly",
_engine_readonly,
), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.sources.upload.ingest.apply_async",
apply_mock,
), ThreadPoolExecutor(max_workers=8) as ex:
responses = list(ex.map(fire, range(8)))
assert all(r.status_code == 200 for r in responses)
# Only one writer wins the claim, so only one apply_async is fired.
assert apply_mock.call_count == 1
# All 8 responses share the same task_id (winner's predetermined id).
ids = {r.json["task_id"] for r in responses}
assert len(ids) == 1
assert "deduplicated" not in ids
def test_empty_header_treated_as_absent(self, app, pg_conn):
from sqlalchemy import text
from application.api.user.sources.upload import UploadFile
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.sources.upload.ingest.apply_async",
apply_mock,
), app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"content"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": ""},
):
from flask import request
request.decoded_token = {"sub": "alice"}
response = UploadFile().post()
assert response.status_code == 200
assert apply_mock.call_count == 1
count = pg_conn.execute(
text("SELECT count(*) FROM task_dedup")
).scalar()
assert count == 0
def test_oversized_header_rejected_with_400(self, app, pg_conn):
from application.api.user.sources.upload import UploadFile
fake_storage = MagicMock()
oversized = "x" * 257
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.sources.upload.ingest.apply_async",
) as mock_apply, app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"content"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": oversized},
):
from flask import request
request.decoded_token = {"sub": "alice"}
response = UploadFile().post()
assert response.status_code == 400
assert mock_apply.call_count == 0
def test_stale_dedup_row_does_not_block_new_work(self, app, pg_conn):
"""Regression for the TTL fail-shut bug: a >24h-old dedup row
must not silently drop a new upload. Pre-fix, the second POST
returned ``task_id="deduplicated"`` and never enqueued.
"""
from sqlalchemy import text
from application.api.user.sources.upload import UploadFile
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.sources.upload.ingest.apply_async",
apply_mock,
):
with app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"content"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "stale-up"},
):
from flask import request
request.decoded_token = {"sub": "alice"}
first = UploadFile().post()
assert first.status_code == 200
first_task_id = first.json["task_id"]
assert first_task_id != "deduplicated"
# Backdate the row past TTL.
pg_conn.execute(
text(
"UPDATE task_dedup SET created_at = "
"clock_timestamp() - make_interval(hours => 25) "
"WHERE idempotency_key = :k"
),
{"k": "alice:stale-up"},
)
with app.test_request_context(
"/api/upload", method="POST",
data={
"user": "alice", "name": "j",
"file": (io.BytesIO(b"content"), "doc.txt"),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "stale-up"},
):
from flask import request
request.decoded_token = {"sub": "alice"}
second = UploadFile().post()
assert second.status_code == 200
assert second.json["task_id"] != "deduplicated"
assert second.json["task_id"] != first_task_id
assert apply_mock.call_count == 2
class TestRemoteIdempotency:
def test_no_header_enqueues_normally(self, app, pg_conn):
from application.api.user.sources.upload import UploadRemote
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.ingest_remote.apply_async",
apply_mock,
), app.test_request_context(
"/api/remote", method="POST",
data={
"user": "u", "source": "github", "name": "g",
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
},
content_type="multipart/form-data",
):
from flask import request
request.decoded_token = {"sub": "u"}
response = UploadRemote().post()
assert response.status_code == 200
assert apply_mock.call_count == 1
def test_header_first_post_records_row(self, app, pg_conn):
from sqlalchemy import text
from application.api.user.sources.upload import UploadRemote
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.ingest_remote.apply_async",
apply_mock,
), app.test_request_context(
"/api/remote", method="POST",
data={
"user": "u", "source": "github", "name": "g",
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "rem-key-1"},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = UploadRemote().post()
assert response.status_code == 200
assert apply_mock.call_count == 1
predetermined_id = apply_mock.call_args.kwargs["task_id"]
# Scoped key: ``"{user}:{key}"``.
row = pg_conn.execute(
text("SELECT task_id, task_name FROM task_dedup WHERE idempotency_key = :k"),
{"k": "u:rem-key-1"},
).fetchone()
assert row is not None
assert row[0] == predetermined_id
assert row[1] == "ingest_remote"
def test_same_header_second_post_returns_cached(self, app, pg_conn):
from application.api.user.sources.upload import UploadRemote
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.ingest_remote.apply_async",
apply_mock,
):
with app.test_request_context(
"/api/remote", method="POST",
data={
"user": "u", "source": "github", "name": "g",
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "rem-rep"},
):
from flask import request
request.decoded_token = {"sub": "u"}
first = UploadRemote().post()
with app.test_request_context(
"/api/remote", method="POST",
data={
"user": "u", "source": "github", "name": "g",
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "rem-rep"},
):
from flask import request
request.decoded_token = {"sub": "u"}
second = UploadRemote().post()
assert first.status_code == 200
assert second.status_code == 200
assert first.json == second.json
assert apply_mock.call_count == 1
def test_oversized_header_rejected_with_400(self, app, pg_conn):
from application.api.user.sources.upload import UploadRemote
oversized = "x" * 257
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.ingest_remote.apply_async",
) as mock_apply, app.test_request_context(
"/api/remote", method="POST",
data={
"user": "u", "source": "github", "name": "g",
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": oversized},
):
from flask import request
request.decoded_token = {"sub": "u"}
response = UploadRemote().post()
assert response.status_code == 400
assert mock_apply.call_count == 0
def test_no_header_returns_source_id_matching_worker_kwarg(
self, app, pg_conn,
):
"""Regression: without an ``Idempotency-Key``, the route must
still return a ``source_id`` AND pass that same id to the worker
as ``source_id`` so SSE envelopes line up with what the
frontend already has. Previously the route omitted ``source_id``
entirely on the no-key path and the worker minted its own
random uuid, breaking push correlation for the default upload
flow.
"""
from application.api.user.sources.upload import UploadRemote
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.ingest_remote.apply_async",
apply_mock,
), app.test_request_context(
"/api/remote", method="POST",
data={
"user": "u", "source": "github", "name": "g",
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
},
content_type="multipart/form-data",
):
from flask import request
request.decoded_token = {"sub": "u"}
response = UploadRemote().post()
assert response.status_code == 200
assert "source_id" in response.json
assert (
apply_mock.call_args.kwargs["kwargs"]["source_id"]
== response.json["source_id"]
)
def test_no_header_connector_returns_source_id_matching_worker_kwarg(
self, app, pg_conn,
):
"""Same regression as above for the connector branch
(``ingest_connector_task``). The connector path took the
no-key gap independently of the plain remote path."""
from application.api.user.sources.upload import UploadRemote
apply_mock = _apply_async_mock()
# Pick any registered connector — the route only branches on
# ``ConnectorCreator.get_supported_connectors()``.
from application.parser.connectors.connector_creator import (
ConnectorCreator,
)
supported = ConnectorCreator.get_supported_connectors()
if not supported:
pytest.skip("no connectors registered in this build")
connector_source = next(iter(supported))
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.ingest_connector_task.apply_async",
apply_mock,
), app.test_request_context(
"/api/remote", method="POST",
data={
"user": "u", "source": connector_source, "name": "g",
"data": json.dumps({
"session_token": "tok",
"file_ids": ["f1"],
}),
},
content_type="multipart/form-data",
):
from flask import request
request.decoded_token = {"sub": "u"}
response = UploadRemote().post()
assert response.status_code == 200
assert "source_id" in response.json
assert (
apply_mock.call_args.kwargs["kwargs"]["source_id"]
== response.json["source_id"]
)
def _seed_source(pg_conn, user="u", **kw):
from application.storage.db.repositories.sources import SourcesRepository
return SourcesRepository(pg_conn).create("manage-src", user_id=user, **kw)
class TestManageSourceFilesIdempotency:
"""Same-key dedup contract for ``ManageSourceFiles.post``: a duplicate
POST must not enqueue a second ``reingest_source_task``. The worker
decorator only deduplicates *post-completion*, so the HTTP handler is
the only place that can serialize concurrent in-flight requests.
"""
def _add_request(self, app, src_id, user, key=None):
kwargs = dict(
data={
"source_id": str(src_id),
"operation": "add",
"file": (io.BytesIO(b"content"), "new.txt"),
},
content_type="multipart/form-data",
)
if key is not None:
kwargs["headers"] = {"Idempotency-Key": key}
return app.test_request_context(
"/api/manage_source_files", method="POST", **kwargs,
)
def test_no_header_enqueues_normally_no_claim_row(self, app, pg_conn):
from sqlalchemy import text
from application.api.user.sources.upload import ManageSourceFiles
user = "alice-mgr-noh"
src = _seed_source(pg_conn, user=user, file_path="/data")
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
), self._add_request(app, src["id"], user):
from flask import request
request.decoded_token = {"sub": user}
response = ManageSourceFiles().post()
assert response.status_code == 200
assert apply_mock.call_count == 1
# No key → predetermined task_id is None, Celery generates one.
assert apply_mock.call_args.kwargs["task_id"] is None
n = pg_conn.execute(
text("SELECT count(*) FROM task_dedup")
).scalar()
assert n == 0
def test_header_records_dedup_row_with_predetermined_id(self, app, pg_conn):
from sqlalchemy import text
from application.api.user.sources.upload import ManageSourceFiles
user = "alice-mgr-rec"
src = _seed_source(pg_conn, user=user, file_path="/data")
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
), self._add_request(app, src["id"], user, key="mgr-key-1"):
from flask import request
request.decoded_token = {"sub": user}
response = ManageSourceFiles().post()
assert response.status_code == 200
assert apply_mock.call_count == 1
predetermined_id = apply_mock.call_args.kwargs["task_id"]
assert predetermined_id is not None
assert response.json["reingest_task_id"] == predetermined_id
row = pg_conn.execute(
text(
"SELECT task_id, task_name, status FROM task_dedup "
"WHERE idempotency_key = :k"
),
{"k": f"{user}:mgr-key-1"},
).fetchone()
assert row is not None
assert row[0] == predetermined_id
assert row[1] == "reingest_source_task"
assert row[2] == "pending"
def test_same_key_second_post_returns_cached(self, app, pg_conn):
from application.api.user.sources.upload import ManageSourceFiles
user = "alice-mgr-rep"
src = _seed_source(pg_conn, user=user, file_path="/data")
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
):
with self._add_request(app, src["id"], user, key="mgr-rep"):
from flask import request
request.decoded_token = {"sub": user}
first = ManageSourceFiles().post()
with self._add_request(app, src["id"], user, key="mgr-rep"):
from flask import request
request.decoded_token = {"sub": user}
second = ManageSourceFiles().post()
assert first.status_code == 200
assert second.status_code == 200
# Loser short-circuits before any storage mutation.
assert apply_mock.call_count == 1
# Loser's response carries the winner's task_id, not the
# original 200-with-added_files payload.
# ``manage_source_files`` aliases ``task_id`` ->
# ``reingest_task_id`` in the cached payload so the dedup
# response shape matches the fresh-request response (FileTree
# keys reingest correlation on ``reingest_task_id`` /
# ``source_id``).
assert second.json["reingest_task_id"] == first.json["reingest_task_id"]
# Cached ``source_id`` must equal the real source row id (not
# the helper's uuid5-of-key) so FileTree's SSE correlation on
# ``event.scope.id === result.source_id`` keeps working.
assert second.json["source_id"] == first.json["source_id"]
assert second.json["source_id"] == str(src["id"])
# Confirm the loser never invoked the file-save path.
assert fake_storage.save_file.call_count == 1
def test_remove_same_key_second_post_returns_real_source_id(
self, app, pg_conn
):
"""Regression: the ``remove`` cached branch used to leave the
helper's synthetic ``source_id`` (uuid5 of the scoped key) in
place. The reingest worker publishes SSE events tagged with the
real source row id, so the cached response had to be patched to
match what the fresh response returns — otherwise FileTree's
SSE correlation silently fails on every idempotent retry and
the user never sees the directory refresh.
"""
from application.api.user.sources.upload import ManageSourceFiles
user = "alice-mgr-rmrep"
src = _seed_source(
pg_conn,
user=user,
file_path="/data",
file_name_map={"a.txt": "a.txt"},
)
fake_storage = MagicMock()
fake_storage.file_exists.return_value = True
apply_mock = _apply_async_mock()
def _do_remove():
return app.test_request_context(
"/api/manage_source_files",
method="POST",
data={
"source_id": str(src["id"]),
"operation": "remove",
"file_paths": json.dumps(["a.txt"]),
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "mgr-rmrep"},
)
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
):
with _do_remove():
from flask import request
request.decoded_token = {"sub": user}
first = ManageSourceFiles().post()
with _do_remove():
from flask import request
request.decoded_token = {"sub": user}
second = ManageSourceFiles().post()
assert first.status_code == 200
assert second.status_code == 200
assert apply_mock.call_count == 1
assert second.json["reingest_task_id"] == first.json["reingest_task_id"]
# The contract under test: cached source_id matches the fresh
# response (the real source row id), not the helper's uuid5.
assert second.json["source_id"] == first.json["source_id"]
assert second.json["source_id"] == str(src["id"])
def test_remove_directory_same_key_second_post_returns_real_source_id(
self, app, pg_conn
):
"""Same regression as the ``remove`` test, for the
``remove_directory`` branch.
"""
from application.api.user.sources.upload import ManageSourceFiles
user = "alice-mgr-rmdir-rep"
src = _seed_source(
pg_conn,
user=user,
file_path="/data",
file_name_map={"sub/a.txt": "a.txt"},
)
fake_storage = MagicMock()
fake_storage.is_directory.return_value = True
fake_storage.remove_directory.return_value = True
apply_mock = _apply_async_mock()
def _do_remove_dir():
return app.test_request_context(
"/api/manage_source_files",
method="POST",
data={
"source_id": str(src["id"]),
"operation": "remove_directory",
"directory_path": "sub",
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "mgr-rmdir-rep"},
)
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
):
with _do_remove_dir():
from flask import request
request.decoded_token = {"sub": user}
first = ManageSourceFiles().post()
with _do_remove_dir():
from flask import request
request.decoded_token = {"sub": user}
second = ManageSourceFiles().post()
assert first.status_code == 200
assert second.status_code == 200
assert apply_mock.call_count == 1
assert second.json["reingest_task_id"] == first.json["reingest_task_id"]
assert second.json["source_id"] == first.json["source_id"]
assert second.json["source_id"] == str(src["id"])
def test_concurrent_same_key_only_one_apply_async(self, app, pg_engine):
"""N parallel same-key POSTs → exactly one apply_async."""
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from application.api.user.sources.upload import ManageSourceFiles
from application.storage.db.repositories.sources import (
SourcesRepository,
)
user = "alice-mgr-race"
with pg_engine.begin() as conn:
src = SourcesRepository(conn).create(
"race-src", user_id=user, file_path="/data",
)
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
@contextmanager
def _engine_session():
with pg_engine.begin() as conn:
yield conn
@contextmanager
def _engine_readonly():
with pg_engine.connect() as conn:
yield conn
def fire(_idx):
# Patches sit outside the thread pool (see below); only the
# per-thread Flask request context is set up inside.
with self._add_request(app, src["id"], user, key="mgr-race"):
from flask import request
request.decoded_token = {"sub": user}
return ManageSourceFiles().post()
# ``unittest.mock.patch`` is not thread-safe; set up the
# module-attribute patches once before fanning out so every
# thread sees the mock instead of racing on save/restore.
with patch(
"application.api.user.sources.upload.db_session",
_engine_session,
), patch(
"application.api.user.sources.upload.db_readonly",
_engine_readonly,
), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
), ThreadPoolExecutor(max_workers=8) as ex:
responses = list(ex.map(fire, range(8)))
assert all(r.status_code == 200 for r in responses)
assert apply_mock.call_count == 1
def test_remove_directory_failure_releases_claim(self, app, pg_conn):
"""When storage.remove_directory returns False the handler must
release the dedup row so a client retry can win the claim. Without
the release, the next retry would silently 200-cache to a task_id
that was never enqueued.
"""
from sqlalchemy import text
from application.api.user.sources.upload import ManageSourceFiles
user = "alice-mgr-rmfail"
src = _seed_source(pg_conn, user=user, file_path="/data")
fake_storage = MagicMock()
fake_storage.is_directory.return_value = True
fake_storage.remove_directory.return_value = False
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
), app.test_request_context(
"/api/manage_source_files", method="POST",
data={
"source_id": str(src["id"]),
"operation": "remove_directory",
"directory_path": "subdir",
},
content_type="multipart/form-data",
headers={"Idempotency-Key": "mgr-rmfail"},
):
from flask import request
request.decoded_token = {"sub": user}
response = ManageSourceFiles().post()
assert response.status_code == 500
assert apply_mock.call_count == 0
# Claim row was released so a retry can re-claim.
n = pg_conn.execute(
text(
"SELECT count(*) FROM task_dedup "
"WHERE idempotency_key = :k AND status = 'pending'"
),
{"k": f"{user}:mgr-rmfail"},
).scalar()
assert n == 0
def test_storage_save_failure_releases_claim(self, app, pg_conn):
"""Regression: an exception from ``storage.save_file`` after the
claim must release the dedup row. Pre-fix the outer ``except``
logged + 500'd without releasing, so a retry within 24h returned
a cached predetermined ``task_id`` for a task that never enqueued.
"""
from sqlalchemy import text
from application.api.user.sources.upload import ManageSourceFiles
user = "alice-mgr-storefail"
src = _seed_source(pg_conn, user=user, file_path="/data")
fake_storage = MagicMock()
fake_storage.save_file.side_effect = RuntimeError("disk full")
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
), self._add_request(app, src["id"], user, key="mgr-storefail"):
from flask import request
request.decoded_token = {"sub": user}
response = ManageSourceFiles().post()
assert response.status_code == 500
assert apply_mock.call_count == 0
n = pg_conn.execute(
text(
"SELECT count(*) FROM task_dedup "
"WHERE idempotency_key = :k"
),
{"k": f"{user}:mgr-storefail"},
).scalar()
assert n == 0
def test_apply_async_failure_releases_claim(self, app, pg_conn):
"""If the broker is unreachable, ``apply_async`` raises *after*
the claim was made. The claim row must be released so a retry
can re-claim with a fresh task_id.
"""
from sqlalchemy import text
from application.api.user.sources.upload import ManageSourceFiles
user = "alice-mgr-brokerdown"
src = _seed_source(pg_conn, user=user, file_path="/data")
fake_storage = MagicMock()
def _broker_down(*args, **kwargs):
raise ConnectionError("broker unreachable")
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
side_effect=_broker_down,
), self._add_request(app, src["id"], user, key="mgr-brokerdown"):
from flask import request
request.decoded_token = {"sub": user}
response = ManageSourceFiles().post()
assert response.status_code == 500
n = pg_conn.execute(
text(
"SELECT count(*) FROM task_dedup "
"WHERE idempotency_key = :k"
),
{"k": f"{user}:mgr-brokerdown"},
).scalar()
assert n == 0
def test_db_update_failure_releases_claim(self, app, pg_conn):
"""A DB blip during the ``file_name_map`` update happens after
storage mutated and the claim was made. Must still release.
"""
from sqlalchemy import text
from application.api.user.sources.upload import ManageSourceFiles
from application.storage.db.repositories import sources as src_module
user = "alice-mgr-dbfail"
src = _seed_source(pg_conn, user=user, file_path="/data")
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
original_update = src_module.SourcesRepository.update
def _explode(self, *args, **kwargs):
raise RuntimeError("db transient")
src_module.SourcesRepository.update = _explode
try:
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
), self._add_request(app, src["id"], user, key="mgr-dbfail"):
from flask import request
request.decoded_token = {"sub": user}
response = ManageSourceFiles().post()
finally:
src_module.SourcesRepository.update = original_update
assert response.status_code == 500
# apply_async never ran because the DB raise pre-empted it.
assert apply_mock.call_count == 0
n = pg_conn.execute(
text(
"SELECT count(*) FROM task_dedup "
"WHERE idempotency_key = :k"
),
{"k": f"{user}:mgr-dbfail"},
).scalar()
assert n == 0
def test_successful_path_keeps_claim_for_worker(self, app, pg_conn):
"""The claim row must persist after a successful ``apply_async``
— the worker owns the predetermined task_id and same-key retries
should resolve to the in-flight task, not re-enqueue.
"""
from sqlalchemy import text
from application.api.user.sources.upload import ManageSourceFiles
user = "alice-mgr-keep"
src = _seed_source(pg_conn, user=user, file_path="/data")
fake_storage = MagicMock()
apply_mock = _apply_async_mock()
with _patch_db(pg_conn), patch(
"application.api.user.sources.upload.StorageCreator.get_storage",
return_value=fake_storage,
), patch(
"application.api.user.tasks.reingest_source_task.apply_async",
apply_mock,
), self._add_request(app, src["id"], user, key="mgr-keep"):
from flask import request
request.decoded_token = {"sub": user}
response = ManageSourceFiles().post()
assert response.status_code == 200
# Row is still there in pending status — worker has not finalised yet.
row = pg_conn.execute(
text(
"SELECT status FROM task_dedup "
"WHERE idempotency_key = :k"
),
{"k": f"{user}:mgr-keep"},
).fetchone()
assert row is not None
assert row[0] == "pending"