mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-09-26 19:36:35 +08:00
fix(runner): normalize PostgreSQL journal timestamps to UTC-naive
Preserve the published timezone-less schema and epoch API contract across run lifecycle, deadlines, leases, heartbeat registry, events, transcripts and retention cutoffs. Exercise the real Host journal with asyncpg under a non-superuser role on metadata and published-migration schemas; 32 PostgreSQL regressions fail before the fix and pass after it.
This commit is contained in:
@@ -232,6 +232,7 @@ Persistence is centered on `pkg/persistence/mgr.py`.
|
||||
|
||||
- SQLite is the default database; PostgreSQL is supported.
|
||||
- Models live under `pkg/entity/persistence/`.
|
||||
- Timezone-less SQL `DateTime` columns store UTC-naive values on both backends. Normalize aware values with `pkg/persistence/datetime_utils.py::as_naive_utc` at the bind boundary, including deadlines, leases, event times, and query/retention cutoffs; do not merely strip an offset. Existing naive rows represent UTC. Restore UTC awareness for application comparisons and epoch serialization. Runner ledger/event/transcript stores follow this existing schema contract without a data migration or changes to journal/tenant authorization.
|
||||
- Fresh schemas are created from current metadata, then Alembic migrations run to head. LangBot 4.x does not upgrade 3.x databases.
|
||||
- New schema changes use Alembic under `pkg/persistence/alembic/versions/`; there is no legacy migration chain in 4.x.
|
||||
|
||||
|
||||
@@ -11,6 +11,8 @@ import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from ...persistence.datetime_utils import as_naive_utc
|
||||
|
||||
from ...entity.persistence.event_log import EventLog
|
||||
|
||||
|
||||
@@ -106,7 +108,7 @@ class EventLogStore:
|
||||
event = EventLog(
|
||||
event_id=event_id,
|
||||
event_type=event_type,
|
||||
event_time=event_time,
|
||||
event_time=as_naive_utc(event_time),
|
||||
source=source,
|
||||
bot_id=bot_id,
|
||||
workspace_id=workspace_id,
|
||||
@@ -123,7 +125,7 @@ class EventLogStore:
|
||||
run_id=run_id,
|
||||
runner_id=runner_id,
|
||||
metadata_json=json.dumps(metadata) if metadata else None,
|
||||
created_at=_utc_now(),
|
||||
created_at=as_naive_utc(_utc_now()),
|
||||
)
|
||||
session.add(event)
|
||||
await session.commit()
|
||||
@@ -279,7 +281,9 @@ class EventLogStore:
|
||||
) -> int:
|
||||
"""Delete EventLog rows created before the supplied timestamp."""
|
||||
async with self._session_factory() as session:
|
||||
result = await session.execute(sqlalchemy.delete(EventLog).where(EventLog.created_at < before))
|
||||
result = await session.execute(
|
||||
sqlalchemy.delete(EventLog).where(EventLog.created_at < as_naive_utc(before))
|
||||
)
|
||||
await session.commit()
|
||||
return result.rowcount or 0
|
||||
|
||||
|
||||
@@ -11,6 +11,8 @@ import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from ...persistence.datetime_utils import as_naive_utc
|
||||
|
||||
from ...entity.persistence.agent_run import AgentRun, AgentRunEvent, AgentRuntime
|
||||
|
||||
|
||||
@@ -45,7 +47,7 @@ def _epoch_to_datetime(value: typing.Any) -> datetime.datetime | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return datetime.datetime.fromtimestamp(float(value), UTC)
|
||||
return as_naive_utc(datetime.datetime.fromtimestamp(float(value), UTC))
|
||||
except (TypeError, ValueError, OSError):
|
||||
return None
|
||||
|
||||
@@ -138,9 +140,9 @@ class RunLedgerStore:
|
||||
queue_name=queue_name,
|
||||
priority=priority,
|
||||
requested_runtime_id=requested_runtime_id,
|
||||
created_at=now,
|
||||
started_at=now if status == 'running' else None,
|
||||
updated_at=now,
|
||||
created_at=as_naive_utc(now),
|
||||
started_at=as_naive_utc(now) if status == 'running' else None,
|
||||
updated_at=as_naive_utc(now),
|
||||
deadline_at=_epoch_to_datetime(deadline_at),
|
||||
authorization_json=_json_dumps(authorization),
|
||||
metadata_json=_json_dumps(metadata),
|
||||
@@ -172,7 +174,7 @@ class RunLedgerStore:
|
||||
sqlalchemy.and_(
|
||||
AgentRun.status == 'claimed',
|
||||
AgentRun.claim_lease_expires_at.is_not(None),
|
||||
AgentRun.claim_lease_expires_at <= now,
|
||||
AgentRun.claim_lease_expires_at <= as_naive_utc(now),
|
||||
),
|
||||
),
|
||||
sqlalchemy.or_(
|
||||
@@ -199,10 +201,10 @@ class RunLedgerStore:
|
||||
run.status = 'claimed'
|
||||
run.claimed_by_runtime_id = runtime_id
|
||||
run.claim_token = uuid.uuid4().hex
|
||||
run.claim_lease_expires_at = lease_expires_at
|
||||
run.claim_lease_expires_at = as_naive_utc(lease_expires_at)
|
||||
run.dispatch_attempts = (run.dispatch_attempts or 0) + 1
|
||||
run.last_claimed_at = now
|
||||
run.updated_at = now
|
||||
run.last_claimed_at = as_naive_utc(now)
|
||||
run.updated_at = as_naive_utc(now)
|
||||
await session.commit()
|
||||
return self._run_to_dict(run, include_claim_token=True)
|
||||
|
||||
@@ -221,8 +223,8 @@ class RunLedgerStore:
|
||||
if run is None or not _claim_is_active(run, runtime_id=runtime_id, claim_token=claim_token, now=now):
|
||||
return None
|
||||
|
||||
run.claim_lease_expires_at = now + datetime.timedelta(seconds=max(int(lease_seconds), 1))
|
||||
run.updated_at = now
|
||||
run.claim_lease_expires_at = as_naive_utc(now + datetime.timedelta(seconds=max(int(lease_seconds), 1)))
|
||||
run.updated_at = as_naive_utc(now)
|
||||
await session.commit()
|
||||
return self._run_to_dict(run)
|
||||
|
||||
@@ -248,9 +250,9 @@ class RunLedgerStore:
|
||||
run.claimed_by_runtime_id = None
|
||||
run.claim_token = None
|
||||
run.claim_lease_expires_at = None
|
||||
run.updated_at = now
|
||||
run.updated_at = as_naive_utc(now)
|
||||
if status in TERMINAL_STATUSES:
|
||||
run.finished_at = run.finished_at or now
|
||||
run.finished_at = run.finished_at or as_naive_utc(now)
|
||||
await session.commit()
|
||||
return self._run_to_dict(run)
|
||||
|
||||
@@ -264,9 +266,7 @@ class RunLedgerStore:
|
||||
) -> list[dict[str, typing.Any]]:
|
||||
"""Release claimed runs whose claim lease has expired."""
|
||||
status = _validate_run_status(status)
|
||||
current_time = now or _utc_now()
|
||||
if current_time.tzinfo is None:
|
||||
current_time = current_time.replace(tzinfo=UTC)
|
||||
current_time = as_naive_utc(now or _utc_now())
|
||||
limit = min(max(int(limit), 1), 500)
|
||||
|
||||
async with self._session_factory() as session:
|
||||
@@ -326,7 +326,7 @@ class RunLedgerStore:
|
||||
type=event_type,
|
||||
data_json=_json_dumps(data or {}),
|
||||
usage_json=_json_dumps(usage),
|
||||
created_at=_utc_now(),
|
||||
created_at=as_naive_utc(_utc_now()),
|
||||
source=source,
|
||||
metadata_json=_json_dumps(metadata),
|
||||
)
|
||||
@@ -360,7 +360,7 @@ class RunLedgerStore:
|
||||
type=event_type,
|
||||
data_json=_json_dumps(data or {}),
|
||||
usage_json=None,
|
||||
created_at=_utc_now(),
|
||||
created_at=as_naive_utc(_utc_now()),
|
||||
source='host',
|
||||
metadata_json=_json_dumps(metadata or {}),
|
||||
)
|
||||
@@ -392,9 +392,9 @@ class RunLedgerStore:
|
||||
run.status = status
|
||||
if status_reason is not None:
|
||||
run.status_reason = status_reason
|
||||
run.updated_at = now
|
||||
run.updated_at = as_naive_utc(now)
|
||||
if status in TERMINAL_STATUSES:
|
||||
run.finished_at = run.finished_at or now
|
||||
run.finished_at = run.finished_at or as_naive_utc(now)
|
||||
run.claimed_by_runtime_id = None
|
||||
run.claim_token = None
|
||||
run.claim_lease_expires_at = None
|
||||
@@ -439,8 +439,8 @@ class RunLedgerStore:
|
||||
run = await self._get_run_row(session, run_id)
|
||||
if run is None:
|
||||
return None
|
||||
run.cancel_requested_at = now
|
||||
run.updated_at = now
|
||||
run.cancel_requested_at = as_naive_utc(now)
|
||||
run.updated_at = as_naive_utc(now)
|
||||
run.status_reason = status_reason or run.status_reason
|
||||
await session.commit()
|
||||
return self._run_to_dict(run)
|
||||
@@ -469,7 +469,7 @@ class RunLedgerStore:
|
||||
async with self._session_factory() as session:
|
||||
runtime = await self._get_runtime_row(session, runtime_id)
|
||||
if runtime is None:
|
||||
runtime = AgentRuntime(runtime_id=runtime_id, created_at=now)
|
||||
runtime = AgentRuntime(runtime_id=runtime_id, created_at=as_naive_utc(now))
|
||||
session.add(runtime)
|
||||
|
||||
runtime.status = status
|
||||
@@ -479,9 +479,11 @@ class RunLedgerStore:
|
||||
runtime.capabilities_json = _json_dumps(capabilities or {})
|
||||
runtime.labels_json = _json_dumps(labels or {})
|
||||
runtime.metadata_json = _json_dumps(metadata or {})
|
||||
runtime.last_heartbeat_at = now
|
||||
runtime.heartbeat_deadline_at = now + datetime.timedelta(seconds=max(int(heartbeat_deadline_seconds), 1))
|
||||
runtime.updated_at = now
|
||||
runtime.last_heartbeat_at = as_naive_utc(now)
|
||||
runtime.heartbeat_deadline_at = as_naive_utc(
|
||||
now + datetime.timedelta(seconds=max(int(heartbeat_deadline_seconds), 1))
|
||||
)
|
||||
runtime.updated_at = as_naive_utc(now)
|
||||
await session.commit()
|
||||
return self._runtime_to_dict(runtime)
|
||||
|
||||
@@ -503,9 +505,11 @@ class RunLedgerStore:
|
||||
return None
|
||||
|
||||
runtime.status = status
|
||||
runtime.last_heartbeat_at = now
|
||||
runtime.heartbeat_deadline_at = now + datetime.timedelta(seconds=max(int(heartbeat_deadline_seconds), 1))
|
||||
runtime.updated_at = now
|
||||
runtime.last_heartbeat_at = as_naive_utc(now)
|
||||
runtime.heartbeat_deadline_at = as_naive_utc(
|
||||
now + datetime.timedelta(seconds=max(int(heartbeat_deadline_seconds), 1))
|
||||
)
|
||||
runtime.updated_at = as_naive_utc(now)
|
||||
if capabilities is not None:
|
||||
runtime.capabilities_json = _json_dumps(capabilities)
|
||||
if labels is not None:
|
||||
@@ -588,9 +592,7 @@ class RunLedgerStore:
|
||||
stale_after_seconds: int | float | None = None,
|
||||
) -> list[dict[str, typing.Any]]:
|
||||
"""Mark runtimes stale when their heartbeat deadline has passed."""
|
||||
current_time = now or _utc_now()
|
||||
if current_time.tzinfo is None:
|
||||
current_time = current_time.replace(tzinfo=UTC)
|
||||
current_time = as_naive_utc(now or _utc_now())
|
||||
stale_conditions: list[typing.Any] = [
|
||||
sqlalchemy.and_(
|
||||
AgentRuntime.heartbeat_deadline_at.is_not(None),
|
||||
|
||||
@@ -11,6 +11,8 @@ import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from ...persistence.datetime_utils import as_naive_utc
|
||||
|
||||
from ...entity.persistence.transcript import Transcript
|
||||
from langbot_plugin.api.entities.builtin.provider import message as provider_message
|
||||
|
||||
@@ -110,7 +112,7 @@ class TranscriptStore:
|
||||
seq=0,
|
||||
run_id=run_id,
|
||||
runner_id=runner_id,
|
||||
created_at=_utc_now(),
|
||||
created_at=as_naive_utc(_utc_now()),
|
||||
metadata_json=json.dumps(metadata) if metadata else None,
|
||||
)
|
||||
session.add(item)
|
||||
@@ -350,7 +352,9 @@ class TranscriptStore:
|
||||
) -> int:
|
||||
"""Delete Transcript rows created before the supplied timestamp."""
|
||||
async with self._session_factory() as session:
|
||||
result = await session.execute(sqlalchemy.delete(Transcript).where(Transcript.created_at < before))
|
||||
result = await session.execute(
|
||||
sqlalchemy.delete(Transcript).where(Transcript.created_at < as_naive_utc(before))
|
||||
)
|
||||
await session.commit()
|
||||
return result.rowcount or 0
|
||||
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
"""Datetime values bound to LangBot's timezone-less SQL columns."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
|
||||
|
||||
def as_naive_utc(value: datetime.datetime | None) -> datetime.datetime | None:
|
||||
"""Keep the existing UTC-naive storage contract on SQLite and PostgreSQL.
|
||||
|
||||
Legacy naive values already represent UTC. Aware values must be converted
|
||||
to UTC *before* dropping tzinfo, including query cutoffs and deadlines.
|
||||
This is a bind-boundary conversion, not a schema or wire-format change.
|
||||
"""
|
||||
if value is None or value.tzinfo is None:
|
||||
return value
|
||||
return value.astimezone(datetime.timezone.utc).replace(tzinfo=None)
|
||||
@@ -0,0 +1,243 @@
|
||||
"""Runner timestamp contract against real asyncpg and published schema DDL.
|
||||
|
||||
Run only on a disposable PostgreSQL via TEST_POSTGRES_URL. Each case owns a
|
||||
unique schema and an unprivileged role; no existing tables or policies change.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
import importlib
|
||||
import os
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from alembic.migration import MigrationContext
|
||||
from alembic.operations import Operations
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.agent.runner import event_log_store, run_ledger_store, transcript_store
|
||||
from langbot.pkg.agent.runner.run_journal import AgentRunJournal
|
||||
from langbot.pkg.entity.persistence.agent_run import AgentRun, AgentRunEvent, AgentRuntime
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.event_log import EventLog
|
||||
from langbot.pkg.entity.persistence.transcript import Transcript
|
||||
|
||||
pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
|
||||
|
||||
UTC = dt.timezone.utc
|
||||
NOW = dt.datetime(2026, 9, 22, 4, 5, 6, 123000, tzinfo=UTC)
|
||||
TABLES = [AgentRun.__table__, AgentRunEvent.__table__, AgentRuntime.__table__, EventLog.__table__, Transcript.__table__]
|
||||
|
||||
|
||||
def _published_schema(connection):
|
||||
with Operations.context(MigrationContext.configure(connection)):
|
||||
for revision in ('58846a8d7a81_add_event_log_and_transcript_tables', '8d3a1f2c4b6e_add_agent_run_ledger'):
|
||||
importlib.import_module(f'langbot.pkg.persistence.alembic.versions.{revision}').upgrade()
|
||||
|
||||
|
||||
@pytest.fixture(params=['metadata', 'published_migrations'])
|
||||
async def pg_engine(request, monkeypatch):
|
||||
url = os.environ.get('TEST_POSTGRES_URL')
|
||||
if not url:
|
||||
pytest.skip('TEST_POSTGRES_URL is required (disposable PostgreSQL)')
|
||||
namespace = f'runner_ts_{uuid.uuid4().hex}'
|
||||
admin = create_async_engine(url)
|
||||
engine = None
|
||||
try:
|
||||
async with admin.begin() as conn:
|
||||
await conn.execute(sa.text(f'CREATE ROLE {namespace} NOLOGIN NOSUPERUSER NOBYPASSRLS'))
|
||||
await conn.execute(sa.text(f'CREATE SCHEMA {namespace} AUTHORIZATION {namespace}'))
|
||||
engine = create_async_engine(
|
||||
url,
|
||||
connect_args={
|
||||
'server_settings': {'search_path': namespace, 'role': namespace, 'timezone': 'Asia/Shanghai'}
|
||||
},
|
||||
)
|
||||
async with engine.begin() as conn:
|
||||
role = (
|
||||
await conn.execute(sa.text('SELECT rolsuper, rolbypassrls FROM pg_roles WHERE rolname = current_user'))
|
||||
).one()
|
||||
assert role == (False, False)
|
||||
if request.param == 'metadata':
|
||||
await conn.run_sync(lambda sync: Base.metadata.create_all(sync, tables=TABLES))
|
||||
else:
|
||||
await conn.run_sync(_published_schema)
|
||||
columns = (
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'SELECT table_name, column_name, data_type FROM information_schema.columns '
|
||||
"WHERE table_schema = current_schema() AND data_type LIKE 'timestamp%'"
|
||||
)
|
||||
)
|
||||
).all()
|
||||
assert columns
|
||||
assert all(column.data_type == 'timestamp without time zone' for column in columns)
|
||||
for module in (run_ledger_store, event_log_store, transcript_store):
|
||||
monkeypatch.setattr(module, '_utc_now', lambda: NOW)
|
||||
yield engine
|
||||
finally:
|
||||
if engine is not None:
|
||||
await engine.dispose()
|
||||
async with admin.begin() as conn:
|
||||
await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS {namespace} CASCADE'))
|
||||
await conn.execute(sa.text(f'DROP ROLE IF EXISTS {namespace}'))
|
||||
await admin.dispose()
|
||||
|
||||
|
||||
async def _legacy_run(engine, *, status='queued'):
|
||||
"""An existing UTC-naive row must remain readable/updatable without DDL."""
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(
|
||||
sa.insert(AgentRun).values(
|
||||
run_id='legacy',
|
||||
runner_id='runner',
|
||||
status=status,
|
||||
created_at=NOW.replace(tzinfo=None),
|
||||
updated_at=(NOW - dt.timedelta(minutes=1)).replace(tzinfo=None),
|
||||
claim_token='old-token' if status == 'claimed' else None,
|
||||
claimed_by_runtime_id='old-runtime' if status == 'claimed' else None,
|
||||
claim_lease_expires_at=(NOW - dt.timedelta(seconds=1)).replace(tzinfo=None)
|
||||
if status == 'claimed'
|
||||
else None,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def test_run_lifecycle_deadline_events_and_stats(pg_engine):
|
||||
store = run_ledger_store.RunLedgerStore(pg_engine)
|
||||
deadline = NOW + dt.timedelta(minutes=5)
|
||||
journal = AgentRunJournal(SimpleNamespace(persistence_mgr=SimpleNamespace(get_db_engine=lambda: pg_engine)))
|
||||
created = await journal.create_run(
|
||||
event=SimpleNamespace(
|
||||
event_id='event',
|
||||
conversation_id='conversation',
|
||||
thread_id=None,
|
||||
workspace_id='workspace-a',
|
||||
bot_id=None,
|
||||
event_type='message.received',
|
||||
source='debug',
|
||||
delivery=SimpleNamespace(model_dump=lambda **kwargs: {}),
|
||||
),
|
||||
binding=SimpleNamespace(binding_id='binding', agent_id='agent', processor_id='agent', processor_type='agent'),
|
||||
descriptor=SimpleNamespace(id='runner'),
|
||||
context={'run_id': 'debug-run', 'runtime': {'deadline_at': deadline.timestamp()}, 'input': {'text': 'test'}},
|
||||
authorization={},
|
||||
)
|
||||
assert created['created_at'] == int(NOW.timestamp())
|
||||
assert created['created_at_ms'] == round(NOW.timestamp() * 1000)
|
||||
assert created['started_at_ms'] == created['created_at_ms']
|
||||
assert created['deadline_at'] == int(deadline.timestamp())
|
||||
assert await store.get_run('debug-run') == created
|
||||
assert (
|
||||
await store.create_run(
|
||||
run_id='debug-run',
|
||||
event_id=None,
|
||||
binding_id=None,
|
||||
runner_id='runner',
|
||||
)
|
||||
) == created
|
||||
event = await store.append_event(
|
||||
run_id='debug-run', sequence=1, event_type='message.completed', data={'text': 'ok'}
|
||||
)
|
||||
assert event['created_at'] == int(NOW.timestamp())
|
||||
assert await store.append_event(run_id='debug-run', sequence=1, event_type='duplicate') == event
|
||||
audit = await store.append_audit_event(run_id='debug-run', event_type='host.cancel_requested')
|
||||
assert audit['sequence'] == 2
|
||||
assert audit['created_at'] == int(NOW.timestamp())
|
||||
cancelled = await store.request_cancel(run_id='debug-run')
|
||||
assert cancelled['cancel_requested_at'] == int(NOW.timestamp())
|
||||
finished = await store.finalize_run(run_id='debug-run', status='completed')
|
||||
assert finished['finished_at_ms'] == round(NOW.timestamp() * 1000)
|
||||
assert await store.get_run('debug-run') == finished
|
||||
events, _, _, has_more = await store.page_run_events(run_id='debug-run')
|
||||
assert [row['sequence'] for row in events] == [1, 2]
|
||||
assert not has_more
|
||||
window = {'start_time': int(NOW.timestamp()) - 1, 'end_time': int(NOW.timestamp()) + 1}
|
||||
assert (await store.get_run_stats(**window))['completed_count'] == 1
|
||||
assert (await store.get_runner_stats(**window))[0]['completed_runs'] == 1
|
||||
async with pg_engine.connect() as conn:
|
||||
raw = (await conn.execute(sa.select(AgentRun.created_at, AgentRun.deadline_at))).one()
|
||||
assert raw == (NOW.replace(tzinfo=None), deadline.replace(tzinfo=None))
|
||||
|
||||
|
||||
async def test_claim_renew_release_existing_naive_row(pg_engine):
|
||||
await _legacy_run(pg_engine)
|
||||
store = run_ledger_store.RunLedgerStore(pg_engine)
|
||||
claimed = await store.claim_next_run(runtime_id='runtime', lease_seconds=30)
|
||||
assert claimed['last_claimed_at'] == int(NOW.timestamp())
|
||||
token = claimed['claim_token']
|
||||
assert await store.validate_active_claim(run_id='legacy', runtime_id='runtime', claim_token=token)
|
||||
assert await store.renew_claim(run_id='legacy', claim_token='wrong') is None
|
||||
renewed = await store.renew_claim(run_id='legacy', claim_token=token, lease_seconds=90)
|
||||
assert renewed['claim_lease_expires_at'] == int(NOW.timestamp()) + 90
|
||||
released = await store.release_claim(run_id='legacy', claim_token=token, status='cancelled')
|
||||
assert released['finished_at'] == int(NOW.timestamp())
|
||||
assert released['claim_lease_expires_at'] is None
|
||||
assert await store.get_run('legacy') == released
|
||||
|
||||
|
||||
async def test_reclaim_expired_naive_lease(pg_engine):
|
||||
await _legacy_run(pg_engine, status='claimed')
|
||||
store = run_ledger_store.RunLedgerStore(pg_engine)
|
||||
assert not await store.validate_active_claim(run_id='legacy', runtime_id='old-runtime', claim_token='old-token')
|
||||
assert await store.release_claim(run_id='legacy', claim_token='old-token') is None
|
||||
claimed = await store.claim_next_run(runtime_id='replacement')
|
||||
assert claimed['claimed_by_runtime_id'] == 'replacement'
|
||||
assert claimed['claim_token'] != 'old-token'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('offset', [None, 0, 8, -7])
|
||||
async def test_expired_claim_cutoff_accepts_naive_and_offset_times(pg_engine, offset):
|
||||
await _legacy_run(pg_engine, status='claimed')
|
||||
store = run_ledger_store.RunLedgerStore(pg_engine)
|
||||
cutoff = NOW.replace(tzinfo=None) if offset is None else NOW.astimezone(dt.timezone(dt.timedelta(hours=offset)))
|
||||
released = await store.release_expired_claims(now=cutoff, status='timeout')
|
||||
assert len(released) == 1
|
||||
assert released[0]['updated_at'] == int(NOW.timestamp())
|
||||
assert released[0]['finished_at'] == int(NOW.timestamp())
|
||||
assert not await store.validate_active_claim(run_id='legacy', runtime_id='old-runtime', claim_token='old-token')
|
||||
|
||||
|
||||
@pytest.mark.parametrize('offset', [None, 8, -7])
|
||||
async def test_runtime_heartbeat_stale_cutoffs_and_stats(pg_engine, offset):
|
||||
store = run_ledger_store.RunLedgerStore(pg_engine)
|
||||
registered = await store.register_runtime(runtime_id='runtime', heartbeat_deadline_seconds=30)
|
||||
assert registered['last_heartbeat_at'] == int(NOW.timestamp())
|
||||
heartbeat = await store.heartbeat_runtime(runtime_id='runtime', heartbeat_deadline_seconds=60)
|
||||
assert heartbeat['heartbeat_deadline_at'] == int(NOW.timestamp()) + 60
|
||||
assert (await store.get_runtime_stats())['avg_heartbeat_age_seconds'] == 0
|
||||
later = NOW + dt.timedelta(seconds=61)
|
||||
cutoff = later.replace(tzinfo=None) if offset is None else later.astimezone(dt.timezone(dt.timedelta(hours=offset)))
|
||||
stale = await store.mark_stale_runtimes(now=cutoff, stale_after_seconds=60)
|
||||
assert len(stale) == 1
|
||||
assert stale[0]['updated_at'] == int(later.timestamp())
|
||||
assert (await store.get_runtime('runtime'))['status'] == 'stale'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('offset', [None, 8, -7])
|
||||
async def test_event_log_timestamp_and_retention(pg_engine, offset):
|
||||
store = event_log_store.EventLogStore(pg_engine)
|
||||
event_time = NOW.replace(tzinfo=None) if offset is None else NOW.astimezone(dt.timezone(dt.timedelta(hours=offset)))
|
||||
await store.append_event(event_id='event', event_type='message', source='test', event_time=event_time)
|
||||
async with pg_engine.connect() as conn:
|
||||
row = (await conn.execute(sa.select(EventLog.event_time, EventLog.created_at))).one()
|
||||
assert row == (NOW.replace(tzinfo=None), NOW.replace(tzinfo=None))
|
||||
assert await store.cleanup_events_older_than(event_time) == 0
|
||||
assert await store.cleanup_events_older_than(event_time + dt.timedelta(seconds=1)) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize('offset', [None, 8, -7])
|
||||
async def test_transcript_timestamp_and_retention(pg_engine, offset):
|
||||
store = transcript_store.TranscriptStore(pg_engine)
|
||||
await store.append_transcript(
|
||||
transcript_id='transcript', event_id='event', conversation_id='conversation', role='assistant'
|
||||
)
|
||||
async with pg_engine.connect() as conn:
|
||||
raw = (await conn.execute(sa.select(Transcript.created_at))).scalar_one()
|
||||
assert raw == NOW.replace(tzinfo=None)
|
||||
cutoff = NOW.replace(tzinfo=None) if offset is None else NOW.astimezone(dt.timezone(dt.timedelta(hours=offset)))
|
||||
assert await store.cleanup_transcripts_older_than(cutoff) == 0
|
||||
assert await store.cleanup_transcripts_older_than(cutoff + dt.timedelta(seconds=1)) == 1
|
||||
Reference in New Issue
Block a user