mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 18:13:03 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
1767 lines
64 KiB
Python
Executable File
1767 lines
64 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""End-to-end verification for the Tier 1 durability work.
|
|
|
|
Two tiers of coverage:
|
|
|
|
**Service-layer scenarios (s1-s10)** — exercise the WAL + reconciler +
|
|
idempotency + token-usage primitives without spinning up Flask, Celery
|
|
workers, or a live LLM:
|
|
|
|
* Service-level calls (e.g. ``ConversationService.save_user_question``)
|
|
drive the same code paths the routes use.
|
|
* Flask ``test_request_context`` lets us hit ``Resource.post()`` directly
|
|
for HTTP idempotency checks.
|
|
* The reconciler is invoked synchronously via ``run_reconciliation()``.
|
|
* Stuck-row scenarios are simulated by backdating ``timestamp`` /
|
|
``attempted_at`` columns so the reconciler treats them as past the
|
|
threshold without having to wait the real 5/15 minutes.
|
|
|
|
**Live scenarios (s11-s13)** — boot the local mock LLM stub from
|
|
``scripts/e2e/mock_llm.py`` as a subprocess and drive real network
|
|
streams + real Celery workers:
|
|
|
|
* s11: real OpenAI-protocol stream through WAL + finalize + token_usage.
|
|
* s12: full ingest of a synthesised markdown file (chunker + embeddings
|
|
+ faiss). The final ``upload_index`` self-call to Flask is patched
|
|
out — everything before that (chunk progress, source id, vector store
|
|
on disk) runs for real.
|
|
* s13: ``kill -9`` a Celery worker subprocess mid-task; verify
|
|
``acks_late`` + ``visibility_timeout`` make the broker redeliver to a
|
|
fresh worker. Uses isolated Redis DBs (11/12) and a unique queue so
|
|
it can't interfere with production work on the same broker.
|
|
|
|
Usage::
|
|
|
|
.venv/bin/python scripts/qa/durability_e2e.py
|
|
.venv/bin/python scripts/qa/durability_e2e.py --list
|
|
.venv/bin/python scripts/qa/durability_e2e.py --only s1,s5
|
|
.venv/bin/python scripts/qa/durability_e2e.py --only s11,s12,s13
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
import traceback
|
|
import uuid
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bootstrap: import docsgpt/* and connect to the configured Postgres.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
# Suppress AUTO_MIGRATE / AUTO_CREATE_DB side-effects from app boot — this
|
|
# script only reads/writes data, the schema is assumed to be at head.
|
|
os.environ.setdefault("AUTO_MIGRATE", "false")
|
|
os.environ.setdefault("AUTO_CREATE_DB", "false")
|
|
|
|
from dotenv import load_dotenv # noqa: E402
|
|
|
|
load_dotenv(ROOT / ".env")
|
|
|
|
from sqlalchemy import create_engine, text # noqa: E402
|
|
|
|
POSTGRES_URI = os.environ.get("POSTGRES_URI", "")
|
|
if not POSTGRES_URI:
|
|
print("POSTGRES_URI not set; cannot run E2E checks.", file=sys.stderr)
|
|
sys.exit(2)
|
|
if POSTGRES_URI.startswith("postgresql://"):
|
|
POSTGRES_URI = POSTGRES_URI.replace("postgresql://", "postgresql+psycopg://", 1)
|
|
|
|
ENGINE = create_engine(POSTGRES_URI, pool_pre_ping=True)
|
|
UNIQUE = f"qa_e2e_{int(time.time())}_{uuid.uuid4().hex[:6]}"
|
|
|
|
|
|
# Quiet noisy library loggers so the script's own output stays readable.
|
|
for name in (
|
|
"application", "docsgpt.api", "docsgpt.storage",
|
|
"docsgpt.usage", "docsgpt.parser",
|
|
"docsgpt.api.user.reconciliation",
|
|
):
|
|
logging.getLogger(name).setLevel(logging.ERROR)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Output helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
GREEN = "\033[32m"
|
|
RED = "\033[31m"
|
|
YELLOW = "\033[33m"
|
|
DIM = "\033[2m"
|
|
RESET = "\033[0m"
|
|
|
|
|
|
def header(msg: str) -> None:
|
|
print(f"\n{'=' * 72}\n{msg}\n{'=' * 72}")
|
|
|
|
|
|
def step(msg: str) -> None:
|
|
print(f" {GREEN}✓{RESET} {msg}")
|
|
|
|
|
|
def info(msg: str) -> None:
|
|
print(f" {DIM}· {msg}{RESET}")
|
|
|
|
|
|
def fail(msg: str) -> None:
|
|
print(f" {RED}✗ {msg}{RESET}")
|
|
|
|
|
|
def warn(msg: str) -> None:
|
|
print(f" {YELLOW}! {msg}{RESET}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Patch helpers — reroute service module-level db_session/db_readonly to
|
|
# OUR engine so all writes land on the same Postgres the script is observing.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@contextmanager
|
|
def patch_db_session(*module_paths: str):
|
|
"""Reroute ``db_session`` / ``db_readonly`` (whichever each module
|
|
actually imports) at our engine. Modules that import only one of the
|
|
two are handled gracefully — patches are skipped for missing attrs.
|
|
"""
|
|
import importlib
|
|
|
|
@contextmanager
|
|
def _ses():
|
|
with ENGINE.begin() as conn:
|
|
yield conn
|
|
|
|
@contextmanager
|
|
def _ro():
|
|
with ENGINE.connect() as conn:
|
|
yield conn
|
|
|
|
patches = []
|
|
for module in module_paths:
|
|
mod = importlib.import_module(module)
|
|
if hasattr(mod, "db_session"):
|
|
patches.append(patch(f"{module}.db_session", _ses))
|
|
if hasattr(mod, "db_readonly"):
|
|
patches.append(patch(f"{module}.db_readonly", _ro))
|
|
for p in patches:
|
|
p.start()
|
|
try:
|
|
yield
|
|
finally:
|
|
for p in patches:
|
|
p.stop()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cleanup helper — run after each scenario regardless of pass/fail.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def cleanup_artifacts(prefix: str = UNIQUE) -> None:
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM stack_logs "
|
|
"WHERE query LIKE 'reconciler_%' "
|
|
"AND timestamp > now() - interval '15 minutes' "
|
|
"AND (user_id LIKE :p OR query LIKE :p)"
|
|
),
|
|
{"p": f"%{prefix}%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM tool_call_attempts WHERE call_id LIKE :p"
|
|
),
|
|
{"p": f"{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM conversation_messages WHERE conversation_id IN ("
|
|
" SELECT id FROM conversations WHERE user_id LIKE :p"
|
|
")"
|
|
),
|
|
{"p": f"{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM pending_tool_state WHERE user_id LIKE :p"
|
|
),
|
|
{"p": f"{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM conversations WHERE user_id LIKE :p"
|
|
),
|
|
{"p": f"{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM token_usage WHERE user_id LIKE :p"
|
|
),
|
|
{"p": f"{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM webhook_dedup WHERE idempotency_key LIKE :p"
|
|
),
|
|
{"p": f"%:{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM task_dedup WHERE idempotency_key LIKE :p"
|
|
),
|
|
{"p": f"%:{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM ingest_chunk_progress WHERE source_id IN ("
|
|
" SELECT id FROM sources WHERE name LIKE :p"
|
|
")"
|
|
),
|
|
{"p": f"{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM agents WHERE name LIKE :p"
|
|
),
|
|
{"p": f"{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM prompts WHERE user_id LIKE :p"
|
|
),
|
|
{"p": f"{prefix}-%"},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM users WHERE user_id LIKE :p"
|
|
),
|
|
{"p": f"{prefix}-%"},
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Scenarios
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def s1_wal_and_reconciler() -> None:
|
|
"""WAL placeholder is reserved before LLM call; reconciler escalates
|
|
a stuck pending row to ``failed`` after 3 attempts and emits a
|
|
structured alert."""
|
|
|
|
from docsgpt.api.answer.services.conversation_service import (
|
|
ConversationService,
|
|
TERMINATED_RESPONSE_PLACEHOLDER,
|
|
)
|
|
from docsgpt.api.user.reconciliation import run_reconciliation
|
|
|
|
user = f"{UNIQUE}-s1"
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": user},
|
|
)
|
|
|
|
with patch_db_session("docsgpt.api.answer.services.conversation_service"):
|
|
res = ConversationService().save_user_question(
|
|
conversation_id=None,
|
|
question=f"{UNIQUE} s1 stuck WAL",
|
|
decoded_token={"sub": user},
|
|
)
|
|
mid, cid = res["message_id"], res["conversation_id"]
|
|
info(f"reserved message_id={mid} conv={cid}")
|
|
|
|
with ENGINE.connect() as c:
|
|
row = c.execute(
|
|
text(
|
|
"SELECT status, response FROM conversation_messages "
|
|
"WHERE id = CAST(:m AS uuid)"
|
|
),
|
|
{"m": mid},
|
|
).fetchone()
|
|
assert row[0] == "pending", f"expected pending, got {row[0]!r}"
|
|
assert row[1] == TERMINATED_RESPONSE_PLACEHOLDER, (
|
|
f"placeholder text mismatch: {row[1]!r}"
|
|
)
|
|
step("placeholder reserved with status='pending' + terminated text")
|
|
|
|
# Backdate timestamp so the reconciler sweep treats it as stuck.
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text(
|
|
"UPDATE conversation_messages "
|
|
"SET timestamp = now() - interval '6 minutes' "
|
|
"WHERE id = CAST(:m AS uuid)"
|
|
),
|
|
{"m": mid},
|
|
)
|
|
|
|
# Reconciler MAX_MESSAGE_RECONCILE_ATTEMPTS = 3 → 3 ticks to escalate.
|
|
for i in range(3):
|
|
summary = run_reconciliation()
|
|
info(f"tick {i + 1}: {summary}")
|
|
|
|
with ENGINE.connect() as c:
|
|
row = c.execute(
|
|
text(
|
|
"SELECT status, message_metadata->>'reconcile_attempts', "
|
|
"message_metadata->>'error' FROM conversation_messages "
|
|
"WHERE id = CAST(:m AS uuid)"
|
|
),
|
|
{"m": mid},
|
|
).fetchone()
|
|
assert row[0] == "failed", f"expected failed, got {row[0]!r}"
|
|
assert int(row[1] or 0) >= 3, f"expected attempts>=3, got {row[1]!r}"
|
|
assert row[2] and "reconciler" in row[2], f"missing reconciler error text: {row[2]!r}"
|
|
step(
|
|
f"reconciler escalated to status='failed' "
|
|
f"after attempts={row[1]} (error: {row[2][:60]}…)"
|
|
)
|
|
|
|
with ENGINE.connect() as c:
|
|
alerts = c.execute(
|
|
text(
|
|
"SELECT count(*) FROM stack_logs "
|
|
"WHERE query = 'reconciler_message_failed' "
|
|
"AND timestamp > now() - interval '10 minutes'"
|
|
),
|
|
).scalar()
|
|
assert alerts >= 1, f"expected stack_logs alert; got {alerts}"
|
|
step(f"alert written to stack_logs (count={alerts})")
|
|
|
|
|
|
def s2_webhook_idempotency_and_scoping() -> None:
|
|
"""Webhook double-POST returns the same task_id; concurrent same-key
|
|
POSTs only enqueue once; two agents sharing a raw key each get their
|
|
own dedup row (cross-tenant scoping)."""
|
|
|
|
from flask import Flask
|
|
|
|
from docsgpt.api.user.agents.webhooks import AgentWebhookListener
|
|
from docsgpt.storage.db.repositories.agents import AgentsRepository
|
|
|
|
user_a = f"{UNIQUE}-s2a"
|
|
user_b = f"{UNIQUE}-s2b"
|
|
with ENGINE.begin() as c:
|
|
for u in (user_a, user_b):
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": u},
|
|
)
|
|
agent_a = AgentsRepository(c).create(
|
|
user_a, f"{UNIQUE}-s2-a", "published",
|
|
incoming_webhook_token=f"{UNIQUE}-tok-a",
|
|
)
|
|
agent_b = AgentsRepository(c).create(
|
|
user_b, f"{UNIQUE}-s2-b", "published",
|
|
incoming_webhook_token=f"{UNIQUE}-tok-b",
|
|
)
|
|
info(f"seeded agents a={agent_a['id']} b={agent_b['id']}")
|
|
|
|
app = Flask(__name__)
|
|
|
|
apply_calls: list[dict] = []
|
|
|
|
def _apply_async_side_effect(*_a, **kw):
|
|
apply_calls.append(kw)
|
|
return MagicMock(id=kw.get("task_id") or "auto")
|
|
|
|
apply_mock = MagicMock(side_effect=_apply_async_side_effect)
|
|
|
|
raw_key = f"{UNIQUE}-key-1"
|
|
|
|
# Sequential dedup: same key twice → same task_id, single apply_async call.
|
|
with patch_db_session("docsgpt.api.user.agents.webhooks"), patch(
|
|
"docsgpt.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
|
apply_mock,
|
|
):
|
|
for _ in range(2):
|
|
with app.test_request_context(
|
|
f"/api/webhooks/agents/{agent_a['incoming_webhook_token']}",
|
|
method="POST",
|
|
json={"event": "x"},
|
|
headers={"Idempotency-Key": raw_key},
|
|
):
|
|
listener = AgentWebhookListener()
|
|
resp = listener.post(
|
|
webhook_token=agent_a["incoming_webhook_token"],
|
|
agent=agent_a,
|
|
agent_id_str=str(agent_a["id"]),
|
|
)
|
|
assert resp.status_code == 200
|
|
last_seq = resp.json["task_id"]
|
|
assert apply_mock.call_count == 1, (
|
|
f"sequential same-key: expected 1 apply_async, got {apply_mock.call_count}"
|
|
)
|
|
step(f"sequential same-key dedup: 1 apply_async, task_id={last_seq}")
|
|
|
|
# Cross-agent same key — must NOT collide. Two distinct apply_async calls.
|
|
apply_mock.reset_mock()
|
|
apply_calls.clear()
|
|
cross_key = f"{UNIQUE}-cross"
|
|
with patch_db_session("docsgpt.api.user.agents.webhooks"), patch(
|
|
"docsgpt.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
|
apply_mock,
|
|
):
|
|
for ag in (agent_a, agent_b):
|
|
with app.test_request_context(
|
|
f"/api/webhooks/agents/{ag['incoming_webhook_token']}",
|
|
method="POST",
|
|
json={"event": "x"},
|
|
headers={"Idempotency-Key": cross_key},
|
|
):
|
|
listener = AgentWebhookListener()
|
|
resp = listener.post(
|
|
webhook_token=ag["incoming_webhook_token"],
|
|
agent=ag,
|
|
agent_id_str=str(ag["id"]),
|
|
)
|
|
assert resp.status_code == 200
|
|
assert apply_mock.call_count == 2, (
|
|
f"cross-agent same-key should NOT collide; got {apply_mock.call_count}"
|
|
)
|
|
with ENGINE.connect() as c:
|
|
rows = c.execute(
|
|
text(
|
|
"SELECT idempotency_key, agent_id FROM webhook_dedup "
|
|
"WHERE idempotency_key LIKE :p ORDER BY idempotency_key"
|
|
),
|
|
{"p": f"%:{cross_key}"},
|
|
).fetchall()
|
|
assert len(rows) == 2, f"expected 2 scoped dedup rows; got {len(rows)}"
|
|
scopes = {str(r[1]) for r in rows}
|
|
assert scopes == {str(agent_a["id"]), str(agent_b["id"])}, (
|
|
f"scopes mismatch: {scopes}"
|
|
)
|
|
step(
|
|
f"cross-agent same key did NOT collide; rows={[r[0] for r in rows]}"
|
|
)
|
|
|
|
|
|
def s3_ingest_deterministic_and_resume() -> None:
|
|
"""Deterministic source_id is stable; chunk-progress checkpoint
|
|
resumes from the next un-embedded chunk."""
|
|
|
|
from docsgpt.parser.embedding_pipeline import _read_resume_index
|
|
from docsgpt.storage.db.repositories.ingest_chunk_progress import (
|
|
IngestChunkProgressRepository,
|
|
)
|
|
from docsgpt.worker import _derive_source_id
|
|
|
|
# Same scoped key → same uuid5; different scope → different uuid5.
|
|
a = _derive_source_id("alice:UPLOAD-1")
|
|
b = _derive_source_id("alice:UPLOAD-1")
|
|
c = _derive_source_id("bob:UPLOAD-1")
|
|
assert a == b, f"deterministic mismatch: {a} vs {b}"
|
|
assert a != c, f"cross-user collision: {a} vs {c}"
|
|
step(f"_derive_source_id: stable for same key, distinct across scopes ({a}, {c})")
|
|
|
|
# No key → uuid4 (random per call).
|
|
r1 = _derive_source_id(None)
|
|
r2 = _derive_source_id(None)
|
|
assert r1 != r2, "uuid4 fallback should differ across calls"
|
|
step("uuid4 fallback when no key supplied")
|
|
|
|
# Checkpoint primitive: seed progress, observe resume index.
|
|
src_id = uuid.uuid4()
|
|
with patch_db_session("docsgpt.parser.embedding_pipeline"):
|
|
with ENGINE.begin() as conn:
|
|
repo = IngestChunkProgressRepository(conn)
|
|
repo.init_progress(str(src_id), total_chunks=10)
|
|
repo.record_chunk(str(src_id), last_index=3, embedded_chunks=4)
|
|
idx = _read_resume_index(str(src_id))
|
|
assert idx == 4, f"expected resume index=4 (next un-embedded); got {idx}"
|
|
step(f"chunk-progress resume index correctly skips already-embedded ({idx}/10)")
|
|
|
|
# Cleanup
|
|
with ENGINE.begin() as conn:
|
|
conn.execute(
|
|
text("DELETE FROM ingest_chunk_progress WHERE source_id = CAST(:s AS uuid)"),
|
|
{"s": str(src_id)},
|
|
)
|
|
|
|
|
|
def s4_token_usage_atomic_rollback() -> None:
|
|
"""Patch ``update_message_by_id`` to raise after the token_usage
|
|
insert; the surrounding transaction must roll back so no orphan
|
|
token_usage row survives."""
|
|
|
|
from docsgpt.api.answer.services.conversation_service import (
|
|
ConversationService,
|
|
)
|
|
from docsgpt.storage.db.repositories import conversations as conv_mod
|
|
from docsgpt.storage.db.repositories.conversations import (
|
|
ConversationsRepository,
|
|
)
|
|
|
|
user = f"{UNIQUE}-s4"
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": user},
|
|
)
|
|
repo = ConversationsRepository(c)
|
|
conv = repo.create(user, f"{UNIQUE}-s4-conv")
|
|
msg = repo.reserve_message(
|
|
str(conv["id"]),
|
|
prompt="q?",
|
|
placeholder_response="...",
|
|
request_id="r1",
|
|
status="streaming",
|
|
)
|
|
mid = str(msg["id"])
|
|
info(f"seeded message_id={mid}")
|
|
|
|
original = conv_mod.ConversationsRepository.update_message_by_id
|
|
|
|
def explode(self, *_a, **_kw): # noqa: ANN001
|
|
raise RuntimeError("intentional fault")
|
|
|
|
try:
|
|
conv_mod.ConversationsRepository.update_message_by_id = explode
|
|
with patch_db_session(
|
|
"docsgpt.api.answer.services.conversation_service",
|
|
), suppress_logging("docsgpt.api.answer.services.conversation_service"):
|
|
try:
|
|
ConversationService().finalize_message(
|
|
mid,
|
|
"answer",
|
|
status="complete",
|
|
token_usage={"prompt_tokens": 100, "generated_tokens": 50},
|
|
decoded_token={"sub": user},
|
|
)
|
|
except RuntimeError as exc:
|
|
assert "intentional fault" in str(exc)
|
|
else:
|
|
raise AssertionError("expected RuntimeError to propagate")
|
|
finally:
|
|
conv_mod.ConversationsRepository.update_message_by_id = original
|
|
|
|
with ENGINE.connect() as c:
|
|
n = c.execute(
|
|
text("SELECT count(*) FROM token_usage WHERE user_id = :u"),
|
|
{"u": user},
|
|
).scalar()
|
|
assert n == 0, f"expected 0 token_usage rows after rollback; got {n}"
|
|
step("rollback held: 0 token_usage rows for failed finalize")
|
|
|
|
|
|
@contextmanager
|
|
def suppress_logging(*names: str):
|
|
saved = []
|
|
for n in names:
|
|
lg = logging.getLogger(n)
|
|
saved.append((lg, lg.disabled, lg.level))
|
|
lg.disabled = True
|
|
try:
|
|
yield
|
|
finally:
|
|
for lg, was_disabled, was_level in saved:
|
|
lg.disabled = was_disabled
|
|
lg.level = was_level
|
|
|
|
|
|
def s5_reconciler_tool_call_sweeps() -> None:
|
|
"""Seed stuck tool_call_attempts rows of each kind; the reconciler
|
|
flips proposed→failed and executed→failed (or compensated when the
|
|
tool exposes a working compensate)."""
|
|
|
|
from docsgpt.api.user.reconciliation import run_reconciliation
|
|
|
|
user = f"{UNIQUE}-s5"
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": user},
|
|
)
|
|
|
|
proposed_id = f"{UNIQUE}-prop-{uuid.uuid4().hex[:6]}"
|
|
executed_id = f"{UNIQUE}-exec-{uuid.uuid4().hex[:6]}"
|
|
|
|
with ENGINE.begin() as c:
|
|
c.execute(text("ALTER TABLE tool_call_attempts DISABLE TRIGGER USER"))
|
|
try:
|
|
c.execute(
|
|
text(
|
|
"INSERT INTO tool_call_attempts "
|
|
"(call_id, tool_name, action_name, arguments, status, "
|
|
" attempted_at, updated_at) "
|
|
"VALUES (:c, 'telegram', 'send_message', '{}'::jsonb, "
|
|
" 'proposed', now() - interval '6 minutes', "
|
|
" now() - interval '6 minutes')"
|
|
),
|
|
{"c": proposed_id},
|
|
)
|
|
c.execute(
|
|
text(
|
|
"INSERT INTO tool_call_attempts "
|
|
"(call_id, tool_name, action_name, arguments, status, "
|
|
" attempted_at, updated_at) "
|
|
"VALUES (:c, 'telegram', 'send_message', '{}'::jsonb, "
|
|
" 'executed', now() - interval '20 minutes', "
|
|
" now() - interval '20 minutes')"
|
|
),
|
|
{"c": executed_id},
|
|
)
|
|
finally:
|
|
c.execute(text("ALTER TABLE tool_call_attempts ENABLE TRIGGER USER"))
|
|
info(f"seeded proposed={proposed_id} executed={executed_id}")
|
|
|
|
summary = run_reconciliation()
|
|
info(f"reconciler summary: {summary}")
|
|
assert summary["tool_calls_failed"] >= 2, (
|
|
f"expected ≥2 tool_calls_failed; got {summary}"
|
|
)
|
|
|
|
with ENGINE.connect() as c:
|
|
rows = c.execute(
|
|
text(
|
|
"SELECT call_id, status, error FROM tool_call_attempts "
|
|
"WHERE call_id IN (:p, :e)"
|
|
),
|
|
{"p": proposed_id, "e": executed_id},
|
|
).fetchall()
|
|
by_id = {r[0]: r for r in rows}
|
|
assert by_id[proposed_id][1] == "failed"
|
|
assert "proposed" in by_id[proposed_id][2]
|
|
assert by_id[executed_id][1] == "failed"
|
|
assert "executed-not-confirmed" in by_id[executed_id][2]
|
|
step("proposed→failed and executed→failed with descriptive error text")
|
|
|
|
with ENGINE.connect() as c:
|
|
n = c.execute(
|
|
text(
|
|
"SELECT count(*) FROM stack_logs "
|
|
"WHERE query IN ('reconciler_tool_call_failed_proposed', "
|
|
"'reconciler_tool_call_failed_executed') "
|
|
"AND timestamp > now() - interval '5 minutes'"
|
|
),
|
|
).scalar()
|
|
assert n >= 2, f"expected ≥2 reconciler tool-call alerts; got {n}"
|
|
step(f"alerts written to stack_logs (count={n})")
|
|
|
|
|
|
def s6_resume_janitor() -> None:
|
|
"""A pending_tool_state row in ``status='resuming'`` whose
|
|
``resumed_at`` is older than the 10-min grace window is reverted to
|
|
``status='pending'`` by the cleanup janitor."""
|
|
|
|
|
|
user = f"{UNIQUE}-s6"
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": user},
|
|
)
|
|
conv = c.execute(
|
|
text(
|
|
"INSERT INTO conversations (user_id, name) VALUES (:u, :n) "
|
|
"RETURNING id"
|
|
),
|
|
{"u": user, "n": f"{UNIQUE}-s6"},
|
|
).fetchone()
|
|
cid = str(conv[0])
|
|
c.execute(
|
|
text(
|
|
"INSERT INTO pending_tool_state "
|
|
"(conversation_id, user_id, messages, pending_tool_calls, "
|
|
" tools_dict, tool_schemas, agent_config, "
|
|
" created_at, expires_at, status, resumed_at) "
|
|
"VALUES (CAST(:c AS uuid), :u, '[]'::jsonb, '[]'::jsonb, "
|
|
" '{}'::jsonb, '[]'::jsonb, '{}'::jsonb, "
|
|
" now(), now() + interval '30 minutes', "
|
|
" 'resuming', now() - interval '11 minutes')"
|
|
),
|
|
{"c": cid, "u": user},
|
|
)
|
|
info(f"seeded pending_tool_state in 'resuming' (resumed_at=-11m) for conv={cid}")
|
|
|
|
from docsgpt.api.user.tasks import cleanup_pending_tool_state
|
|
|
|
res = cleanup_pending_tool_state.run()
|
|
info(f"janitor result: {res}")
|
|
assert res["reverted"] >= 1, f"expected reverted>=1; got {res}"
|
|
|
|
with ENGINE.connect() as c:
|
|
row = c.execute(
|
|
text(
|
|
"SELECT status, resumed_at FROM pending_tool_state "
|
|
"WHERE conversation_id = CAST(:c AS uuid) AND user_id = :u"
|
|
),
|
|
{"c": cid, "u": user},
|
|
).fetchone()
|
|
assert row[0] == "pending", f"expected status=pending after revert; got {row[0]!r}"
|
|
assert row[1] is None, f"expected resumed_at=NULL after revert; got {row[1]!r}"
|
|
step("janitor reverted stale 'resuming' → 'pending' and cleared resumed_at")
|
|
|
|
|
|
def s7_token_usage_attribution() -> None:
|
|
"""``finalize_message`` reads from ``agent.llm.token_usage`` (the LLM
|
|
the stream actually mutated); side-channel LLMs marked with
|
|
``_persist_token_usage_inline`` self-write a row; the canary fires
|
|
when ``status='complete'`` arrives with zero counts."""
|
|
|
|
from docsgpt.api.answer.services.conversation_service import (
|
|
ConversationService,
|
|
)
|
|
from docsgpt.storage.db.repositories.conversations import (
|
|
ConversationsRepository,
|
|
)
|
|
from docsgpt.usage import _maybe_persist_inline
|
|
|
|
user = f"{UNIQUE}-s7"
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": user},
|
|
)
|
|
repo = ConversationsRepository(c)
|
|
conv = repo.create(user, f"{UNIQUE}-s7-conv")
|
|
msg = repo.reserve_message(
|
|
str(conv["id"]),
|
|
prompt="q?",
|
|
placeholder_response="...",
|
|
request_id="r1",
|
|
status="streaming",
|
|
)
|
|
mid_real = str(msg["id"])
|
|
|
|
msg2 = repo.reserve_message(
|
|
str(conv["id"]),
|
|
prompt="q?",
|
|
placeholder_response="...",
|
|
request_id="r2",
|
|
status="streaming",
|
|
)
|
|
mid_zero = str(msg2["id"])
|
|
|
|
# 1. Real counts → row written, source='agent_stream'.
|
|
with patch_db_session("docsgpt.api.answer.services.conversation_service"):
|
|
ok = ConversationService().finalize_message(
|
|
mid_real,
|
|
"answer",
|
|
status="complete",
|
|
token_usage={"prompt_tokens": 3000, "generated_tokens": 800},
|
|
decoded_token={"sub": user},
|
|
)
|
|
assert ok, "finalize_message returned False"
|
|
|
|
with ENGINE.connect() as c:
|
|
row = c.execute(
|
|
text(
|
|
"SELECT prompt_tokens, generated_tokens, source "
|
|
"FROM token_usage WHERE user_id = :u "
|
|
"ORDER BY timestamp DESC LIMIT 1"
|
|
),
|
|
{"u": user},
|
|
).fetchone()
|
|
assert row == (3000, 800, "agent_stream"), f"row mismatch: {row}"
|
|
step(f"primary stream row written: {row}")
|
|
|
|
# 2. Zero-count complete fires the canary alert.
|
|
canary_logger = logging.getLogger(
|
|
"docsgpt.api.answer.services.conversation_service"
|
|
)
|
|
seen: list[logging.LogRecord] = []
|
|
|
|
class _Catch(logging.Handler):
|
|
def emit(self, record):
|
|
seen.append(record)
|
|
|
|
handler = _Catch(level=logging.WARNING)
|
|
canary_logger.addHandler(handler)
|
|
prior_disabled = canary_logger.disabled
|
|
prior_level = canary_logger.level
|
|
canary_logger.disabled = False
|
|
canary_logger.setLevel(logging.WARNING)
|
|
try:
|
|
with patch_db_session(
|
|
"docsgpt.api.answer.services.conversation_service",
|
|
):
|
|
ConversationService().finalize_message(
|
|
mid_zero,
|
|
"answer",
|
|
status="complete",
|
|
token_usage={"prompt_tokens": 0, "generated_tokens": 0},
|
|
decoded_token={"sub": user},
|
|
)
|
|
finally:
|
|
canary_logger.removeHandler(handler)
|
|
canary_logger.disabled = prior_disabled
|
|
canary_logger.setLevel(prior_level)
|
|
|
|
assert any(
|
|
getattr(r, "alert", None) == "token_usage_zero_on_complete" for r in seen
|
|
), (
|
|
"expected structured warning with alert=token_usage_zero_on_complete; "
|
|
f"saw alerts={[getattr(r, 'alert', None) for r in seen]}"
|
|
)
|
|
step("canary fired on zero-count 'complete' finalize")
|
|
|
|
# 3. Side-channel inline persist via _maybe_persist_inline.
|
|
class _SideLLM:
|
|
decoded_token = {"sub": user}
|
|
user_api_key = None
|
|
agent_id = None
|
|
_persist_token_usage_inline = True
|
|
_token_usage_source = "compression"
|
|
|
|
with patch_db_session("docsgpt.usage"):
|
|
_maybe_persist_inline(
|
|
_SideLLM(),
|
|
{"prompt_tokens": 250, "generated_tokens": 80},
|
|
)
|
|
with ENGINE.connect() as c:
|
|
side = c.execute(
|
|
text(
|
|
"SELECT prompt_tokens, generated_tokens, source "
|
|
"FROM token_usage WHERE user_id = :u AND source = 'compression'"
|
|
),
|
|
{"u": user},
|
|
).fetchone()
|
|
assert side == (250, 80, "compression"), f"side-channel row mismatch: {side}"
|
|
step(f"side-channel inline persist row written: {side}")
|
|
|
|
|
|
def s8_regenerate_replaces() -> None:
|
|
"""``save_user_question(index=N)`` truncates messages at and after
|
|
position N before reserving the placeholder, so the new row lands at
|
|
position=N (replacing the old) rather than appending at the end."""
|
|
|
|
from docsgpt.api.answer.services.conversation_service import (
|
|
ConversationService,
|
|
)
|
|
from docsgpt.storage.db.repositories.conversations import (
|
|
ConversationsRepository,
|
|
)
|
|
|
|
user = f"{UNIQUE}-s8"
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": user},
|
|
)
|
|
repo = ConversationsRepository(c)
|
|
conv = repo.create(user, f"{UNIQUE}-s8-conv")
|
|
cid = str(conv["id"])
|
|
for i in range(5):
|
|
repo.append_message(
|
|
cid,
|
|
{
|
|
"prompt": f"q{i}",
|
|
"response": f"a{i}",
|
|
"thought": "",
|
|
"sources": [],
|
|
"tool_calls": [],
|
|
"metadata": {},
|
|
},
|
|
)
|
|
before = [m["position"] for m in repo.get_messages(cid)]
|
|
assert before == [0, 1, 2, 3, 4], f"seed wrong: {before}"
|
|
info(f"seeded conv={cid} positions={before}")
|
|
|
|
with patch_db_session("docsgpt.api.answer.services.conversation_service"):
|
|
result = ConversationService().save_user_question(
|
|
conversation_id=cid,
|
|
question=f"{UNIQUE}-regen",
|
|
decoded_token={"sub": user},
|
|
index=3,
|
|
)
|
|
info(f"regenerated at index=3 → message_id={result['message_id']}")
|
|
|
|
with ENGINE.connect() as c:
|
|
repo = ConversationsRepository(c)
|
|
msgs = repo.get_messages(cid)
|
|
positions = [m["position"] for m in msgs]
|
|
assert positions == [0, 1, 2, 3], f"expected [0,1,2,3]; got {positions}"
|
|
regen = next(m for m in msgs if m["position"] == 3)
|
|
assert regen["prompt"] == f"{UNIQUE}-regen", f"prompt mismatch: {regen['prompt']!r}"
|
|
assert regen["status"] == "pending", f"status mismatch: {regen['status']!r}"
|
|
assert not any(m["response"] == "a3" for m in msgs), "old answer at index=3 not gone"
|
|
assert not any(m["prompt"] == "q4" for m in msgs), "tail not truncated"
|
|
step(f"regenerate replaced position 3 and truncated tail (positions={positions})")
|
|
|
|
|
|
def s9_idempotency_concurrent_claim() -> None:
|
|
"""N concurrent ``claim_task`` calls against the same scoped key:
|
|
exactly one writer wins (returns a row), the rest get None.
|
|
Confirms the ON CONFLICT DO NOTHING race semantics for the HTTP
|
|
claim-first pattern."""
|
|
|
|
from docsgpt.storage.db.repositories.idempotency import (
|
|
IdempotencyRepository,
|
|
)
|
|
|
|
key = f"alice:{UNIQUE}-claim-race"
|
|
n = 20
|
|
|
|
def attempt(_idx: int) -> bool:
|
|
# Each thread checks out its own engine connection so we don't
|
|
# serialize at the SQLAlchemy Connection layer.
|
|
with ENGINE.begin() as conn:
|
|
row = IdempotencyRepository(conn).claim_task(
|
|
key=key, task_name="ingest", task_id=f"task-{_idx}",
|
|
)
|
|
return row is not None
|
|
|
|
with ThreadPoolExecutor(max_workers=n) as ex:
|
|
winners = sum(1 for r in ex.map(attempt, range(n)) if r)
|
|
|
|
with ENGINE.connect() as conn:
|
|
rows = conn.execute(
|
|
text("SELECT count(*) FROM task_dedup WHERE idempotency_key = :k"),
|
|
{"k": key},
|
|
).scalar()
|
|
assert winners == 1, f"expected exactly 1 winner; got {winners}"
|
|
assert rows == 1, f"expected exactly 1 dedup row; got {rows}"
|
|
step(f"{n} concurrent claim_task → 1 winner, 1 row")
|
|
|
|
with ENGINE.begin() as conn:
|
|
conn.execute(
|
|
text("DELETE FROM task_dedup WHERE idempotency_key = :k"),
|
|
{"k": key},
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Live-LLM scenarios: bootstrap helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
import shutil # noqa: E402
|
|
import signal # noqa: E402
|
|
import socket # noqa: E402
|
|
import subprocess # noqa: E402
|
|
from typing import Optional # noqa: E402
|
|
|
|
|
|
MOCK_LLM_HOST = "127.0.0.1"
|
|
MOCK_LLM_PORT = 7899
|
|
MOCK_LLM_BASE_URL = f"http://{MOCK_LLM_HOST}:{MOCK_LLM_PORT}/v1"
|
|
|
|
|
|
def _port_open(host: str, port: int, timeout: float = 0.4) -> bool:
|
|
try:
|
|
with socket.create_connection((host, port), timeout=timeout):
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _wait_for(predicate, timeout_s: float = 15.0, interval_s: float = 0.2) -> bool:
|
|
deadline = time.monotonic() + timeout_s
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(interval_s)
|
|
return False
|
|
|
|
|
|
class MockLLMService:
|
|
"""Owns a subprocess running ``scripts/e2e/mock_llm.py``. Idempotent —
|
|
if the port is already open, treats that as "already running" and
|
|
skips bootstrap (so a developer can run their own ``up.sh`` and have
|
|
this script reuse it)."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
chunk_delay_ms: int = 0,
|
|
total_delay_ms: int = 0,
|
|
) -> None:
|
|
self._proc: Optional[subprocess.Popen] = None
|
|
self._reused_existing = False
|
|
self._chunk_delay_ms = chunk_delay_ms
|
|
self._total_delay_ms = total_delay_ms
|
|
self._log_path = (
|
|
ROOT / ".e2e-tmp" / "logs" / f"mock-llm-{UNIQUE}.log"
|
|
)
|
|
|
|
def start(self) -> None:
|
|
if _port_open(MOCK_LLM_HOST, MOCK_LLM_PORT):
|
|
info(f"mock LLM already running at {MOCK_LLM_BASE_URL} — reusing")
|
|
self._reused_existing = True
|
|
return
|
|
self._log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
env = os.environ.copy()
|
|
env["MOCK_LLM_HOST"] = MOCK_LLM_HOST
|
|
env["MOCK_LLM_PORT"] = str(MOCK_LLM_PORT)
|
|
if self._chunk_delay_ms:
|
|
env["MOCK_LLM_FORCE_STREAM_CHUNK_DELAY_MS"] = str(self._chunk_delay_ms)
|
|
if self._total_delay_ms:
|
|
env["MOCK_LLM_FORCE_TOTAL_DELAY_MS"] = str(self._total_delay_ms)
|
|
env["PYTHONUNBUFFERED"] = "1"
|
|
self._proc = subprocess.Popen( # noqa: S603
|
|
[sys.executable, str(ROOT / "scripts" / "e2e" / "mock_llm.py")],
|
|
stdout=open(self._log_path, "w"),
|
|
stderr=subprocess.STDOUT,
|
|
env=env,
|
|
)
|
|
if not _wait_for(
|
|
lambda: _port_open(MOCK_LLM_HOST, MOCK_LLM_PORT), timeout_s=10
|
|
):
|
|
raise RuntimeError(
|
|
f"mock LLM did not open port {MOCK_LLM_PORT} within 10s; "
|
|
f"see {self._log_path}"
|
|
)
|
|
info(f"mock LLM started (pid={self._proc.pid}); log: {self._log_path}")
|
|
|
|
def stop(self) -> None:
|
|
if self._proc is None or self._reused_existing:
|
|
return
|
|
try:
|
|
self._proc.send_signal(signal.SIGTERM)
|
|
self._proc.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
self._proc.kill()
|
|
self._proc.wait(timeout=2)
|
|
finally:
|
|
self._proc = None
|
|
|
|
|
|
def _build_openai_llm():
|
|
"""Construct an OpenAI-provider LLM pointing at the mock LLM. Skips
|
|
the BYOM resolution path that LLMCreator usually goes through, since
|
|
we just want a direct openai-client wrapper for the mock."""
|
|
from docsgpt.llm.openai import OpenAILLM
|
|
|
|
return OpenAILLM(
|
|
api_key="mock-key",
|
|
user_api_key=None,
|
|
base_url=MOCK_LLM_BASE_URL,
|
|
model_id="gpt-4o-mini",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# s11: live LLM stream end-to-end
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def s11_live_llm_stream_through_wal() -> None:
|
|
"""Reserve a WAL placeholder, drive a real OpenAI-protocol stream
|
|
against the mock LLM, then ``finalize_message`` with the resulting
|
|
``llm.token_usage``. Verifies:
|
|
|
|
* The OpenAI client + ``stream_token_usage`` decorator accumulate
|
|
non-zero counts on the LLM instance during a real network stream.
|
|
* ``finalize_message`` writes a ``token_usage`` row with
|
|
``source='agent_stream'`` and the recorded counts.
|
|
* The conversation message is left at ``status='complete'`` with the
|
|
streamed response captured.
|
|
"""
|
|
|
|
from docsgpt.api.answer.services.conversation_service import (
|
|
ConversationService,
|
|
)
|
|
|
|
user = f"{UNIQUE}-s11"
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": user},
|
|
)
|
|
|
|
with patch_db_session("docsgpt.api.answer.services.conversation_service"):
|
|
res = ConversationService().save_user_question(
|
|
conversation_id=None,
|
|
question=f"{UNIQUE} say hello",
|
|
decoded_token={"sub": user},
|
|
)
|
|
mid = res["message_id"]
|
|
info(f"reserved message_id={mid}")
|
|
|
|
llm = _build_openai_llm()
|
|
chunks = list(
|
|
llm.gen_stream(
|
|
model="gpt-4o-mini",
|
|
messages=[
|
|
{"role": "user", "content": f"{UNIQUE} say hello"},
|
|
],
|
|
stream=True,
|
|
tools=None,
|
|
)
|
|
)
|
|
response_text = "".join(c for c in chunks if isinstance(c, str))
|
|
assert response_text, f"no streamed content received; chunks={chunks!r}"
|
|
assert llm.token_usage["prompt_tokens"] > 0, (
|
|
f"expected prompt_tokens > 0; got {llm.token_usage}"
|
|
)
|
|
assert llm.token_usage["generated_tokens"] > 0, (
|
|
f"expected generated_tokens > 0; got {llm.token_usage}"
|
|
)
|
|
step(
|
|
f"streamed {len(chunks)} chunks; tokens={llm.token_usage}; "
|
|
f"response[:40]={response_text[:40]!r}"
|
|
)
|
|
|
|
with patch_db_session("docsgpt.api.answer.services.conversation_service"):
|
|
ok = ConversationService().finalize_message(
|
|
mid,
|
|
response_text,
|
|
status="complete",
|
|
token_usage=llm.token_usage,
|
|
decoded_token={"sub": user},
|
|
)
|
|
assert ok, "finalize_message returned False"
|
|
|
|
with ENGINE.connect() as c:
|
|
msg_row = c.execute(
|
|
text(
|
|
"SELECT status, response FROM conversation_messages "
|
|
"WHERE id = CAST(:m AS uuid)"
|
|
),
|
|
{"m": mid},
|
|
).fetchone()
|
|
assert msg_row[0] == "complete", f"expected complete; got {msg_row[0]!r}"
|
|
assert msg_row[1] == response_text, "response mismatch"
|
|
step("WAL row finalised at status='complete' with the streamed text")
|
|
|
|
with ENGINE.connect() as c:
|
|
usage = c.execute(
|
|
text(
|
|
"SELECT prompt_tokens, generated_tokens, source "
|
|
"FROM token_usage WHERE user_id = :u "
|
|
"ORDER BY timestamp DESC LIMIT 1"
|
|
),
|
|
{"u": user},
|
|
).fetchone()
|
|
assert usage is not None, "no token_usage row written"
|
|
assert usage[0] == llm.token_usage["prompt_tokens"], (
|
|
f"prompt mismatch: row={usage[0]} llm={llm.token_usage['prompt_tokens']}"
|
|
)
|
|
assert usage[1] == llm.token_usage["generated_tokens"], (
|
|
f"generated mismatch: row={usage[1]} llm={llm.token_usage['generated_tokens']}"
|
|
)
|
|
assert usage[2] == "agent_stream", f"source mismatch: {usage[2]!r}"
|
|
step(f"token_usage row written: {usage}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# s12: full embedding loop with a mock markdown document
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def s12_full_embedding_loop() -> None:
|
|
"""Drive ``ingest_worker`` end-to-end against a tiny markdown file:
|
|
chunker runs, the mock LLM answers each ``/v1/embeddings`` call,
|
|
chunks land in a faiss index, and the row in ``ingest_chunk_progress``
|
|
reflects ``embedded_chunks == total_chunks``.
|
|
|
|
Note: we do NOT go through Flask/Celery here — the worker entrypoint
|
|
is just a regular Python function; ``self`` is a MagicMock with the
|
|
minimum attrs the worker reads."""
|
|
|
|
from docsgpt.core.settings import settings
|
|
|
|
user = f"{UNIQUE}-s12"
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": user},
|
|
)
|
|
|
|
# ``LocalStorage`` enforces a path-traversal guard against the project
|
|
# root, so the markdown file has to live UNDER ``<root>/<UPLOAD_FOLDER>``.
|
|
# We use the configured upload folder (relative path) and create a
|
|
# uniquely-named subdir we can clean up afterwards.
|
|
upload_folder_rel = settings.UPLOAD_FOLDER # e.g. "inputs"
|
|
upload_root = ROOT / upload_folder_rel
|
|
upload_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
safe_user = user.replace("/", "_")
|
|
job_dir_name = f"qa-md-{UNIQUE}"
|
|
file_dir_abs = upload_root / safe_user / job_dir_name
|
|
file_dir_abs.mkdir(parents=True, exist_ok=True)
|
|
md_path = file_dir_abs / "qa.md"
|
|
# The FAISS path seeds the index with ``docs[0]`` *before* the loop
|
|
# that records per-chunk progress. To exercise both the seed AND at
|
|
# least one progress-recorded chunk, the markdown needs to be long
|
|
# enough to chunk into ≥2 pieces under the default token settings
|
|
# (~1000 tokens / chunk). We use a synthesised body of clearly-
|
|
# separated paragraphs to give the chunker something realistic to
|
|
# split on.
|
|
paragraph = (
|
|
"DocsGPT durability is verified by an E2E script that drives "
|
|
"the chunker, the embedding pipeline, and the vector store on "
|
|
"a real Postgres + mock-LLM stack. This paragraph exists to "
|
|
"give the chunker enough text to split into multiple chunks "
|
|
"so the per-chunk progress recorder fires at least once after "
|
|
"the FAISS seed step. "
|
|
)
|
|
body = (paragraph * 80) + "\n\n" # ~80 paragraphs of ~50 tokens each
|
|
md_path.write_text(
|
|
"# QA E2E Document\n\n"
|
|
+ "## Section A\n\n" + body
|
|
+ "## Section B\n\n" + body
|
|
+ "## Section C\n\n" + body,
|
|
encoding="utf-8",
|
|
)
|
|
|
|
# Path passed downstream is the relative form the upload route would
|
|
# produce (``LocalStorage`` joins this against project root for I/O).
|
|
file_path_rel = f"{upload_folder_rel}/{safe_user}/{job_dir_name}"
|
|
|
|
# Point embeddings at the mock LLM via env+settings, but leave
|
|
# UPLOAD_FOLDER alone so LocalStorage's cached base_dir still points
|
|
# at the project root.
|
|
saved = {
|
|
"EMBEDDINGS_BASE_URL": settings.EMBEDDINGS_BASE_URL,
|
|
"EMBEDDINGS_KEY": settings.EMBEDDINGS_KEY,
|
|
"EMBEDDINGS_NAME": settings.EMBEDDINGS_NAME,
|
|
"VECTOR_STORE": settings.VECTOR_STORE,
|
|
}
|
|
settings.EMBEDDINGS_BASE_URL = MOCK_LLM_BASE_URL
|
|
settings.EMBEDDINGS_KEY = "mock-key"
|
|
settings.EMBEDDINGS_NAME = "openai_text-embedding-3-small"
|
|
settings.VECTOR_STORE = "faiss"
|
|
|
|
idempotency_key = f"{user}:{UNIQUE}-s12-key"
|
|
from docsgpt.worker import _derive_source_id
|
|
|
|
expected_source_id = str(_derive_source_id(idempotency_key))
|
|
|
|
fake_self = MagicMock()
|
|
fake_self.update_state = MagicMock()
|
|
fake_self.request = MagicMock()
|
|
fake_self.request.id = f"{UNIQUE}-s12-task"
|
|
|
|
# The final ``upload_index`` step in ingest_worker self-calls Flask to
|
|
# register the index with the API. We're not running Flask here, so
|
|
# patch it to a no-op — everything we want to assert (chunker,
|
|
# embedding pipeline, ``ingest_chunk_progress`` rows, vector store
|
|
# files written to disk) happens BEFORE that registration call.
|
|
upload_index_calls: list = []
|
|
|
|
def _fake_upload_index(full_path, file_data):
|
|
upload_index_calls.append({"full_path": full_path, "file_data": file_data})
|
|
# Sanity-check the faiss artefacts exist before we'd POST them.
|
|
assert os.path.exists(full_path), f"vector store dir missing: {full_path}"
|
|
|
|
try:
|
|
with patch_db_session(
|
|
"docsgpt.parser.embedding_pipeline",
|
|
"docsgpt.worker",
|
|
"docsgpt.api.answer.services.conversation_service",
|
|
), patch("docsgpt.worker.upload_index", _fake_upload_index):
|
|
from docsgpt.worker import ingest_worker
|
|
|
|
resp = ingest_worker(
|
|
fake_self,
|
|
directory=upload_folder_rel,
|
|
formats=[".md"],
|
|
job_name=job_dir_name,
|
|
file_path=file_path_rel,
|
|
filename=job_dir_name,
|
|
user=user,
|
|
file_name_map={"qa.md": "qa.md"},
|
|
idempotency_key=idempotency_key,
|
|
)
|
|
info(f"ingest_worker returned: {str(resp)[:80]}")
|
|
finally:
|
|
for k, v in saved.items():
|
|
setattr(settings, k, v)
|
|
|
|
assert len(upload_index_calls) == 1, (
|
|
f"expected 1 upload_index call; got {len(upload_index_calls)}"
|
|
)
|
|
captured = upload_index_calls[0]["file_data"]
|
|
assert captured["id"] == expected_source_id, (
|
|
f"upload_index id mismatch: {captured['id']} vs {expected_source_id}"
|
|
)
|
|
step(
|
|
f"deterministic source_id derived from scoped key: {expected_source_id} "
|
|
f"(would-have-been-POSTed to /api/upload_index)"
|
|
)
|
|
|
|
with ENGINE.connect() as c:
|
|
prog = c.execute(
|
|
text(
|
|
"SELECT total_chunks, embedded_chunks, last_index "
|
|
"FROM ingest_chunk_progress WHERE source_id = CAST(:s AS uuid)"
|
|
),
|
|
{"s": expected_source_id},
|
|
).fetchone()
|
|
assert prog is not None, "no ingest_chunk_progress row for source"
|
|
assert prog[0] >= 1, f"expected at least 1 chunk; got total={prog[0]}"
|
|
assert prog[1] == prog[0], (
|
|
f"expected embedded == total; got embedded={prog[1]} total={prog[0]}"
|
|
)
|
|
step(
|
|
f"ingest_chunk_progress: embedded={prog[1]}/{prog[0]} "
|
|
f"last_index={prog[2]}"
|
|
)
|
|
|
|
# Cleanup
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text(
|
|
"DELETE FROM ingest_chunk_progress WHERE source_id = CAST(:s AS uuid)"
|
|
),
|
|
{"s": expected_source_id},
|
|
)
|
|
shutil.rmtree(file_dir_abs.parent, ignore_errors=True)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# s13: SIGKILL Celery worker mid-task; verify acks_late redelivers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class CeleryWorkerService:
|
|
"""Lightweight Celery worker subprocess for the chaos test. Uses
|
|
isolated Redis DBs (11/12) so it can't grab DocsGPT production work
|
|
by accident, and a unique queue name so two consecutive runs don't
|
|
inherit each other's unacked messages."""
|
|
|
|
_log_seq = 0
|
|
|
|
def __init__(self, queue: str, label: str = "celery") -> None:
|
|
self._proc: Optional[subprocess.Popen] = None
|
|
self._queue = queue
|
|
CeleryWorkerService._log_seq += 1
|
|
self._log_path = (
|
|
ROOT / ".e2e-tmp" / "logs"
|
|
/ f"{label}-{UNIQUE}-{CeleryWorkerService._log_seq}.log"
|
|
)
|
|
|
|
@property
|
|
def pid(self) -> Optional[int]:
|
|
return self._proc.pid if self._proc else None
|
|
|
|
def start(self, env_overrides: dict) -> None:
|
|
self._log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
env = os.environ.copy()
|
|
env.update(env_overrides)
|
|
env["PYTHONUNBUFFERED"] = "1"
|
|
self._proc = subprocess.Popen( # noqa: S603
|
|
[
|
|
sys.executable, "-m", "celery",
|
|
"-A", "docsgpt.app.celery", "worker",
|
|
"-l", "INFO", "--pool=solo",
|
|
"-Q", self._queue,
|
|
"-n", f"qa-{UNIQUE}@%h",
|
|
"--without-gossip", "--without-mingle", "--without-heartbeat",
|
|
],
|
|
stdout=open(self._log_path, "w"),
|
|
stderr=subprocess.STDOUT,
|
|
env=env,
|
|
cwd=str(ROOT),
|
|
)
|
|
|
|
def _ready() -> bool:
|
|
try:
|
|
txt = self._log_path.read_text()
|
|
except OSError:
|
|
return False
|
|
return f"qa-{UNIQUE}@" in txt and "ready." in txt
|
|
|
|
if not _wait_for(_ready, timeout_s=20):
|
|
raise RuntimeError(
|
|
f"celery worker not ready within 20s; see {self._log_path}"
|
|
)
|
|
info(f"celery worker started pid={self._proc.pid} log={self._log_path}")
|
|
|
|
def kill(self, sig: int = signal.SIGKILL) -> None:
|
|
if self._proc is None:
|
|
return
|
|
try:
|
|
os.kill(self._proc.pid, sig)
|
|
except ProcessLookupError:
|
|
pass
|
|
self._proc.wait(timeout=5)
|
|
info(f"celery worker pid={self._proc.pid} terminated by signal {sig}")
|
|
self._proc = None
|
|
|
|
def stop(self) -> None:
|
|
if self._proc is None:
|
|
return
|
|
try:
|
|
self._proc.send_signal(signal.SIGTERM)
|
|
self._proc.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
self._proc.kill()
|
|
self._proc.wait(timeout=2)
|
|
self._proc = None
|
|
|
|
|
|
def s13_sigkill_redelivers_via_acks_late() -> None:
|
|
"""Real-process chaos test:
|
|
|
|
1. Start a Celery worker in a subprocess (isolated Redis DBs + queue).
|
|
2. Dispatch ``process_agent_webhook`` configured to hit a slow mock
|
|
LLM stream (``MOCK_LLM_FORCE_STREAM_CHUNK_DELAY_MS=400``).
|
|
3. ``kill -9`` the worker mid-task before the LLM finishes.
|
|
4. Restart the worker with the same broker.
|
|
5. Watch the broker re-deliver the same task; the second run completes
|
|
the LLM stream and exits cleanly.
|
|
"""
|
|
|
|
import redis # type: ignore[import-not-found]
|
|
|
|
queue = f"docsgpt-chaos-{UNIQUE}"
|
|
broker_db = 11
|
|
backend_db = 12
|
|
broker_url = f"redis://127.0.0.1:6379/{broker_db}"
|
|
backend_url = f"redis://127.0.0.1:6379/{backend_db}"
|
|
|
|
# Make sure these Redis DBs are clean of leftovers from prior runs.
|
|
r = redis.Redis(host="127.0.0.1", port=6379, db=broker_db)
|
|
r.flushdb()
|
|
redis.Redis(host="127.0.0.1", port=6379, db=backend_db).flushdb()
|
|
# Bust the stream-response cache (DB 0) too — without this, a prior
|
|
# run's cached response satisfies the stream instantly and the mock
|
|
# LLM's chunk delay never gets exercised, so SIGKILL fires after the
|
|
# task already completed.
|
|
redis.Redis(host="127.0.0.1", port=6379, db=0).flushdb()
|
|
|
|
# The mock LLM running on the main port (started by the runner) has
|
|
# no chunk delay. s13 needs a slow stream so SIGKILL lands while the
|
|
# worker is still consuming chunks. Stop the existing mock and start
|
|
# a delayed one in its place; the delay env reaches mock_llm.py via
|
|
# ``MOCK_LLM_FORCE_STREAM_CHUNK_DELAY_MS``.
|
|
delayed_mock = MockLLMService(chunk_delay_ms=400)
|
|
# Force-restart even if the existing one is reused: kill via port.
|
|
if _port_open(MOCK_LLM_HOST, MOCK_LLM_PORT):
|
|
# The runner's mock is owned by ``main()``; we need to take it
|
|
# offline temporarily so our delayed copy can bind the port.
|
|
# Send SIGTERM to whoever owns the port via lsof (best-effort).
|
|
try:
|
|
pids = subprocess.check_output(
|
|
["lsof", "-tiTCP:7899", "-sTCP:LISTEN"], text=True
|
|
).split()
|
|
for pid in pids:
|
|
os.kill(int(pid), signal.SIGTERM)
|
|
time.sleep(0.3)
|
|
except (subprocess.CalledProcessError, FileNotFoundError, ProcessLookupError):
|
|
pass
|
|
delayed_mock.start()
|
|
|
|
# We need an agent (with key + webhook token + prompt) to drive the
|
|
# webhook task. The mock LLM has no idea about retrieval, so the
|
|
# 'classic' agent_type is fine.
|
|
user = f"{UNIQUE}-s13"
|
|
with ENGINE.begin() as c:
|
|
c.execute(
|
|
text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
|
|
{"u": user},
|
|
)
|
|
prompt_row = c.execute(
|
|
text(
|
|
"INSERT INTO prompts (user_id, name, content) "
|
|
"VALUES (:u, :n, :c) RETURNING id"
|
|
),
|
|
{"u": user, "n": f"{UNIQUE}-s13", "c": "you are a tester"},
|
|
).fetchone()
|
|
prompt_id = str(prompt_row[0])
|
|
agent_id = str(uuid.uuid4())
|
|
agent_key = f"{UNIQUE}-akey-{uuid.uuid4().hex[:8]}"
|
|
c.execute(
|
|
text(
|
|
"INSERT INTO agents (id, user_id, name, status, retriever, "
|
|
"prompt_id, tools, agent_type, key, chunks) "
|
|
"VALUES (CAST(:id AS uuid), :u, :n, 'published', 'classic', "
|
|
"CAST(:p AS uuid), '[]'::jsonb, 'classic', :k, '2')"
|
|
),
|
|
{
|
|
"id": agent_id, "u": user, "n": f"{UNIQUE}-s13",
|
|
"p": prompt_id, "k": agent_key,
|
|
},
|
|
)
|
|
info(f"seeded agent={agent_id} user={user}")
|
|
|
|
env_overrides = {
|
|
"POSTGRES_URI": os.environ["POSTGRES_URI"],
|
|
"CELERY_BROKER_URL": broker_url,
|
|
"CELERY_RESULT_BACKEND": backend_url,
|
|
"LLM_PROVIDER": "openai",
|
|
"LLM_NAME": "gpt-4o-mini",
|
|
"API_KEY": "mock-key",
|
|
"OPENAI_API_KEY": "mock-key",
|
|
"OPENAI_BASE_URL": MOCK_LLM_BASE_URL,
|
|
"EMBEDDINGS_BASE_URL": MOCK_LLM_BASE_URL,
|
|
"EMBEDDINGS_KEY": "mock-key",
|
|
"EMBEDDINGS_NAME": "openai_text-embedding-3-small",
|
|
"VECTOR_STORE": "faiss",
|
|
"MOCK_LLM_FORCE_STREAM_CHUNK_DELAY_MS": "400",
|
|
"AUTO_MIGRATE": "false",
|
|
"AUTO_CREATE_DB": "false",
|
|
# Short visibility window so the broker re-queues the SIGKILLed
|
|
# task within seconds rather than the 1-hour production default.
|
|
# Has to be set on BOTH workers — the unack entry's age is
|
|
# checked by whichever worker scans next.
|
|
"CELERY_VISIBILITY_TIMEOUT": "5",
|
|
}
|
|
|
|
worker = CeleryWorkerService(queue=queue)
|
|
worker.start(env_overrides)
|
|
|
|
# Dispatch via celery's send_task using the same broker URL the worker
|
|
# listens on. The DocsGPT celery instance has already loaded
|
|
# ``docsgpt.celeryconfig`` (broker = production .env DB 0); since
|
|
# Celery's conf re-reads ``CELERY_BROKER_URL`` from the env at publish
|
|
# time, we override the env for the *whole* dispatch + wait block.
|
|
from celery import Celery
|
|
|
|
saved_broker = os.environ.get("CELERY_BROKER_URL")
|
|
saved_backend = os.environ.get("CELERY_RESULT_BACKEND")
|
|
os.environ["CELERY_BROKER_URL"] = broker_url
|
|
os.environ["CELERY_RESULT_BACKEND"] = backend_url
|
|
|
|
try:
|
|
chaos_app = Celery("chaos", broker=broker_url, backend=backend_url)
|
|
chaos_app.conf.broker_url = broker_url
|
|
chaos_app.conf.result_backend = backend_url
|
|
chaos_app.conf.task_default_queue = queue
|
|
chaos_app.conf.task_default_exchange = queue
|
|
chaos_app.conf.task_default_routing_key = queue
|
|
chaos_app.conf.task_acks_late = True
|
|
chaos_app.conf.task_reject_on_worker_lost = True
|
|
info(
|
|
f"chaos_app broker={chaos_app.conf.broker_url} "
|
|
f"queue={chaos_app.conf.task_default_queue}"
|
|
)
|
|
|
|
idempotency_key = f"{agent_id}:{UNIQUE}-s13-key"
|
|
task_id = str(uuid.uuid4())
|
|
info(f"dispatching task_id={task_id} key={idempotency_key}")
|
|
chaos_app.send_task(
|
|
"docsgpt.api.user.tasks.process_agent_webhook",
|
|
kwargs={
|
|
"agent_id": agent_id,
|
|
"payload": {"event": f"{UNIQUE}-s13"},
|
|
"idempotency_key": idempotency_key,
|
|
},
|
|
queue=queue,
|
|
task_id=task_id,
|
|
)
|
|
|
|
# Sanity: the task should be visible in the broker for at least a
|
|
# brief moment before the worker pulls it. If it never appears,
|
|
# the publish missed the queue entirely.
|
|
queue_seen = _wait_for(
|
|
lambda: r.llen(queue) >= 1, timeout_s=2.0, interval_s=0.05,
|
|
)
|
|
info(
|
|
f"after send_task: queue len={r.llen(queue)} "
|
|
f"(visible-briefly={queue_seen}) "
|
|
f"unacked={r.hlen('unacked') if r.exists('unacked') else 0}"
|
|
)
|
|
|
|
# Wait until the worker actually starts running the task. The
|
|
# mock LLM streams 5 chunks at 400ms each (= 2s), giving us a
|
|
# window to SIGKILL while the task is still in flight.
|
|
def _task_received() -> bool:
|
|
try:
|
|
txt = worker._log_path.read_text() # noqa: SLF001
|
|
except OSError:
|
|
return False
|
|
return task_id in txt
|
|
|
|
started = _wait_for(_task_received, timeout_s=20)
|
|
if not started:
|
|
try:
|
|
tail = worker._log_path.read_text()[-2000:] # noqa: SLF001
|
|
warn(f"worker log tail:\n{tail}")
|
|
except OSError:
|
|
pass
|
|
qlen = r.llen(queue)
|
|
unacked_now = r.hlen("unacked") if r.exists("unacked") else 0
|
|
warn(f"broker state: queue len={qlen} unacked={unacked_now}")
|
|
worker.stop()
|
|
raise AssertionError(
|
|
f"task did not start on worker 1 within 20s — "
|
|
f"see {worker._log_path}"
|
|
)
|
|
info("task picked up by worker 1; killing mid-flight")
|
|
sigkill_pid = worker.pid
|
|
worker.kill(signal.SIGKILL)
|
|
info(f"SIGKILLed worker pid={sigkill_pid}")
|
|
|
|
unacked_count = r.hlen("unacked") if r.exists("unacked") else 0
|
|
info(f"broker 'unacked' hash size after SIGKILL: {unacked_count}")
|
|
|
|
# Restart the worker — acks_late + reject_on_worker_lost should
|
|
# make the broker redeliver the same task to it.
|
|
worker2 = CeleryWorkerService(queue=queue)
|
|
worker2.start(env_overrides)
|
|
|
|
def _task_completed() -> bool:
|
|
try:
|
|
txt = worker2._log_path.read_text() # noqa: SLF001
|
|
except OSError:
|
|
return False
|
|
return (
|
|
task_id in txt
|
|
and ("succeeded in" in txt or "Task succeeded" in txt)
|
|
)
|
|
|
|
completed = _wait_for(_task_completed, timeout_s=60)
|
|
worker2.stop()
|
|
assert completed, (
|
|
f"redelivered task did not complete within 60s; "
|
|
f"see {worker2._log_path}"
|
|
)
|
|
step(f"task {task_id} redelivered after SIGKILL and completed cleanly")
|
|
finally:
|
|
if saved_broker is not None:
|
|
os.environ["CELERY_BROKER_URL"] = saved_broker
|
|
else:
|
|
os.environ.pop("CELERY_BROKER_URL", None)
|
|
if saved_backend is not None:
|
|
os.environ["CELERY_RESULT_BACKEND"] = saved_backend
|
|
else:
|
|
os.environ.pop("CELERY_RESULT_BACKEND", None)
|
|
|
|
# Cleanup
|
|
delayed_mock.stop()
|
|
redis.Redis(host="127.0.0.1", port=6379, db=broker_db).flushdb()
|
|
redis.Redis(host="127.0.0.1", port=6379, db=backend_db).flushdb()
|
|
with ENGINE.begin() as c:
|
|
c.execute(text("DELETE FROM agents WHERE id = CAST(:a AS uuid)"), {"a": agent_id})
|
|
c.execute(text("DELETE FROM prompts WHERE id = CAST(:p AS uuid)"), {"p": prompt_id})
|
|
|
|
|
|
def s10_celery_default_queue() -> None:
|
|
"""Sanity: the broker queue is project-scoped so a sibling worker on
|
|
the same Redis can't grab DocsGPT tasks."""
|
|
|
|
from docsgpt import celeryconfig
|
|
|
|
assert celeryconfig.task_default_queue == "docsgpt", (
|
|
f"expected default queue 'docsgpt'; got {celeryconfig.task_default_queue!r}"
|
|
)
|
|
assert celeryconfig.task_acks_late is True
|
|
assert celeryconfig.task_reject_on_worker_lost is True
|
|
assert celeryconfig.broker_transport_options.get("visibility_timeout") == 3600, (
|
|
f"visibility_timeout mismatch: "
|
|
f"{celeryconfig.broker_transport_options}"
|
|
)
|
|
step("celeryconfig: docsgpt queue, acks_late, reject_on_lost, visibility=3600")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Runner
|
|
# ---------------------------------------------------------------------------
|
|
|
|
SCENARIOS = [
|
|
("s1", s1_wal_and_reconciler,
|
|
"WAL placeholder + reconciler escalation"),
|
|
("s2", s2_webhook_idempotency_and_scoping,
|
|
"webhook idempotency + cross-agent scope"),
|
|
("s3", s3_ingest_deterministic_and_resume,
|
|
"deterministic source_id + chunk-progress resume"),
|
|
("s4", s4_token_usage_atomic_rollback,
|
|
"token_usage rolls back atomically with message update"),
|
|
("s5", s5_reconciler_tool_call_sweeps,
|
|
"reconciler escalates stuck tool_call_attempts"),
|
|
("s6", s6_resume_janitor,
|
|
"pending_tool_state janitor reverts stale 'resuming'"),
|
|
("s7", s7_token_usage_attribution,
|
|
"token_usage attribution + canary + side-channel inline persist"),
|
|
("s8", s8_regenerate_replaces,
|
|
"regenerate at index truncates tail and replaces"),
|
|
("s9", s9_idempotency_concurrent_claim,
|
|
"concurrent claim_task race: exactly one winner"),
|
|
("s10", s10_celery_default_queue,
|
|
"celeryconfig sanity: docsgpt queue + durability defaults"),
|
|
("s11", s11_live_llm_stream_through_wal,
|
|
"LIVE: real OpenAI-protocol stream → WAL → finalize → token_usage"),
|
|
("s12", s12_full_embedding_loop,
|
|
"LIVE: ingest a markdown file end-to-end (chunker + embeddings + faiss)"),
|
|
("s13", s13_sigkill_redelivers_via_acks_late,
|
|
"LIVE: SIGKILL Celery worker mid-task; verify acks_late redelivers"),
|
|
]
|
|
|
|
|
|
# Scenarios that need the mock LLM running. The runner starts/stops it
|
|
# automatically when any of these are selected.
|
|
_LIVE_SCENARIOS = {"s11", "s12", "s13"}
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--list", action="store_true", help="list scenarios and exit")
|
|
parser.add_argument(
|
|
"--only",
|
|
help="comma-separated scenario keys to run (e.g. s1,s5)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if args.list:
|
|
for key, _, desc in SCENARIOS:
|
|
print(f" {key}: {desc}")
|
|
return 0
|
|
|
|
selected = SCENARIOS
|
|
if args.only:
|
|
wanted = {s.strip() for s in args.only.split(",") if s.strip()}
|
|
selected = [s for s in SCENARIOS if s[0] in wanted]
|
|
missing = wanted - {s[0] for s in selected}
|
|
if missing:
|
|
print(f"unknown scenarios: {sorted(missing)}", file=sys.stderr)
|
|
return 2
|
|
|
|
print(f"{DIM}prefix: {UNIQUE}{RESET}")
|
|
print(f"{DIM}engine: {POSTGRES_URI[:60]}…{RESET}")
|
|
|
|
selected_keys = {k for k, _, _ in selected}
|
|
needs_mock_llm = bool(selected_keys & _LIVE_SCENARIOS)
|
|
mock_llm: Optional[MockLLMService] = None
|
|
if needs_mock_llm:
|
|
mock_llm = MockLLMService()
|
|
try:
|
|
mock_llm.start()
|
|
except Exception as exc: # noqa: BLE001
|
|
print(f" {RED}failed to start mock LLM: {exc}{RESET}", file=sys.stderr)
|
|
return 2
|
|
|
|
passes = 0
|
|
fails: list[tuple[str, str]] = []
|
|
|
|
try:
|
|
for key, fn, desc in selected:
|
|
header(f"[{key}] {desc}")
|
|
try:
|
|
fn()
|
|
passes += 1
|
|
except AssertionError as exc:
|
|
fail(str(exc))
|
|
fails.append((key, str(exc)))
|
|
except Exception as exc: # noqa: BLE001
|
|
fail(f"{type(exc).__name__}: {exc}")
|
|
traceback.print_exc()
|
|
fails.append((key, f"{type(exc).__name__}: {exc}"))
|
|
finally:
|
|
try:
|
|
cleanup_artifacts(UNIQUE)
|
|
except Exception as exc: # noqa: BLE001
|
|
warn(f"cleanup raised: {exc}")
|
|
finally:
|
|
if mock_llm is not None:
|
|
mock_llm.stop()
|
|
|
|
print()
|
|
header("summary")
|
|
print(f" {GREEN}{passes} passed{RESET}, {RED}{len(fails)} failed{RESET}")
|
|
if fails:
|
|
print()
|
|
for key, why in fails:
|
|
print(f" {RED}{key}{RESET}: {why}")
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|