diff --git a/haiku_rag_slim/haiku/rag/ingester/api/__init__.py b/haiku_rag_slim/haiku/rag/ingester/api/__init__.py new file mode 100644 index 00000000..f01618fd --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/api/__init__.py @@ -0,0 +1,3 @@ +from haiku.rag.ingester.api.server import APIState, build_app + +__all__ = ["APIState", "build_app"] diff --git a/haiku_rag_slim/haiku/rag/ingester/api/auth.py b/haiku_rag_slim/haiku/rag/ingester/api/auth.py new file mode 100644 index 00000000..97341cd7 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/api/auth.py @@ -0,0 +1,22 @@ +from fastapi import Depends, HTTPException, Request, status +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer + +_bearer = HTTPBearer(auto_error=False) + + +async def require_auth( + request: Request, + creds: HTTPAuthorizationCredentials | None = Depends(_bearer), +) -> None: + """Bearer-token gate. When the app has no auth_token configured, all + requests are allowed (and a warning was logged at startup). Otherwise + the Authorization header's bearer must match exactly.""" + expected = getattr(request.app.state, "auth_token", None) + if expected is None: + return + if creds is None or creds.credentials != expected: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Unauthorized", + headers={"WWW-Authenticate": "Bearer"}, + ) diff --git a/haiku_rag_slim/haiku/rag/ingester/api/routes/__init__.py b/haiku_rag_slim/haiku/rag/ingester/api/routes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/haiku_rag_slim/haiku/rag/ingester/api/routes/dlq.py b/haiku_rag_slim/haiku/rag/ingester/api/routes/dlq.py new file mode 100644 index 00000000..9cef98e0 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/api/routes/dlq.py @@ -0,0 +1,32 @@ +from fastapi import APIRouter, Depends, HTTPException, status + +from haiku.rag.ingester.api.server import APIState, get_state +from haiku.rag.ingester.queue.models import Job, JobStatus + +router = APIRouter(prefix="/dlq", tags=["dlq"]) + + +@router.get("", response_model=list[Job]) +async def list_dlq( + source_id: str | None = None, + limit: int = 50, + offset: int = 0, + state: APIState = Depends(get_state), +) -> list[Job]: + """Jobs that exhausted retries or hit a permanent error.""" + return await state.job_repo.list_jobs( + status=JobStatus.DEAD, source_id=source_id, limit=limit, offset=offset + ) + + +@router.post("/{job_id}/retry", response_model=Job) +async def retry_from_dlq(job_id: str, state: APIState = Depends(get_state)) -> Job: + """Convenience alias for /jobs/{id}/retry — same effect, separate + endpoint so operator tooling can wire DLQ rescues without coupling + to the generic /jobs path.""" + try: + return await state.job_repo.retry(job_id) + except KeyError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="job not found" + ) from exc diff --git a/haiku_rag_slim/haiku/rag/ingester/api/routes/health.py b/haiku_rag_slim/haiku/rag/ingester/api/routes/health.py new file mode 100644 index 00000000..573dfb20 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/api/routes/health.py @@ -0,0 +1,23 @@ +from fastapi import APIRouter, Depends + +from haiku.rag.ingester.api.schemas import HealthResponse +from haiku.rag.ingester.api.server import APIState, get_state + +router = APIRouter() + + +@router.get("/health", response_model=HealthResponse) +async def health(state: APIState = Depends(get_state)) -> HealthResponse: + """Liveness signal + queue/worker overview. Unauthenticated so load + balancers and uptime monitors can hit it without a token.""" + counts = await state.job_repo.counts_by_status() + worker_count = ( + state.config.ingester.workers.worker_count if state.pool is not None else 0 + ) + poller_count = len(state.pollers.pollers) if state.pollers is not None else 0 + return HealthResponse( + status="ok", + queue_counts=counts, + worker_count=worker_count, + poller_count=poller_count, + ) diff --git a/haiku_rag_slim/haiku/rag/ingester/api/routes/jobs.py b/haiku_rag_slim/haiku/rag/ingester/api/routes/jobs.py new file mode 100644 index 00000000..7db80b58 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/api/routes/jobs.py @@ -0,0 +1,55 @@ +from fastapi import APIRouter, Depends, HTTPException, status + +from haiku.rag.ingester.api.schemas import CancelResponse +from haiku.rag.ingester.api.server import APIState, get_state +from haiku.rag.ingester.queue.models import Job, JobStatus + +router = APIRouter(prefix="/jobs", tags=["jobs"]) + + +@router.get("", response_model=list[Job]) +async def list_jobs( + status: JobStatus | None = None, + source_id: str | None = None, + uri: str | None = None, + limit: int = 50, + offset: int = 0, + state: APIState = Depends(get_state), +) -> list[Job]: + return await state.job_repo.list_jobs( + status=status, source_id=source_id, uri=uri, limit=limit, offset=offset + ) + + +@router.get("/{job_id}", response_model=Job) +async def get_job(job_id: str, state: APIState = Depends(get_state)) -> Job: + job = await state.job_repo.get_job(job_id) + if job is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="job not found" + ) + return job + + +@router.post("/{job_id}/retry", response_model=Job) +async def retry_job(job_id: str, state: APIState = Depends(get_state)) -> Job: + """Force a dead or queued job back to queued with attempts=0.""" + try: + return await state.job_repo.retry(job_id) + except KeyError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="job not found" + ) from exc + + +@router.delete("/{job_id}", response_model=CancelResponse) +async def cancel_job( + job_id: str, state: APIState = Depends(get_state) +) -> CancelResponse: + cancelled = await state.job_repo.cancel(job_id) + if not cancelled: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="job not found or already terminal", + ) + return CancelResponse(job_id=job_id, cancelled=True) diff --git a/haiku_rag_slim/haiku/rag/ingester/api/routes/sources.py b/haiku_rag_slim/haiku/rag/ingester/api/routes/sources.py new file mode 100644 index 00000000..838e974c --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/api/routes/sources.py @@ -0,0 +1,44 @@ +from fastapi import APIRouter, Depends, HTTPException, status + +from haiku.rag.ingester.api.schemas import RefreshResponse, SourceSummary +from haiku.rag.ingester.api.server import APIState, get_state + +router = APIRouter(prefix="/sources", tags=["sources"]) + + +@router.get("", response_model=list[SourceSummary]) +async def list_sources( + state: APIState = Depends(get_state), +) -> list[SourceSummary]: + if state.pollers is None: + return [] + summaries: list[SourceSummary] = [] + for poller in state.pollers.pollers: + summaries.append( + SourceSummary( + source_id=poller.source_id, + type=type(poller.config).__name__, + last_polled_at=poller.last_polled_at, + circuit_breaker_open=poller._breaker.is_open, + ) + ) + return summaries + + +@router.post("/{source_id}/refresh", response_model=RefreshResponse) +async def refresh_source( + source_id: str, state: APIState = Depends(get_state) +) -> RefreshResponse: + """Out-of-band poll: forces an immediate `discover()` sweep.""" + if state.pollers is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="pollers not running", + ) + for poller in state.pollers.pollers: + if poller.source_id == source_id: + ok = await poller._sweep_once() + return RefreshResponse(source_id=source_id, refreshed=ok) + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="source not found" + ) diff --git a/haiku_rag_slim/haiku/rag/ingester/api/schemas.py b/haiku_rag_slim/haiku/rag/ingester/api/schemas.py new file mode 100644 index 00000000..c57742dc --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/api/schemas.py @@ -0,0 +1,27 @@ +from datetime import datetime + +from pydantic import BaseModel + + +class HealthResponse(BaseModel): + status: str + queue_counts: dict[str, int] + worker_count: int + poller_count: int + + +class SourceSummary(BaseModel): + source_id: str + type: str + last_polled_at: datetime | None + circuit_breaker_open: bool + + +class RefreshResponse(BaseModel): + source_id: str + refreshed: bool + + +class CancelResponse(BaseModel): + job_id: str + cancelled: bool diff --git a/haiku_rag_slim/haiku/rag/ingester/api/server.py b/haiku_rag_slim/haiku/rag/ingester/api/server.py new file mode 100644 index 00000000..6fd301a1 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/ingester/api/server.py @@ -0,0 +1,52 @@ +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from fastapi import Depends, FastAPI, Request + +from haiku.rag.ingester.api.auth import require_auth + +if TYPE_CHECKING: + from haiku.rag.config import AppConfig + from haiku.rag.ingester.pollers.manager import PollerManager + from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo + from haiku.rag.ingester.workers.pool import WorkerPool + + +@dataclass +class APIState: + """Everything the API handlers need to read or act on. Pool/pollers are + optional so the FastAPI app can be tested in isolation.""" + + config: "AppConfig" + job_repo: "JobRepo" + sync_repo: "SyncStateRepo" + pool: "WorkerPool | None" = None + pollers: "PollerManager | None" = None + + +def get_state(request: Request) -> APIState: + return request.app.state.api_state + + +def build_app( + state: APIState, + *, + auth_token: str | None = None, +) -> FastAPI: + """Construct the ingester's FastAPI control plane.""" + from haiku.rag.ingester.api.routes import dlq, health, jobs, sources + + app = FastAPI( + title="haiku-ingester", + description="Control plane for the haiku.rag production ingester.", + version="1", + ) + app.state.api_state = state + app.state.auth_token = auth_token + + auth_dep = [Depends(require_auth)] + app.include_router(health.router) # /health is unauthenticated by design + app.include_router(jobs.router, dependencies=auth_dep) + app.include_router(sources.router, dependencies=auth_dep) + app.include_router(dlq.router, dependencies=auth_dep) + return app diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index 0c2b7f9e..16447b17 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -42,69 +42,108 @@ class IngesterApp: ingester_cfg = self._config.ingester self._queue_conn = await open_queue(ingester_cfg.queue.path) - self._jobs = JobRepo(self._queue_conn) - self._sync = SyncStateRepo(self._queue_conn) + try: + self._jobs = JobRepo(self._queue_conn) + self._sync = SyncStateRepo(self._queue_conn) - supported_extensions = get_converter(self._config).supported_extensions - retry = RetryPolicy( - max_attempts=ingester_cfg.workers.retry.max_attempts, - base_delay_s=ingester_cfg.workers.retry.base_delay_s, - max_delay_s=ingester_cfg.workers.retry.max_delay_s, - jitter=ingester_cfg.workers.retry.jitter, - ) - - async with HaikuRAG(self._db_path, config=self._config) as client: - self._client = client - self._pool = WorkerPool( - client=client, - job_repo=self._jobs, - sync_repo=self._sync, - worker_count=ingester_cfg.workers.worker_count, - max_concurrent=ingester_cfg.workers.max_concurrent, - retry_policy=retry, - poll_idle_interval_s=ingester_cfg.workers.poll_idle_interval_s, - claim_timeout_s=ingester_cfg.workers.claim_timeout_s, - reaper_interval_s=ingester_cfg.workers.reaper_interval_s, - ) - self._pollers = PollerManager( - configs=ingester_cfg.sources, - job_repo=self._jobs, - sync_repo=self._sync, - supported_extensions=supported_extensions, - default_max_attempts=ingester_cfg.workers.retry.max_attempts, + supported_extensions = get_converter(self._config).supported_extensions + retry = RetryPolicy( + max_attempts=ingester_cfg.workers.retry.max_attempts, + base_delay_s=ingester_cfg.workers.retry.base_delay_s, + max_delay_s=ingester_cfg.workers.retry.max_delay_s, + jitter=ingester_cfg.workers.retry.jitter, ) - stop_event = asyncio.Event() - loop = asyncio.get_running_loop() - for sig in (signal.SIGINT, signal.SIGTERM): - try: - loop.add_signal_handler(sig, stop_event.set) - except NotImplementedError: - # Windows; signal handlers unavailable in asyncio. - pass - - await self._pool.start() - await self._pollers.start() - logger.info( - "Ingester running: %d worker(s), %d source(s)", - ingester_cfg.workers.worker_count, - len(ingester_cfg.sources), - ) - - if api: - # HTTP control plane lands in a follow-up; for now this branch - # is a no-op so callers can still pass api=True without error. - logger.info( - "HTTP API not yet implemented; running pollers + workers only" + async with HaikuRAG(self._db_path, config=self._config) as client: + self._client = client + self._pool = WorkerPool( + client=client, + job_repo=self._jobs, + sync_repo=self._sync, + worker_count=ingester_cfg.workers.worker_count, + max_concurrent=ingester_cfg.workers.max_concurrent, + retry_policy=retry, + poll_idle_interval_s=ingester_cfg.workers.poll_idle_interval_s, + claim_timeout_s=ingester_cfg.workers.claim_timeout_s, + reaper_interval_s=ingester_cfg.workers.reaper_interval_s, + ) + self._pollers = PollerManager( + configs=ingester_cfg.sources, + job_repo=self._jobs, + sync_repo=self._sync, + supported_extensions=supported_extensions, + default_max_attempts=ingester_cfg.workers.retry.max_attempts, ) - try: - await stop_event.wait() - finally: - logger.info("Shutting down ingester") - await self._pollers.stop() - await self._pool.stop() + stop_event = asyncio.Event() + loop = asyncio.get_running_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, stop_event.set) + except NotImplementedError: + # Windows; signal handlers unavailable in asyncio. + pass - if self._queue_conn is not None: - await self._queue_conn.close() - self._queue_conn = None + await self._pool.start() + await self._pollers.start() + logger.info( + "Ingester running: %d worker(s), %d source(s)", + ingester_cfg.workers.worker_count, + len(ingester_cfg.sources), + ) + + api_task, api_server = await self._maybe_start_api(api) + + try: + await stop_event.wait() + finally: + logger.info("Shutting down ingester") + if api_server is not None: + api_server.should_exit = True + if api_task is not None: + await asyncio.gather(api_task, return_exceptions=True) + await self._pollers.stop() + await self._pool.stop() + finally: + # Close the queue connection unconditionally. aiosqlite runs the + # underlying sqlite3 in a background thread; leaving it open holds + # the event loop alive and blocks process exit on early failures + # (e.g. HaikuRAG raising MigrationRequiredError). + if self._queue_conn is not None: + await self._queue_conn.close() + self._queue_conn = None + + async def _maybe_start_api(self, api: bool): + """Spin up the FastAPI control plane on an asyncio task. Returns + (task, server) or (None, None) when the API is disabled.""" + ingester_cfg = self._config.ingester + if not (api and ingester_cfg.api.enabled): + return None, None + + import uvicorn + + from haiku.rag.ingester.api.server import APIState, build_app + + assert self._jobs is not None and self._sync is not None + state = APIState( + config=self._config, + job_repo=self._jobs, + sync_repo=self._sync, + pool=self._pool, + pollers=self._pollers, + ) + if ingester_cfg.api.auth_token is None: + logger.warning("API auth_token is unset — control plane is unauthenticated") + app = build_app(state, auth_token=ingester_cfg.api.auth_token) + config = uvicorn.Config( + app, + host=ingester_cfg.api.host, + port=ingester_cfg.api.port, + log_level="info", + lifespan="off", + ) + server = uvicorn.Server(config) + logger.info( + "API listening on %s:%d", ingester_cfg.api.host, ingester_cfg.api.port + ) + return asyncio.create_task(server.serve()), server diff --git a/haiku_rag_slim/haiku/rag/ingester/cli.py b/haiku_rag_slim/haiku/rag/ingester/cli.py index 66c8f752..059b3e08 100644 --- a/haiku_rag_slim/haiku/rag/ingester/cli.py +++ b/haiku_rag_slim/haiku/rag/ingester/cli.py @@ -1,4 +1,5 @@ import asyncio +import sys import uuid from datetime import UTC, datetime from pathlib import Path @@ -21,20 +22,51 @@ from haiku.rag.ingester.exceptions import PermanentError, TransientError # noqa from haiku.rag.ingester.queue.migrations import open_queue # noqa: E402 from haiku.rag.ingester.queue.models import Job, JobOp, JobStatus # noqa: E402 from haiku.rag.ingester.workers.pipeline import run_job # noqa: E402 +from haiku.rag.store.exceptions import ( # noqa: E402 + MigrationRequiredError, + ReadOnlyError, +) -cli = typer.Typer( +_cli = typer.Typer( name="haiku-ingester", no_args_is_help=True, pretty_exceptions_show_locals=False, help="Production ingester for haiku.rag.", ) + +def _configure_logfire() -> None: + """Logfire emits spans only when LOGFIRE_TOKEN is set; otherwise it + stays silent (no warning either way). Matches the haiku-rag CLI.""" + try: + import logfire + + is_production = get_config().environment != "development" + logfire.configure( + send_to_logfire="if-token-present", + console=False if is_production else None, + ) + logfire.instrument_pydantic_ai() + except Exception: # pragma: no cover + pass + + +def cli() -> None: + """Entry point that translates store-state errors into a clean exit.""" + _configure_logfire() + try: + _cli() + except (MigrationRequiredError, ReadOnlyError) as e: + typer.echo(f"Error: {e}", err=True) + sys.exit(1) + + queue_cli = typer.Typer( name="queue", no_args_is_help=True, help="Operate the ingester's SQLite job queue.", ) -cli.add_typer(queue_cli) +_cli.add_typer(queue_cli) def _load_config_with_override(config_path: Path | None) -> AppConfig: @@ -102,7 +134,7 @@ def _resolve_db_path(config: AppConfig, override: Path | None) -> Path: return override or (config.storage.data_dir / "haiku.rag.lancedb") -@cli.command("serve") +@_cli.command("serve") def serve( config: Path | None = typer.Option( None, "--config", "-c", help="Path to haiku.rag.yaml." @@ -126,7 +158,7 @@ def serve( asyncio.run(app.serve(api=not no_api)) -@cli.command("run-once") +@_cli.command("run-once") def run_once( uri: str = typer.Argument(..., help="URI to ingest (file://, http(s)://, s3://)."), config: Path | None = typer.Option( diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/factory.py b/haiku_rag_slim/haiku/rag/ingester/pollers/factory.py index d50eaec8..d4e37bc2 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/factory.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/factory.py @@ -24,6 +24,7 @@ def build_source( ignore_patterns=cfg.ignore_patterns or None, include_patterns=cfg.include_patterns or None, supported_extensions=supported_extensions, + source_id=cfg.id, ) if isinstance(cfg, HTTPSourceConfig): if cfg.id is None: diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py index 1835187b..91878a42 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py @@ -33,12 +33,13 @@ class FSSource: ignore_patterns: list[str] | None = None, include_patterns: list[str] | None = None, supported_extensions: list[str] | None = None, + source_id: str | None = None, ) -> None: # Resolve so symlinks and relative paths collapse to one canonical - # source_id. The queue uses source_id as a foreign key — two paths - # for the same root would mean duplicate sync_state rows. + # root. The queue uses source_id as a foreign key — two paths for + # the same root would mean duplicate sync_state rows. self.root = Path(root).resolve() - self.source_id = f"fs:{self.root}" + self.source_id = source_id or f"fs:{self.root}" self.supported_extensions = ( list(supported_extensions) if supported_extensions is not None diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py index 6b6c3402..ebf5d728 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py @@ -2,6 +2,7 @@ import asyncio from typing import TYPE_CHECKING import httpx +import logfire from pydantic import BaseModel from haiku.rag.ingester.exceptions import PermanentError, TransientError @@ -67,28 +68,6 @@ def _classify(exc: BaseException) -> Exception: return TransientError(f"unexpected: {exc!r}") -def _logfire_span_or_null(name: str, **attrs): - """logfire is available via pydantic-ai but may be disabled; use the - no-op API surface so tests don't need it configured.""" - try: - import logfire - - return logfire.span(name, **attrs) - except ImportError: # pragma: no cover - logfire ships with pydantic-ai - - class _Null: - def __enter__(self): - return self - - def __exit__(self, *_): - return False - - def set_attribute(self, *_args, **_kwargs): - pass - - return _Null() - - async def run_job(client: "HaikuRAG", job: Job) -> JobResult: """Execute the work described by `job`. Raises PermanentError or TransientError; the worker uses that to decide dead vs retry.""" @@ -96,7 +75,7 @@ async def run_job(client: "HaikuRAG", job: Job) -> JobResult: storage_options = extra.get("storage_options") user_metadata = extra.get("metadata", {}) - with _logfire_span_or_null( + with logfire.span( "ingester.job", job_id=job.id, source_id=job.source_id, diff --git a/pyproject.toml b/pyproject.toml index 8c985a1b..c6e2d560 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,8 +70,8 @@ members = ["haiku_rag_slim", "evaluations"] [dependency-groups] dev = [ "haiku.rag-evals", + "haiku.rag-slim[ingester]", "datasets>=4.8.4", - "obstore>=0.9,<0.10", "zensical", "pre-commit>=4.5.1", "pydantic-ai-slim[anthropic]", diff --git a/tests/ingester/test_api.py b/tests/ingester/test_api.py new file mode 100644 index 00000000..5d61b770 --- /dev/null +++ b/tests/ingester/test_api.py @@ -0,0 +1,340 @@ +from datetime import UTC, datetime + +import aiosqlite +import httpx +import pytest +from httpx import ASGITransport + +from haiku.rag.config import AppConfig +from haiku.rag.ingester.api.server import APIState, build_app +from haiku.rag.ingester.queue.migrations import apply_migrations +from haiku.rag.ingester.queue.models import JobOp, JobStatus +from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo +from haiku.rag.ingester.sources.base import ( + FetchResult, + SourceEvent, + SourceEventKind, +) + + +@pytest.fixture +async def conn(tmp_path): + path = tmp_path / "queue.db" + connection = await aiosqlite.connect(str(path)) + connection.row_factory = aiosqlite.Row + await apply_migrations(connection) + yield connection + await connection.close() + + +@pytest.fixture +def jobs(conn): + return JobRepo(conn) + + +@pytest.fixture +def sync(conn): + return SyncStateRepo(conn) + + +@pytest.fixture +def state(jobs, sync): + return APIState( + config=AppConfig(), + job_repo=jobs, + sync_repo=sync, + ) + + +def _client(state, *, auth_token: str | None = None) -> httpx.AsyncClient: + app = build_app(state, auth_token=auth_token) + return httpx.AsyncClient( + transport=ASGITransport(app=app), base_url="http://testserver" + ) + + +# --- /health --- + + +@pytest.mark.asyncio +async def test_health_ok_with_counts(state, jobs): + await jobs.enqueue("src", "u1", JobOp.UPSERT) + j2 = await jobs.enqueue("src", "u2", JobOp.UPSERT) + assert j2 is not None + await jobs.mark_dead(j2.id, "boom") + + async with _client(state) as client: + resp = await client.get("/health") + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "ok" + assert body["queue_counts"] == {"queued": 1, "dead": 1} + assert body["worker_count"] == 0 # pool not attached in the test state + assert body["poller_count"] == 0 + + +@pytest.mark.asyncio +async def test_health_skips_auth(state): + async with _client(state, auth_token="secret") as client: + resp = await client.get("/health") + assert resp.status_code == 200 + + +# --- auth --- + + +@pytest.mark.asyncio +async def test_protected_endpoint_rejects_without_token(state): + async with _client(state, auth_token="secret") as client: + resp = await client.get("/jobs") + assert resp.status_code == 401 + + +@pytest.mark.asyncio +async def test_protected_endpoint_rejects_wrong_token(state): + async with _client(state, auth_token="secret") as client: + resp = await client.get("/jobs", headers={"Authorization": "Bearer nope"}) + assert resp.status_code == 401 + + +@pytest.mark.asyncio +async def test_protected_endpoint_accepts_correct_token(state): + async with _client(state, auth_token="secret") as client: + resp = await client.get("/jobs", headers={"Authorization": "Bearer secret"}) + assert resp.status_code == 200 + + +@pytest.mark.asyncio +async def test_no_auth_token_allows_everything(state, jobs): + async with _client(state, auth_token=None) as client: + assert (await client.get("/jobs")).status_code == 200 + assert (await client.get("/health")).status_code == 200 + + +# --- /jobs --- + + +@pytest.mark.asyncio +async def test_list_jobs_returns_recent_first(state, jobs): + j1 = await jobs.enqueue("a", "u1", JobOp.UPSERT) + j2 = await jobs.enqueue("b", "u2", JobOp.UPSERT) + assert j1 is not None and j2 is not None + + async with _client(state) as client: + resp = await client.get("/jobs") + assert resp.status_code == 200 + payload = resp.json() + assert [j["id"] for j in payload] == [j2.id, j1.id] + + +@pytest.mark.asyncio +async def test_list_jobs_filters_by_source_and_status(state, jobs): + await jobs.enqueue("a", "u", JobOp.UPSERT) + j = await jobs.enqueue("b", "u", JobOp.UPSERT) + assert j is not None + await jobs.mark_dead(j.id, "err") + + async with _client(state) as client: + resp = await client.get("/jobs?source_id=b&status=dead") + payload = resp.json() + assert len(payload) == 1 + assert payload[0]["id"] == j.id + + +@pytest.mark.asyncio +async def test_get_job_returns_record(state, jobs): + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + async with _client(state) as client: + resp = await client.get(f"/jobs/{job.id}") + assert resp.status_code == 200 + assert resp.json()["id"] == job.id + + +@pytest.mark.asyncio +async def test_get_job_404(state): + async with _client(state) as client: + resp = await client.get("/jobs/nope") + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_retry_revives_dead_job(state, jobs): + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + await jobs.mark_dead(job.id, "err") + + async with _client(state) as client: + resp = await client.post(f"/jobs/{job.id}/retry") + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == JobStatus.QUEUED.value + assert body["attempts"] == 0 + + +@pytest.mark.asyncio +async def test_retry_404(state): + async with _client(state) as client: + resp = await client.post("/jobs/missing/retry") + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_cancel_queued_job(state, jobs): + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + + async with _client(state) as client: + resp = await client.delete(f"/jobs/{job.id}") + assert resp.status_code == 200 + assert resp.json() == {"job_id": job.id, "cancelled": True} + assert await jobs.get_job(job.id) is None + + +@pytest.mark.asyncio +async def test_cancel_succeeded_returns_404(state, jobs): + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_succeeded(claimed.id) + + async with _client(state) as client: + resp = await client.delete(f"/jobs/{job.id}") + assert resp.status_code == 404 + + +# --- /dlq --- + + +@pytest.mark.asyncio +async def test_dlq_lists_dead_jobs_only(state, jobs): + j1 = await jobs.enqueue("src", "u1", JobOp.UPSERT) + j2 = await jobs.enqueue("src", "u2", JobOp.UPSERT) + assert j1 is not None and j2 is not None + await jobs.mark_dead(j2.id, "err") + + async with _client(state) as client: + resp = await client.get("/dlq") + payload = resp.json() + assert len(payload) == 1 + assert payload[0]["id"] == j2.id + + +@pytest.mark.asyncio +async def test_dlq_retry_resurrects(state, jobs): + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + await jobs.mark_dead(job.id, "err") + + async with _client(state) as client: + resp = await client.post(f"/dlq/{job.id}/retry") + assert resp.status_code == 200 + assert resp.json()["status"] == JobStatus.QUEUED.value + + +# --- /sources --- + + +class _StubSource: + def __init__(self, source_id, sweeps=()): + self.source_id = source_id + self._sweeps = list(sweeps) + + def supports(self, uri): # pragma: no cover + return True + + async def head(self, uri): # pragma: no cover + return None + + async def fetch(self, uri) -> FetchResult: # pragma: no cover + raise NotImplementedError + + async def discover(self, since=None): + events = self._sweeps.pop(0) if self._sweeps else [] + for event in events: + yield event + + +def _build_pollers_state(tmp_path, jobs, sync, source_id: str = "local"): + """Build an APIState with a real PollerManager containing one FS poller.""" + from haiku.rag.config import FSSourceConfig + from haiku.rag.ingester.pollers.manager import PollerManager + + cfg = FSSourceConfig(type="fs", id=source_id, root=tmp_path) + manager = PollerManager(configs=[cfg], job_repo=jobs, sync_repo=sync) + # Build pollers without starting tasks — we want to inspect/refresh directly. + manager._pollers = manager.build_pollers() + state = APIState( + config=AppConfig(), + job_repo=jobs, + sync_repo=sync, + pollers=manager, + ) + return state, manager + + +@pytest.mark.asyncio +async def test_sources_empty_when_no_pollers(state): + async with _client(state) as client: + resp = await client.get("/sources") + assert resp.status_code == 200 + assert resp.json() == [] + + +@pytest.mark.asyncio +async def test_sources_lists_configured(tmp_path, jobs, sync): + state, _ = _build_pollers_state(tmp_path, jobs, sync) + async with _client(state) as client: + resp = await client.get("/sources") + payload = resp.json() + assert len(payload) == 1 + assert payload[0]["source_id"] == "local" + assert payload[0]["type"] == "FSSourceConfig" + assert payload[0]["circuit_breaker_open"] is False + + +@pytest.mark.asyncio +async def test_source_refresh_triggers_sweep(tmp_path, jobs, sync): + state, manager = _build_pollers_state(tmp_path, jobs, sync) + # Replace the real source with a stub that records the sweep + emits an event. + poller = manager.pollers[0] + poller.source = _StubSource( + poller.source_id, + [ + [ + SourceEvent( + source_id=poller.source_id, + uri="file:///x.md", + kind=SourceEventKind.UPSERT, + revision="v1", + discovered_at=datetime.now(UTC), + ) + ] + ], + ) + + async with _client(state) as client: + resp = await client.post(f"/sources/{poller.source_id}/refresh") + assert resp.status_code == 200 + body = resp.json() + assert body["refreshed"] is True + assert body["source_id"] == poller.source_id + + queued = await jobs.list_jobs(source_id=poller.source_id) + assert len(queued) == 1 + + +@pytest.mark.asyncio +async def test_source_refresh_unknown_id_404(tmp_path, jobs, sync): + state, _ = _build_pollers_state(tmp_path, jobs, sync) + async with _client(state) as client: + resp = await client.post("/sources/missing/refresh") + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_source_refresh_503_when_pollers_absent(state): + async with _client(state) as client: + resp = await client.post("/sources/anything/refresh") + assert resp.status_code == 503 diff --git a/tests/ingester/test_pollers.py b/tests/ingester/test_pollers.py index fee76170..de36adb9 100644 --- a/tests/ingester/test_pollers.py +++ b/tests/ingester/test_pollers.py @@ -326,13 +326,13 @@ async def test_fs_poller_enqueues_initial_files(tmp_path, jobs, sync): try: # Wait for the initial sweep to land jobs. for _ in range(40): - queued = await jobs.list_jobs(source_id=f"fs:{tmp_path.resolve()}") + queued = await jobs.list_jobs(source_id="local") if len(queued) == 2: break await asyncio.sleep(0.05) finally: await manager.stop() - queued = await jobs.list_jobs(source_id=f"fs:{tmp_path.resolve()}") + queued = await jobs.list_jobs(source_id="local") assert {Path(j.uri).name for j in queued} == {"a.md", "b.md"} assert all(j.status is JobStatus.QUEUED for j in queued) diff --git a/uv.lock b/uv.lock index e1c6df7c..315266e6 100644 --- a/uv.lock +++ b/uv.lock @@ -1481,7 +1481,7 @@ tui = [ dev = [ { name = "datasets" }, { name = "haiku-rag-evals" }, - { name = "obstore" }, + { name = "haiku-rag-slim", extra = ["ingester"] }, { name = "pre-commit" }, { name = "pydantic-ai-slim", extra = ["anthropic", "bedrock", "google", "groq"] }, { name = "pytest" }, @@ -1508,7 +1508,7 @@ provides-extras = ["tui", "s3", "cross-encoder", "ingester"] dev = [ { name = "datasets", specifier = ">=4.8.4" }, { name = "haiku-rag-evals", editable = "evaluations" }, - { name = "obstore", specifier = ">=0.9,<0.10" }, + { name = "haiku-rag-slim", extras = ["ingester"], editable = "haiku_rag_slim" }, { name = "pre-commit", specifier = ">=4.5.1" }, { name = "pydantic-ai-slim", extras = ["anthropic"] }, { name = "pydantic-ai-slim", extras = ["bedrock"] },