tiennm99/DocsGPT
1
0 Fork 0
mirror of https://github.com/tiennm99/DocsGPT.git synced 2026-10-15 14:16:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Feat notification system (#2472)

Browse Source 
* 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
This commit is contained in:
Alex authored and GitHub committed 2026-05-15 12:23:31 +01:00
1 parent 4d6f360e3a
commit e351f45d88
135 files changed
+14394 -1253
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/api/answer/services/test_conversation_service.py
+21 -6
View File
@@ -449,6 +449,9 @@ class TestFinalizeMessage:
question="q",
decoded_token={"sub": user},
)
from application.storage.db.repositories.conversations import (
MessageUpdateOutcome,
)
assert svc.finalize_message(
res["message_id"],
"real answer",
@@ -458,7 +461,7 @@ class TestFinalizeMessage:
model_id="gpt-4",
metadata={"foo": "bar"},
status="complete",
) is True
) is MessageUpdateOutcome.UPDATED
msgs = ConversationsRepository(pg_conn).get_messages(
res["conversation_id"],
@@ -488,12 +491,15 @@ class TestFinalizeMessage:
decoded_token={"sub": user},
)
err = RuntimeError("provider down")
from application.storage.db.repositories.conversations import (
MessageUpdateOutcome,
)
assert svc.finalize_message(
res["message_id"],
"fallback text",
status="failed",
error=err,
) is True
) is MessageUpdateOutcome.UPDATED
msgs = ConversationsRepository(pg_conn).get_messages(
res["conversation_id"],
@@ -526,9 +532,12 @@ class TestFinalizeMessage:
),
{"cid": "c1", "mid": res["message_id"]},
)
from application.storage.db.repositories.conversations import (
MessageUpdateOutcome,
)
assert svc.finalize_message(
res["message_id"], "ans", status="complete",
) is True
) is MessageUpdateOutcome.UPDATED
status = pg_conn.execute(
sql_text("SELECT status FROM tool_call_attempts WHERE call_id = :cid"),
@@ -536,16 +545,19 @@ class TestFinalizeMessage:
).scalar()
assert status == "confirmed"
def test_finalize_returns_false_for_unknown_message(self, pg_conn):
def test_finalize_returns_not_found_for_unknown_message(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.storage.db.repositories.conversations import (
MessageUpdateOutcome,
)
with _patch_db(pg_conn):
assert ConversationService().finalize_message(
"00000000-0000-0000-0000-000000000000",
"x",
status="complete",
) is False
) is MessageUpdateOutcome.NOT_FOUND
def test_finalize_rolls_back_tool_call_confirm_on_message_update_failure(
self, pg_conn
@@ -651,6 +663,9 @@ class TestFinalizeMessage:
question="long question that becomes the fallback name",
decoded_token={"sub": user},
)
from application.storage.db.repositories.conversations import (
MessageUpdateOutcome,
)
assert svc.finalize_message(
res["message_id"],
"answer",
@@ -664,7 +679,7 @@ class TestFinalizeMessage:
"long question that becomes the fallback name"[:50]
),
},
) is True
) is MessageUpdateOutcome.UPDATED
repo = ConversationsRepository(pg_conn)
conv = repo.get_any(res["conversation_id"], user)
tests/api/answer/test_base_routes.py
+566 -2
View File
@@ -343,6 +343,50 @@ class TestProcessResponseStreamExtended:
result = resource.process_response_stream(iter(stream))
assert result["thought"] == "thinking..."
def test_handles_id_prefixed_chunks(self, mock_mongo_db, flask_app):
"""``complete_stream`` emits ``id: <seq>\\n`` before each
``data:`` line so reconnects can resume. The non-streaming
``/api/answer`` consumer must skip the ``id:`` header (and the
informational ``message_id`` event) without breaking JSON
decoding.
"""
from application.api.answer.routes.base import BaseAnswerResource
with flask_app.app_context():
resource = BaseAnswerResource()
stream = [
'id: 0\n'
f'data: {json.dumps({"type": "message_id", "message_id": "abc"})}\n\n',
'id: 1\n'
f'data: {json.dumps({"type": "answer", "answer": "Hello "})}\n\n',
'id: 2\n'
f'data: {json.dumps({"type": "answer", "answer": "world"})}\n\n',
'id: 3\n'
f'data: {json.dumps({"type": "id", "id": "conv-1"})}\n\n',
'id: 4\n'
f'data: {json.dumps({"type": "end"})}\n\n',
]
result = resource.process_response_stream(iter(stream))
assert result["error"] is None
assert result["answer"] == "Hello world"
assert result["conversation_id"] == "conv-1"
def test_skips_keepalive_comment_lines(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
with flask_app.app_context():
resource = BaseAnswerResource()
stream = [
': keepalive\n\n',
'id: 0\n'
f'data: {json.dumps({"type": "answer", "answer": "ok"})}\n\n',
'id: 1\n'
f'data: {json.dumps({"type": "end"})}\n\n',
]
result = resource.process_response_stream(iter(stream))
assert result["answer"] == "ok"
assert result["error"] is None
@pytest.mark.unit
class TestCheckUsageStringBooleans:
@@ -593,6 +637,251 @@ class TestCompleteStreamGeneratorExit:
next(gen)
gen.close() # Should not crash even with save error
def test_generator_exit_before_any_response_journals_error_not_end(
self, mock_mongo_db, flask_app,
):
"""A client disconnect right after the early ``message_id`` frame
leaves ``response_full`` empty, so finalize never runs. The abort
handler must journal ``error`` (not ``end``) and flip the row to
``failed`` — otherwise a reconnecting client sees a terminal
``end`` for a row whose DB status is still non-terminal and the
UI parks on a blank successful answer.
"""
from application.api.answer.routes.base import BaseAnswerResource
with flask_app.app_context():
resource = BaseAnswerResource()
mock_agent = MagicMock()
# An agent that's primed but never yields anything before the
# client disconnects — keeps ``response_full`` empty.
def gen_never_yields():
if False:
yield # pragma: no cover
return
mock_agent.gen.return_value = gen_never_yields()
mock_agent.compression_metadata = None
mock_agent.compression_saved = False
mock_agent.tool_calls = []
resource.conversation_service = MagicMock()
resource.conversation_service.save_user_question.return_value = {
"conversation_id": "conv1",
"message_id": "msg1",
"request_id": "req1",
}
journaled: list[tuple] = []
def _capture_record(message_id, sequence_no, event_type, payload):
journaled.append((message_id, sequence_no, event_type, payload))
return True
with patch(
"application.api.answer.routes.base.record_event",
side_effect=_capture_record,
):
gen = resource.complete_stream(
question="Q",
agent=mock_agent,
conversation_id="conv1",
user_api_key=None,
decoded_token={"sub": "u"},
should_save_conversation=True,
model_id="gpt-4",
)
# Drain the early ``message_id`` event, then close before
# the agent yields anything.
next(gen)
gen.close()
# The early message_id frame got journaled (seq 0); the abort
# handler must follow with an ``error`` event (NOT ``end``).
terminal_events = [
(et, pl) for (_, _, et, pl) in journaled if et in ("end", "error")
]
assert len(terminal_events) == 1, (
f"expected exactly one terminal journal write, got {terminal_events}"
)
assert terminal_events[0][0] == "error", (
f"expected ``error`` terminal but got ``end``: {terminal_events}"
)
payload = terminal_events[0][1]
assert payload.get("type") == "error"
assert payload.get("code") == "client_disconnect"
# And the DB row should have been flipped to ``failed`` via
# finalize_message. The mocked service records the call.
finalize_calls = (
resource.conversation_service.finalize_message.call_args_list
)
assert len(finalize_calls) == 1
assert finalize_calls[0].kwargs.get("status") == "failed"
def test_generator_exit_after_response_still_journals_end(
self, mock_mongo_db, flask_app,
):
"""Regression guard: a disconnect AFTER partial response was
produced and ``finalize_message`` succeeded must still journal
``end`` (the row matches ``complete``). Only the empty-response
branch flips to ``error``.
"""
from application.api.answer.routes.base import BaseAnswerResource
from application.storage.db.repositories.conversations import (
MessageUpdateOutcome,
)
with flask_app.app_context():
resource = BaseAnswerResource()
mock_agent = MagicMock()
def gen_answers():
yield {"answer": "partial"}
mock_agent.gen.return_value = gen_answers()
mock_agent.compression_metadata = None
mock_agent.compression_saved = False
mock_agent.tool_calls = []
resource.conversation_service = MagicMock()
resource.conversation_service.finalize_message.return_value = (
MessageUpdateOutcome.UPDATED
)
resource.conversation_service.save_user_question.return_value = {
"conversation_id": "conv1",
"message_id": "msg1",
"request_id": "req1",
}
journaled: list[tuple] = []
def _capture_record(message_id, sequence_no, event_type, payload):
journaled.append((message_id, sequence_no, event_type, payload))
return True
with patch(
"application.api.answer.routes.base.record_event",
side_effect=_capture_record,
):
gen = resource.complete_stream(
question="Q",
agent=mock_agent,
conversation_id="conv1",
user_api_key=None,
decoded_token={"sub": "u"},
should_save_conversation=True,
model_id="gpt-4",
)
next(gen) # message_id frame
next(gen) # answer frame (consumes ``partial``)
gen.close()
terminal_events = [
(et, pl) for (_, _, et, pl) in journaled if et in ("end", "error")
]
assert terminal_events and terminal_events[0][0] == "end", (
f"finalize succeeded but abort journaled {terminal_events}"
)
def test_generator_exit_after_normal_finalize_already_complete_journals_end(
self, mock_mongo_db, flask_app,
):
"""Regression for the race where the normal-path finalize wins
against a client disconnect.
Trace: agent finishes, ``complete_stream`` runs the normal-path
``finalize_message`` at base.py:632 and flips the row to
``complete``. The client TCP-resets before the ``end`` frame can
be journaled. The GeneratorExit handler calls ``finalize_message``
again — and the repository, gated by ``only_if_non_terminal``,
reports ``ALREADY_COMPLETE`` because the row is already at the
target state. The abort handler must journal ``end``, not
``error``: the DB says ``complete``, the reconnecting client
must see the same.
Without the fix the abort handler treated ``ALREADY_COMPLETE``
as a failure and journaled ``error`` — a reconnect would then
replay a successful completion as a failed answer.
"""
from application.api.answer.routes.base import BaseAnswerResource
from application.storage.db.repositories.conversations import (
MessageUpdateOutcome,
)
with flask_app.app_context():
resource = BaseAnswerResource()
mock_agent = MagicMock()
def gen_answers():
yield {"answer": "partial"}
mock_agent.gen.return_value = gen_answers()
mock_agent.compression_metadata = None
mock_agent.compression_saved = False
mock_agent.tool_calls = []
resource.conversation_service = MagicMock()
# Normal-path call returns UPDATED; abort-handler call sees
# the row already terminal and returns ALREADY_COMPLETE.
resource.conversation_service.finalize_message.side_effect = [
MessageUpdateOutcome.UPDATED,
MessageUpdateOutcome.ALREADY_COMPLETE,
]
resource.conversation_service.save_user_question.return_value = {
"conversation_id": "conv1",
"message_id": "msg1",
"request_id": "req1",
}
journaled: list[tuple] = []
def _capture_record(message_id, sequence_no, event_type, payload):
journaled.append((message_id, sequence_no, event_type, payload))
return True
with patch(
"application.api.answer.routes.base.record_event",
side_effect=_capture_record,
):
gen = resource.complete_stream(
question="Q",
agent=mock_agent,
conversation_id="conv1",
user_api_key=None,
decoded_token={"sub": "u"},
should_save_conversation=True,
model_id="gpt-4",
)
next(gen) # message_id frame
next(gen) # answer frame
# Pull the ``id`` frame so the normal-path finalize at
# line 632 actually runs (it sits between the answer
# frame yield and the id frame yield).
next(gen) # id frame — runs normal-path finalize first
gen.close() # GeneratorExit at the id frame yield
terminal_events = [
(et, pl) for (_, _, et, pl) in journaled if et in ("end", "error")
]
assert len(terminal_events) == 1, (
f"expected exactly one terminal journal write, got "
f"{terminal_events}"
)
assert terminal_events[0][0] == "end", (
f"row was already ``complete`` but abort journaled "
f"{terminal_events[0]} — reconnect would surface a "
f"successful answer as failed"
)
# Both finalize_message calls were made: the normal-path
# one and the abort-handler one. Asserting on side_effect
# consumption ensures the test really exercised both
# branches.
assert (
resource.conversation_service.finalize_message.call_count == 2
)
@contextmanager
def _patch_db_session(conn):
@@ -606,10 +895,34 @@ def _patch_db_session(conn):
), patch(
"application.api.answer.services.conversation_service.db_readonly",
_yield,
), patch(
# ``record_event`` opens its own short-lived ``db_session`` for
# cross-connection visibility. In tests we route it back to the
# same ``pg_conn`` so the journal write can see the message row
# the conversation_service just wrote in this transaction.
"application.streaming.message_journal.db_session",
_yield,
), patch(
# ``complete_stream`` reads ``latest_sequence_no`` via
# ``db_readonly`` to seed continuation runs. Same patch reason
# as the journal — keep the read on the same pg_conn so it sees
# uncommitted writes from this transaction.
"application.api.answer.routes.base.db_readonly",
_yield,
):
yield
def _extract_sse_data(chunk: str) -> str:
"""Pull the ``data:`` payload from an SSE record, ignoring any
``id:`` header introduced by the journal wiring.
"""
for line in chunk.split("\n"):
if line.startswith("data:"):
return line[len("data:") :].lstrip()
return ""
@pytest.mark.unit
class TestCompleteStreamWalAcceptance:
"""Acceptance for the WAL pre-persist behaviour: when the LLM raises
@@ -659,6 +972,158 @@ class TestCompleteStreamWalAcceptance:
assert "RuntimeError" in msgs[0]["metadata"]["error"]
assert "LLM upstream failed" in msgs[0]["metadata"]["error"]
def test_tool_approval_event_only_fires_when_state_saved(
self, pg_conn, flask_app,
):
"""A `tool.approval.required` notification with no resumable
``pending_tool_state`` row would deep-link the user to a 404.
Gate the publish on save_state actually committing.
"""
from application.api.answer.routes.base import BaseAnswerResource
with flask_app.app_context():
resource = BaseAnswerResource()
mock_agent = MagicMock()
# Drive the agent into the paused branch with a single
# ``tool_calls_pending`` event.
mock_agent.gen.return_value = iter(
[
{
"type": "tool_calls_pending",
"data": {"pending_tool_calls": [{"call_id": "c1"}]},
}
]
)
mock_agent._pending_continuation = {
"messages": [],
"tools_dict": {},
"pending_tool_calls": [{"call_id": "c1"}],
}
mock_agent.tool_calls = []
mock_agent.compression_metadata = None
mock_agent.compression_saved = False
published: list = []
def _capture(*args, **kwargs):
published.append(args)
with _patch_db_session(pg_conn), patch(
"application.api.answer.routes.base.publish_user_event",
side_effect=_capture,
), patch(
"application.api.answer.services.continuation_service."
"ContinuationService.save_state",
side_effect=RuntimeError("PG outage"),
):
list(
resource.complete_stream(
question="run my tool",
agent=mock_agent,
conversation_id=None,
user_api_key=None,
decoded_token={"sub": "u-tap"},
should_save_conversation=True,
model_id="gpt-4",
)
)
# No tool.approval.required publish when save_state failed.
event_types = [a[1] for a in published if len(a) >= 2]
assert "tool.approval.required" not in event_types
def test_continuation_seeds_sequence_no_from_journal_high_water_mark(
self, pg_conn, flask_app,
):
"""A resumed (tool-actions) stream must continue numbering past
the original run's max ``sequence_no``. Otherwise the second
invocation collides on the duplicate-PK and silently drops
every journal write past the resume point.
We pre-seed the journal with synthetic rows simulating an
original run, then invoke ``complete_stream`` with
``_continuation`` set and assert seq numbering picks up past
the high-water mark.
"""
import uuid as _uuid
from sqlalchemy import text as sql_text
from application.api.answer.routes.base import BaseAnswerResource
from application.storage.db.repositories.message_events import (
MessageEventsRepository,
)
# Pre-seed: insert a parent conversation + message + a few
# journal rows so latest_sequence_no returns 7.
user_id = "u-resume"
conv_id = _uuid.uuid4()
message_id = _uuid.uuid4()
pg_conn.execute(
sql_text("INSERT INTO users (user_id) VALUES (:u)"),
{"u": user_id},
)
pg_conn.execute(
sql_text(
"INSERT INTO conversations (id, user_id, name) "
"VALUES (:id, :u, 'pre-seed')"
),
{"id": conv_id, "u": user_id},
)
pg_conn.execute(
sql_text(
"INSERT INTO conversation_messages "
"(id, conversation_id, user_id, position, prompt) "
"VALUES (:id, :c, :u, 0, 'q')"
),
{"id": message_id, "c": conv_id, "u": user_id},
)
repo = MessageEventsRepository(pg_conn)
for seq in range(8): # rows 0..7
repo.record(str(message_id), seq, "answer", {"chunk": str(seq)})
original_max = repo.latest_sequence_no(str(message_id))
assert original_max == 7
with flask_app.app_context():
resource = BaseAnswerResource()
cont_agent = MagicMock()
cont_agent.gen_continuation.return_value = iter(
[{"answer": "d"}, {"answer": "e"}]
)
cont_agent.tool_calls = []
cont_agent.compression_metadata = None
cont_agent.compression_saved = False
cont_agent.tool_executor = None
with _patch_db_session(pg_conn):
list(
resource.complete_stream(
question="",
agent=cont_agent,
conversation_id=str(conv_id),
user_api_key=None,
decoded_token={"sub": user_id},
should_save_conversation=True,
model_id="gpt-4",
_continuation={
"messages": [],
"tools_dict": {},
"pending_tool_calls": [],
"tool_actions": [],
"reserved_message_id": str(message_id),
"request_id": "req-resume",
},
)
)
new_max = MessageEventsRepository(pg_conn).latest_sequence_no(
str(message_id)
)
# Continuation extended the journal — the new high-water mark
# is strictly greater than the seeded ``original_max=7``,
# confirming the allocator picked up past the resume point.
assert new_max is not None and new_max > original_max
def test_request_id_consistent_across_sse_event_and_wal_row(
self, pg_conn, flask_app,
):
@@ -694,9 +1159,9 @@ class TestCompleteStreamWalAcceptance:
)
sse_events = [
json.loads(s.replace("data: ", "").strip())
json.loads(_extract_sse_data(s))
for s in stream
if s.startswith("data: ")
if "data:" in s
]
early_events = [e for e in sse_events if e.get("type") == "message_id"]
assert len(early_events) == 1
@@ -712,3 +1177,102 @@ class TestCompleteStreamWalAcceptance:
msgs = ConversationsRepository(pg_conn).get_messages(str(convs[0][0]))
assert len(msgs) == 1
assert msgs[0]["request_id"] == sse_request_id
@pytest.mark.unit
class TestStreamingHeartbeatSeed:
"""Regression guard: when the row flips to ``streaming`` we must seed
``last_heartbeat_at`` so the watchdog doesn't fall back to
``timestamp`` (creation time) on slow LLM cold-starts (>idle_secs).
"""
def test_heartbeat_seeded_on_first_chunk(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
with flask_app.app_context():
resource = BaseAnswerResource()
mock_agent = MagicMock()
mock_agent.gen.return_value = iter(
[{"answer": "a"}, {"answer": "b"}]
)
mock_agent.compression_metadata = None
mock_agent.compression_saved = False
mock_agent.tool_calls = []
mock_agent.tool_executor = None
resource.conversation_service = MagicMock()
resource.conversation_service.save_conversation.return_value = "conv1"
resource.conversation_service.save_user_question.return_value = {
"conversation_id": "conv1",
"message_id": "msg1",
"request_id": "req1",
}
list(
resource.complete_stream(
question="Q",
agent=mock_agent,
conversation_id="conv1",
user_api_key=None,
decoded_token={"sub": "u"},
should_save_conversation=True,
model_id="gpt-4",
)
)
# update_message_status flips the row to ``streaming`` exactly
# once (idempotent via streaming_marked).
status_calls = [
c for c in resource.conversation_service
.update_message_status.call_args_list
if c.args[1] == "streaming"
]
assert len(status_calls) == 1
assert status_calls[0].args[0] == "msg1"
# heartbeat seed runs exactly once at the same flip — multiple
# chunks don't re-stamp inside this window (the throttled
# _heartbeat_streaming path is gated by STREAM_HEARTBEAT_INTERVAL).
hb_calls = (
resource.conversation_service.heartbeat_message.call_args_list
)
assert len(hb_calls) == 1
assert hb_calls[0].args[0] == "msg1"
def test_heartbeat_seed_skipped_without_reserved_message_id(
self, mock_mongo_db, flask_app,
):
"""No DB-backed message row → no heartbeat call (and no error)."""
from application.api.answer.routes.base import BaseAnswerResource
with flask_app.app_context():
resource = BaseAnswerResource()
mock_agent = MagicMock()
mock_agent.gen.return_value = iter([{"answer": "a"}])
mock_agent.compression_metadata = None
mock_agent.compression_saved = False
mock_agent.tool_calls = []
mock_agent.tool_executor = None
resource.conversation_service = MagicMock()
resource.conversation_service.save_conversation.return_value = "conv1"
# save_user_question returns no message_id → reservation absent
resource.conversation_service.save_user_question.return_value = {
"conversation_id": "conv1",
"message_id": None,
"request_id": "req1",
}
list(
resource.complete_stream(
question="Q",
agent=mock_agent,
conversation_id="conv1",
user_api_key=None,
decoded_token={"sub": "u"},
should_save_conversation=True,
model_id="gpt-4",
)
)
resource.conversation_service.heartbeat_message.assert_not_called()
tests/api/answer/test_snapshot_tail_integration.py
+233
View File
@@ -0,0 +1,233 @@
"""Integration tests for the end-to-end snapshot+tail handoff.
Exercises the publisher → journal → reconnect endpoint round-trip
without mocking the journal layer, so a regression in any of:
- complete_stream's _emit closure
- record_event's commit-per-call contract
- build_message_event_stream's snapshot-from-DB path
- the reconnect route's auth + ownership gates
- message_events repo SQL
would surface here as a failed integration assertion.
"""
from __future__ import annotations
import uuid as _uuid
from contextlib import contextmanager
from unittest.mock import patch
import pytest
from sqlalchemy import text as sql_text
@contextmanager
def _patch_journal_session(conn):
@contextmanager
def _yield():
yield conn
with patch(
"application.streaming.message_journal.db_session", _yield
), patch(
"application.streaming.event_replay.db_readonly", _yield
):
yield
def _seed_message(conn, user_id: str | None = None):
user_id = user_id or f"u-{_uuid.uuid4().hex[:6]}"
conv_id = _uuid.uuid4()
msg_id = _uuid.uuid4()
conn.execute(sql_text("INSERT INTO users (user_id) VALUES (:u)"), {"u": user_id})
conn.execute(
sql_text(
"INSERT INTO conversations (id, user_id, name) VALUES (:id, :u, 't')"
),
{"id": conv_id, "u": user_id},
)
conn.execute(
sql_text(
"INSERT INTO conversation_messages (id, conversation_id, user_id, position) "
"VALUES (:id, :c, :u, 0)"
),
{"id": msg_id, "c": conv_id, "u": user_id},
)
return user_id, str(msg_id)
@pytest.mark.integration
class TestSnapshotPlusTailRoundTrip:
def test_record_event_then_snapshot_returns_what_was_journaled(
self, pg_conn,
):
"""End-to-end of the journal half: ``record_event`` writes through
a real ``MessageEventsRepository``, ``build_message_event_stream``
reads the snapshot back via the same repo + a stub Topic.
"""
from application.streaming import event_replay
from application.streaming.message_journal import record_event
_, message_id = _seed_message(pg_conn)
with _patch_journal_session(pg_conn):
# Three events stamp the journal.
record_event(message_id, 0, "answer", {"type": "answer", "answer": "A"})
record_event(message_id, 1, "answer", {"type": "answer", "answer": "B"})
record_event(message_id, 2, "end", {"type": "end"})
# Reconnect path: subscribe yields nothing (Redis-down
# branch); the post-loop fallback runs the snapshot read
# synchronously and yields the journal contents.
def _empty_subscribe(self, on_subscribe=None, poll_timeout=1.0):
return
yield # pragma: no cover
with patch.object(
event_replay.Topic,
"subscribe",
_empty_subscribe,
create=False,
):
gen = event_replay.build_message_event_stream(
message_id,
last_event_id=None,
keepalive_seconds=0.05,
poll_timeout_seconds=0.01,
)
out = list(gen)
# Prelude + 3 snapshot frames.
assert out[0] == ": connected\n\n"
assert "id: 0" in out[1] and '"answer": "A"' in out[1]
assert "id: 1" in out[2] and '"answer": "B"' in out[2]
assert "id: 2" in out[3] and '"type": "end"' in out[3]
def test_snapshot_resumes_past_last_event_id(self, pg_conn):
from application.streaming import event_replay
from application.streaming.message_journal import record_event
_, message_id = _seed_message(pg_conn)
with _patch_journal_session(pg_conn):
for seq in range(5):
record_event(
message_id, seq, "answer", {"type": "answer", "answer": str(seq)}
)
def _empty_subscribe(self, on_subscribe=None, poll_timeout=1.0):
return
yield # pragma: no cover
with patch.object(
event_replay.Topic,
"subscribe",
_empty_subscribe,
create=False,
):
# Client says it has seen up through seq=2; expect 3 + 4.
out = list(
event_replay.build_message_event_stream(
message_id,
last_event_id=2,
keepalive_seconds=0.05,
poll_timeout_seconds=0.01,
)
)
ids_seen = [line for line in out if line.startswith("id: ")]
# Multi-line records: extract the id integers we delivered.
emitted = sorted(
int(line.split(": ", 1)[1].split("\n")[0])
for line in out
if line.startswith("id: ")
)
# Filter for non-negative (the snapshot-failure synthetic uses -1).
emitted = [e for e in emitted if e >= 0]
assert emitted == [3, 4]
assert ids_seen # sanity
def test_reconnect_route_round_trip(self, pg_conn, flask_app):
"""``/api/messages/<id>/events`` returns the journaled events
for an authenticated owner.
"""
from flask import Flask, request
from application.api.answer.routes.messages import messages_bp
from application.streaming.message_journal import record_event
# Build a fresh Flask app routing to the reconnect blueprint
# plus a tiny auth shim that injects the test user.
user_id, message_id = _seed_message(pg_conn)
app = Flask(__name__)
app.register_blueprint(messages_bp)
app.config["TESTING"] = True
@app.before_request
def _shim_auth():
request.decoded_token = {"sub": user_id}
with _patch_journal_session(pg_conn):
record_event(message_id, 0, "answer", {"type": "answer", "answer": "x"})
record_event(message_id, 1, "end", {"type": "end"})
from application.streaming import event_replay
def _empty_subscribe(self, on_subscribe=None, poll_timeout=1.0):
return
yield # pragma: no cover
with patch.object(
event_replay.Topic,
"subscribe",
_empty_subscribe,
create=False,
), patch(
"application.api.answer.routes.messages.db_readonly"
) as ro:
ro.return_value.__enter__.return_value = pg_conn
with app.test_client() as c:
r = c.get(f"/api/messages/{message_id}/events")
assert r.status_code == 200
body = b""
for chunk in r.iter_encoded():
body += chunk
if body.count(b"\n\n") >= 4:
break
r.close()
# Both journaled events present in the response.
text = body.decode("utf-8")
assert ": connected" in text
assert '"answer": "x"' in text
assert '"type": "end"' in text
# The seq lines are correct.
assert "id: 0" in text and "id: 1" in text
def test_reconnect_rejects_non_owner(self, pg_conn, flask_app):
from flask import Flask, request
from application.api.answer.routes.messages import messages_bp
user_id, message_id = _seed_message(pg_conn)
app = Flask(__name__)
app.register_blueprint(messages_bp)
@app.before_request
def _shim_auth():
request.decoded_token = {"sub": "different-user"}
# Make the ownership check use the test connection.
with patch(
"application.api.answer.routes.messages.db_readonly"
) as ro:
from contextlib import contextmanager as _cm
@_cm
def _yield():
yield pg_conn
ro.side_effect = lambda: _yield()
with app.test_client() as c:
r = c.get(f"/api/messages/{message_id}/events")
assert r.status_code == 404
tests/api/test_events_routes.py
+649
View File
@@ -0,0 +1,649 @@
"""Tests for application/api/events/routes.py — the SSE endpoint.
The SSE generator runs in a separate thread under the WSGI test client;
we drive it with mocked Redis (the ``pubsub.get_message`` and ``xrange``
sequences) and read the response body until we have enough records to
assert on, then close the response to terminate the generator.
"""
from __future__ import annotations
import json
import threading
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask, request
def _make_app():
"""Mount the events blueprint on a bare Flask app + JWT shim.
The shim mimics ``application/app.py`` populating
``request.decoded_token`` so the SSE handler's auth gate sees a
user-id without requiring the full app stack.
"""
from application.api.events.routes import events
app = Flask(__name__)
app.register_blueprint(events)
app.config["TESTING"] = True
@app.before_request
def _shim_auth(): # noqa: D401
header = request.headers.get("X-Test-Sub")
request.decoded_token = {"sub": header} if header else None
return app
class _FakePubSub:
"""Minimal Redis pub/sub stand-in for the SSE handler.
``messages`` is a list of message dicts the generator should see in
order. After exhausting it, ``get_message`` returns ``None`` (poll
timeout) so the generator stays alive emitting keepalives until the
test closes the response.
"""
def __init__(self, messages: list[dict[str, Any]]):
self._messages = list(messages)
self.subscribed: list[str] = []
self.unsubscribed: list[str] = []
self.closed = False
self._lock = threading.Lock()
def subscribe(self, name: str):
self.subscribed.append(name)
def unsubscribe(self, name: str):
self.unsubscribed.append(name)
def close(self):
self.closed = True
def get_message(self, timeout: float = 0):
with self._lock:
if self._messages:
return self._messages.pop(0)
return None
def _drain_until(response, predicate, max_chunks: int = 200) -> bytes:
"""Consume the streamed response until ``predicate(buf)`` is true.
Returns the accumulated bytes. Closes the response so the generator
exits cleanly via GeneratorExit.
"""
buf = b""
iterator = response.iter_encoded()
for _ in range(max_chunks):
try:
chunk = next(iterator)
except StopIteration:
break
if not chunk:
continue
buf += chunk
if predicate(buf):
break
response.close()
return buf
# ── auth gate ───────────────────────────────────────────────────────────
class TestAuthGate:
def test_rejects_when_no_decoded_token(self):
app = _make_app()
with app.test_client() as c:
r = c.get("/api/events")
assert r.status_code == 401
def test_rejects_when_decoded_token_missing_sub(self):
from application.api.events import routes as events_module
app = _make_app()
# Clear the shim's behavior — supply a decoded_token without sub.
@app.before_request
def _override():
request.decoded_token = {"email": "x@y.z"}
with patch.object(events_module, "get_redis_instance", return_value=None):
with app.test_client() as c:
r = c.get("/api/events")
assert r.status_code == 401
# ── streaming response shape ────────────────────────────────────────────
class TestStreamShape:
def test_returns_event_stream_mimetype_and_no_buffering_header(self):
from application.api.events import routes as events_module
app = _make_app()
with patch.object(events_module, "get_redis_instance", return_value=None):
with app.test_client() as c:
r = c.get("/api/events", headers={"X-Test-Sub": "alice"})
assert r.status_code == 200
assert r.mimetype == "text/event-stream"
assert r.headers.get("Cache-Control") == "no-store"
assert r.headers.get("X-Accel-Buffering") == "no"
# Drain enough to see the prelude comment then close.
body = _drain_until(r, lambda b: b": connected" in b)
assert b": connected" in body
def test_emits_push_disabled_when_setting_off(self):
from application.api.events import routes as events_module
app = _make_app()
with patch.object(events_module, "get_redis_instance", return_value=None), \
patch.object(events_module.settings, "ENABLE_SSE_PUSH", False):
with app.test_client() as c:
r = c.get("/api/events", headers={"X-Test-Sub": "alice"})
body = _drain_until(r, lambda b: b": push_disabled" in b)
assert b": push_disabled" in body
assert b": connected" in body # prelude still emitted
# ── concurrency cap ─────────────────────────────────────────────────────
class TestConcurrencyCap:
def test_returns_429_when_user_over_cap(self):
from application.api.events import routes as events_module
app = _make_app()
redis_client = MagicMock()
# First INCR returns 9 (over cap of 8).
redis_client.incr.return_value = 9
with patch.object(events_module, "get_redis_instance", return_value=redis_client), \
patch.object(events_module.settings, "SSE_MAX_CONCURRENT_PER_USER", 8):
with app.test_client() as c:
r = c.get("/api/events", headers={"X-Test-Sub": "alice"})
assert r.status_code == 429
# DECR fired to release the over-cap increment.
redis_client.decr.assert_called_once_with("user:alice:sse_count")
def test_skips_cap_when_zero_disabled(self):
from application.api.events import routes as events_module
app = _make_app()
redis_client = MagicMock()
with patch.object(events_module, "get_redis_instance", return_value=redis_client), \
patch.object(events_module.settings, "SSE_MAX_CONCURRENT_PER_USER", 0), \
patch.object(events_module, "Topic") as mock_topic_cls:
mock_topic = MagicMock()
mock_topic.subscribe.return_value = iter([])
mock_topic_cls.return_value = mock_topic
redis_client.xinfo_stream.side_effect = Exception("no stream")
redis_client.xrange.return_value = []
with app.test_client() as c:
r = c.get("/api/events", headers={"X-Test-Sub": "alice"})
assert r.status_code == 200
# Concurrency counter not touched when cap is 0. The
# replay-budget INCR is unrelated and may still fire.
incr_keys = [
call.args[0] for call in redis_client.incr.call_args_list
]
assert "user:alice:sse_count" not in incr_keys
_drain_until(r, lambda b: b": connected" in b)
# ── replay + live tail ──────────────────────────────────────────────────
class TestReplayAndTail:
def test_replay_yields_xrange_entries_with_injected_id(self):
from application.api.events import routes as events_module
app = _make_app()
redis_client = MagicMock()
redis_client.incr.return_value = 1
# Empty stream (no truncation).
redis_client.xinfo_stream.side_effect = Exception("nope")
# XRANGE returns one stored envelope (without ``id``); the route
# injects the entry id on the way out.
stored_event = json.dumps(
{
"type": "source.ingest.progress",
"ts": "2026-04-28T00:00:00.000Z",
"user_id": "alice",
"topic": "user:alice",
"scope": {"kind": "source", "id": "src-1"},
"payload": {"current": 25, "total": 100},
}
).encode()
redis_client.xrange.return_value = [
(b"1735682400000-0", {b"event": stored_event}),
]
# Topic.subscribe yields an immediate timeout so the generator
# keeps running long enough to flush replay; subsequent calls
# also return None.
from application.api.events.routes import _SSE_LINE_SPLIT # noqa: F401
# Fake the broadcast Topic to invoke on_subscribe immediately
# then yield None ticks until close.
def _fake_subscribe(self, on_subscribe=None, poll_timeout=1.0):
if on_subscribe is not None:
on_subscribe()
while True:
yield None
with patch.object(events_module, "get_redis_instance", return_value=redis_client), \
patch.object(
events_module.Topic, "subscribe", _fake_subscribe, create=False
):
with app.test_client() as c:
r = c.get(
"/api/events",
headers={"X-Test-Sub": "alice", "Last-Event-ID": "1735682300000-0"},
)
body = _drain_until(
r,
lambda b: b'"current": 25' in b or b'"current":25' in b,
max_chunks=80,
)
# Replay yields the entry id as the SSE id field.
assert b"id: 1735682400000-0" in body
# Envelope was rewritten to include the injected id.
assert b'"id": "1735682400000-0"' in body or b'"id":"1735682400000-0"' in body
# The connect log fires before replay.
assert b": connected" in body
def test_snapshot_flushed_when_subscribe_dies_after_callback(self):
"""Regression: if ``on_subscribe`` populated ``replay_lines`` but
``Topic.subscribe`` exits before yielding once (transient Redis
hiccup between SUBSCRIBE-ack and the first poll), the snapshot
must still reach the client. Prior to the fix the in-loop flush
was the only flush, so the backlog was silently dropped.
"""
from application.api.events import routes as events_module
app = _make_app()
redis_client = MagicMock()
redis_client.incr.return_value = 1
redis_client.xinfo_stream.side_effect = Exception("nope")
stored_event = json.dumps(
{
"type": "notification",
"payload": {"text": "from snapshot"},
}
).encode()
redis_client.xrange.return_value = [
(b"1735682400000-0", {b"event": stored_event}),
]
# Mimic the broadcast_channel race: SUBSCRIBE acks, on_subscribe
# runs, then the next get_message raises and the generator
# returns without ever yielding.
def _subscribe_dies_after_callback(
self, on_subscribe=None, poll_timeout=1.0
):
if on_subscribe is not None:
on_subscribe()
return
yield # pragma: no cover (make the function a generator)
with patch.object(events_module, "get_redis_instance", return_value=redis_client), \
patch.object(
events_module.Topic,
"subscribe",
_subscribe_dies_after_callback,
create=False,
):
with app.test_client() as c:
r = c.get(
"/api/events",
headers={
"X-Test-Sub": "alice",
"Last-Event-ID": "1735682300000-0",
},
)
body = _drain_until(
r,
lambda b: b"from snapshot" in b,
max_chunks=80,
)
# Snapshot frame must have been flushed via the post-loop
# safety net even though Topic.subscribe exited before
# the in-loop flush could fire.
assert b"id: 1735682400000-0" in body
assert b"from snapshot" in body
# XRANGE was issued exactly once (no double-flush).
redis_client.xrange.assert_called_once()
def test_invalid_last_event_id_emits_truncation_notice(self):
from application.api.events import routes as events_module
app = _make_app()
redis_client = MagicMock()
redis_client.incr.return_value = 1
redis_client.xinfo_stream.return_value = {"first-entry": [b"1-0", []]}
redis_client.xrange.return_value = []
def _fake_subscribe(self, on_subscribe=None, poll_timeout=1.0):
if on_subscribe is not None:
on_subscribe()
while True:
yield None
with patch.object(events_module, "get_redis_instance", return_value=redis_client), \
patch.object(events_module.Topic, "subscribe", _fake_subscribe, create=False):
with app.test_client() as c:
r = c.get(
"/api/events",
headers={"X-Test-Sub": "alice", "Last-Event-ID": "definitely-not-an-id"},
)
body = _drain_until(
r, lambda b: b"backlog.truncated" in b, max_chunks=80
)
assert b"backlog.truncated" in body
def test_live_tail_rejects_malformed_event_id_for_dedupe(self):
"""A pub/sub envelope carrying a non-Redis-Streams ``id`` must not
seed the dedup floor. Otherwise an adversarial or buggy publisher
could ship ``id="9999999999999-9"`` (lex-greater than any real
id) and pin every subsequent legitimate event below the floor,
silently dropping the user's notifications.
The event itself should still be delivered to the client — we
just refuse to use the bogus id for ordering, so it ships
without an SSE ``id:`` header and ``max_replayed_id`` stays put.
"""
from application.api.events import routes as events_module
app = _make_app()
redis_client = MagicMock()
redis_client.incr.return_value = 1
redis_client.xinfo_stream.side_effect = Exception("nope")
# Snapshot covers ids up to 1735682400000-0; max_replayed_id
# becomes that value after the in-loop flush.
replay_event = json.dumps({
"type": "source.ingest.progress",
"payload": {"step": "replay"},
}).encode()
redis_client.xrange.return_value = [
(b"1735682400000-0", {b"event": replay_event}),
]
live_bogus = json.dumps({
"id": "definitely-not-an-id",
"type": "source.ingest.completed",
"payload": {"step": "live-bogus"},
})
live_real = json.dumps({
"id": "1735682500000-0",
"type": "source.ingest.completed",
"payload": {"step": "live-real"},
})
def _fake_subscribe(self, on_subscribe=None, poll_timeout=1.0):
# ``Topic.subscribe`` already unpacks redis-py pubsub dicts
# and yields the raw ``data`` bytes (or ``None`` on poll
# timeout). Mirror that contract.
if on_subscribe is not None:
on_subscribe()
yield live_bogus.encode()
yield live_real.encode()
while True:
yield None
with patch.object(
events_module, "get_redis_instance", return_value=redis_client
), patch.object(
events_module.Topic, "subscribe", _fake_subscribe, create=False
):
with app.test_client() as c:
r = c.get(
"/api/events",
headers={
"X-Test-Sub": "alice",
"Last-Event-ID": "1735682300000-0",
},
)
body = _drain_until(
r, lambda b: b"live-real" in b, max_chunks=80
)
# Live-real arrived (its id is strictly greater than the
# replayed snapshot's id), with its valid id surfaced as
# the SSE ``id:`` header so the frontend can advance.
assert b"live-real" in body
assert b"id: 1735682500000-0" in body
# The bogus-id event was still delivered to the client,
# but no ``id: definitely-not-an-id`` line was emitted —
# the malformed id never reached the SSE wire and so
# could not pin the dedup floor.
assert b"live-bogus" in body
assert b"id: definitely-not-an-id" not in body
# ── format helpers (already covered in test_events_substrate but
# duplicated here as a smoke for the route's surface) ─────────────────
class TestReplayRateLimit:
"""Enumeration defenses on the per-user backlog."""
def test_allow_replay_returns_true_when_budget_disabled(self):
from application.api.events.routes import _allow_replay
with patch("application.api.events.routes.settings") as mock_settings:
mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 0
mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60
assert _allow_replay(MagicMock(), "alice", "1735682400000-0") is True
def test_allow_replay_returns_true_when_redis_unavailable(self):
from application.api.events.routes import _allow_replay
with patch("application.api.events.routes.settings") as mock_settings:
mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 5
mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60
assert _allow_replay(None, "alice", "1735682400000-0") is True
def test_allow_replay_skips_incr_when_no_cursor_and_empty_backlog(self):
"""Fresh client with no cursor and an empty user stream cannot
do snapshot work — INCR'ing the counter would needlessly
burn budget. Catches the React-StrictMode dev-burst case where
double-mounted components would otherwise 429 in 5 connects.
"""
from application.api.events.routes import _allow_replay
with patch("application.api.events.routes.settings") as mock_settings:
mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 3
mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60
redis = MagicMock()
redis.xlen.return_value = 0
# 5 connects in a row, all with no cursor — none consume
# budget because the backlog is empty.
for _ in range(5):
assert _allow_replay(redis, "alice", None) is True
redis.xlen.assert_called()
redis.incr.assert_not_called()
def test_allow_replay_incrs_when_no_cursor_but_backlog_present(self):
"""A no-cursor connect against a non-empty backlog *will* do
snapshot work, so it consumes budget normally.
"""
from application.api.events.routes import _allow_replay
with patch("application.api.events.routes.settings") as mock_settings:
mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 5
mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60
redis = MagicMock()
redis.xlen.return_value = 42
redis.incr.return_value = 1
assert _allow_replay(redis, "alice", None) is True
redis.incr.assert_called_once()
def test_allow_replay_passes_until_budget_exhausted(self):
from application.api.events.routes import _allow_replay
with patch("application.api.events.routes.settings") as mock_settings:
mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 3
mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60
redis = MagicMock()
counter = {"v": 0}
def _incr(_key):
counter["v"] += 1
return counter["v"]
redis.incr.side_effect = _incr
# Cursor set → XLEN short-circuit doesn't fire, INCR always runs.
cursor = "1735682400000-0"
# First three pass.
assert _allow_replay(redis, "alice", cursor) is True
assert _allow_replay(redis, "alice", cursor) is True
assert _allow_replay(redis, "alice", cursor) is True
# Fourth refused.
assert _allow_replay(redis, "alice", cursor) is False
# TTL re-seeded on every successful INCR so a transient
# EXPIRE failure on the seeding call can't wedge the key.
assert redis.expire.call_count == 4
for call in redis.expire.call_args_list:
assert call.args[1] == 60
def test_allow_replay_fail_open_on_redis_error(self):
from application.api.events.routes import _allow_replay
with patch("application.api.events.routes.settings") as mock_settings:
mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 5
mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60
redis = MagicMock()
redis.incr.side_effect = Exception("redis down")
assert _allow_replay(redis, "alice", "1735682400000-0") is True
def test_allow_replay_recovers_when_seeding_expire_raises(self):
"""Regression: INCR=1 then EXPIRE raising must not wedge the key.
Earlier code only called EXPIRE when ``used == 1``. If that EXPIRE
raised, the counter persisted with no TTL and every subsequent
call hit ``used > 1`` without re-seeding — locking the user out
until an operator DEL'd the key. The fix calls EXPIRE on every
successful INCR so the next call still re-seeds the TTL.
"""
from application.api.events.routes import _allow_replay
with patch("application.api.events.routes.settings") as mock_settings:
mock_settings.EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW = 5
mock_settings.EVENTS_REPLAY_BUDGET_WINDOW_SECONDS = 60
redis = MagicMock()
counter = {"v": 0}
def _incr(_key):
counter["v"] += 1
return counter["v"]
redis.incr.side_effect = _incr
# First EXPIRE raises (the seeding call that would have
# wedged the key under the old gated logic). Second EXPIRE
# succeeds — and crucially, must still run.
redis.expire.side_effect = [Exception("expire blip"), True]
cursor = "1735682400000-0"
# First call: INCR=1 succeeds, EXPIRE raises -> outer except
# returns True (fail-open) for this call.
assert _allow_replay(redis, "alice", cursor) is True
# Second call: INCR=2, EXPIRE succeeds -> still under budget,
# and the TTL is now seeded (no permanent lockout).
assert _allow_replay(redis, "alice", cursor) is True
assert redis.expire.call_count == 2
# Both EXPIRE calls used the configured window.
for call in redis.expire.call_args_list:
assert call.args[1] == 60
def test_replay_backlog_passes_count_to_xrange(self):
from application.api.events.routes import _replay_backlog
redis = MagicMock()
redis.xrange.return_value = []
# Drain the iterator so xrange is actually called.
list(_replay_backlog(redis, "alice", None, 200))
redis.xrange.assert_called_once()
kwargs = redis.xrange.call_args.kwargs
assert kwargs.get("count") == 200
def test_returns_429_when_replay_budget_exhausted(self):
"""Route refuses the connection rather than serving live tail
only. Earlier behavior silently skipped replay and let the
client advance ``lastEventId`` via id-bearing live frames,
permanently stranding the un-replayed window. The 429 keeps
the cursor pinned so the next reconnect (after the budget
window slides) can replay normally.
"""
from application.api.events import routes as events_module
app = _make_app()
redis_client = MagicMock()
def _incr(key):
if key == "user:alice:sse_count":
return 1
# Budget counter: report over-limit.
return 31
redis_client.incr.side_effect = _incr
with patch.object(
events_module, "get_redis_instance", return_value=redis_client
), patch.object(
events_module.settings,
"EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW",
30,
):
with app.test_client() as c:
r = c.get(
"/api/events",
headers={
"X-Test-Sub": "alice",
"Last-Event-ID": "1735682300000-0",
},
)
assert r.status_code == 429
# Concurrency slot is released so a budget-denied request
# doesn't permanently consume a connection from the cap.
redis_client.decr.assert_called_once_with("user:alice:sse_count")
class TestFormatHelpers:
def test_format_sse_two_terminating_newlines(self):
from application.api.events.routes import _format_sse
out = _format_sse("hello", event_id="1-0")
assert out.endswith("\n\n")
# Exactly one ``id:`` and one ``data:``.
lines = out.rstrip("\n").split("\n")
assert lines == ["id: 1-0", "data: hello"]
@pytest.mark.parametrize(
"candidate, expected",
[
("1234", "1234"),
("1234-5", "1234-5"),
(" 1234-0 ", "1234-0"),
(None, None),
("", None),
(" ", None),
("nope", None),
("1234-foo", None),
],
)
def test_normalize_last_event_id(self, candidate, expected):
from application.api.events.routes import _normalize_last_event_id
assert _normalize_last_event_id(candidate) == expected
tests/api/test_message_stream_routes.py
+220
View File
@@ -0,0 +1,220 @@
"""Tests for ``application/api/answer/routes/messages.py``.
Reconnect endpoint: GET /api/messages/<id>/events. Auth gate, ownership
gate, malformed-id rejection, Last-Event-ID normalisation, and a smoke
test that the SSE response shape matches the user-events endpoint.
"""
from __future__ import annotations
from unittest.mock import patch
from flask import Flask, request
from application.api.answer.routes.messages import (
_MESSAGE_ID_RE,
_normalise_last_event_id,
messages_bp,
)
def _make_app(decoded_token=None):
app = Flask(__name__)
app.register_blueprint(messages_bp)
app.config["TESTING"] = True
@app.before_request
def _shim_auth():
request.decoded_token = decoded_token
return app
VALID_UUID = "67d65e8f-e7fb-4df1-9e6e-99ea6c830206"
class TestNormaliseLastEventId:
def test_none_passthrough(self):
assert _normalise_last_event_id(None) is None
def test_empty_string(self):
assert _normalise_last_event_id("") is None
def test_whitespace_only(self):
assert _normalise_last_event_id(" ") is None
def test_valid_int(self):
assert _normalise_last_event_id("42") == 42
def test_stripped_whitespace(self):
assert _normalise_last_event_id(" 7 ") == 7
def test_zero_is_valid(self):
assert _normalise_last_event_id("0") == 0
def test_negative_rejected(self):
# We expose only non-negative cursors; -1 is reserved for the
# snapshot-failure synthetic terminal event and shouldn't
# round-trip back.
assert _normalise_last_event_id("-1") is None
def test_non_numeric_rejected(self):
for bad in ("foo", "1.5", "1e3", "abc-123", "null"):
assert _normalise_last_event_id(bad) is None, bad
class TestMessageIdRegex:
def test_canonical_uuid_accepted(self):
assert _MESSAGE_ID_RE.match(VALID_UUID)
def test_uppercase_uuid_accepted(self):
assert _MESSAGE_ID_RE.match(VALID_UUID.upper())
def test_no_dashes_rejected(self):
assert not _MESSAGE_ID_RE.match(VALID_UUID.replace("-", ""))
def test_legacy_mongo_id_rejected(self):
# 24-char hex with no dashes — a Mongo objectid-shaped string
# that happened to leak through somewhere.
assert not _MESSAGE_ID_RE.match("507f1f77bcf86cd799439011")
class TestAuthGate:
def test_401_when_no_decoded_token(self):
app = _make_app(decoded_token=None)
with app.test_client() as c:
r = c.get(f"/api/messages/{VALID_UUID}/events")
assert r.status_code == 401
def test_401_when_decoded_token_missing_sub(self):
app = _make_app(decoded_token={"email": "x@y"})
with app.test_client() as c:
r = c.get(f"/api/messages/{VALID_UUID}/events")
assert r.status_code == 401
class TestMessageIdValidation:
def test_400_on_malformed_id(self):
app = _make_app(decoded_token={"sub": "alice"})
with app.test_client() as c:
r = c.get("/api/messages/not-a-uuid/events")
assert r.status_code == 400
class TestOwnershipGate:
def test_404_when_user_does_not_own_message(self):
from application.api.answer.routes import messages as messages_module
app = _make_app(decoded_token={"sub": "alice"})
with patch.object(
messages_module, "_user_owns_message", return_value=False
):
with app.test_client() as c:
r = c.get(f"/api/messages/{VALID_UUID}/events")
assert r.status_code == 404
def test_200_when_user_owns_message(self):
from application.api.answer.routes import messages as messages_module
app = _make_app(decoded_token={"sub": "alice"})
# Have build_message_event_stream yield just the prelude then
# exit so the test can drain the response without blocking on
# a live pubsub subscription.
def _fake_builder(message_id, last_event_id=None, **kwargs):
yield ": connected\n\n"
with patch.object(
messages_module, "_user_owns_message", return_value=True
), patch.object(
messages_module, "build_message_event_stream", _fake_builder
):
with app.test_client() as c:
r = c.get(f"/api/messages/{VALID_UUID}/events")
assert r.status_code == 200
assert r.mimetype == "text/event-stream"
assert r.headers.get("Cache-Control") == "no-store"
assert r.headers.get("X-Accel-Buffering") == "no"
body = b""
for chunk in r.iter_encoded():
body += chunk
if b": connected" in body:
break
r.close()
assert b": connected" in body
class TestLastEventIdParsing:
def test_header_passes_through_to_builder(self):
from application.api.answer.routes import messages as messages_module
captured = {}
def _fake_builder(message_id, last_event_id=None, **kwargs):
captured["message_id"] = message_id
captured["last_event_id"] = last_event_id
yield ": connected\n\n"
app = _make_app(decoded_token={"sub": "alice"})
with patch.object(
messages_module, "_user_owns_message", return_value=True
), patch.object(
messages_module, "build_message_event_stream", _fake_builder
):
with app.test_client() as c:
r = c.get(
f"/api/messages/{VALID_UUID}/events",
headers={"Last-Event-ID": "12"},
)
# Drain a tick.
next(iter(r.iter_encoded()), None)
r.close()
assert captured["message_id"] == VALID_UUID
assert captured["last_event_id"] == 12
def test_query_param_fallback(self):
from application.api.answer.routes import messages as messages_module
captured = {}
def _fake_builder(message_id, last_event_id=None, **kwargs):
captured["last_event_id"] = last_event_id
yield ": connected\n\n"
app = _make_app(decoded_token={"sub": "alice"})
with patch.object(
messages_module, "_user_owns_message", return_value=True
), patch.object(
messages_module, "build_message_event_stream", _fake_builder
):
with app.test_client() as c:
r = c.get(
f"/api/messages/{VALID_UUID}/events?last_event_id=5"
)
next(iter(r.iter_encoded()), None)
r.close()
assert captured["last_event_id"] == 5
def test_invalid_cursor_normalised_to_none(self):
from application.api.answer.routes import messages as messages_module
captured = {}
def _fake_builder(message_id, last_event_id=None, **kwargs):
captured["last_event_id"] = last_event_id
yield ": connected\n\n"
app = _make_app(decoded_token={"sub": "alice"})
with patch.object(
messages_module, "_user_owns_message", return_value=True
), patch.object(
messages_module, "build_message_event_stream", _fake_builder
):
with app.test_client() as c:
r = c.get(
f"/api/messages/{VALID_UUID}/events",
headers={"Last-Event-ID": "definitely-not-a-number"},
)
next(iter(r.iter_encoded()), None)
r.close()
assert captured["last_event_id"] is None
tests/api/user/sources/test_upload_idempotency.py
+209 -1
View File
@@ -543,6 +543,84 @@ class TestRemoteIdempotency:
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
@@ -673,10 +751,140 @@ class TestManageSourceFilesIdempotency:
assert apply_mock.call_count == 1
# Loser's response carries the winner's task_id, not the
# original 200-with-added_files payload.
assert second.json["task_id"] == first.json["reingest_task_id"]
# ``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
tests/api/user/test_conversations.py
+54
View File
@@ -506,6 +506,60 @@ class TestGetMessageTail:
assert response.status_code == 404
def test_streaming_row_returns_partial_from_journal(self, app, pg_conn):
"""Mid-stream rows must rebuild from message_events, not return the placeholder."""
from application.api.user.conversations.routes import GetMessageTail
from application.storage.db.repositories.message_events import (
MessageEventsRepository,
)
owner = "user-tail-partial"
_, msg_id = self._seed_in_flight_message(pg_conn, owner)
events_repo = MessageEventsRepository(pg_conn)
events_repo.record(msg_id, 0, "message_id", {"type": "message_id"})
events_repo.record(msg_id, 1, "answer", {"type": "answer", "answer": "Hello"})
events_repo.record(msg_id, 2, "answer", {"type": "answer", "answer": ", world"})
events_repo.record(
msg_id, 3, "source", {"type": "source", "source": [{"id": "s1"}]}
)
with _patch_conversations_db(pg_conn), app.test_request_context(
f"/api/messages/{msg_id}/tail"
):
from flask import request
request.decoded_token = {"sub": owner}
response = GetMessageTail().get(msg_id)
assert response.status_code == 200
assert response.json["status"] == "streaming"
assert response.json["response"] == "Hello, world"
assert response.json["sources"] == [{"id": "s1"}]
assert "terminated prior to completion" not in (
response.json["response"] or ""
)
def test_streaming_row_with_empty_journal_returns_empty_response(
self, app, pg_conn
):
"""Empty journal returns empty response, not the placeholder."""
from application.api.user.conversations.routes import GetMessageTail
owner = "user-tail-empty"
_, msg_id = self._seed_in_flight_message(pg_conn, owner)
with _patch_conversations_db(pg_conn), app.test_request_context(
f"/api/messages/{msg_id}/tail"
):
from flask import request
request.decoded_token = {"sub": owner}
response = GetMessageTail().get(msg_id)
assert response.status_code == 200
assert response.json["status"] == "streaming"
assert response.json["response"] == ""
class TestUpdateConversationNameHappy:
def test_returns_401_unauthenticated(self, app):
tests/api/user/test_tasks.py
+97 -20
View File
@@ -33,7 +33,7 @@ class TestIngestTask:
mock_worker.assert_called_once_with(
ANY, "dir", ["pdf"], "job1", "/path", "file.pdf", "user1",
file_name_map=None, idempotency_key=None,
file_name_map=None, idempotency_key=None, source_id=None,
)
assert result == {"status": "ok"}
@@ -50,7 +50,7 @@ class TestIngestTask:
mock_worker.assert_called_once_with(
ANY, "dir", ["pdf"], "job1", "/path", "file.pdf", "user1",
file_name_map=name_map, idempotency_key=None,
file_name_map=name_map, idempotency_key=None, source_id=None,
)
@@ -66,7 +66,7 @@ class TestIngestRemoteTask:
mock_worker.assert_called_once_with(
ANY, {"url": "http://x"}, "job1", "user1", "web",
idempotency_key=None,
idempotency_key=None, source_id=None,
)
assert result == {"status": "ok"}
@@ -169,6 +169,7 @@ class TestIngestConnectorTask:
doc_id=None,
sync_frequency="never",
idempotency_key=None,
source_id=None,
)
assert result == {"status": "ok"}
@@ -207,6 +208,7 @@ class TestIngestConnectorTask:
doc_id="doc1",
sync_frequency="daily",
idempotency_key=None,
source_id=None,
)
assert result == {"status": "ok"}
@@ -220,7 +222,7 @@ class TestSetupPeriodicTasks:
setup_periodic_tasks(sender)
assert sender.add_periodic_task.call_count == 7
assert sender.add_periodic_task.call_count == 8
calls = sender.add_periodic_task.call_args_list
@@ -241,6 +243,9 @@ class TestSetupPeriodicTasks:
assert calls[5][1].get("name") == "reconciliation"
# version-check (every 7h)
assert calls[6][0][0] == timedelta(hours=7)
# message_events retention sweep (24h)
assert calls[7][0][0] == timedelta(hours=24)
assert calls[7][1].get("name") == "cleanup-message-events"
class TestMcpOauthTask:
@@ -257,20 +262,6 @@ class TestMcpOauthTask:
assert result == {"url": "http://auth"}
class TestMcpOauthStatusTask:
@pytest.mark.unit
@patch("application.api.user.tasks.mcp_oauth_status")
def test_calls_mcp_oauth_status(self, mock_worker):
from application.api.user.tasks import mcp_oauth_status_task
mock_worker.return_value = {"status": "authorized"}
result = mcp_oauth_status_task("task123")
mock_worker.assert_called_once_with(ANY, "task123")
assert result == {"status": "authorized"}
class TestDurableTaskRetryPolicy:
"""The long-running tasks share a uniform retry policy."""
@@ -302,7 +293,6 @@ class TestDurableTaskRetryPolicy:
"schedule_syncs",
"sync_source",
"mcp_oauth_task",
"mcp_oauth_status_task",
"cleanup_pending_tool_state",
"reconciliation_task",
"version_check_task",
@@ -438,6 +428,93 @@ class TestCleanupPendingToolState:
}
class TestCleanupMessageEventsTask:
"""Retention janitor delegates to MessageEventsRepository.cleanup_older_than."""
@pytest.mark.unit
def test_skips_when_postgres_uri_missing(self, monkeypatch):
from application.api.user.tasks import cleanup_message_events
from application.core.settings import settings
monkeypatch.setattr(settings, "POSTGRES_URI", None, raising=False)
result = cleanup_message_events.run()
assert result == {"deleted": 0, "skipped": "POSTGRES_URI not set"}
@pytest.mark.unit
def test_deletes_rows_past_retention_window(self, pg_conn, monkeypatch):
import uuid
from sqlalchemy import text as _text
from application.api.user.tasks import cleanup_message_events
from application.core.settings import settings
from application.storage.db.repositories.message_events import (
MessageEventsRepository,
)
# Seed parent rows so the FK on message_events holds.
user_id = f"user-{uuid.uuid4().hex[:8]}"
conv_id = uuid.uuid4()
msg_id = uuid.uuid4()
pg_conn.execute(
_text("INSERT INTO users (user_id) VALUES (:u)"),
{"u": user_id},
)
pg_conn.execute(
_text(
"INSERT INTO conversations (id, user_id, name) "
"VALUES (:id, :u, 'test')"
),
{"id": conv_id, "u": user_id},
)
pg_conn.execute(
_text(
"INSERT INTO conversation_messages (id, conversation_id, "
"user_id, position) VALUES (:id, :c, :u, 0)"
),
{"id": msg_id, "c": conv_id, "u": user_id},
)
repo = MessageEventsRepository(pg_conn)
repo.record(str(msg_id), 0, "answer", {"chunk": "stale"})
repo.record(str(msg_id), 1, "answer", {"chunk": "fresh"})
# Backdate seq=0 past the default 14-day retention so the
# janitor catches it; seq=1 stays at "now" and must survive.
pg_conn.execute(
_text(
"UPDATE message_events SET created_at = now() - interval '20 days' "
"WHERE message_id = CAST(:id AS uuid) AND sequence_no = 0"
),
{"id": str(msg_id)},
)
monkeypatch.setattr(
settings, "POSTGRES_URI", "postgresql://stub", raising=False
)
@contextmanager
def _fake_begin():
yield pg_conn
fake_engine = MagicMock()
fake_engine.begin = _fake_begin
with patch(
"application.storage.db.engine.get_engine",
return_value=fake_engine,
):
result = cleanup_message_events.run()
assert result == {
"deleted": 1,
"ttl_days": settings.MESSAGE_EVENTS_RETENTION_DAYS,
}
# Only the fresh row survives.
rows = repo.read_after(str(msg_id))
assert [r["sequence_no"] for r in rows] == [1]
class TestIngestIdempotency:
"""Same short-circuit applies to the ingest task path."""
@@ -449,7 +526,7 @@ class TestIngestIdempotency:
def _fake_worker(self, directory, formats, job_name, file_path,
filename, user, file_name_map=None,
idempotency_key=None):
idempotency_key=None, source_id=None):
worker_calls.append(filename)
return {"status": "ok", "directory": directory}
tests/api/user/test_tools_mcp_pg.py
-41
View File
@@ -363,47 +363,6 @@ class TestMCPServerSave:
assert response.status_code in (200, 201)
class TestMCPOAuthStatus:
def test_returns_pending_when_no_data(self, app):
from application.api.user.tools.mcp import MCPOAuthStatus
fake_redis = MagicMock()
fake_redis.get.return_value = None
with patch(
"application.api.user.tools.mcp.get_redis_instance",
return_value=fake_redis,
), app.test_request_context("/api/mcp_oauth_status/t1"):
response = MCPOAuthStatus().get("t1")
assert response.status_code == 200
assert response.json["status"] == "pending"
def test_returns_status_from_redis(self, app):
from application.api.user.tools.mcp import MCPOAuthStatus
fake_redis = MagicMock()
fake_redis.get.return_value = '{"status": "completed", "tools": [{"name": "t1", "description": "d"}]}'
with patch(
"application.api.user.tools.mcp.get_redis_instance",
return_value=fake_redis,
), app.test_request_context("/api/mcp_oauth_status/t1"):
response = MCPOAuthStatus().get("t1")
assert response.status_code == 200
assert response.json["status"] == "completed"
assert response.json["tools"][0]["name"] == "t1"
def test_redis_error_returns_500(self, app):
from application.api.user.tools.mcp import MCPOAuthStatus
with patch(
"application.api.user.tools.mcp.get_redis_instance",
side_effect=RuntimeError("boom"),
), app.test_request_context("/api/mcp_oauth_status/t1"):
response = MCPOAuthStatus().get("t1")
assert response.status_code == 500
class TestMCPOAuthCallback:
def test_error_param_redirects_error(self, app):
from application.api.user.tools.mcp import MCPOAuthCallback
tests/api/v1/test_routes_extended.py
+69
View File
@@ -513,4 +513,73 @@ class TestChatCompletionsHappyPath:
assert resp.status_code == 200
assert resp.mimetype == "text/event-stream"
def test_stream_handles_id_prefixed_chunks(self, pg_conn):
"""``complete_stream`` emits ``id: <seq>\\n`` before each
``data:`` line. The v1 streaming consumer must skip the id
header and the informational ``message_id`` event, not silently
drop every chunk.
"""
app = _build_app()
def _fake_translate(data, api_key):
return {"question": "hi"}
fake_processor = MagicMock()
fake_processor.decoded_token = {"sub": "u"}
fake_processor.conversation_id = "conv-1"
fake_processor.agent_config = {}
fake_processor.agent_id = None
fake_processor.model_id = "m"
def _fake_helper_complete_stream(**kw):
# Mirror the new wire format: id-prefixed records, plus
# the informational message_id event the v1 path doesn't
# have an analog for.
yield 'id: 0\ndata: {"type": "message_id", "message_id": "m-1"}\n\n'
yield 'id: 1\ndata: {"type": "id", "id": "conv-1"}\n\n'
yield 'id: 2\ndata: {"type": "answer", "answer": "hi"}\n\n'
translated_chunks: list = []
def _fake_translate_stream_event(event_data, completion_id, model_name):
translated_chunks.append(event_data)
return ['data: x\n\n']
fake_helper = MagicMock()
fake_helper.check_usage.return_value = None
fake_helper.complete_stream.side_effect = _fake_helper_complete_stream
with _patch_v1_db(pg_conn), patch(
"application.api.v1.routes.translate_request",
side_effect=_fake_translate,
), patch(
"application.api.v1.routes.StreamProcessor",
return_value=fake_processor,
), patch(
"application.api.v1.routes._V1AnswerHelper",
return_value=fake_helper,
), patch(
"application.api.v1.routes.translate_stream_event",
side_effect=_fake_translate_stream_event,
):
with app.test_client() as c:
resp = c.post(
"/v1/chat/completions",
headers={"Authorization": "Bearer x"},
json={
"messages": [{"role": "user", "content": "Hi"}],
"stream": True,
},
)
# Drain the response so the generator runs to completion.
list(resp.iter_encoded())
assert resp.status_code == 200
# message_id event is skipped (no v1 analog); id + answer are
# decoded and forwarded to the translator.
types_translated = [c.get("type") for c in translated_chunks]
assert "message_id" not in types_translated
assert "id" in types_translated
assert "answer" in types_translated