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

This commit is contained in:
wehub-resource-sync
2026-07-13 13:28:29 +08:00
commit fed8b2eed7
1531 changed files with 1107494 additions and 0 deletions
View File
@@ -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
+587
View File
@@ -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)
+439
View File
@@ -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