YoVinchen
ecdb10c37a
Some checks failed
verify / verify (push) Has been cancelled
Separate local parsing from model indexing, bind review decisions to immutable manifests, persist vectors behind active profiles, and expose retrieval, chat, evaluation, and document workflows through the React workbench. Constraint: Live Bailian authentication currently fails for all three configured capabilities Rejected: Direct upload-to-embedding flow | bypasses local review and manifest binding Confidence: high Scope-risk: broad Directive: Keep private-data deployment blocked until authentication, RBAC, and separate database roles land Tested: make verify; fresh and replay Docker document smoke; worker recovery smoke; frozen synthetic evaluation; migration 0003-0004 roundtrip Not-tested: Successful live Bailian calls, OCR, real multi-user authorization
360 lines
10 KiB
Python
360 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import signal
|
|
import threading
|
|
import uuid
|
|
from collections.abc import Sequence
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import cast
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from app.core.config import Settings
|
|
from app.persistence.job_queue import (
|
|
BackgroundJob,
|
|
JobLease,
|
|
JobState,
|
|
LeaseHeartbeat,
|
|
LeaseLostError,
|
|
)
|
|
from app.worker import Worker, WorkerConfig, build_default_handlers, install_signal_handlers
|
|
|
|
NOW = datetime(2026, 7, 13, 8, 0, tzinfo=UTC)
|
|
JOB_ID = uuid.UUID("40000000-0000-0000-0000-000000000004")
|
|
RESOURCE_ID = uuid.UUID("50000000-0000-0000-0000-000000000005")
|
|
LEASE_TOKEN = uuid.UUID("60000000-0000-0000-0000-000000000006")
|
|
|
|
|
|
def _job(job_type: str = "KNOWN_JOB") -> BackgroundJob:
|
|
lease = JobLease(JOB_ID, "worker-a", LEASE_TOKEN)
|
|
return BackgroundJob(
|
|
id=JOB_ID,
|
|
job_type=job_type,
|
|
required_capability="embedding",
|
|
resource_type="document_version",
|
|
resource_id=RESOURCE_ID,
|
|
idempotency_key="job:one",
|
|
payload={"resource_id": str(RESOURCE_ID)},
|
|
stage="PROCESSING",
|
|
progress=0,
|
|
priority=0,
|
|
attempt=1,
|
|
max_attempts=3,
|
|
run_after=NOW,
|
|
lease_until=NOW + timedelta(seconds=60),
|
|
created_at=NOW,
|
|
updated_at=NOW,
|
|
lease=lease,
|
|
)
|
|
|
|
|
|
def _state(status: str) -> JobState:
|
|
return JobState(
|
|
job_id=JOB_ID,
|
|
status=status,
|
|
attempt=1,
|
|
max_attempts=3,
|
|
finished_at=NOW if status in {"SUCCEEDED", "FAILED"} else None,
|
|
)
|
|
|
|
|
|
class FakeQueue:
|
|
def __init__(self, job: BackgroundJob | None = None) -> None:
|
|
self.job = job
|
|
self.claim_count = 0
|
|
self.claim_called = threading.Event()
|
|
self.in_claim = False
|
|
self.heartbeat_count = 0
|
|
self.complete_leases: list[JobLease] = []
|
|
self.failures: list[tuple[JobLease, str, str, int]] = []
|
|
self.reap_count = 0
|
|
self.heartbeat_error: Exception | None = None
|
|
self.complete_error: Exception | None = None
|
|
|
|
def claim(
|
|
self,
|
|
*,
|
|
worker_id: str,
|
|
worker_capabilities: Sequence[str],
|
|
lease_seconds: int,
|
|
) -> BackgroundJob | None:
|
|
assert worker_id == "worker-a"
|
|
assert tuple(worker_capabilities) == ("embedding",)
|
|
assert lease_seconds == 1
|
|
self.claim_count += 1
|
|
self.claim_called.set()
|
|
self.in_claim = True
|
|
try:
|
|
result = self.job
|
|
self.job = None
|
|
return result
|
|
finally:
|
|
self.in_claim = False
|
|
|
|
def heartbeat(self, lease: JobLease, *, lease_seconds: int) -> LeaseHeartbeat:
|
|
assert lease == JobLease(JOB_ID, "worker-a", LEASE_TOKEN)
|
|
assert lease_seconds == 1
|
|
self.heartbeat_count += 1
|
|
if self.heartbeat_error is not None:
|
|
raise self.heartbeat_error
|
|
return LeaseHeartbeat(JOB_ID, NOW + timedelta(seconds=1))
|
|
|
|
def complete(self, lease: JobLease) -> JobState:
|
|
self.complete_leases.append(lease)
|
|
if self.complete_error is not None:
|
|
raise self.complete_error
|
|
return _state("SUCCEEDED")
|
|
|
|
def fail_or_retry(
|
|
self,
|
|
lease: JobLease,
|
|
*,
|
|
error_code: str,
|
|
error_message: str,
|
|
retry_delay_seconds: int,
|
|
) -> JobState:
|
|
self.failures.append((lease, error_code, error_message, retry_delay_seconds))
|
|
return _state("QUEUED")
|
|
|
|
def reap_expired(
|
|
self,
|
|
*,
|
|
lock_key: int,
|
|
batch_size: int = 100,
|
|
) -> tuple[JobState, ...]:
|
|
assert isinstance(lock_key, int)
|
|
assert batch_size == 10
|
|
self.reap_count += 1
|
|
return ()
|
|
|
|
|
|
def _config(
|
|
*,
|
|
capabilities: tuple[str, ...] = ("embedding",),
|
|
heartbeat_seconds: float = 0.01,
|
|
poll_seconds: float = 0.01,
|
|
reaper_batch_size: int = 10,
|
|
) -> WorkerConfig:
|
|
return WorkerConfig(
|
|
worker_id="worker-a",
|
|
capabilities=capabilities,
|
|
lease_seconds=1,
|
|
heartbeat_seconds=heartbeat_seconds,
|
|
poll_seconds=poll_seconds,
|
|
retry_delay_seconds=7,
|
|
reaper_interval_seconds=30.0,
|
|
reaper_batch_size=reaper_batch_size,
|
|
reaper_lock_key=42,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handler_runs_after_claim_transaction_and_completes_with_same_lease() -> None:
|
|
queue = FakeQueue(_job())
|
|
observed_payload: dict[str, object] = {}
|
|
|
|
async def handler(job: BackgroundJob) -> None:
|
|
assert queue.in_claim is False
|
|
observed_payload.update(job.payload)
|
|
|
|
worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})
|
|
|
|
worked = await worker.run_once()
|
|
|
|
assert worked is True
|
|
assert observed_payload == {"resource_id": str(RESOURCE_ID)}
|
|
assert queue.complete_leases == [JobLease(JOB_ID, "worker-a", LEASE_TOKEN)]
|
|
assert queue.failures == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_long_handler_is_heartbeated_before_completion() -> None:
|
|
queue = FakeQueue(_job())
|
|
|
|
async def handler(_job: BackgroundJob) -> None:
|
|
await asyncio.sleep(0.035)
|
|
|
|
worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})
|
|
|
|
await worker.run_once()
|
|
|
|
assert queue.heartbeat_count >= 2
|
|
assert queue.complete_leases == [JobLease(JOB_ID, "worker-a", LEASE_TOKEN)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heartbeat_lease_loss_cancels_handler_without_terminal_write() -> None:
|
|
queue = FakeQueue(_job())
|
|
queue.heartbeat_error = LeaseLostError("lease moved")
|
|
handler_cancelled = asyncio.Event()
|
|
|
|
async def handler(_job: BackgroundJob) -> None:
|
|
try:
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
handler_cancelled.set()
|
|
|
|
worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})
|
|
|
|
await worker.run_once()
|
|
|
|
assert handler_cancelled.is_set()
|
|
assert queue.complete_leases == []
|
|
assert queue.failures == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unknown_job_is_safely_failed_without_handler_execution() -> None:
|
|
job = _job("UNREGISTERED_JOB")
|
|
queue = FakeQueue(job)
|
|
worker = Worker(queue, _config(), handlers={})
|
|
|
|
await worker.run_once()
|
|
|
|
assert queue.complete_leases == []
|
|
assert queue.failures == [
|
|
(
|
|
job.lease,
|
|
"UNKNOWN_JOB_TYPE",
|
|
"No registered handler exists for this job type.",
|
|
7,
|
|
)
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handler_exception_is_redacted_before_queue_failure(
|
|
caplog: pytest.LogCaptureFixture,
|
|
) -> None:
|
|
job = _job()
|
|
queue = FakeQueue(job)
|
|
|
|
async def handler(_job: BackgroundJob) -> None:
|
|
raise RuntimeError("private-document-text database-password")
|
|
|
|
worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})
|
|
|
|
caplog.set_level(logging.ERROR, logger="geological_rag.worker")
|
|
await worker.run_once()
|
|
|
|
assert len(queue.failures) == 1
|
|
lease, code, message, delay = queue.failures[0]
|
|
assert lease == job.lease
|
|
assert code == "JOB_HANDLER_FAILED"
|
|
assert message == "Registered job handler failed."
|
|
assert "private-document-text" not in message
|
|
assert delay == 7
|
|
assert [getattr(record, "error_type", None) for record in caplog.records] == ["RuntimeError"]
|
|
assert "private-document-text" not in caplog.text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_completion_fence_rejection_does_not_retry_completed_handler() -> None:
|
|
queue = FakeQueue(_job())
|
|
queue.complete_error = LeaseLostError("lease moved")
|
|
calls = 0
|
|
|
|
async def handler(_job: BackgroundJob) -> None:
|
|
nonlocal calls
|
|
calls += 1
|
|
|
|
worker = Worker(queue, _config(), handlers={"KNOWN_JOB": handler})
|
|
|
|
await worker.run_once()
|
|
|
|
assert calls == 1
|
|
assert queue.failures == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reaper_is_rate_limited_by_monotonic_schedule() -> None:
|
|
queue = FakeQueue()
|
|
clock_value = 100.0
|
|
worker = Worker(
|
|
queue,
|
|
_config(),
|
|
handlers={},
|
|
monotonic=lambda: clock_value,
|
|
)
|
|
|
|
assert await worker.run_once() is False
|
|
assert await worker.run_once() is False
|
|
assert queue.reap_count == 1
|
|
|
|
clock_value = 131.0
|
|
assert await worker.run_once() is False
|
|
assert queue.reap_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stop_wakes_poll_and_prevents_new_claims() -> None:
|
|
queue = FakeQueue()
|
|
worker = Worker(queue, _config(poll_seconds=30.0), handlers={})
|
|
task = asyncio.create_task(worker.run())
|
|
|
|
assert await asyncio.to_thread(queue.claim_called.wait, 0.2)
|
|
worker.request_stop()
|
|
await asyncio.wait_for(task, timeout=0.2)
|
|
|
|
assert worker.stopping is True
|
|
assert queue.claim_count == 1
|
|
|
|
|
|
def test_sigterm_callback_requests_graceful_stop() -> None:
|
|
worker = Worker(FakeQueue(), _config(), handlers={})
|
|
loop = MagicMock(spec=asyncio.AbstractEventLoop)
|
|
|
|
installed = install_signal_handlers(worker, loop)
|
|
|
|
assert signal.SIGTERM in installed
|
|
sigterm_call = next(
|
|
call for call in loop.add_signal_handler.call_args_list if call.args[0] == signal.SIGTERM
|
|
)
|
|
sigterm_call.args[1]()
|
|
assert worker.stopping is True
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("field", "value"),
|
|
[
|
|
("capabilities", ()),
|
|
("capabilities", ("embedding", "embedding")),
|
|
("heartbeat_seconds", 0.34),
|
|
("heartbeat_seconds", 1.0),
|
|
("reaper_batch_size", 0),
|
|
],
|
|
)
|
|
def test_worker_config_rejects_unsafe_lease_or_routing_values(
|
|
field: str,
|
|
value: object,
|
|
) -> None:
|
|
with pytest.raises(ValueError):
|
|
if field == "capabilities":
|
|
assert isinstance(value, tuple)
|
|
_config(capabilities=cast(tuple[str, ...], value))
|
|
elif field == "heartbeat_seconds":
|
|
assert isinstance(value, float)
|
|
_config(heartbeat_seconds=value)
|
|
else:
|
|
assert isinstance(value, int)
|
|
_config(reaper_batch_size=value)
|
|
|
|
|
|
def test_default_handler_registry_is_capability_isolated(tmp_path: Path) -> None:
|
|
settings = Settings(upload_root=tmp_path.resolve())
|
|
|
|
local_handlers = build_default_handlers(settings, ("document_parse",))
|
|
|
|
assert set(local_handlers) == {"PARSE_DOCUMENT"}
|
|
|
|
|
|
def test_default_handler_registry_rejects_unimplemented_capabilities(tmp_path: Path) -> None:
|
|
settings = Settings(upload_root=tmp_path.resolve())
|
|
|
|
with pytest.raises(RuntimeError, match="unsupported worker capabilities: evaluation"):
|
|
build_default_handlers(settings, ("evaluation",))
|