chore: import upstream snapshot with attribution
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,599 @@
|
||||
"""Streaming-heartbeat liveness for reasoning models that "think" before
|
||||
answering (``complete_stream`` in ``application/api/answer/routes/base.py``).
|
||||
|
||||
Background
|
||||
----------
|
||||
The reconciler (``ReconciliationRepository.find_and_lock_stuck_messages``)
|
||||
fails a ``pending``/``streaming`` row whose effective freshness —
|
||||
``GREATEST(timestamp, message_metadata.last_heartbeat_at)`` — is older than
|
||||
5 minutes. ``complete_stream`` keeps that freshness up via a streaming
|
||||
heartbeat.
|
||||
|
||||
The bug these tests pin: the heartbeat pump used to early-return until the row
|
||||
had been flipped to ``streaming``, and the row is only flipped on the first
|
||||
``answer``/``sources`` chunk. A reasoning model (e.g. ``reasoning_effort:
|
||||
high``) that streams only ``thought`` chunks for minutes before its first
|
||||
answer token therefore left the row ``pending`` with a *frozen* heartbeat, so
|
||||
the reconciler would falsely fail a live request.
|
||||
|
||||
The fix makes the heartbeat a true "agent is producing output / is alive"
|
||||
signal: it pumps whenever a ``reserved_message_id`` exists (regardless of the
|
||||
``streaming`` status flip), and is seeded once at generation start. The
|
||||
``pending → streaming`` status transition itself is unchanged (still driven by
|
||||
the first ``answer``/``sources`` chunk).
|
||||
|
||||
Residual (documented, intentionally not covered here): a model that emits *no*
|
||||
chunks at all — not even ``thought`` — for >5 min would still go stale, since
|
||||
the pump only ticks when a chunk flows. Covering a fully-silent stream would
|
||||
need a background-thread heartbeat or a higher threshold; both are out of
|
||||
scope for this surgical fix.
|
||||
|
||||
These tests drive the real ``complete_stream`` against an ephemeral Postgres
|
||||
database (``pg_engine``) and reuse the wiring/fakes style of
|
||||
``tests/api/v1/test_v1_tool_pause_finalization.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
from application.api.answer.services.conversation_service import ConversationService
|
||||
from application.storage.db.repositories.conversations import ConversationsRepository
|
||||
from application.storage.db.repositories.reconciliation import (
|
||||
ReconciliationRepository,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fakes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeLLM:
|
||||
"""Minimal LLM stand-in so ``complete_stream`` can stamp ``_request_id``."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._request_id: Optional[str] = None
|
||||
self.model_id = "gpt-5"
|
||||
|
||||
|
||||
class _ThoughtOnlyAgent:
|
||||
"""Agent that streams only ``thought`` chunks and never answers.
|
||||
|
||||
Models a reasoning model mid-"thinking" phase: the stream is alive and
|
||||
producing output, but no ``answer``/``sources`` chunk has arrived yet, so
|
||||
the row is still ``pending`` (``streaming`` not marked).
|
||||
"""
|
||||
|
||||
def __init__(self, n_thoughts: int = 4) -> None:
|
||||
self.llm = _FakeLLM()
|
||||
self._n_thoughts = n_thoughts
|
||||
|
||||
def gen(self, query: str = ""):
|
||||
for i in range(self._n_thoughts):
|
||||
yield {"thought": f"reasoning step {i}..."}
|
||||
|
||||
|
||||
class _ThoughtThenAnswerAgent:
|
||||
"""Agent that thinks (``thought`` chunks) then emits a final ``answer``.
|
||||
|
||||
Used for the regression case: the row must still flip to ``streaming`` on
|
||||
the first ``answer`` chunk and finalize ``complete``.
|
||||
"""
|
||||
|
||||
ANSWER_TEXT = "The answer is 42."
|
||||
|
||||
def __init__(self, n_thoughts: int = 3) -> None:
|
||||
self.llm = _FakeLLM()
|
||||
self._n_thoughts = n_thoughts
|
||||
|
||||
def gen(self, query: str = ""):
|
||||
for i in range(self._n_thoughts):
|
||||
yield {"thought": f"reasoning step {i}..."}
|
||||
yield {"answer": self.ANSWER_TEXT}
|
||||
|
||||
|
||||
class _NoopJournalWriter:
|
||||
"""No-op journal writer so the test exercises DB row state, not the
|
||||
message-events journal / Redis broadcast (orthogonal to heartbeats)."""
|
||||
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
pass
|
||||
|
||||
def record(self, *args, **kwargs) -> None:
|
||||
pass
|
||||
|
||||
def flush(self) -> None:
|
||||
pass
|
||||
|
||||
def close(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test harness
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _seed_user(conn, user_id: str) -> None:
|
||||
conn.execute(
|
||||
text(
|
||||
"INSERT INTO users (user_id) VALUES (:u) "
|
||||
"ON CONFLICT (user_id) DO NOTHING"
|
||||
),
|
||||
{"u": user_id},
|
||||
)
|
||||
|
||||
|
||||
def _row(conn, message_id: str) -> Dict[str, Any]:
|
||||
"""Fetch ``(status, message_metadata)`` for a reserved row."""
|
||||
r = conn.execute(
|
||||
text(
|
||||
"SELECT status, message_metadata FROM conversation_messages "
|
||||
"WHERE id = CAST(:m AS uuid)"
|
||||
),
|
||||
{"m": message_id},
|
||||
).fetchone()
|
||||
return {"status": r[0], "metadata": r[1]} if r is not None else {}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _wire_db(engine, monkeypatch):
|
||||
"""Point conversation_service / continuation_service / base at ``engine``.
|
||||
|
||||
Each helper opens its own short-lived connection (matching production),
|
||||
so we hand out fresh connections from the same ephemeral engine and swap
|
||||
the journal writer for a no-op. Mirrors the helper in
|
||||
``test_v1_tool_pause_finalization.py``.
|
||||
"""
|
||||
from application.api.answer.services import conversation_service as conv_mod
|
||||
from application.api.answer.services import continuation_service as cont_mod
|
||||
from application.api.answer.routes import base as base_mod
|
||||
|
||||
@contextmanager
|
||||
def _session():
|
||||
conn = engine.connect()
|
||||
txn = conn.begin()
|
||||
try:
|
||||
yield conn
|
||||
txn.commit()
|
||||
except Exception:
|
||||
txn.rollback()
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
@contextmanager
|
||||
def _readonly():
|
||||
conn = engine.connect()
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
monkeypatch.setattr(conv_mod, "db_session", _session)
|
||||
monkeypatch.setattr(conv_mod, "db_readonly", _readonly)
|
||||
monkeypatch.setattr(cont_mod, "db_session", _session)
|
||||
monkeypatch.setattr(cont_mod, "db_readonly", _readonly)
|
||||
monkeypatch.setattr(base_mod, "db_session", _session)
|
||||
monkeypatch.setattr(base_mod, "db_readonly", _readonly)
|
||||
monkeypatch.setattr(base_mod, "BatchedJournalWriter", _NoopJournalWriter)
|
||||
monkeypatch.setattr(base_mod, "record_event", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(base_mod, "publish_user_event", lambda *a, **kw: None)
|
||||
yield
|
||||
|
||||
|
||||
class _FakeMonotonic:
|
||||
"""Deterministic ``time.monotonic`` that jumps forward on every call.
|
||||
|
||||
``complete_stream`` reads ``time.monotonic()`` once per loop iteration
|
||||
(via ``_heartbeat_streaming``) plus at seed/mark points. Advancing by more
|
||||
than ``STREAM_HEARTBEAT_INTERVAL`` (60s) on each read guarantees the
|
||||
interval gate fires every iteration, so a thought-only stream attempts a
|
||||
heartbeat on every ``thought`` chunk.
|
||||
"""
|
||||
|
||||
def __init__(self, step: float = 120.0) -> None:
|
||||
self._t = 0.0
|
||||
self._step = step
|
||||
|
||||
def __call__(self) -> float:
|
||||
self._t += self._step
|
||||
return self._t
|
||||
|
||||
|
||||
def _make_base() -> BaseAnswerResource:
|
||||
base = BaseAnswerResource.__new__(BaseAnswerResource)
|
||||
base.default_model_id = "gpt-5"
|
||||
base.conversation_service = ConversationService()
|
||||
return base
|
||||
|
||||
|
||||
def _drain(gen) -> List[str]:
|
||||
return list(gen)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Pump-during-reasoning: a thought-only stream heartbeats while pending
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestHeartbeatPumpsDuringReasoning:
|
||||
"""A reasoning model that only streams ``thought`` chunks (no answer yet)
|
||||
must still keep its ``last_heartbeat_at`` fresh while the row is ``pending``
|
||||
— otherwise the reconciler falsely fails a live request.
|
||||
|
||||
Pre-fix the pump early-returns while ``streaming`` is unmarked, so a
|
||||
thought-only phase produces NO heartbeat and these assertions fail.
|
||||
"""
|
||||
|
||||
def test_thought_only_stream_heartbeats_while_pending(
|
||||
self, pg_engine, monkeypatch
|
||||
):
|
||||
from application.api.answer.routes import base as base_mod
|
||||
|
||||
user_id = f"user-{uuid.uuid4().hex[:8]}"
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, user_id)
|
||||
|
||||
# Count heartbeat invocations directly on the service so the test is
|
||||
# explicit that the pump fired during the thought-only phase.
|
||||
heartbeat_calls: List[str] = []
|
||||
real_heartbeat = ConversationService.heartbeat_message
|
||||
|
||||
def _counting_heartbeat(self, message_id: str) -> bool:
|
||||
heartbeat_calls.append(message_id)
|
||||
return real_heartbeat(self, message_id)
|
||||
|
||||
with _wire_db(pg_engine, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
ConversationService, "heartbeat_message", _counting_heartbeat
|
||||
)
|
||||
# Make every loop iteration cross the heartbeat interval.
|
||||
monkeypatch.setattr(base_mod.time, "monotonic", _FakeMonotonic())
|
||||
|
||||
base = _make_base()
|
||||
agent = _ThoughtOnlyAgent(n_thoughts=4)
|
||||
frames = _drain(
|
||||
base.complete_stream(
|
||||
question="hard reasoning question",
|
||||
agent=agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": user_id},
|
||||
should_persist=True,
|
||||
model_id="gpt-5",
|
||||
)
|
||||
)
|
||||
|
||||
# Sanity: the stream only ever emitted thoughts (plus framing/end),
|
||||
# never an answer — so ``streaming`` was never marked during the loop.
|
||||
joined = "\n".join(frames)
|
||||
assert '"type": "thought"' in joined
|
||||
assert '"type": "answer"' not in joined
|
||||
|
||||
# Resolve the reserved row.
|
||||
with pg_engine.connect() as conn:
|
||||
msg_id = conn.execute(
|
||||
text(
|
||||
"SELECT cm.id FROM conversation_messages cm "
|
||||
"JOIN conversations c ON c.id = cm.conversation_id "
|
||||
"WHERE c.user_id = :u ORDER BY cm.timestamp DESC LIMIT 1"
|
||||
),
|
||||
{"u": user_id},
|
||||
).scalar()
|
||||
assert msg_id is not None
|
||||
row = _row(conn, str(msg_id))
|
||||
|
||||
# Core assertion: the heartbeat fired for this row during a thought-only
|
||||
# phase. Pre-fix the pump is gated behind ``streaming_marked`` and never
|
||||
# runs, so ``heartbeat_calls`` is empty.
|
||||
assert str(msg_id) in heartbeat_calls, (
|
||||
"heartbeat_message was never called for the pending reasoning row; "
|
||||
"the heartbeat pump is still gated behind the streaming flip"
|
||||
)
|
||||
# And it is observable on the row: a non-null last_heartbeat_at.
|
||||
assert row["metadata"] is not None
|
||||
assert row["metadata"].get("last_heartbeat_at") is not None
|
||||
|
||||
def test_heartbeat_advances_last_heartbeat_at_while_row_stays_pending(
|
||||
self, pg_engine, monkeypatch
|
||||
):
|
||||
"""End-of-stream the answer arrives and the row finalizes, so to observe
|
||||
the *pending* heartbeat we snapshot the row mid-stream: after the first
|
||||
``thought`` chunk the row must be ``pending`` with a fresh
|
||||
``last_heartbeat_at`` already stamped (seed-at-start), and a subsequent
|
||||
``thought`` must bump it further — all before any ``answer``.
|
||||
"""
|
||||
from application.api.answer.routes import base as base_mod
|
||||
|
||||
user_id = f"user-{uuid.uuid4().hex[:8]}"
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, user_id)
|
||||
|
||||
observed: List[Dict[str, Any]] = []
|
||||
|
||||
class _ObservingAgent:
|
||||
"""Yields thoughts; after each, the harness inspects the DB row."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.llm = _FakeLLM()
|
||||
|
||||
def gen(self, query: str = ""):
|
||||
# message_id is surfaced before this generator runs, so by the
|
||||
# time we yield, the reserved row already exists. We can't read
|
||||
# it from inside gen() (no msg id handle), so just emit; the
|
||||
# snapshots happen in the consuming loop below via a wrapper.
|
||||
for i in range(3):
|
||||
yield {"thought": f"step {i}"}
|
||||
|
||||
with _wire_db(pg_engine, monkeypatch):
|
||||
monkeypatch.setattr(base_mod.time, "monotonic", _FakeMonotonic())
|
||||
base = _make_base()
|
||||
agent = _ObservingAgent()
|
||||
|
||||
gen = base.complete_stream(
|
||||
question="reason about it",
|
||||
agent=agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": user_id},
|
||||
should_persist=True,
|
||||
model_id="gpt-5",
|
||||
)
|
||||
|
||||
# Pull the first frame (the ``message_id`` event) to learn the row.
|
||||
first = next(gen)
|
||||
assert "message_id" in first
|
||||
msg_id = None
|
||||
with pg_engine.connect() as conn:
|
||||
msg_id = conn.execute(
|
||||
text(
|
||||
"SELECT cm.id FROM conversation_messages cm "
|
||||
"JOIN conversations c ON c.id = cm.conversation_id "
|
||||
"WHERE c.user_id = :u ORDER BY cm.timestamp DESC LIMIT 1"
|
||||
),
|
||||
{"u": user_id},
|
||||
).scalar()
|
||||
assert msg_id is not None
|
||||
|
||||
# Drain the rest, snapshotting the row's heartbeat after each frame.
|
||||
for _ in gen:
|
||||
with pg_engine.connect() as conn:
|
||||
observed.append(_row(conn, str(msg_id)))
|
||||
|
||||
# While the stream was thought-only the row stayed ``pending`` and the
|
||||
# heartbeat advanced (was non-null on each snapshot).
|
||||
pending_snaps = [s for s in observed if s.get("status") == "pending"]
|
||||
assert pending_snaps, "expected at least one pending snapshot mid-stream"
|
||||
for snap in pending_snaps:
|
||||
assert snap["metadata"] is not None
|
||||
assert snap["metadata"].get("last_heartbeat_at") is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Seed-at-start: the reserved row has a heartbeat from generation start
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestHeartbeatSeededAtGenerationStart:
|
||||
"""Right after ``complete_stream`` reserves the row and begins consuming the
|
||||
generator (before any ``answer``), the row must already carry a non-null
|
||||
``last_heartbeat_at`` — so even a model that takes a while to emit its first
|
||||
chunk is covered from t=0, not only from the first interval tick.
|
||||
"""
|
||||
|
||||
def test_reserved_row_has_heartbeat_before_first_answer(
|
||||
self, pg_engine, monkeypatch
|
||||
):
|
||||
# This case keeps real ``time.monotonic`` so it asserts the *seed*
|
||||
# (which runs before the loop, with no time advance), not an
|
||||
# interval-driven pump — so no ``base.time`` monkeypatch is needed.
|
||||
user_id = f"user-{uuid.uuid4().hex[:8]}"
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, user_id)
|
||||
|
||||
captured: Dict[str, Any] = {}
|
||||
|
||||
class _SlowFirstChunkAgent:
|
||||
"""Yields a single thought then stops — emulates a model that has
|
||||
produced no answer yet by the time we inspect the row."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.llm = _FakeLLM()
|
||||
|
||||
def gen(self, query: str = ""):
|
||||
# Inspect the reserved row at the very first generator step,
|
||||
# before any answer/sources chunk has been processed.
|
||||
with pg_engine.connect() as conn:
|
||||
mid = conn.execute(
|
||||
text(
|
||||
"SELECT cm.id FROM conversation_messages cm "
|
||||
"JOIN conversations c ON c.id = cm.conversation_id "
|
||||
"WHERE c.user_id = :u "
|
||||
"ORDER BY cm.timestamp DESC LIMIT 1"
|
||||
),
|
||||
{"u": user_id},
|
||||
).scalar()
|
||||
captured["row"] = _row(conn, str(mid)) if mid else {}
|
||||
yield {"thought": "starting to reason"}
|
||||
|
||||
with _wire_db(pg_engine, monkeypatch):
|
||||
# Keep real monotonic time so this asserts the *seed*, not an
|
||||
# interval-driven pump (no time advance happens before gen runs).
|
||||
base = _make_base()
|
||||
agent = _SlowFirstChunkAgent()
|
||||
_drain(
|
||||
base.complete_stream(
|
||||
question="seed check",
|
||||
agent=agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": user_id},
|
||||
should_persist=True,
|
||||
model_id="gpt-5",
|
||||
)
|
||||
)
|
||||
|
||||
row = captured.get("row") or {}
|
||||
# The seed runs before the loop, so the row is still ``pending`` and
|
||||
# already has a heartbeat stamped.
|
||||
assert row.get("status") == "pending"
|
||||
assert row.get("metadata") is not None
|
||||
assert row["metadata"].get("last_heartbeat_at") is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. End-to-end reconciler guarantee: a fresh-heartbeat pending row is safe
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestReconcilerProtectsFreshlyHeartbeatedPendingRow:
|
||||
"""The whole point of the fix: a ``pending`` row older than 5 min by
|
||||
``timestamp`` but with a fresh ``last_heartbeat_at`` (the reasoning pump)
|
||||
must NOT be swept by ``find_and_lock_stuck_messages``; a stale-heartbeat
|
||||
``pending`` row still IS swept.
|
||||
|
||||
This mirrors the ``streaming`` cases in
|
||||
``tests/storage/db/repositories/test_reconciliation.py`` but for the
|
||||
``pending`` (still-reasoning) status the fix is about.
|
||||
"""
|
||||
|
||||
def _seed_pending_with_heartbeat(
|
||||
self, conn, *, hb_minutes_ago: float, ts_minutes_ago: int = 20,
|
||||
) -> str:
|
||||
conv = ConversationsRepository(conn).create("u-recon", "reasoning recon")
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO conversation_messages (
|
||||
conversation_id, position, prompt, response, status,
|
||||
user_id, timestamp, message_metadata
|
||||
)
|
||||
VALUES (
|
||||
CAST(:cid AS uuid), 0, 'p', '', 'pending', 'u-recon',
|
||||
clock_timestamp() - make_interval(mins => :ts),
|
||||
jsonb_build_object(
|
||||
'last_heartbeat_at',
|
||||
to_jsonb(
|
||||
clock_timestamp()
|
||||
- make_interval(secs => :hb_secs)
|
||||
)
|
||||
)
|
||||
)
|
||||
RETURNING id
|
||||
"""
|
||||
),
|
||||
{
|
||||
"cid": conv["id"],
|
||||
"ts": ts_minutes_ago,
|
||||
"hb_secs": int(hb_minutes_ago * 60),
|
||||
},
|
||||
).fetchone()
|
||||
return str(row[0])
|
||||
|
||||
def test_fresh_heartbeat_pending_row_is_not_swept(self, pg_conn):
|
||||
# timestamp 20 min old (well past the 5-min age gate) but the reasoning
|
||||
# pump heartbeat is only 30s old → the row is live, must be skipped.
|
||||
msg_id = self._seed_pending_with_heartbeat(
|
||||
pg_conn, hb_minutes_ago=0.5, ts_minutes_ago=20
|
||||
)
|
||||
rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages()
|
||||
assert all(str(r["id"]) != msg_id for r in rows), (
|
||||
"a still-reasoning pending row with a fresh heartbeat was "
|
||||
"incorrectly selected for reconciliation"
|
||||
)
|
||||
|
||||
def test_stale_heartbeat_pending_row_is_swept(self, pg_conn):
|
||||
# Both timestamp and heartbeat are older than 5 min → genuinely stuck.
|
||||
msg_id = self._seed_pending_with_heartbeat(
|
||||
pg_conn, hb_minutes_ago=10, ts_minutes_ago=20
|
||||
)
|
||||
rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages()
|
||||
assert any(str(r["id"]) == msg_id for r in rows), (
|
||||
"a pending row stale by both timestamp and heartbeat should be "
|
||||
"selected for reconciliation"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Regression: a normal answer turn still marks streaming + heartbeats
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestNormalAnswerTurnUnchanged:
|
||||
"""A normal turn that streams an ``answer`` must still flip the row to
|
||||
``streaming`` on the first answer chunk (status semantics unchanged) and
|
||||
finalize ``complete``. The heartbeat liveness change must not regress this.
|
||||
"""
|
||||
|
||||
def test_answer_turn_marks_streaming_and_finalizes_complete(
|
||||
self, pg_engine, monkeypatch
|
||||
):
|
||||
from application.api.answer.routes import base as base_mod
|
||||
|
||||
user_id = f"user-{uuid.uuid4().hex[:8]}"
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, user_id)
|
||||
|
||||
# Observe the status transition: snapshot after the first answer frame.
|
||||
statuses_seen: List[str] = []
|
||||
real_update = ConversationService.update_message_status
|
||||
|
||||
def _recording_update(self, message_id: str, status: str) -> bool:
|
||||
statuses_seen.append(status)
|
||||
return real_update(self, message_id, status)
|
||||
|
||||
with _wire_db(pg_engine, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
ConversationService, "update_message_status", _recording_update
|
||||
)
|
||||
monkeypatch.setattr(base_mod.time, "monotonic", _FakeMonotonic())
|
||||
base = _make_base()
|
||||
agent = _ThoughtThenAnswerAgent(n_thoughts=2)
|
||||
frames = _drain(
|
||||
base.complete_stream(
|
||||
question="normal question",
|
||||
agent=agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": user_id},
|
||||
should_persist=True,
|
||||
model_id="gpt-5",
|
||||
)
|
||||
)
|
||||
|
||||
joined = "\n".join(frames)
|
||||
assert '"type": "answer"' in joined
|
||||
assert '"type": "end"' in joined
|
||||
|
||||
# The first answer chunk marked the row ``streaming`` exactly once.
|
||||
assert statuses_seen == ["streaming"]
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
text(
|
||||
"SELECT status, response, message_metadata "
|
||||
"FROM conversation_messages cm "
|
||||
"JOIN conversations c ON c.id = cm.conversation_id "
|
||||
"WHERE c.user_id = :u ORDER BY cm.timestamp DESC LIMIT 1"
|
||||
),
|
||||
{"u": user_id},
|
||||
).fetchone()
|
||||
|
||||
# Final state: terminal ``complete`` with the answer, and a heartbeat
|
||||
# was stamped along the way (liveness intact).
|
||||
assert row[0] == "complete"
|
||||
assert row[1] == _ThoughtThenAnswerAgent.ANSWER_TEXT
|
||||
assert row[2] is not None
|
||||
assert row[2].get("last_heartbeat_at") is not None
|
||||
@@ -0,0 +1,587 @@
|
||||
"""Extended unit tests for application/api/v1/routes.py.
|
||||
|
||||
Covers:
|
||||
- _extract_bearer_token helper
|
||||
- _lookup_agent helper (happy path + exception)
|
||||
- _get_model_name helper
|
||||
- /v1/chat/completions: auth error, missing messages, translate error,
|
||||
non-stream success, stream success, ValueError, generic Exception,
|
||||
tool_actions continuation path (missing conversation_id),
|
||||
usage_error path
|
||||
- /v1/models: missing auth, mongo exception path, with createdAt timestamp
|
||||
"""
|
||||
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
from application.api.v1.routes import (
|
||||
_extract_bearer_token,
|
||||
_get_model_name,
|
||||
v1_bp,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeCollection:
|
||||
pass
|
||||
|
||||
def __init__(self, docs):
|
||||
self.docs = list(docs)
|
||||
|
||||
def find_one(self, query):
|
||||
for doc in self.docs:
|
||||
if all(doc.get(k) == v for k, v in query.items()):
|
||||
return doc
|
||||
return None
|
||||
|
||||
def find(self, query):
|
||||
return [d for d in self.docs if all(d.get(k) == v for k, v in query.items())]
|
||||
|
||||
|
||||
def _build_app():
|
||||
app = Flask(__name__)
|
||||
app.register_blueprint(v1_bp)
|
||||
return app
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _extract_bearer_token
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestExtractBearerToken:
|
||||
pass
|
||||
|
||||
def test_returns_token_from_bearer_header(self):
|
||||
app = _build_app()
|
||||
with app.test_request_context(headers={"Authorization": "Bearer my-api-key"}):
|
||||
assert _extract_bearer_token() == "my-api-key"
|
||||
|
||||
def test_returns_none_when_no_authorization_header(self):
|
||||
app = _build_app()
|
||||
with app.test_request_context():
|
||||
assert _extract_bearer_token() is None
|
||||
|
||||
def test_returns_none_when_not_bearer_scheme(self):
|
||||
app = _build_app()
|
||||
with app.test_request_context(headers={"Authorization": "Token my-api-key"}):
|
||||
assert _extract_bearer_token() is None
|
||||
|
||||
def test_strips_whitespace(self):
|
||||
app = _build_app()
|
||||
with app.test_request_context(headers={"Authorization": "Bearer spaced-key "}):
|
||||
assert _extract_bearer_token() == "spaced-key"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _lookup_agent
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestLookupAgent:
|
||||
pass
|
||||
|
||||
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _get_model_name
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGetModelName:
|
||||
pass
|
||||
|
||||
def test_returns_agent_name_when_agent_has_name(self):
|
||||
assert _get_model_name({"name": "My Agent"}, "api-key") == "My Agent"
|
||||
|
||||
def test_falls_back_to_api_key_when_agent_has_no_name(self):
|
||||
assert _get_model_name({"user": "u"}, "api-key") == "api-key"
|
||||
|
||||
def test_returns_api_key_when_no_agent(self):
|
||||
assert _get_model_name(None, "api-key") == "api-key"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# /v1/chat/completions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestChatCompletions:
|
||||
pass
|
||||
|
||||
def _make_mongo(self, key="key-1", user="user-1"):
|
||||
agent = {"_id": "agent-1", "key": key, "user": user}
|
||||
fake_col = _FakeCollection([agent])
|
||||
return {"testdb": {"agents": fake_col}}
|
||||
|
||||
def _patch_mongo(self, monkeypatch, mongo=None):
|
||||
if mongo is None:
|
||||
mongo = self._make_mongo()
|
||||
monkeypatch.setattr("application.api.v1.routes.MongoDB.get_client", lambda: mongo)
|
||||
monkeypatch.setattr("application.api.v1.routes.settings.MONGO_DB_NAME", "testdb")
|
||||
|
||||
def test_missing_auth_returns_401(self):
|
||||
app = _build_app()
|
||||
with app.test_client() as client:
|
||||
resp = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"messages": [{"role": "user", "content": "Hi"}]},
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
assert resp.get_json()["error"]["type"] == "auth_error"
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# /v1/models — additional edge cases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestListModelsExtra:
|
||||
pass
|
||||
|
||||
def test_missing_auth_header_returns_401(self):
|
||||
app = _build_app()
|
||||
with app.test_client() as client:
|
||||
resp = client.get("/v1/models")
|
||||
assert resp.status_code == 401
|
||||
assert resp.get_json()["error"]["type"] == "auth_error"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests using the ephemeral pg_conn fixture against real PG
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_v1_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch("application.api.v1.routes.db_readonly", _yield):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestLookupAgentHappy:
|
||||
def test_returns_agent_for_valid_key(self, pg_conn):
|
||||
from application.api.v1.routes import _lookup_agent
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
|
||||
AgentsRepository(pg_conn).create("u1", "Test", "published", key="k-ok")
|
||||
|
||||
with _patch_v1_db(pg_conn):
|
||||
got = _lookup_agent("k-ok")
|
||||
assert got is not None
|
||||
assert got["key"] == "k-ok"
|
||||
|
||||
def test_returns_none_when_not_found(self, pg_conn):
|
||||
from application.api.v1.routes import _lookup_agent
|
||||
|
||||
with _patch_v1_db(pg_conn):
|
||||
got = _lookup_agent("nope")
|
||||
assert got is None
|
||||
|
||||
def test_returns_none_on_exception(self):
|
||||
from application.api.v1.routes import _lookup_agent
|
||||
|
||||
@contextmanager
|
||||
def _broken():
|
||||
raise RuntimeError("db down")
|
||||
yield
|
||||
|
||||
with patch("application.api.v1.routes.db_readonly", _broken):
|
||||
got = _lookup_agent("k")
|
||||
assert got is None
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestListModelsPgConn:
|
||||
def test_invalid_key_returns_401(self, pg_conn):
|
||||
app = _build_app()
|
||||
with _patch_v1_db(pg_conn):
|
||||
with app.test_client() as c:
|
||||
resp = c.get(
|
||||
"/v1/models", headers={"Authorization": "Bearer bad"}
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_returns_agent_for_valid_key(self, pg_conn):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
|
||||
app = _build_app()
|
||||
repo = AgentsRepository(pg_conn)
|
||||
repo.create("u-m", "A1", "published", key="models-key")
|
||||
repo.create("u-m", "A2", "published", key="models-key-2")
|
||||
|
||||
with _patch_v1_db(pg_conn):
|
||||
with app.test_client() as c:
|
||||
resp = c.get(
|
||||
"/v1/models",
|
||||
headers={"Authorization": "Bearer models-key"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.get_json()
|
||||
assert data["object"] == "list"
|
||||
names = {m["name"] for m in data["data"]}
|
||||
assert names == {"A1"}
|
||||
|
||||
def test_db_error_returns_500(self):
|
||||
app = _build_app()
|
||||
|
||||
@contextmanager
|
||||
def _broken():
|
||||
raise RuntimeError("boom")
|
||||
yield
|
||||
|
||||
with patch("application.api.v1.routes.db_readonly", _broken):
|
||||
with app.test_client() as c:
|
||||
resp = c.get(
|
||||
"/v1/models",
|
||||
headers={"Authorization": "Bearer any"},
|
||||
)
|
||||
assert resp.status_code == 500
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestChatCompletionsHappyPath:
|
||||
"""Tests that reach the request-translate / processor code paths."""
|
||||
|
||||
def test_missing_messages_returns_400(self, pg_conn):
|
||||
app = _build_app()
|
||||
with _patch_v1_db(pg_conn):
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": "Bearer x"},
|
||||
json={},
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
assert resp.get_json()["error"]["type"] == "invalid_request"
|
||||
|
||||
def test_translate_error_returns_400(self, pg_conn):
|
||||
app = _build_app()
|
||||
with _patch_v1_db(pg_conn), patch(
|
||||
"application.api.v1.routes.translate_request",
|
||||
side_effect=ValueError("bad"),
|
||||
):
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": "Bearer x"},
|
||||
json={"messages": [{"role": "user", "content": "Hi"}]},
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_tool_actions_missing_conversation_id_returns_400(self, pg_conn):
|
||||
app = _build_app()
|
||||
|
||||
def _fake_translate(data, api_key):
|
||||
return {
|
||||
"question": "",
|
||||
"tool_actions": [{"id": "t1", "result": "r"}],
|
||||
# missing conversation_id
|
||||
}
|
||||
|
||||
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=MagicMock(decoded_token={"sub": "u"}),
|
||||
):
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": "Bearer x"},
|
||||
json={"messages": [{"role": "user", "content": "Hi"}]},
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_processor_value_error_returns_400(self, pg_conn):
|
||||
app = _build_app()
|
||||
|
||||
def _fake_translate(data, api_key):
|
||||
return {"question": "q"}
|
||||
|
||||
fake_processor = MagicMock()
|
||||
fake_processor.build_agent.side_effect = ValueError("boom")
|
||||
|
||||
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,
|
||||
):
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": "Bearer x"},
|
||||
json={"messages": [{"role": "user", "content": "Hi"}]},
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_processor_generic_exception_returns_500(self, pg_conn):
|
||||
app = _build_app()
|
||||
|
||||
def _fake_translate(data, api_key):
|
||||
return {"question": "q"}
|
||||
|
||||
fake_processor = MagicMock()
|
||||
fake_processor.build_agent.side_effect = RuntimeError("boom")
|
||||
|
||||
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,
|
||||
):
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": "Bearer x"},
|
||||
json={"messages": [{"role": "user", "content": "Hi"}]},
|
||||
)
|
||||
assert resp.status_code == 500
|
||||
|
||||
def test_non_stream_success_path(self, pg_conn):
|
||||
app = _build_app()
|
||||
|
||||
def _fake_translate(data, api_key):
|
||||
return {"question": "hi", "save_conversation": False}
|
||||
|
||||
fake_processor = MagicMock()
|
||||
fake_processor.decoded_token = {"sub": "u"}
|
||||
fake_processor.conversation_id = None
|
||||
fake_processor.agent_config = {"user_api_key": "k"}
|
||||
fake_processor.agent_id = None
|
||||
fake_processor.model_id = "m"
|
||||
fake_processor.build_agent.return_value = MagicMock()
|
||||
|
||||
def _fake_stream(*a, **kw):
|
||||
yield 'data: {"type": "end"}'
|
||||
|
||||
fake_helper = MagicMock()
|
||||
fake_helper.check_usage.return_value = None
|
||||
fake_helper.complete_stream.side_effect = _fake_stream
|
||||
fake_helper.process_response_stream.return_value = {
|
||||
"error": None,
|
||||
"conversation_id": "conv-1",
|
||||
"answer": "hello",
|
||||
"sources": [],
|
||||
"tool_calls": [],
|
||||
"thought": "",
|
||||
"extra": {},
|
||||
}
|
||||
|
||||
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_response",
|
||||
return_value={"id": "x", "choices": []},
|
||||
):
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": "Bearer x"},
|
||||
json={"messages": [{"role": "user", "content": "Hi"}]},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_non_stream_error_in_result_returns_500(self, pg_conn):
|
||||
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 = None
|
||||
fake_processor.agent_config = {}
|
||||
fake_processor.agent_id = None
|
||||
fake_processor.model_id = "m"
|
||||
fake_helper = MagicMock()
|
||||
fake_helper.check_usage.return_value = None
|
||||
fake_helper.complete_stream.return_value = iter([])
|
||||
fake_helper.process_response_stream.return_value = {
|
||||
"error": "something went wrong",
|
||||
"conversation_id": None,
|
||||
"answer": None,
|
||||
"sources": None,
|
||||
"tool_calls": None,
|
||||
"thought": None,
|
||||
"extra": None,
|
||||
}
|
||||
|
||||
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,
|
||||
):
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": "Bearer x"},
|
||||
json={"messages": [{"role": "user", "content": "Hi"}]},
|
||||
)
|
||||
assert resp.status_code == 500
|
||||
|
||||
def test_stream_success_returns_sse(self, pg_conn):
|
||||
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):
|
||||
yield 'data: {"type": "id", "id": "conv-1"}'
|
||||
yield 'data: {"type": "answer", "answer": "hi"}'
|
||||
|
||||
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",
|
||||
return_value=["data: chunk\n\n"],
|
||||
):
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": "Bearer x"},
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Hi"}],
|
||||
"stream": True,
|
||||
},
|
||||
)
|
||||
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, strip_reasoning_leak=False
|
||||
):
|
||||
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
|
||||
|
||||
|
||||
@@ -0,0 +1,235 @@
|
||||
"""Foreign-``conversation_id`` ownership gate on the ``/v1`` continuation path.
|
||||
|
||||
A ``/v1`` tool continuation persists its final answer into a client-supplied
|
||||
``conversation_id``. Ownership IS enforced deep in the stack:
|
||||
|
||||
route → ``build_continuation_from_messages`` → ``build_agent("")``
|
||||
→ ``initialize`` → ``_load_conversation_history``
|
||||
→ ``get_conversation(conversation_id, owner)`` ← owner-scoped read
|
||||
→ ``None`` for a foreign conversation
|
||||
→ ``ValueError("Conversation not found or unauthorized")``
|
||||
→ route's ``except ValueError`` → HTTP 400
|
||||
|
||||
…but the existing E2E suite mocks ``build_agent``, so this security gate is
|
||||
never actually exercised. This test seeds a conversation owned by user A and an
|
||||
agent owned by user B, then POSTs a continuation with B's api_key targeting A's
|
||||
conversation. The ownership read runs against the **real** DB (not mocked): we
|
||||
only short-circuit ``StreamProcessor.initialize`` to the
|
||||
``_load_conversation_history`` step so the test doesn't have to stand up the
|
||||
full agent/model/source/retriever machinery to reach the gate.
|
||||
|
||||
Asserted: the route returns 400 AND no assistant turn is appended to A's
|
||||
conversation (the foreign write is rejected before persistence).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.answer.services import stream_processor as sp_mod
|
||||
from application.api.v1.routes import v1_bp
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.conversations import ConversationsRepository
|
||||
|
||||
# Reuse the route-level DB wiring + seed helpers from the tool-pause suite.
|
||||
from tests.api.v1.test_v1_tool_pause_finalization import (
|
||||
_seed_user,
|
||||
_wire_v1_route_db,
|
||||
)
|
||||
|
||||
|
||||
CLIENT_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather for a city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _seed_conversation(conn, user_id: str, name: str) -> str:
|
||||
"""Create a conversation owned by ``user_id`` with one finalized turn."""
|
||||
repo = ConversationsRepository(conn)
|
||||
conv = repo.create(user_id, name)
|
||||
conv_id = str(conv["id"])
|
||||
repo.reserve_message(
|
||||
conv_id,
|
||||
prompt="hello",
|
||||
placeholder_response="hi there",
|
||||
status="complete",
|
||||
)
|
||||
return conv_id
|
||||
|
||||
|
||||
def _seed_agent(conn, user_id: str, key: str) -> None:
|
||||
AgentsRepository(conn).create(user_id, "Weather Agent", "published", key=key)
|
||||
|
||||
|
||||
def _build_app() -> Flask:
|
||||
app = Flask(__name__)
|
||||
app.register_blueprint(v1_bp)
|
||||
return app
|
||||
|
||||
|
||||
def _row_count(conn, conv_id: str) -> int:
|
||||
return conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM conversation_messages "
|
||||
"WHERE conversation_id = CAST(:c AS uuid)"
|
||||
),
|
||||
{"c": conv_id},
|
||||
).scalar()
|
||||
|
||||
|
||||
def _only_run_history_load(self) -> None:
|
||||
"""Stand-in for ``StreamProcessor.initialize`` that runs ONLY the
|
||||
ownership-checked history load.
|
||||
|
||||
The real ownership gate (``_load_conversation_history`` →
|
||||
``get_conversation(conversation_id, owner)``) is NOT mocked — it runs
|
||||
against the real DB. We skip the surrounding agent/model/source/retriever
|
||||
configuration only so the test can reach the gate without provisioning a
|
||||
full agent runtime; the ``ValueError`` raised by the gate for a foreign
|
||||
conversation propagates exactly as in production.
|
||||
"""
|
||||
self._load_conversation_history()
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestForeignConversationOwnership:
|
||||
"""A continuation targeting another user's conversation is rejected at the
|
||||
ownership gate (400) and never appends a turn to that conversation."""
|
||||
|
||||
PENDING_MESSAGES: List[Dict[str, Any]] = [
|
||||
{"role": "user", "content": "weather in SF?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_abc",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": json.dumps({"city": "SF"}),
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_abc", "content": "72F sunny"},
|
||||
]
|
||||
|
||||
def test_continuation_into_foreign_conversation_is_rejected(
|
||||
self, pg_engine, monkeypatch
|
||||
):
|
||||
user_a = f"userA-{uuid.uuid4().hex[:8]}" # owns the conversation
|
||||
user_b = f"userB-{uuid.uuid4().hex[:8]}" # owns the agent (the caller)
|
||||
api_key = f"key-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, user_a)
|
||||
_seed_user(conn, user_b)
|
||||
# Agent owned by B → the v1 route resolves decoded_token={"sub": B}.
|
||||
_seed_agent(conn, user_b, api_key)
|
||||
# Conversation owned by A, with one existing complete turn.
|
||||
conv_id = _seed_conversation(conn, user_a, "A's private chat")
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
rows_before = _row_count(conn, conv_id)
|
||||
assert rows_before == 1 # the seeded turn
|
||||
|
||||
app = _build_app()
|
||||
|
||||
# Drive the REAL ownership read; only short-circuit the surrounding
|
||||
# initialize steps so we reach the gate without a full agent runtime.
|
||||
with _wire_v1_route_db(pg_engine, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
sp_mod.StreamProcessor, "initialize", _only_run_history_load
|
||||
)
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
json={
|
||||
# B (the agent owner) tries to write into A's conversation.
|
||||
"conversation_id": conv_id,
|
||||
"messages": self.PENDING_MESSAGES,
|
||||
"tools": [CLIENT_TOOL],
|
||||
"docsgpt": {"save_conversation": True},
|
||||
},
|
||||
)
|
||||
|
||||
# The ownership gate raised ValueError → route returns 400.
|
||||
assert resp.status_code == 400, resp.get_data(as_text=True)
|
||||
body = resp.get_json()
|
||||
assert body["error"]["type"] == "invalid_request"
|
||||
|
||||
# No assistant turn was appended to A's conversation — the foreign
|
||||
# write was rejected before any persistence.
|
||||
with pg_engine.connect() as conn:
|
||||
rows_after = _row_count(conn, conv_id)
|
||||
assert rows_after == rows_before
|
||||
|
||||
def test_owner_can_continue_into_own_conversation(self, pg_engine, monkeypatch):
|
||||
"""Control: the SAME continuation into the caller's OWN conversation
|
||||
passes the ownership gate (no 400 from the gate). Proves the 400 above
|
||||
is the ownership check firing, not an unrelated error in the path.
|
||||
"""
|
||||
owner = f"owner-{uuid.uuid4().hex[:8]}"
|
||||
api_key = f"key-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, owner)
|
||||
_seed_agent(conn, owner, api_key)
|
||||
# Conversation owned by the agent owner (the caller).
|
||||
conv_id = _seed_conversation(conn, owner, "owner's chat")
|
||||
|
||||
app = _build_app()
|
||||
|
||||
# Capture whether the gate passed: a flag the stubbed initialize sets
|
||||
# only if ``_load_conversation_history`` returned without raising.
|
||||
gate = {"passed": False}
|
||||
|
||||
def _history_then_flag(self):
|
||||
self._load_conversation_history()
|
||||
gate["passed"] = True
|
||||
# Stop here: returning lets build_continuation_from_messages proceed
|
||||
# to create_agent, which needs a full runtime we deliberately skip.
|
||||
# Raise a sentinel the test recognises, distinct from the ownership
|
||||
# ValueError, so we can assert the gate itself did not reject us.
|
||||
raise RuntimeError("__gate_passed_sentinel__")
|
||||
|
||||
with _wire_v1_route_db(pg_engine, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
sp_mod.StreamProcessor, "initialize", _history_then_flag
|
||||
)
|
||||
with app.test_client() as c:
|
||||
resp = c.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
json={
|
||||
"conversation_id": conv_id,
|
||||
"messages": (
|
||||
TestForeignConversationOwnership.PENDING_MESSAGES
|
||||
),
|
||||
"tools": [CLIENT_TOOL],
|
||||
"docsgpt": {"save_conversation": True},
|
||||
},
|
||||
)
|
||||
|
||||
# The ownership read accepted the owner (gate passed); the request then
|
||||
# fails on the deliberately-skipped runtime (500), NOT on a 400 from the
|
||||
# gate. This isolates the ownership decision from the rest of the path.
|
||||
assert gate["passed"] is True
|
||||
assert resp.status_code == 500, resp.get_data(as_text=True)
|
||||
@@ -0,0 +1,439 @@
|
||||
"""Layer-1 idempotency for ``/v1/chat/completions`` (non-streaming, OpenAI-style).
|
||||
|
||||
Dropping the native ``resume_from_tool_actions`` path removed its
|
||||
``mark_resuming`` guard, so a duplicated/retried non-streaming POST could
|
||||
re-run the agent → a duplicate answer row + double token billing. The fix
|
||||
honors a client-supplied ``Idempotency-Key`` header: a retry returns the
|
||||
*stored first response* instead of re-running.
|
||||
|
||||
These tests pin both layers against an ephemeral Postgres (``pg_engine``):
|
||||
|
||||
- the ``application.api.v1.idempotency`` helper contract directly
|
||||
(completed → cached, fresh-pending → 409, stale-pending → re-claim,
|
||||
non-2xx → not cached); and
|
||||
- the real ``/v1/chat/completions`` route end-to-end: the same key twice
|
||||
returns the cached body and does NOT re-run the agent or append a second
|
||||
conversation row; no key / different keys run independently; a 5xx first
|
||||
response is not cached so a retry re-runs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Dict, List
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.v1 import idempotency as v1_idem
|
||||
from application.api.v1.idempotency import (
|
||||
STALE_PENDING_SECONDS,
|
||||
TASK_NAME,
|
||||
claim_or_replay,
|
||||
finalize,
|
||||
release,
|
||||
scoped_key,
|
||||
)
|
||||
from application.api.v1.routes import v1_bp
|
||||
|
||||
# Reuse the route-level DB wiring + fake answering agent from the tool-pause
|
||||
# suite so a real two-POST round-trip runs against the ephemeral Postgres.
|
||||
from tests.api.v1.test_v1_tool_pause_finalization import (
|
||||
_seed_agent,
|
||||
_seed_user,
|
||||
_wire_v1_route_db,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper-contract tests (drive the real task_dedup row directly)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_response(app: Flask, body: Dict[str, Any], status: int):
|
||||
"""Build a real Flask Response inside an app context (for ``finalize``)."""
|
||||
from flask import jsonify, make_response
|
||||
|
||||
with app.test_request_context():
|
||||
return make_response(jsonify(body), status)
|
||||
|
||||
|
||||
def _wire_idem_db(engine, monkeypatch) -> None:
|
||||
"""Point the v1 idempotency helper's ``db_session``/``db_readonly`` at
|
||||
the ephemeral engine (each call opens its own short-lived connection)."""
|
||||
|
||||
@contextmanager
|
||||
def _session():
|
||||
conn = engine.connect()
|
||||
txn = conn.begin()
|
||||
try:
|
||||
yield conn
|
||||
txn.commit()
|
||||
except Exception:
|
||||
txn.rollback()
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
@contextmanager
|
||||
def _readonly():
|
||||
conn = engine.connect()
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
monkeypatch.setattr(v1_idem, "db_session", _session)
|
||||
monkeypatch.setattr(v1_idem, "db_readonly", _readonly)
|
||||
|
||||
|
||||
def _backdate_pending(engine, key: str, secs_ago: int) -> None:
|
||||
with engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(
|
||||
"UPDATE task_dedup "
|
||||
"SET created_at = clock_timestamp() - make_interval(secs => :s) "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"s": secs_ago, "k": key},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestIdempotencyHelperContract:
|
||||
"""The three-outcome claim-before-process contract, against a real row."""
|
||||
|
||||
def test_scoped_key_namespaces_by_agent(self):
|
||||
assert scoped_key("abc", "agent-1") == "agent-1:abc"
|
||||
# Two agents with the same key value never collide.
|
||||
assert scoped_key("abc", "agent-1") != scoped_key("abc", "agent-2")
|
||||
# Opt-in: missing either component disables idempotency.
|
||||
assert scoped_key(None, "agent-1") is None
|
||||
assert scoped_key("abc", None) is None
|
||||
|
||||
def test_first_claim_then_completed_replays_cached(self, pg_engine, monkeypatch):
|
||||
# ``claim_or_replay`` builds Flask responses (``jsonify``) for the
|
||||
# replay/409 paths, so it runs inside an app context — exactly as it
|
||||
# does in the real route.
|
||||
app = Flask(__name__)
|
||||
_wire_idem_db(pg_engine, monkeypatch)
|
||||
key = f"agent:{uuid.uuid4().hex}"
|
||||
|
||||
with app.app_context():
|
||||
claimed, replay = claim_or_replay(key)
|
||||
assert claimed is True
|
||||
assert replay is None
|
||||
|
||||
# Finalize a 200 with a body, then a retry replays it byte-for-byte
|
||||
# without re-claiming.
|
||||
body = {"id": "chatcmpl-1", "choices": [{"x": 1}]}
|
||||
finalize(key, _make_response(app, body, 200))
|
||||
|
||||
claimed2, replay2 = claim_or_replay(key)
|
||||
assert claimed2 is False
|
||||
assert replay2 is not None
|
||||
assert replay2.status_code == 200
|
||||
assert replay2.get_json() == body
|
||||
|
||||
def test_fresh_pending_returns_409(self, pg_engine, monkeypatch):
|
||||
app = Flask(__name__)
|
||||
_wire_idem_db(pg_engine, monkeypatch)
|
||||
key = f"agent:{uuid.uuid4().hex}"
|
||||
|
||||
with app.app_context():
|
||||
claimed, _ = claim_or_replay(key)
|
||||
assert claimed is True # first request in flight, never finalized
|
||||
|
||||
# A concurrent same-key request sees a fresh pending row → 409.
|
||||
claimed2, replay2 = claim_or_replay(key)
|
||||
assert claimed2 is False
|
||||
assert replay2 is not None
|
||||
assert replay2.status_code == 409
|
||||
assert replay2.get_json()["error"]["type"] == "idempotency_conflict"
|
||||
|
||||
def test_stale_pending_is_reclaimed(self, pg_engine, monkeypatch):
|
||||
app = Flask(__name__)
|
||||
_wire_idem_db(pg_engine, monkeypatch)
|
||||
key = f"agent:{uuid.uuid4().hex}"
|
||||
|
||||
with app.app_context():
|
||||
claimed, _ = claim_or_replay(key)
|
||||
assert claimed is True
|
||||
# The original request "died" before finalize; age the claim past
|
||||
# the safety window so a retry may re-claim instead of 409-ing
|
||||
# forever.
|
||||
_backdate_pending(pg_engine, key, secs_ago=STALE_PENDING_SECONDS + 5)
|
||||
|
||||
claimed2, replay2 = claim_or_replay(key)
|
||||
assert claimed2 is True
|
||||
assert replay2 is None
|
||||
|
||||
def test_non_2xx_response_is_not_cached(self, pg_engine, monkeypatch):
|
||||
app = Flask(__name__)
|
||||
_wire_idem_db(pg_engine, monkeypatch)
|
||||
key = f"agent:{uuid.uuid4().hex}"
|
||||
|
||||
with app.app_context():
|
||||
claimed, _ = claim_or_replay(key)
|
||||
assert claimed is True
|
||||
# A 5xx must release the claim (not cache the error).
|
||||
finalize(key, _make_response(app, {"error": {"message": "boom"}}, 500))
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
text("SELECT * FROM task_dedup WHERE idempotency_key = :k"),
|
||||
{"k": key},
|
||||
).fetchone()
|
||||
assert row is None # released → a genuine retry can re-claim
|
||||
|
||||
def test_release_drops_pending_claim(self, pg_engine, monkeypatch):
|
||||
app = Flask(__name__)
|
||||
_wire_idem_db(pg_engine, monkeypatch)
|
||||
key = f"agent:{uuid.uuid4().hex}"
|
||||
with app.app_context():
|
||||
claim_or_replay(key)
|
||||
release(key)
|
||||
with pg_engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
text("SELECT * FROM task_dedup WHERE idempotency_key = :k"),
|
||||
{"k": key},
|
||||
).fetchone()
|
||||
assert row is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Route-level tests — same key twice does NOT re-run the agent
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _CountingAnswerAgent:
|
||||
"""Agent that answers once per ``gen`` and records how often it ran.
|
||||
|
||||
A class-level counter survives across the fresh instances handed back by
|
||||
the patched ``build_agent``, so the route-level test can assert the agent
|
||||
ran exactly once across two same-key POSTs.
|
||||
"""
|
||||
|
||||
ANSWER_TEXT = "The answer is 42."
|
||||
gen_calls = 0
|
||||
|
||||
def __init__(self) -> None:
|
||||
from tests.api.v1.test_v1_tool_pause_finalization import (
|
||||
_FakeClientToolExecutor,
|
||||
_FakeLLM,
|
||||
)
|
||||
|
||||
self.llm = _FakeLLM()
|
||||
self.tool_executor = _FakeClientToolExecutor()
|
||||
self.conversation_id = None
|
||||
self.initial_user_id = None
|
||||
self.tools: List[Dict[str, Any]] = []
|
||||
|
||||
def gen(self, query: str = ""):
|
||||
# Write the base-class counter explicitly so subclasses (e.g. a
|
||||
# failing variant) share one run-count rather than shadowing it.
|
||||
_CountingAnswerAgent.gen_calls += 1
|
||||
yield {"answer": self.ANSWER_TEXT}
|
||||
|
||||
|
||||
def _build_app() -> Flask:
|
||||
app = Flask(__name__)
|
||||
app.register_blueprint(v1_bp)
|
||||
return app
|
||||
|
||||
|
||||
def _post_chat(client, body, api_key, idem_key=None):
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
if idem_key is not None:
|
||||
headers["Idempotency-Key"] = idem_key
|
||||
return client.post("/v1/chat/completions", headers=headers, json=body)
|
||||
|
||||
|
||||
def _conv_count(conn, user_id: str) -> int:
|
||||
return conn.execute(
|
||||
text("SELECT count(*) FROM conversations WHERE user_id = :u"),
|
||||
{"u": user_id},
|
||||
).scalar()
|
||||
|
||||
|
||||
def _row_count(conn, user_id: str) -> int:
|
||||
return conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM conversation_messages cm "
|
||||
"JOIN conversations c ON c.id = cm.conversation_id "
|
||||
"WHERE c.user_id = :u"
|
||||
),
|
||||
{"u": user_id},
|
||||
).scalar()
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestV1IdempotencyRoute:
|
||||
"""Drive the real ``/v1/chat/completions`` route with/without a key."""
|
||||
|
||||
QUESTION = {"messages": [{"role": "user", "content": "what is the answer?"}],
|
||||
"docsgpt": {"save_conversation": True}}
|
||||
|
||||
def _fake_build_agent(self, question): # noqa: ARG002
|
||||
return _CountingAnswerAgent()
|
||||
|
||||
def test_same_key_twice_replays_without_rerunning_agent(
|
||||
self, pg_engine, monkeypatch
|
||||
):
|
||||
user_id = f"user-{uuid.uuid4().hex[:8]}"
|
||||
api_key = f"key-{uuid.uuid4().hex[:8]}"
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, user_id)
|
||||
_seed_agent(conn, user_id, api_key)
|
||||
|
||||
app = _build_app()
|
||||
idem_key = f"idem-{uuid.uuid4().hex}"
|
||||
_CountingAnswerAgent.gen_calls = 0
|
||||
|
||||
with _wire_v1_route_db(pg_engine, monkeypatch), patch(
|
||||
"application.api.answer.services.stream_processor.StreamProcessor"
|
||||
".build_agent",
|
||||
self._fake_build_agent,
|
||||
):
|
||||
_wire_idem_db(pg_engine, monkeypatch)
|
||||
with app.test_client() as c:
|
||||
resp1 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key)
|
||||
resp2 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key)
|
||||
|
||||
assert resp1.status_code == 200, resp1.get_data(as_text=True)
|
||||
assert resp2.status_code == 200, resp2.get_data(as_text=True)
|
||||
body1, body2 = resp1.get_json(), resp2.get_json()
|
||||
# The retry returns the cached first response byte-for-byte.
|
||||
assert body2 == body1
|
||||
assert body2["choices"][0]["message"]["content"] == (
|
||||
_CountingAnswerAgent.ANSWER_TEXT
|
||||
)
|
||||
|
||||
# Real dedup: the agent ran exactly once across the two POSTs...
|
||||
assert _CountingAnswerAgent.gen_calls == 1
|
||||
# ...and no duplicate conversation / answer row was appended.
|
||||
with pg_engine.connect() as conn:
|
||||
assert _conv_count(conn, user_id) == 1
|
||||
# One conversation, one terminal answer row (not two).
|
||||
assert _row_count(conn, user_id) == 1
|
||||
|
||||
def test_no_key_runs_each_time(self, pg_engine, monkeypatch):
|
||||
"""Opt-in: without a key, behavior is today's — each POST runs."""
|
||||
user_id = f"user-{uuid.uuid4().hex[:8]}"
|
||||
api_key = f"key-{uuid.uuid4().hex[:8]}"
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, user_id)
|
||||
_seed_agent(conn, user_id, api_key)
|
||||
|
||||
app = _build_app()
|
||||
_CountingAnswerAgent.gen_calls = 0
|
||||
|
||||
with _wire_v1_route_db(pg_engine, monkeypatch), patch(
|
||||
"application.api.answer.services.stream_processor.StreamProcessor"
|
||||
".build_agent",
|
||||
self._fake_build_agent,
|
||||
):
|
||||
_wire_idem_db(pg_engine, monkeypatch)
|
||||
with app.test_client() as c:
|
||||
resp1 = _post_chat(c, self.QUESTION, api_key)
|
||||
resp2 = _post_chat(c, self.QUESTION, api_key)
|
||||
|
||||
assert resp1.status_code == 200
|
||||
assert resp2.status_code == 200
|
||||
# Two independent runs → two conversations.
|
||||
assert _CountingAnswerAgent.gen_calls == 2
|
||||
with pg_engine.connect() as conn:
|
||||
assert _conv_count(conn, user_id) == 2
|
||||
|
||||
def test_different_keys_run_independently(self, pg_engine, monkeypatch):
|
||||
user_id = f"user-{uuid.uuid4().hex[:8]}"
|
||||
api_key = f"key-{uuid.uuid4().hex[:8]}"
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, user_id)
|
||||
_seed_agent(conn, user_id, api_key)
|
||||
|
||||
app = _build_app()
|
||||
_CountingAnswerAgent.gen_calls = 0
|
||||
|
||||
with _wire_v1_route_db(pg_engine, monkeypatch), patch(
|
||||
"application.api.answer.services.stream_processor.StreamProcessor"
|
||||
".build_agent",
|
||||
self._fake_build_agent,
|
||||
):
|
||||
_wire_idem_db(pg_engine, monkeypatch)
|
||||
with app.test_client() as c:
|
||||
_post_chat(c, self.QUESTION, api_key, idem_key="k-A")
|
||||
_post_chat(c, self.QUESTION, api_key, idem_key="k-B")
|
||||
|
||||
assert _CountingAnswerAgent.gen_calls == 2
|
||||
with pg_engine.connect() as conn:
|
||||
assert _conv_count(conn, user_id) == 2
|
||||
|
||||
def test_5xx_first_response_is_not_cached_and_retry_reruns(
|
||||
self, pg_engine, monkeypatch
|
||||
):
|
||||
"""A failed first response must NOT be cached: a retry re-runs."""
|
||||
user_id = f"user-{uuid.uuid4().hex[:8]}"
|
||||
api_key = f"key-{uuid.uuid4().hex[:8]}"
|
||||
with pg_engine.begin() as conn:
|
||||
_seed_user(conn, user_id)
|
||||
_seed_agent(conn, user_id, api_key)
|
||||
|
||||
app = _build_app()
|
||||
idem_key = f"idem-{uuid.uuid4().hex}"
|
||||
|
||||
class _BoomAgent(_CountingAnswerAgent):
|
||||
def gen(self, query: str = ""):
|
||||
_CountingAnswerAgent.gen_calls += 1
|
||||
raise RuntimeError("upstream exploded")
|
||||
yield # pragma: no cover - keeps gen a generator
|
||||
|
||||
boom_state = {"fail": True}
|
||||
|
||||
def _build(self, question): # noqa: ARG001
|
||||
return _BoomAgent() if boom_state["fail"] else _CountingAnswerAgent()
|
||||
|
||||
_CountingAnswerAgent.gen_calls = 0
|
||||
|
||||
with _wire_v1_route_db(pg_engine, monkeypatch), patch(
|
||||
"application.api.answer.services.stream_processor.StreamProcessor"
|
||||
".build_agent",
|
||||
_build,
|
||||
):
|
||||
_wire_idem_db(pg_engine, monkeypatch)
|
||||
with app.test_client() as c:
|
||||
resp1 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key)
|
||||
# First attempt blew up (500) and released the claim. The key
|
||||
# is agent-id scoped at the route, so confirm release via a
|
||||
# table-wide count for this task_name (no live claim left).
|
||||
assert resp1.status_code == 500, resp1.get_data(as_text=True)
|
||||
with pg_engine.connect() as conn:
|
||||
rows = conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM task_dedup "
|
||||
"WHERE task_name = :n"
|
||||
),
|
||||
{"n": TASK_NAME},
|
||||
).scalar()
|
||||
assert rows == 0 # error was not cached; claim released
|
||||
|
||||
# Retry: the agent now answers; idempotency did NOT cache the error.
|
||||
boom_state["fail"] = False
|
||||
resp2 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key)
|
||||
|
||||
assert resp2.status_code == 200, resp2.get_data(as_text=True)
|
||||
# Agent ran twice total: the failing first + the successful retry.
|
||||
assert _CountingAnswerAgent.gen_calls == 2
|
||||
# The successful retry IS now cached.
|
||||
with pg_engine.connect() as conn:
|
||||
completed = conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM task_dedup "
|
||||
"WHERE task_name = :n AND status = 'completed'"
|
||||
),
|
||||
{"n": TASK_NAME},
|
||||
).scalar()
|
||||
assert completed == 1
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user