Files
wehub-resource-sync fed8b2eed7
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
chore: import upstream snapshot with attribution
2026-07-13 13:28:29 +08:00

888 lines
35 KiB
Python

"""Tests for ConversationsRepository against a real Postgres instance."""
from __future__ import annotations
import uuid
from datetime import datetime, timezone
from sqlalchemy import text
from application.storage.db.repositories.conversations import (
ConversationsRepository,
MessageUpdateOutcome,
)
def _repo(conn) -> ConversationsRepository:
return ConversationsRepository(conn)
# ------------------------------------------------------------------
# Conversation CRUD
# ------------------------------------------------------------------
class TestCreate:
def test_creates_conversation(self, pg_conn):
repo = _repo(pg_conn)
doc = repo.create("user-1", "My Chat")
assert doc["user_id"] == "user-1"
assert doc["name"] == "My Chat"
assert doc["id"] is not None
assert doc["_id"] == doc["id"]
def test_create_with_agent(self, pg_conn):
from application.storage.db.repositories.agents import AgentsRepository
agent_repo = AgentsRepository(pg_conn)
agent = agent_repo.create("user-1", "a", "active")
repo = _repo(pg_conn)
doc = repo.create(
"user-1", "Chat",
agent_id=agent["id"],
api_key="ak-123",
is_shared_usage=True,
shared_token="tok-abc",
)
assert str(doc["agent_id"]) == agent["id"]
assert doc["api_key"] == "ak-123"
assert doc["is_shared_usage"] is True
assert doc["shared_token"] == "tok-abc"
class TestGet:
def test_get_owned(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "c")
fetched = repo.get(created["id"], "user-1")
assert fetched["id"] == created["id"]
def test_get_nonexistent(self, pg_conn):
repo = _repo(pg_conn)
assert repo.get("00000000-0000-0000-0000-000000000000", "u") is None
def test_get_wrong_user(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "c")
assert repo.get(created["id"], "user-other") is None
class TestListForUser:
def test_lists_own_conversations(self, pg_conn):
repo = _repo(pg_conn)
repo.create("alice", "c1")
repo.create("alice", "c2")
repo.create("bob", "c3")
results = repo.list_for_user("alice")
assert len(results) == 2
assert all(r["user_id"] == "alice" for r in results)
def test_excludes_hidden_conversations(self, pg_conn):
repo = _repo(pg_conn)
repo.create("alice", "normal")
repo.create("alice", "hidden-one", visibility="hidden")
results = repo.list_for_user("alice")
assert len(results) == 1
assert results[0]["name"] == "normal"
def test_lists_api_key_conversation_when_listed(self, pg_conn):
# Visibility is explicit now: an api_key conversation created as
# ``listed`` surfaces (the old api_key/agent_id heuristic is gone).
repo = _repo(pg_conn)
repo.create("alice", "api-listed", api_key="key-1", visibility="listed")
results = repo.list_for_user("alice")
assert len(results) == 1
assert results[0]["name"] == "api-listed"
def test_create_defaults_to_listed(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("alice", "default-vis")
assert conv["visibility"] == "listed"
class TestRename:
def test_renames(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "old")
assert repo.rename(created["id"], "user-1", "new") is True
fetched = repo.get(created["id"], "user-1")
assert fetched["name"] == "new"
def test_rename_wrong_user(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "old")
assert repo.rename(created["id"], "user-other", "new") is False
class TestDelete:
def test_deletes(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "c")
assert repo.delete(created["id"], "user-1") is True
assert repo.get(created["id"], "user-1") is None
def test_delete_wrong_user(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "c")
assert repo.delete(created["id"], "user-other") is False
def test_delete_cascades_messages(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "hi", "response": "hello"})
repo.delete(conv["id"], "user-1")
assert repo.get_messages(conv["id"]) == []
def test_delete_all_for_user(self, pg_conn):
repo = _repo(pg_conn)
repo.create("user-1", "c1")
repo.create("user-1", "c2")
repo.create("user-2", "c3")
count = repo.delete_all_for_user("user-1")
assert count == 2
assert repo.list_for_user("user-1") == []
assert len(repo.list_for_user("user-2")) == 1
# ------------------------------------------------------------------
# Messages
# ------------------------------------------------------------------
class TestAppendMessage:
def test_append_first_message(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
msg = repo.append_message(conv["id"], {
"prompt": "hello",
"response": "hi there",
"model_id": "gpt-4",
})
assert msg["position"] == 0
assert msg["prompt"] == "hello"
assert msg["response"] == "hi there"
assert msg["model_id"] == "gpt-4"
def test_append_increments_position(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
m0 = repo.append_message(conv["id"], {"prompt": "q1", "response": "a1"})
m1 = repo.append_message(conv["id"], {"prompt": "q2", "response": "a2"})
assert m0["position"] == 0
assert m1["position"] == 1
def test_append_with_sources_and_tools(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
msg = repo.append_message(conv["id"], {
"prompt": "q",
"response": "a",
"sources": [{"title": "doc1"}],
"tool_calls": [{"name": "search", "args": {}}],
"metadata": {"search_query": "rewritten"},
})
assert msg["sources"] == [{"title": "doc1"}]
assert msg["tool_calls"] == [{"name": "search", "args": {}}]
assert msg["metadata"] == {"search_query": "rewritten"}
def test_append_preserves_explicit_timestamp(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
ts = datetime.now(timezone.utc)
msg = repo.append_message(conv["id"], {
"prompt": "q",
"response": "a",
"timestamp": ts,
})
# ``row_to_dict`` coerces datetimes to ISO strings at the SELECT
# boundary; round-trip via ``fromisoformat`` to compare values.
assert datetime.fromisoformat(msg["timestamp"]) == ts
class TestGetMessages:
def test_returns_ordered_messages(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q1", "response": "a1"})
repo.append_message(conv["id"], {"prompt": "q2", "response": "a2"})
msgs = repo.get_messages(conv["id"])
assert len(msgs) == 2
assert msgs[0]["position"] == 0
assert msgs[1]["position"] == 1
def test_get_message_at(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q1", "response": "a1"})
repo.append_message(conv["id"], {"prompt": "q2", "response": "a2"})
msg = repo.get_message_at(conv["id"], 1)
assert msg["prompt"] == "q2"
def test_get_message_at_nonexistent(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
assert repo.get_message_at(conv["id"], 99) is None
class TestFirstNSnapshot:
"""``message_in_first_n`` / ``first_n_message_ids`` scope a share snapshot
to the first ``first_n`` messages by position (matches ``messages[:first_n]``)."""
def test_message_in_first_n_membership_by_position(self, pg_conn):
repo = _repo(pg_conn)
cid = repo.create("user-1", "c")["id"]
m0 = repo.append_message(cid, {"prompt": "q0", "response": "a0"})
m1 = repo.append_message(cid, {"prompt": "q1", "response": "a1"})
m2 = repo.append_message(cid, {"prompt": "q2", "response": "a2"})
assert repo.message_in_first_n(cid, m0["id"], 2) is True
assert repo.message_in_first_n(cid, m1["id"], 2) is True
assert repo.message_in_first_n(cid, m2["id"], 2) is False
def test_message_in_first_n_zero_is_empty(self, pg_conn):
repo = _repo(pg_conn)
cid = repo.create("user-1", "c")["id"]
m0 = repo.append_message(cid, {"prompt": "q", "response": "a"})
assert repo.message_in_first_n(cid, m0["id"], 0) is False
def test_message_in_first_n_unknown_and_null_ids(self, pg_conn):
repo = _repo(pg_conn)
cid = repo.create("user-1", "c")["id"]
repo.append_message(cid, {"prompt": "q", "response": "a"})
assert repo.message_in_first_n(cid, str(uuid.uuid4()), 5) is False
assert repo.message_in_first_n(cid, None, 5) is False
# A non-UUID id is shape-gated, never poisoning the enclosing txn.
assert repo.message_in_first_n(cid, "not-a-uuid", 5) is False
assert pg_conn.execute(text("SELECT 1")).scalar() == 1
def test_first_n_message_ids_returns_snapshot_set(self, pg_conn):
repo = _repo(pg_conn)
cid = repo.create("user-1", "c")["id"]
m0 = repo.append_message(cid, {"prompt": "q0", "response": "a0"})
m1 = repo.append_message(cid, {"prompt": "q1", "response": "a1"})
m2 = repo.append_message(cid, {"prompt": "q2", "response": "a2"})
assert repo.first_n_message_ids(cid, 2) == {str(m0["id"]), str(m1["id"])}
assert str(m2["id"]) not in repo.first_n_message_ids(cid, 2)
assert repo.first_n_message_ids(cid, 0) == set()
def test_first_n_message_ids_rejects_non_uuid(self, pg_conn):
repo = _repo(pg_conn)
assert repo.first_n_message_ids("not-a-uuid", 5) == set()
assert pg_conn.execute(text("SELECT 1")).scalar() == 1
class TestUpdateMessageAt:
def test_updates_response(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q", "response": "old"})
assert repo.update_message_at(conv["id"], 0, {"response": "new"}) is True
msg = repo.get_message_at(conv["id"], 0)
assert msg["response"] == "new"
def test_update_disallowed_field(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q", "response": "a"})
assert repo.update_message_at(conv["id"], 0, {"id": "bad"}) is False
def test_updates_explicit_timestamp(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q", "response": "old"})
ts = datetime.now(timezone.utc)
assert repo.update_message_at(
conv["id"], 0, {"response": "new", "timestamp": ts},
) is True
msg = repo.get_message_at(conv["id"], 0)
assert msg["response"] == "new"
assert datetime.fromisoformat(msg["timestamp"]) == ts
class TestTruncateAfter:
def test_truncates_messages(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
for i in range(5):
repo.append_message(conv["id"], {"prompt": f"q{i}", "response": f"a{i}"})
deleted = repo.truncate_after(conv["id"], 2)
assert deleted == 2
msgs = repo.get_messages(conv["id"])
assert len(msgs) == 3
assert [m["position"] for m in msgs] == [0, 1, 2]
class TestSetFeedback:
def test_set_feedback(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q", "response": "a"})
assert repo.set_feedback(conv["id"], 0, {"text": "thumbs_up"}) is True
msg = repo.get_message_at(conv["id"], 0)
assert msg["feedback"] == {"text": "thumbs_up"}
def test_unset_feedback(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q", "response": "a"})
repo.set_feedback(conv["id"], 0, {"text": "thumbs_up"})
assert repo.set_feedback(conv["id"], 0, None) is True
msg = repo.get_message_at(conv["id"], 0)
assert msg["feedback"] is None
def test_set_feedback_nonexistent_position(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
assert repo.set_feedback(conv["id"], 99, {"text": "x"}) is False
class TestMessageCount:
def test_counts_messages(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
assert repo.message_count(conv["id"]) == 0
repo.append_message(conv["id"], {"prompt": "q", "response": "a"})
assert repo.message_count(conv["id"]) == 1
class TestCompressionMetadata:
def test_set_compression_metadata(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
meta = {"is_compressed": True, "last_compression_at": "2026-01-01T00:00:00Z"}
assert repo.update_compression_metadata(conv["id"], "user-1", meta) is True
fetched = repo.get(conv["id"], "user-1")
assert fetched["compression_metadata"]["is_compressed"] is True
def test_set_compression_flags_preserves_points(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.update_compression_metadata(conv["id"], "user-1", {
"is_compressed": False,
"compression_points": [{"summary": "earlier"}],
})
assert repo.set_compression_flags(
conv["id"], is_compressed=True, last_compression_at="2026-01-02",
) is True
fetched = repo.get(conv["id"], "user-1")
assert fetched["compression_metadata"]["is_compressed"] is True
assert fetched["compression_metadata"]["last_compression_at"] == "2026-01-02"
assert fetched["compression_metadata"]["compression_points"] == [
{"summary": "earlier"}
]
def test_append_compression_point_slices_to_max(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
for i in range(5):
assert repo.append_compression_point(
conv["id"], {"summary": f"p{i}"}, max_points=3,
) is True
fetched = repo.get(conv["id"], "user-1")
points = fetched["compression_metadata"]["compression_points"]
assert [p["summary"] for p in points] == ["p2", "p3", "p4"]
class TestResolveAgentRef:
"""The repo must translate Mongo ObjectId-shaped ``agent_id`` values
to Postgres UUIDs on ``create`` so that dual-write from the
ObjectId-era conversation service doesn't silently lose rows."""
def test_create_translates_objectid_agent_id(self, pg_conn):
from application.storage.db.repositories.agents import AgentsRepository
agent_repo = AgentsRepository(pg_conn)
legacy_oid = "507f1f77bcf86cd799439099"
agent = agent_repo.create(
"user-1", "a", "active", legacy_mongo_id=legacy_oid,
)
repo = _repo(pg_conn)
conv = repo.create("user-1", "chat", agent_id=legacy_oid)
assert str(conv["agent_id"]) == agent["id"]
def test_create_passes_through_uuid_agent_id(self, pg_conn):
from application.storage.db.repositories.agents import AgentsRepository
agent_repo = AgentsRepository(pg_conn)
agent = agent_repo.create("user-1", "a", "active")
repo = _repo(pg_conn)
conv = repo.create("user-1", "chat", agent_id=agent["id"])
assert str(conv["agent_id"]) == agent["id"]
def test_create_drops_unknown_objectid_agent_id(self, pg_conn):
# Unknown legacy id resolves to None — the conversation row still
# inserts (dual_write stays quiet) but without an agent FK.
repo = _repo(pg_conn)
conv = repo.create(
"user-1", "chat", agent_id="507f1f77bcf86cd7994390aa",
)
assert conv["agent_id"] is None
class TestResolveAttachmentRefs:
"""``append_message`` and ``update_message_at`` must translate Mongo
ObjectId attachment ids to PG UUIDs via ``attachments.legacy_mongo_id``.
Without this, the ``uuid[]`` cast raises and dual_write drops the
whole message."""
def _create_attachment(self, pg_conn, legacy: str) -> str:
from application.storage.db.repositories.attachments import (
AttachmentsRepository,
)
att = AttachmentsRepository(pg_conn).create(
"user-1", "a.txt", "/tmp/a.txt", legacy_mongo_id=legacy,
)
return att["id"]
def test_append_translates_objectid_attachments(self, pg_conn):
att_uuid = self._create_attachment(
pg_conn, "507f1f77bcf86cd799439011",
)
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
msg = repo.append_message(conv["id"], {
"prompt": "q", "response": "a",
"attachments": ["507f1f77bcf86cd799439011"],
})
assert [str(a) for a in msg["attachments"]] == [att_uuid]
def test_append_passes_through_uuid_attachments(self, pg_conn):
att_uuid = self._create_attachment(
pg_conn, "507f1f77bcf86cd799439022",
)
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
msg = repo.append_message(conv["id"], {
"prompt": "q", "response": "a",
"attachments": [att_uuid],
})
assert [str(a) for a in msg["attachments"]] == [att_uuid]
def test_append_drops_unknown_objectid_attachments(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
# Unknown legacy id is silently dropped; message still inserts.
msg = repo.append_message(conv["id"], {
"prompt": "q", "response": "a",
"attachments": ["507f1f77bcf86cd7994390bb"],
})
assert list(msg["attachments"] or []) == []
def test_update_translates_objectid_attachments(self, pg_conn):
att_uuid = self._create_attachment(
pg_conn, "507f1f77bcf86cd799439033",
)
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q", "response": "a"})
assert repo.update_message_at(
conv["id"], 0,
{"attachments": ["507f1f77bcf86cd799439033"]},
) is True
msg = repo.get_message_at(conv["id"], 0)
assert [str(a) for a in msg["attachments"]] == [att_uuid]
class TestUpdateMessageFeedback:
"""``feedback`` / ``feedback_timestamp`` must be in the update whitelist
so continuation-flow re-appends don't silently strip them."""
def test_update_sets_feedback(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q", "response": "a"})
assert repo.update_message_at(
conv["id"], 0, {"feedback": {"text": "thumbs_up"}},
) is True
msg = repo.get_message_at(conv["id"], 0)
assert msg["feedback"] == {"text": "thumbs_up"}
def test_update_clears_feedback(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q", "response": "a"})
repo.set_feedback(conv["id"], 0, {"text": "thumbs_up"})
assert repo.update_message_at(
conv["id"], 0, {"feedback": None},
) is True
msg = repo.get_message_at(conv["id"], 0)
assert msg["feedback"] is None
class TestReserveAndFinalizeMessage:
"""Pre-persist (WAL) + finalisation primitives used by save_user_question /
finalize_message in the answer-streaming path."""
def test_reserve_message_inserts_pending_row(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-r", "c")
msg = repo.reserve_message(
conv["id"],
prompt="q1",
placeholder_response="placeholder",
request_id="req-1",
)
assert msg["position"] == 0
assert msg["status"] == "pending"
assert msg["request_id"] == "req-1"
assert msg["prompt"] == "q1"
assert msg["response"] == "placeholder"
def test_reserve_message_allocates_next_position(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-r", "c")
repo.append_message(conv["id"], {"prompt": "q0", "response": "a0"})
msg = repo.reserve_message(
conv["id"], prompt="q1", placeholder_response="ph",
)
assert msg["position"] == 1
def test_update_message_by_id_updates_response_and_status(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-r", "c")
msg = repo.reserve_message(
conv["id"], prompt="q", placeholder_response="ph",
)
outcome = repo.update_message_by_id(
msg["id"], {"response": "real answer", "status": "complete"},
)
assert outcome is MessageUpdateOutcome.UPDATED
refreshed = repo.get_message_at(conv["id"], 0)
assert refreshed["response"] == "real answer"
assert refreshed["status"] == "complete"
def test_update_message_by_id_writes_metadata_error(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-r", "c")
msg = repo.reserve_message(
conv["id"], prompt="q", placeholder_response="ph",
)
repo.update_message_by_id(
msg["id"],
{"status": "failed", "metadata": {"error": "RuntimeError: boom"}},
)
refreshed = repo.get_message_at(conv["id"], 0)
assert refreshed["status"] == "failed"
assert refreshed["metadata"]["error"] == "RuntimeError: boom"
def test_update_message_by_id_rejects_non_uuid(self, pg_conn):
repo = _repo(pg_conn)
assert (
repo.update_message_by_id("not-a-uuid", {"status": "complete"})
is MessageUpdateOutcome.INVALID
)
def test_update_message_by_id_distinguishes_already_complete(self, pg_conn):
"""When the row is already ``complete``, a subsequent
``only_if_non_terminal=True`` finalize must report
``ALREADY_COMPLETE`` — not ``UPDATED`` and not the generic
not-found case. The SSE abort handler relies on this to
journal ``end`` instead of a spurious ``error`` when the
normal-path finalize wins the race against a client
disconnect.
"""
repo = _repo(pg_conn)
conv = repo.create("user-r", "c")
msg = repo.reserve_message(
conv["id"], prompt="q", placeholder_response="ph",
)
first = repo.update_message_by_id(
msg["id"], {"response": "ok", "status": "complete"},
only_if_non_terminal=True,
)
assert first is MessageUpdateOutcome.UPDATED
second = repo.update_message_by_id(
msg["id"], {"response": "ok again", "status": "complete"},
only_if_non_terminal=True,
)
assert second is MessageUpdateOutcome.ALREADY_COMPLETE
# The second attempt must NOT have overwritten anything.
refreshed = repo.get_message_at(conv["id"], 0)
assert refreshed["response"] == "ok"
def test_update_message_by_id_distinguishes_already_failed(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-r", "c")
msg = repo.reserve_message(
conv["id"], prompt="q", placeholder_response="ph",
)
assert repo.update_message_by_id(
msg["id"], {"response": "boom", "status": "failed"},
only_if_non_terminal=True,
) is MessageUpdateOutcome.UPDATED
assert repo.update_message_by_id(
msg["id"], {"response": "late", "status": "complete"},
only_if_non_terminal=True,
) is MessageUpdateOutcome.ALREADY_FAILED
def test_update_message_by_id_unknown_uuid_is_not_found(self, pg_conn):
repo = _repo(pg_conn)
assert repo.update_message_by_id(
"00000000-0000-0000-0000-000000000000",
{"status": "complete"},
only_if_non_terminal=True,
) is MessageUpdateOutcome.NOT_FOUND
def test_update_message_status_only(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-r", "c")
msg = repo.reserve_message(
conv["id"], prompt="q", placeholder_response="ph",
)
assert repo.update_message_status(msg["id"], "streaming") is True
refreshed = repo.get_message_at(conv["id"], 0)
assert refreshed["status"] == "streaming"
def test_confirm_executed_tool_calls_flips_status(self, pg_conn):
from sqlalchemy import text as sql_text
repo = _repo(pg_conn)
conv = repo.create("user-r", "c")
msg = repo.reserve_message(
conv["id"], prompt="q", placeholder_response="ph",
)
# Insert two rows: one 'executed' (should flip), one 'proposed' (no-op).
pg_conn.execute(
sql_text(
"INSERT INTO tool_call_attempts "
"(call_id, message_id, tool_name, action_name, arguments, status) "
"VALUES (:cid, CAST(:mid AS uuid), 't', 'a', '{}'::jsonb, :status)"
),
[
{"cid": "c-exec", "mid": msg["id"], "status": "executed"},
{"cid": "c-prop", "mid": msg["id"], "status": "proposed"},
],
)
flipped = repo.confirm_executed_tool_calls(msg["id"])
assert flipped == 1
rows = pg_conn.execute(
sql_text(
"SELECT call_id, status FROM tool_call_attempts "
"WHERE message_id = CAST(:mid AS uuid) ORDER BY call_id"
),
{"mid": msg["id"]},
).fetchall()
as_dict = {r[0]: r[1] for r in rows}
assert as_dict == {"c-exec": "confirmed", "c-prop": "proposed"}
def test_confirm_executed_tool_calls_no_rows_is_zero(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-r", "c")
msg = repo.reserve_message(
conv["id"], prompt="q", placeholder_response="ph",
)
# No tool_call_attempts inserted — no-op
assert repo.confirm_executed_tool_calls(msg["id"]) == 0
class TestConcurrentAppend:
"""Two threads appending to the same conversation must not race on
``position``. The migration plan explicitly calls this out as the
single trickiest invariant, so we exercise it directly with two
parallel connections."""
def test_concurrent_appends_get_distinct_positions(self, pg_engine, pg_conn):
import threading
# Arrange — one conversation, created inside the outer test txn so
# it disappears on teardown even if the workers somehow commit.
# We commit it explicitly so the workers' separate sessions see it.
repo_setup = _repo(pg_conn)
conv = repo_setup.create("user-concurrent", "c")
pg_conn.commit()
try:
errors: list[BaseException] = []
def worker() -> None:
try:
with pg_engine.begin() as worker_conn:
ConversationsRepository(worker_conn).append_message(
conv["id"], {"prompt": "q", "response": "a"},
)
except BaseException as e: # noqa: BLE001
errors.append(e)
threads = [threading.Thread(target=worker) for _ in range(2)]
for t in threads:
t.start()
for t in threads:
t.join()
assert errors == [], f"worker threads errored: {errors}"
# Assert — the parent row-lock in append_message must have
# serialised the two inserts so they land at positions {0, 1}.
with pg_engine.connect() as verify_conn:
msgs = ConversationsRepository(verify_conn).get_messages(conv["id"])
positions = sorted(m["position"] for m in msgs)
assert positions == [0, 1], (
f"concurrent appends raced; got positions {positions}"
)
finally:
# Clean up — the conversation was committed, so the transaction
# rollback won't drop it.
with pg_engine.begin() as cleanup_conn:
ConversationsRepository(cleanup_conn).delete(
conv["id"], "user-concurrent"
)
ConversationsRepository(cleanup_conn).delete(
conv["id"], "user-concurrent"
)
class TestSharedWith:
def test_add_shared_user_by_uuid(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("owner", "c")
assert repo.add_shared_user(conv["id"], "bob") is True
fetched = repo.get(conv["id"], "bob")
assert fetched is not None
assert "bob" in fetched["shared_with"]
def test_add_shared_user_is_idempotent(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("owner", "c")
assert repo.add_shared_user(conv["id"], "bob") is True
# Second call is a no-op (mirrors Mongo $addToSet semantics).
assert repo.add_shared_user(conv["id"], "bob") is False
fetched = repo.get(conv["id"], "bob")
assert fetched["shared_with"].count("bob") == 1
def test_add_shared_user_by_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create(
"owner", "c", legacy_mongo_id="507f1f77bcf86cd799439abc"
)
assert repo.add_shared_user("507f1f77bcf86cd799439abc", "bob") is True
fetched = repo.get(conv["id"], "bob")
assert "bob" in fetched["shared_with"]
def test_add_shared_user_empty_user_returns_false(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("owner", "c")
assert repo.add_shared_user(conv["id"], "") is False
def test_remove_shared_user(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("owner", "c")
repo.add_shared_user(conv["id"], "bob")
repo.add_shared_user(conv["id"], "carol")
assert repo.remove_shared_user(conv["id"], "bob") is True
fetched = repo.get(conv["id"], "carol")
assert fetched["shared_with"] == ["carol"]
def test_remove_missing_user_returns_false(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("owner", "c")
assert repo.remove_shared_user(conv["id"], "bob") is False
def test_remove_shared_user_by_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create(
"owner", "c", legacy_mongo_id="507f1f77bcf86cd799439def"
)
repo.add_shared_user("507f1f77bcf86cd799439def", "bob")
assert repo.remove_shared_user("507f1f77bcf86cd799439def", "bob") is True
fetched = repo.get(conv["id"], "owner")
assert fetched["shared_with"] == []
class TestUuidShapeGate:
"""Regression: a non-UUID conversation id (e.g. a legacy Mongo
ObjectId still embedded in old client-side state) must never reach
``CAST(:id AS uuid)``. The cast raises ``InvalidTextRepresentation``
on the server and **aborts the enclosing Postgres transaction**,
making every subsequent query on the same connection fail. These
tests pin the conservative behaviour: return False/None/0 for
non-UUID input and leave the transaction usable."""
@staticmethod
def _assert_txn_alive(conn) -> None:
"""Subsequent trivial query must succeed — proves the txn wasn't
poisoned. This is the load-bearing assertion; "returns False"
alone is insufficient since the old code did exactly that while
leaving the txn dead."""
from sqlalchemy import text as _text
assert conn.execute(_text("SELECT 1")).scalar() == 1
def test_rename_rejects_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
assert repo.rename("507f1f77bcf86cd799439011", "user-1", "new") is False
self._assert_txn_alive(pg_conn)
def test_rename_uuid_path_still_works(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "old")
assert repo.rename(created["id"], "user-1", "new") is True
def test_delete_rejects_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
assert repo.delete("507f1f77bcf86cd799439011", "user-1") is False
self._assert_txn_alive(pg_conn)
def test_delete_uuid_path_still_works(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "c")
assert repo.delete(created["id"], "user-1") is True
def test_set_feedback_rejects_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
assert repo.set_feedback(
"507f1f77bcf86cd799439011", 0, {"text": "x"},
) is False
self._assert_txn_alive(pg_conn)
def test_truncate_after_rejects_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
assert repo.truncate_after("507f1f77bcf86cd799439011", 0) == 0
self._assert_txn_alive(pg_conn)
def test_set_shared_token_rejects_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
assert repo.set_shared_token(
"507f1f77bcf86cd799439011", "user-1", "tok",
) is False
self._assert_txn_alive(pg_conn)
def test_update_compression_metadata_rejects_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
assert repo.update_compression_metadata(
"507f1f77bcf86cd799439011", "user-1", {"is_compressed": True},
) is False
self._assert_txn_alive(pg_conn)
def test_set_compression_flags_rejects_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
assert repo.set_compression_flags(
"507f1f77bcf86cd799439011",
is_compressed=True,
last_compression_at="2026-01-01",
) is False
self._assert_txn_alive(pg_conn)
def test_append_compression_point_rejects_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
assert repo.append_compression_point(
"507f1f77bcf86cd799439011", {"summary": "x"}, max_points=3,
) is False
self._assert_txn_alive(pg_conn)
def test_get_message_at_rejects_legacy_id(self, pg_conn):
repo = _repo(pg_conn)
assert repo.get_message_at("507f1f77bcf86cd799439011", 0) is None
self._assert_txn_alive(pg_conn)
def test_rename_rejects_garbage(self, pg_conn):
repo = _repo(pg_conn)
assert repo.rename("not-an-id", "user-1", "new") is False
self._assert_txn_alive(pg_conn)
def test_get_message_at_uuid_path_still_works(self, pg_conn):
repo = _repo(pg_conn)
conv = repo.create("user-1", "c")
repo.append_message(conv["id"], {"prompt": "q", "response": "a"})
msg = repo.get_message_at(conv["id"], 0)
assert msg is not None
assert msg["prompt"] == "q"