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:
RockChinQ
2026-09-21 17:05:40 +00:00
parent a8d291e581
commit 08a6ed0fdd
6 changed files with 307 additions and 36 deletions
+1
View File
@@ -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