wehub-resource-sync
416 lines
16 KiB
Python
416 lines
16 KiB
Python
"""Tests for the session epoch counter.
|
|
|
|
Covers schema migration of the ``epoch`` column, NULL-row repair on
|
|
connect, reset-increments-epoch, stale-write rejection (single + concurrent),
|
|
SCHEMA_VERSION visibility, and epoch injection into event payloads.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from opensquilla.compat import aiosqlite
|
|
from opensquilla.session.models import SessionNode, TranscriptEntry
|
|
from opensquilla.session.storage import SCHEMA_VERSION, SessionStorage, StaleEpochError
|
|
|
|
# ── fixtures ────────────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def storage():
|
|
s = SessionStorage(":memory:")
|
|
await s.connect()
|
|
yield s
|
|
await s.close()
|
|
|
|
|
|
async def _make_session(storage: SessionStorage, key: str = "agent:main:test") -> SessionNode:
|
|
node = SessionNode(session_key=key, session_id="sid-" + key.split(":")[-1])
|
|
await storage.upsert_session(node)
|
|
return node
|
|
|
|
|
|
# ── SCHEMA_VERSION constant ─────────────────────────────────────────────────
|
|
|
|
|
|
def test_schema_version_constant_visible():
|
|
"""SCHEMA_VERSION is importable and >= 2 (epoch migration bumped it)."""
|
|
assert isinstance(SCHEMA_VERSION, int)
|
|
assert SCHEMA_VERSION >= 2
|
|
|
|
|
|
# ── epoch column default ────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_epoch_column_exists_with_default_zero(storage):
|
|
"""New databases have the epoch column defaulting to 0."""
|
|
await _make_session(storage)
|
|
epoch = await storage.get_epoch("agent:main:test")
|
|
assert epoch == 0
|
|
|
|
|
|
# ── old DB migrates ─────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_old_db_migrates():
|
|
"""A pre-epoch database (no epoch column) is migrated transparently on connect."""
|
|
conn = await aiosqlite.connect(":memory:")
|
|
conn.row_factory = aiosqlite.Row
|
|
await conn.execute("PRAGMA journal_mode=WAL")
|
|
# Create old-style sessions table without epoch column
|
|
await conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS sessions (
|
|
session_key TEXT PRIMARY KEY,
|
|
session_id TEXT NOT NULL,
|
|
created_at INTEGER NOT NULL DEFAULT 0,
|
|
updated_at INTEGER NOT NULL DEFAULT 0,
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens_fresh INTEGER NOT NULL DEFAULT 0,
|
|
estimated_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
total_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
billed_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
estimated_cost_component_usd REAL NOT NULL DEFAULT 0.0,
|
|
cost_source TEXT NOT NULL DEFAULT 'none',
|
|
missing_cost_entries INTEGER NOT NULL DEFAULT 0,
|
|
cache_read INTEGER NOT NULL DEFAULT 0,
|
|
cache_write INTEGER NOT NULL DEFAULT 0,
|
|
compaction_count INTEGER NOT NULL DEFAULT 0,
|
|
forked_from_parent INTEGER NOT NULL DEFAULT 0,
|
|
spawn_depth INTEGER NOT NULL DEFAULT 0,
|
|
status TEXT NOT NULL DEFAULT 'running',
|
|
chat_type TEXT NOT NULL DEFAULT 'unknown',
|
|
fast_mode INTEGER NOT NULL DEFAULT 0,
|
|
send_policy TEXT NOT NULL DEFAULT 'allow',
|
|
queue_mode TEXT NOT NULL DEFAULT 'steer',
|
|
agent_id TEXT NOT NULL DEFAULT 'main',
|
|
schema_version INTEGER NOT NULL DEFAULT 1
|
|
)
|
|
"""
|
|
)
|
|
# Insert a row without epoch
|
|
await conn.execute(
|
|
"INSERT INTO sessions (session_key, session_id, created_at, updated_at) "
|
|
"VALUES (?, ?, 0, 0)",
|
|
("agent:main:old", "sid-old"),
|
|
)
|
|
await conn.commit()
|
|
await conn.close()
|
|
|
|
# Now open via SessionStorage — migration must add epoch column
|
|
# We can't easily pass an existing connection, so use a temp file.
|
|
import os
|
|
import tempfile
|
|
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
|
|
db_path = f.name
|
|
try:
|
|
# Build the old-style DB on disk
|
|
conn2 = await aiosqlite.connect(db_path)
|
|
await conn2.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS sessions (
|
|
session_key TEXT PRIMARY KEY,
|
|
session_id TEXT NOT NULL,
|
|
created_at INTEGER NOT NULL DEFAULT 0,
|
|
updated_at INTEGER NOT NULL DEFAULT 0,
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens_fresh INTEGER NOT NULL DEFAULT 0,
|
|
estimated_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
total_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
billed_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
estimated_cost_component_usd REAL NOT NULL DEFAULT 0.0,
|
|
cost_source TEXT NOT NULL DEFAULT 'none',
|
|
missing_cost_entries INTEGER NOT NULL DEFAULT 0,
|
|
cache_read INTEGER NOT NULL DEFAULT 0,
|
|
cache_write INTEGER NOT NULL DEFAULT 0,
|
|
compaction_count INTEGER NOT NULL DEFAULT 0,
|
|
forked_from_parent INTEGER NOT NULL DEFAULT 0,
|
|
spawn_depth INTEGER NOT NULL DEFAULT 0,
|
|
status TEXT NOT NULL DEFAULT 'running',
|
|
chat_type TEXT NOT NULL DEFAULT 'unknown',
|
|
fast_mode INTEGER NOT NULL DEFAULT 0,
|
|
send_policy TEXT NOT NULL DEFAULT 'allow',
|
|
queue_mode TEXT NOT NULL DEFAULT 'steer',
|
|
agent_id TEXT NOT NULL DEFAULT 'main',
|
|
schema_version INTEGER NOT NULL DEFAULT 1
|
|
)
|
|
"""
|
|
)
|
|
await conn2.execute(
|
|
"INSERT INTO sessions (session_key, session_id, created_at, updated_at) "
|
|
"VALUES (?, ?, 0, 0)",
|
|
("agent:main:old", "sid-old"),
|
|
)
|
|
await conn2.commit()
|
|
await conn2.close()
|
|
|
|
# Open with SessionStorage — triggers migration
|
|
s = SessionStorage(db_path)
|
|
await s.connect()
|
|
try:
|
|
# epoch column must exist and default to 0
|
|
epoch = await s.get_epoch("agent:main:old")
|
|
assert epoch == 0, f"Expected 0, got {epoch}"
|
|
finally:
|
|
await s.close()
|
|
finally:
|
|
os.unlink(db_path)
|
|
|
|
|
|
# ── reset increments epoch ──────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_increments_epoch(storage):
|
|
"""increment_epoch raises epoch by 1 on each call."""
|
|
await _make_session(storage)
|
|
key = "agent:main:test"
|
|
|
|
assert await storage.get_epoch(key) == 0
|
|
|
|
e1 = await storage.increment_epoch(key)
|
|
assert e1 == 1
|
|
|
|
e2 = await storage.increment_epoch(key)
|
|
assert e2 == 2
|
|
|
|
e3 = await storage.increment_epoch(key)
|
|
assert e3 == 3
|
|
|
|
|
|
# ── startup NULL-row guard ──────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_null_epoch_rows_zeroed_on_connect():
|
|
"""If a row somehow has NULL epoch, _migrate_epoch_column zeros it."""
|
|
import os
|
|
import tempfile
|
|
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
|
|
db_path = f.name
|
|
try:
|
|
# Build a DB that already has the epoch column but with a NULL value
|
|
conn = await aiosqlite.connect(db_path)
|
|
await conn.execute(
|
|
"""
|
|
CREATE TABLE sessions (
|
|
session_key TEXT PRIMARY KEY,
|
|
session_id TEXT NOT NULL,
|
|
created_at INTEGER NOT NULL DEFAULT 0,
|
|
updated_at INTEGER NOT NULL DEFAULT 0,
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens_fresh INTEGER NOT NULL DEFAULT 0,
|
|
estimated_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
total_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
billed_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
estimated_cost_component_usd REAL NOT NULL DEFAULT 0.0,
|
|
cost_source TEXT NOT NULL DEFAULT 'none',
|
|
missing_cost_entries INTEGER NOT NULL DEFAULT 0,
|
|
cache_read INTEGER NOT NULL DEFAULT 0,
|
|
cache_write INTEGER NOT NULL DEFAULT 0,
|
|
compaction_count INTEGER NOT NULL DEFAULT 0,
|
|
forked_from_parent INTEGER NOT NULL DEFAULT 0,
|
|
spawn_depth INTEGER NOT NULL DEFAULT 0,
|
|
status TEXT NOT NULL DEFAULT 'running',
|
|
chat_type TEXT NOT NULL DEFAULT 'unknown',
|
|
fast_mode INTEGER NOT NULL DEFAULT 0,
|
|
send_policy TEXT NOT NULL DEFAULT 'allow',
|
|
queue_mode TEXT NOT NULL DEFAULT 'steer',
|
|
agent_id TEXT NOT NULL DEFAULT 'main',
|
|
schema_version INTEGER NOT NULL DEFAULT 1,
|
|
epoch INTEGER
|
|
)
|
|
"""
|
|
)
|
|
# Insert with explicit NULL epoch
|
|
await conn.execute(
|
|
"INSERT INTO sessions (session_key, session_id, created_at, updated_at, epoch) "
|
|
"VALUES (?, ?, 0, 0, NULL)",
|
|
("agent:main:nullrow", "sid-null"),
|
|
)
|
|
await conn.commit()
|
|
await conn.close()
|
|
|
|
s = SessionStorage(db_path)
|
|
await s.connect()
|
|
try:
|
|
epoch = await s.get_epoch("agent:main:nullrow")
|
|
assert epoch == 0
|
|
finally:
|
|
await s.close()
|
|
finally:
|
|
os.unlink(db_path)
|
|
|
|
|
|
# ── stale write rejected ────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stale_write_rejected(storage):
|
|
"""A write using a stale epoch (pre-reset) returns 0 rows affected.
|
|
|
|
Simulates a write that targets epoch=0 but the session has been reset
|
|
(epoch=1). The WHERE epoch = ? clause must match 0 rows.
|
|
"""
|
|
await _make_session(storage)
|
|
key = "agent:main:test"
|
|
|
|
# Advance epoch to 1 (one reset)
|
|
await storage.increment_epoch(key)
|
|
current_epoch = await storage.get_epoch(key)
|
|
assert current_epoch == 1
|
|
|
|
# Simulate a stale write: writer holds old_epoch=0 but session is at 1.
|
|
old_epoch = 0
|
|
async with storage.conn.execute(
|
|
"UPDATE sessions SET status = 'done' WHERE session_key = ? AND epoch = ?",
|
|
(key, old_epoch),
|
|
) as cur:
|
|
rows_affected = cur.rowcount
|
|
|
|
assert rows_affected == 0, "Stale write must not modify rows"
|
|
|
|
# Session status must remain unchanged
|
|
node = await storage.get_session(key)
|
|
assert node is not None
|
|
assert node.status == "running"
|
|
|
|
|
|
# ── concurrent stale writes all rejected ────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_stale_writes_all_rejected(storage):
|
|
"""Three concurrent append_transcript_entry calls with stale epoch are all rejected.
|
|
|
|
Models: turn A, B, C are in-flight capturing epoch=0; a reset fires (epoch→1);
|
|
all three must raise StaleEpochError via the real storage write path.
|
|
"""
|
|
node = await _make_session(storage)
|
|
key = "agent:main:test"
|
|
|
|
# Snapshot the epoch all three in-flight writers captured
|
|
captured_epoch = await storage.get_epoch(key) # 0
|
|
|
|
# A reset fires — epoch is now 1
|
|
await storage.increment_epoch(key)
|
|
post_reset_epoch = await storage.get_epoch(key)
|
|
assert post_reset_epoch == 1
|
|
|
|
def _make_entry(label: str) -> TranscriptEntry:
|
|
return TranscriptEntry(
|
|
session_id=node.session_id,
|
|
session_key=key,
|
|
role="assistant",
|
|
content=f"stale content from {label}",
|
|
)
|
|
|
|
async def _stale_write(label: str) -> bool:
|
|
"""Returns True if StaleEpochError was raised (write correctly rejected)."""
|
|
try:
|
|
await storage.append_transcript_entry(
|
|
_make_entry(label), expected_epoch=captured_epoch
|
|
)
|
|
return False # write was NOT rejected — bad
|
|
except StaleEpochError:
|
|
return True # correctly rejected
|
|
|
|
results = await asyncio.gather(
|
|
_stale_write("write-A"),
|
|
_stale_write("write-B"),
|
|
_stale_write("write-C"),
|
|
)
|
|
|
|
assert results == [True, True, True], (
|
|
f"All stale writes must be rejected; got {results}"
|
|
)
|
|
|
|
# No transcript entries must have been written
|
|
entries = await storage.get_transcript(node.session_id)
|
|
assert entries == [], f"No stale entries should persist; got {entries}"
|
|
|
|
|
|
# ── epoch in event payload ──────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_epoch_in_event_payload(storage):
|
|
"""_emit_to_subscribers injects 'epoch' into session.event.* payloads.
|
|
|
|
Uses a minimal fake ctx with real SessionStorage so that get_epoch()
|
|
returns the live value and confirms it appears in the payload sent to
|
|
the WS connection.
|
|
"""
|
|
from types import SimpleNamespace
|
|
|
|
from opensquilla.gateway.rpc_sessions import _emit_to_subscribers
|
|
|
|
node = await _make_session(storage)
|
|
key = node.session_key
|
|
|
|
# Advance epoch to 3 so we can assert the injected value is non-zero.
|
|
for _ in range(3):
|
|
await storage.increment_epoch(key)
|
|
expected_epoch = await storage.get_epoch(key)
|
|
assert expected_epoch == 3
|
|
|
|
# Capture sent payloads via a fake connection.
|
|
sent: list[tuple[str, dict]] = []
|
|
|
|
class FakeConn:
|
|
async def send_event(self, event_name: str, payload: dict) -> None:
|
|
sent.append((event_name, payload))
|
|
|
|
fake_conn = FakeConn()
|
|
|
|
class FakeRegistry:
|
|
def get(self, conn_id: str):
|
|
return fake_conn
|
|
|
|
class FakeSubMgr:
|
|
def get_message_subscribers(self, session_key: str):
|
|
return {"conn-1"}
|
|
|
|
def get_session_subscribers(self):
|
|
return set()
|
|
|
|
class FakeSessionMgr:
|
|
_storage = storage
|
|
|
|
fake_ctx = SimpleNamespace(
|
|
subscription_manager=FakeSubMgr(),
|
|
session_manager=FakeSessionMgr(),
|
|
)
|
|
|
|
import opensquilla.gateway.websocket as _ws_module
|
|
original_get_registry = _ws_module.get_registry
|
|
|
|
def _patched_registry():
|
|
return FakeRegistry()
|
|
|
|
_ws_module.get_registry = _patched_registry
|
|
try:
|
|
await _emit_to_subscribers(
|
|
fake_ctx, key, "session.event.text_delta", {"text": "hello"}
|
|
)
|
|
finally:
|
|
_ws_module.get_registry = original_get_registry
|
|
|
|
assert sent, "Expected at least one event to be delivered"
|
|
_event_name, payload = sent[0]
|
|
assert "epoch" in payload, f"epoch must be injected into payload; got {payload}"
|
|
assert payload["epoch"] == expected_epoch, (
|
|
f"epoch in payload ({payload['epoch']}) must match storage epoch ({expected_epoch})"
|
|
)
|