diff --git a/docs/adapters/importer-base.md b/docs/adapters/importer-base.md new file mode 100644 index 00000000..a70947d7 --- /dev/null +++ b/docs/adapters/importer-base.md @@ -0,0 +1,282 @@ +# Importer Base — Shared Hardening Helpers + +`layerlens.instrument.adapters._base.importer` provides a small, reusable +hardening surface for any adapter that **imports traces or records from +an upstream system** — Langfuse today, planned ServiceNow / Zendesk / +HubSpot tomorrow, plus anything else that fits the pattern of "list → +fetch → map → ingest". + +The contract was extracted from the production-hardened Agentforce +importer (`frameworks/agentforce/auth.py` + +`frameworks/agentforce/importer.py`) so subsequent importers inherit +the same semantics by default rather than re-deriving them from scratch +(and missing edge cases). + +This document is the contract for **future importer adapters**: when +you write the next one, build on the helpers below. + +--- + +## What's in the box + +| Symbol | Purpose | +| ---------------------------- | ------------------------------------------------------------------------ | +| `validate_id` / `require_valid_id` | Regex-anchored ID validation BEFORE query interpolation | +| `ID_PATTERN_*` constants | Pre-compiled patterns for UUID, slug, integer, Salesforce, date, ts | +| `parse_rate_limit_headers` | Normalised view of `X-RateLimit-*` and `Retry-After` | +| `RateLimitInfo` | Immutable dataclass — `.is_throttled`, `.sleep_seconds()`, `.usage_ratio`| +| `paginate` / `apaginate` | Sync iterator + async generator over cursor-paged endpoints | +| `batched_in` | Chunk a list of IDs for `WHERE x IN (…)` queries | +| `retry_with_backoff` / `aretry_with_backoff` | Decorrelated full-jitter backoff, rate-limit aware | +| `RetryableHTTPError` | Carries optional `RateLimitInfo` so the retry loop can sleep until reset | +| `BaseImporter` | Abstract orchestrator: `fetch_records` → `process_record` with retry | +| `ImporterResult` | Standard counters: imported / skipped / failed / quarantined | + +--- + +## When to use it + +Use the importer base when **every** statement below is true: + +1. The adapter pulls data from an upstream HTTP API (or any cursored, + rate-limited source — gRPC, GraphQL, message queue with ACK + semantics). +2. The adapter does NOT instrument a running framework (those use + `BaseAdapter` directly; importers compose `BaseImporter` *and* + `BaseAdapter` — see `frameworks/langfuse/lifecycle.py`). +3. Records have stable, validatable IDs that are interpolated into + subsequent fetch URLs / queries. + +Don't use it for runtime-wrapping adapters (LangGraph callbacks, +LangChain handlers, Bedrock client decorators) — those have their own +contract via `_base/adapter.py` and `_base/capture.py`. + +--- + +## Minimum viable importer + +```python +from collections.abc import Iterator +from typing import Any + +from layerlens.instrument.adapters._base.importer import ( + ID_PATTERN_UUID, + BaseImporter, + paginate, + require_valid_id, +) + + +class MyImporter(BaseImporter[MyClient, MyState]): + SOURCE = "myservice" + + def fetch_records(self) -> Iterator[dict[str, Any]]: + # paginate handles cursor termination, empty-data short-circuit, + # cursor-loop detection, and the max_pages safety bound. + yield from paginate( + lambda cursor: self._client.list_records(cursor=cursor), + cursor_field="next_cursor", + data_field="records", + ) + + def process_record(self, record: dict[str, Any]) -> bool: + # Validate BEFORE any interpolation into a query string. + rid = require_valid_id(record["id"], ID_PATTERN_UUID, label="record_id") + full = self._client.get_record(rid) # safe: rid passed validation + events = self._mapper.map(full) + for event in events: + self._stratix.emit(event["type"], event["payload"]) + return True + + def record_id(self, record: dict[str, Any]) -> str: + return record.get("id", "") +``` + +`run()` then drives the loop: + +```python +result = MyImporter(client, state).run(dry_run=False) +print(f"imported={result.imported_count} failed={result.failed_count}") +``` + +--- + +## Hardening rules + +These are the rules every importer adapter MUST follow. They map 1:1 to +the supra-spec features that the audit at +`A:/tmp/adapter-cross-pollination-audit.md` (§2.7-2.9) identified as +critical for importer adapters. + +### 1. Always validate IDs before query interpolation + +Any identifier that came from the upstream is untrusted until it has +been matched against an anchored regex. + +```python +# WRONG — direct interpolation; vulnerable to injection. +url = f"/api/records/{record_id}" + +# RIGHT — validate against the canonical pattern first. +if not validate_id(record_id, ID_PATTERN_UUID): + logger.warning("Skipping malformed record id: %r", record_id) + return None +url = f"/api/records/{record_id}" +``` + +The five built-in patterns cover the common cases: + +| Constant | Matches | +| ------------------------- | ------------------------------------------------------ | +| `ID_PATTERN_UUID` | 8-4-4-4-12 hex (with or without hyphens, any case) | +| `ID_PATTERN_SLUG` | URL slugs — letters / digits / hyphens / underscores | +| `ID_PATTERN_INTEGER` | Decimal integers (signed, up to 19 digits) | +| `ID_PATTERN_SALESFORCE` | Salesforce 15-or-18-char alphanumeric record IDs | +| `ID_PATTERN_DATE` | `YYYY-MM-DD` | +| `ID_PATTERN_TIMESTAMP` | ISO 8601 datetime with optional offset / fractional s | + +If your upstream uses a different shape, compile and pin your own: + +```python +import re +ID_PATTERN_MYAPP = re.compile(r"^MA-[A-Z0-9]{12}$") +``` + +### 2. Honour rate-limit headers + +Don't retry blindly into a 429 — sleep until the explicit reset. + +```python +from layerlens.instrument.adapters._base.importer import ( + RetryableHTTPError, + parse_rate_limit_headers, +) + +try: + response = http.get(url) +except HTTPError as exc: + rl = parse_rate_limit_headers(exc.response.headers) + if exc.code == 429 or exc.code >= 500: + # retry_with_backoff will sleep until rl.reset_at if rl is set, + # else fall back to its computed exponential backoff. + raise RetryableHTTPError("transient", status_code=exc.code, rate_limit=rl) + raise # 4xx other than 429 is terminal +``` + +Surface the most recent observation via a `last_rate_limit` property +so dashboards can warn at 80 % usage *before* the upstream throttles +(see `LangfuseAPIClient._warn_if_throttle_imminent`). + +### 3. Use `paginate` / `apaginate` — don't roll your own + +The shared helpers handle: + +* Cursor termination on `null` / empty / missing. +* **Empty-data short-circuit** — providers occasionally return + `next_cursor` on an empty final page; rolling your own loop will + loop forever. +* **Cursor-loop guard** — if the provider buggily returns the same + cursor twice, terminate rather than iterate forever. +* `max_pages` ceiling — defaults to 10 000, raise `RuntimeError` when + exceeded so runaway pagination becomes a visible failure rather than + an OOM. + +```python +yield from paginate( + lambda cursor: client.list(cursor=cursor), + cursor_field="next_cursor", + data_field="data", + max_pages=10_000, +) +``` + +### 4. Use `retry_with_backoff` for transient failures + +Real retry policy, not a sleep-and-pray loop: + +* **Decorrelated full-jitter** sampling per the AWS Architecture Blog + (2015) — best at avoiding synchronised retry storms when many + importer instances fail at once. +* **Rate-limit-aware** — if the inner function raises + `RetryableHTTPError(rate_limit=…)`, the loop sleeps until that + explicit deadline (capped by `max_delay`) instead of using its own + computed backoff. +* **Caps** — `max_delay` is enforced even when `Retry-After` is + unreasonable; `max_retries` is enforced and re-raises the original + exception type so callers see the same exception they always do. + +```python +result = retry_with_backoff( + lambda: client.fetch_one(record_id), + max_retries=5, + base_delay=1.0, + max_delay=30.0, + jitter=True, +) +``` + +### 5. Batch IDs through `WHERE x IN (…)` queries + +Most upstream APIs cap the number of identifiers in a single `IN` +clause. Use `batched_in` rather than ad-hoc slicing: + +```python +from layerlens.instrument.adapters._base.importer import batched_in + +for batch in batched_in(parent_ids, batch_size=200): # SOQL limit + safe = [require_valid_id(p, ID_PATTERN_SALESFORCE) for p in batch] + query = f"SELECT Id FROM Child WHERE ParentId IN ('{chr(39).join(safe)}')" + yield from connection.query(query) +``` + +### 6. Quarantine after repeated failures + +`BaseImporter.run()` will quarantine a record after +`quarantine_threshold` (default 3) consecutive failures, so a +poison-pill record doesn't keep failing forever and consuming retries. + +If your state object exposes `is_quarantined(record_id) -> bool` and +`record_failure(record_id, max_failures: int) -> bool`, `BaseImporter` +will use them automatically. The Langfuse `SyncState` is the canonical +example; see `frameworks/langfuse/config.py`. + +--- + +## Reference implementation + +The Langfuse adapter is the first consumer of every helper in this +module: + +* `frameworks/langfuse/client.py` — uses `retry_with_backoff` + + `parse_rate_limit_headers` + `paginate` + `RetryableHTTPError`. +* `frameworks/langfuse/importer.py` — uses `validate_id` against + `ID_PATTERN_UUID` + a second `retry_with_backoff` layer for + per-trace fetches. +* `frameworks/langfuse/exporter.py` — uses `validate_id` + + `retry_with_backoff` for batch ingestion pushes. + +Use it as the template for the next importer. + +--- + +## Acceptance checklist for new importers + +Before merging an importer adapter PR, verify all of: + +- [ ] Every upstream-supplied ID is validated via `validate_id` / + `require_valid_id` before being interpolated into ANY query + string. +- [ ] `parse_rate_limit_headers` is called on every successful AND + failed response; results are surfaced via a public property. +- [ ] All cursor-paged endpoints use `paginate` / `apaginate` (no + hand-rolled `while True` loops). +- [ ] All transient failures (429, 5xx, network errors) are wrapped in + `RetryableHTTPError` and driven through `retry_with_backoff`. +- [ ] Terminal failures (4xx other than 429) raise the adapter's own + `*APIError` and do **not** retry. +- [ ] `BaseImporter.run()` is the single entry point — no bespoke loops. +- [ ] Tests exercise: invalid-ID rejection, transient-retry recovery, + terminal-4xx no-retry, rate-limit-aware sleep, quarantine + promotion. +- [ ] `mypy --strict` clean on the new module. +- [ ] `ruff check` clean. diff --git a/src/layerlens/instrument/adapters/_base/importer.py b/src/layerlens/instrument/adapters/_base/importer.py new file mode 100644 index 00000000..9ecf3742 --- /dev/null +++ b/src/layerlens/instrument/adapters/_base/importer.py @@ -0,0 +1,969 @@ +"""Shared base for trace-importer adapters. + +Provides the abstract :class:`BaseImporter` class plus reusable helpers +that any importer adapter (Langfuse, future ServiceNow, Zendesk, +HubSpot, etc.) can compose to deliver commercial-grade behaviour: + +* **Regex ID validation** — :func:`validate_id` with ready-to-use + patterns for UUID, slug, integer, and Salesforce-style IDs. Importers + building queries from upstream-supplied IDs MUST validate before + interpolation to block injection. +* **Rate-limit header parsing** — :func:`parse_rate_limit_headers` + understands the de-facto standard (``X-RateLimit-Remaining``, + ``X-RateLimit-Reset``, ``Retry-After``) and surfaces a structured + :class:`RateLimitInfo` so callers can sleep until reset rather than + retrying blindly into a 429 storm. +* **Cursor-based pagination** — :func:`paginate` (sync) and + :func:`apaginate` (async generator) drive any "fetch-page → next + cursor → repeat" upstream into a flat iterator of records, batching + parent IDs through ``WHERE x IN (…)``-style clauses while honouring + upstream batch limits. +* **Retry with exponential backoff + jitter** — :func:`retry_with_backoff` + implements decorrelated full-jitter backoff (AWS Architecture Blog, + *Exponential Backoff and Jitter*, 2015) and respects rate-limit + headers when the failure is HTTP 429. + +The pattern is lifted from the Agentforce importer +(``frameworks/agentforce/auth.py`` + ``frameworks/agentforce/importer.py``) +where it has been hardened against real Salesforce production loads. +This module factors the cross-cutting bits out so subsequent importer +adapters (next planned: Langfuse refactor in this PR; future: +ServiceNow, Zendesk, HubSpot) inherit the same semantics by default. + +Threading model: helpers are thread-safe. :class:`BaseImporter` +subclasses MUST treat instance state as single-writer per ``run()`` +call (no concurrent ``run()`` invocations on the same instance). +""" + +from __future__ import annotations + +import re +import abc +import time +import random +import asyncio +import logging +from typing import ( + Any, + Dict, + List, + Generic, + TypeVar, + Callable, + Iterator, + Optional, + Awaitable, + AsyncIterator, +) +from datetime import datetime, timezone +from dataclasses import field, dataclass + +UTC = timezone.utc # Python 3.11+ has datetime.UTC; alias for 3.9/3.10 compat. + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# ID validation patterns +# --------------------------------------------------------------------------- + +#: Canonical 8-4-4-4-12 hex UUID (with or without hyphens, case-insensitive). +ID_PATTERN_UUID: re.Pattern[str] = re.compile( + r"^[0-9a-fA-F]{8}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{4}-?[0-9a-fA-F]{12}$" +) + +#: Lowercase URL slug — letters, digits, hyphens, underscores. Length 1-256. +ID_PATTERN_SLUG: re.Pattern[str] = re.compile(r"^[a-zA-Z0-9_\-]{1,256}$") + +#: Decimal integer (positive or negative). Up to 19 digits (signed 64-bit max). +ID_PATTERN_INTEGER: re.Pattern[str] = re.compile(r"^-?\d{1,19}$") + +#: Salesforce record ID — 15 case-sensitive or 18 case-insensitive alphanumeric. +ID_PATTERN_SALESFORCE: re.Pattern[str] = re.compile(r"^[a-zA-Z0-9]{15}(?:[a-zA-Z0-9]{3})?$") + +#: ISO 8601 date (YYYY-MM-DD). +ID_PATTERN_DATE: re.Pattern[str] = re.compile(r"^\d{4}-\d{2}-\d{2}$") + +#: ISO 8601 timestamp (date + 'T' + time, optional offset/zulu). +ID_PATTERN_TIMESTAMP: re.Pattern[str] = re.compile( + r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:?\d{2})?$" +) + + +def validate_id(value: Any, pattern: re.Pattern[str]) -> bool: + """Return ``True`` iff ``value`` is a string fully matching ``pattern``. + + Use this BEFORE interpolating an upstream-supplied identifier into + any query string (SOQL, SQL, GraphQL, JSON-Path, URL path segment). + Validation against a compiled regex is the difference between an + importer that rejects malformed input cleanly and one that issues + an injection-vulnerable query upstream. + + Args: + value: Candidate identifier. Non-strings always return ``False`` + (typed as :class:`Any` so the function is safe to invoke on + untrusted upstream payloads where the field type is not + statically guaranteed). + pattern: Compiled regex; must contain ``^…$`` anchors. Use one of + the ``ID_PATTERN_*`` constants in this module or supply a + tighter project-specific pattern. + + Returns: + ``True`` when ``value`` is a string and the pattern matches the + entire string. ``False`` otherwise — never raises. + """ + if not isinstance(value, str): + return False + return pattern.match(value) is not None + + +def require_valid_id(value: Any, pattern: re.Pattern[str], *, label: str = "id") -> str: + """Validate and return ``value`` or raise :class:`ValueError`. + + Strict variant of :func:`validate_id` for code paths where an + invalid id is a hard error (e.g. SOQL/SQL interpolation). The + ``label`` argument is included in the error message so callers can + locate the offending field at the boundary. + + Args: + value: Candidate identifier. + pattern: Compiled regex; see :func:`validate_id`. + label: Field name used in the error message. + + Returns: + The original ``value``, unchanged, when valid. + + Raises: + ValueError: When ``value`` is not a string or the pattern does + not match. + """ + if not validate_id(value, pattern): + raise ValueError(f"Invalid {label} format: {value!r}") + # validate_id has confirmed isinstance(value, str); narrow for mypy. + assert isinstance(value, str) # noqa: S101 — tightening type after validate. + return value + + +# --------------------------------------------------------------------------- +# Rate-limit header parsing +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class RateLimitInfo: + """Parsed view of a response's rate-limit headers. + + Headers vary between providers, so :func:`parse_rate_limit_headers` + normalises the common shapes: + + * ``X-RateLimit-Remaining`` — request budget left in the current window. + * ``X-RateLimit-Limit`` — total budget per window (when advertised). + * ``X-RateLimit-Reset`` — unix timestamp (seconds) when the budget + replenishes. Some providers also send the value in seconds-from-now; + we treat values lower than 10 years' worth of seconds as relative. + * ``Retry-After`` — RFC 7231 §7.1.3. Either delta-seconds or an + HTTP-date. Surfaced as :attr:`retry_after_s`. + + Attributes: + remaining: Budget left in the current window, or ``None`` when + the header was absent or unparseable. + limit: Total budget per window, or ``None``. + reset_at: Absolute timestamp (UTC) when the budget replenishes, + or ``None``. + retry_after_s: Seconds the caller SHOULD sleep before retrying + (Retry-After header), or ``None``. + usage_ratio: ``(limit - remaining) / limit`` when both are + known, else ``None``. ``0.8`` means 80 % of the budget has + been consumed. + """ + + remaining: Optional[int] = None + limit: Optional[int] = None + reset_at: Optional[datetime] = None + retry_after_s: Optional[float] = None + usage_ratio: Optional[float] = None + + @property + def is_throttled(self) -> bool: + """True when the upstream has explicitly asked us to back off. + + This is the case when either ``Retry-After`` was sent (response + was 429 / 503) or ``remaining`` is zero. + """ + if self.retry_after_s is not None and self.retry_after_s > 0: + return True + if self.remaining is not None and self.remaining <= 0: + return True + return False + + def sleep_seconds(self) -> float: + """Recommended sleep duration before the next attempt. + + Prefers ``Retry-After`` when present; otherwise falls back to + seconds-until-reset; otherwise zero. + """ + if self.retry_after_s is not None and self.retry_after_s > 0: + return float(self.retry_after_s) + if self.reset_at is not None: + now = datetime.now(UTC) + delta = (self.reset_at - now).total_seconds() + if delta > 0: + return float(delta) + return 0.0 + + +# Ten years of seconds — anything smaller than this is treated as a +# relative seconds-from-now value rather than an absolute unix epoch. +# Picked deliberately to be large enough that 1970-style epoch values +# above year ~1980 are still parsed as absolute. +_RESET_RELATIVE_THRESHOLD_S = 10 * 365 * 24 * 3600 + + +def _parse_int(value: Optional[str]) -> Optional[int]: + if value is None: + return None + try: + return int(value.strip()) + except (ValueError, AttributeError): + return None + + +def _parse_retry_after(value: Optional[str]) -> Optional[float]: + """Parse RFC 7231 Retry-After (delta-seconds OR HTTP-date).""" + if value is None: + return None + raw = value.strip() + if not raw: + return None + try: + return float(raw) + except ValueError: + pass + # Fallback: HTTP-date. Use email.utils which understands RFC 1123. + from email.utils import parsedate_to_datetime + + try: + dt = parsedate_to_datetime(raw) + except (TypeError, ValueError): + return None + if dt.tzinfo is None: + dt = dt.replace(tzinfo=UTC) + delta = (dt - datetime.now(UTC)).total_seconds() + return max(0.0, delta) + + +def _parse_reset(value: Optional[str]) -> Optional[datetime]: + """Parse a rate-limit reset header. + + Accepts an absolute unix timestamp (seconds since epoch) OR a + seconds-from-now value. Heuristic: values below + ``_RESET_RELATIVE_THRESHOLD_S`` are treated as relative. + """ + parsed = _parse_int(value) + if parsed is None: + return None + now = datetime.now(UTC) + if parsed < _RESET_RELATIVE_THRESHOLD_S: + return now.fromtimestamp(now.timestamp() + parsed, UTC) + return datetime.fromtimestamp(float(parsed), UTC) + + +def _normalise_headers(headers: Any) -> Dict[str, str]: + """Lower-case the header keys for case-insensitive lookup. + + Accepts dict-likes (including ``http.client.HTTPMessage``, + ``requests.structures.CaseInsensitiveDict``, plain dicts). + """ + out: Dict[str, str] = {} + if headers is None: + return out + if hasattr(headers, "items"): + items = headers.items() + else: + try: + items = list(headers) + except TypeError: + return out + for k, v in items: + if k is None: + continue + out[str(k).lower()] = "" if v is None else str(v) + return out + + +def parse_rate_limit_headers(headers: Any) -> RateLimitInfo: + """Parse common rate-limit headers from an HTTP response. + + Looks for (case-insensitively): + + * ``X-RateLimit-Remaining`` / ``RateLimit-Remaining`` + * ``X-RateLimit-Limit`` / ``RateLimit-Limit`` + * ``X-RateLimit-Reset`` / ``RateLimit-Reset`` + * ``Retry-After`` + + Args: + headers: Mapping of header names to values. Accepts plain dicts, + ``requests.Response.headers``, ``http.client.HTTPMessage``, + and any object exposing ``items()``. + + Returns: + A :class:`RateLimitInfo`. All fields default to ``None`` when + the corresponding header is absent or unparseable — this never + raises. + """ + h = _normalise_headers(headers) + + remaining = _parse_int(h.get("x-ratelimit-remaining") or h.get("ratelimit-remaining")) + limit = _parse_int(h.get("x-ratelimit-limit") or h.get("ratelimit-limit")) + reset_at = _parse_reset(h.get("x-ratelimit-reset") or h.get("ratelimit-reset")) + retry_after_s = _parse_retry_after(h.get("retry-after")) + + usage_ratio: Optional[float] = None + if limit is not None and limit > 0 and remaining is not None: + usage_ratio = max(0.0, min(1.0, (limit - remaining) / limit)) + + return RateLimitInfo( + remaining=remaining, + limit=limit, + reset_at=reset_at, + retry_after_s=retry_after_s, + usage_ratio=usage_ratio, + ) + + +# --------------------------------------------------------------------------- +# Pagination +# --------------------------------------------------------------------------- + +T = TypeVar("T") + + +def paginate( + fetch_fn: Callable[[Optional[str]], Dict[str, Any]], + cursor_field: str = "next_cursor", + data_field: str = "data", + max_pages: int = 10_000, +) -> Iterator[T]: + """Sync cursor-based pagination iterator. + + Calls ``fetch_fn(cursor)`` repeatedly until the response no longer + advances the cursor, yielding records from ``response[data_field]`` + one at a time so callers can stream large result sets without + materialising every page in memory. + + The cursor is read from ``response[cursor_field]``. The loop + terminates when: + + * ``response[cursor_field]`` is missing, ``None`` or an empty string + (provider has no further pages), OR + * ``response[data_field]`` is empty (defensive — some providers send + ``next_cursor`` even when no records remain), OR + * ``max_pages`` iterations have been performed (safety bound — set + conservatively for unattended runs). + + Args: + fetch_fn: Callable invoked once per page. First call receives + ``None``; subsequent calls receive the cursor returned by + the previous response. Must return a mapping with at least + ``data_field`` and (optionally) ``cursor_field`` keys. + cursor_field: Key on the response dict that holds the next-page + cursor. Default ``"next_cursor"`` matches Langfuse; + providers that use ``"nextRecordsUrl"`` (Salesforce) or + ``"page_token"`` (Google) supply their own value. + data_field: Key on the response dict that holds the list of + records to yield. Default ``"data"``. + max_pages: Hard cap on the number of pages fetched. Reaching + this raises :class:`RuntimeError` rather than silently + truncating — operators should be alerted when an importer + chases a runaway cursor. + + Yields: + Each record from ``response[data_field]`` across all pages, in + the order returned by the upstream. + + Raises: + RuntimeError: If ``max_pages`` is reached before the upstream + stops paginating. + """ + cursor: Optional[str] = None + pages_seen = 0 + + while pages_seen < max_pages: + response = fetch_fn(cursor) + pages_seen += 1 + + records = response.get(data_field) or [] + for record in records: + yield record + + next_cursor = response.get(cursor_field) + # Empty data ALWAYS terminates — providers occasionally send + # ``next_cursor`` on a final empty page. + if not records: + return + if not next_cursor: + return + if next_cursor == cursor: + # Defensive guard against cursor loops (provider bug). + logger.warning( + "paginate: cursor did not advance (%r) — terminating to avoid loop", + cursor, + ) + return + cursor = str(next_cursor) + + raise RuntimeError( + f"paginate: max_pages={max_pages} reached without exhausting cursor — " + "raise the bound or investigate runaway upstream pagination" + ) + + +async def apaginate( + fetch_fn: Callable[[Optional[str]], Awaitable[Dict[str, Any]]], + cursor_field: str = "next_cursor", + data_field: str = "data", + max_pages: int = 10_000, +) -> AsyncIterator[T]: + """Async generator counterpart to :func:`paginate`. + + Identical contract to the sync variant — see :func:`paginate` — but + ``fetch_fn`` is awaited instead of called synchronously. Use this + when the upstream client is built on ``httpx.AsyncClient`` / + ``aiohttp``. + + Args: + fetch_fn: Coroutine function invoked once per page. See + :func:`paginate`. + cursor_field: See :func:`paginate`. + data_field: See :func:`paginate`. + max_pages: See :func:`paginate`. + + Yields: + Each record from ``response[data_field]`` across all pages. + + Raises: + RuntimeError: As :func:`paginate`. + """ + cursor: Optional[str] = None + pages_seen = 0 + + while pages_seen < max_pages: + response = await fetch_fn(cursor) + pages_seen += 1 + + records = response.get(data_field) or [] + for record in records: + yield record + + next_cursor = response.get(cursor_field) + if not records: + return + if not next_cursor: + return + if next_cursor == cursor: + logger.warning( + "apaginate: cursor did not advance (%r) — terminating to avoid loop", + cursor, + ) + return + cursor = str(next_cursor) + + raise RuntimeError( + f"apaginate: max_pages={max_pages} reached without exhausting cursor — " + "raise the bound or investigate runaway upstream pagination" + ) + + +def batched_in(values: List[str], batch_size: int = 100) -> Iterator[List[str]]: + """Yield ``values`` in fixed-size batches for ``WHERE … IN (…)`` use. + + Most upstream APIs cap the number of identifiers permitted in a + single ``IN`` clause (Salesforce SOQL: 200; PostgreSQL: 32 767; + Langfuse: per-endpoint). Use this to chunk parent IDs before + issuing a related-records query. + + Args: + values: Identifiers to chunk. Empty input yields nothing. + batch_size: Maximum batch size. Must be positive; otherwise + raises :class:`ValueError`. + + Yields: + Sub-lists of ``values`` of length ``<= batch_size``. The final + batch may be shorter than ``batch_size``. + + Raises: + ValueError: When ``batch_size <= 0``. + """ + if batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {batch_size}") + for i in range(0, len(values), batch_size): + yield values[i : i + batch_size] + + +# --------------------------------------------------------------------------- +# Retry with backoff + jitter +# --------------------------------------------------------------------------- + + +class RetryableHTTPError(Exception): + """Marker exception that carries optional rate-limit context. + + Importers should raise this from the inner ``fn`` passed to + :func:`retry_with_backoff` when an HTTP error is transient (429, + 5xx). The ``rate_limit`` attribute, when present, lets the retry + loop sleep until the explicit reset deadline rather than using its + computed backoff — which prevents thundering-herd behaviour against + a clearly-throttled upstream. + + Args: + message: Human-readable description. + status_code: Original HTTP status code, if known. + rate_limit: :class:`RateLimitInfo` parsed from the response + headers, if available. + """ + + def __init__( + self, + message: str, + *, + status_code: Optional[int] = None, + rate_limit: Optional[RateLimitInfo] = None, + ) -> None: + super().__init__(message) + self.status_code = status_code + self.rate_limit = rate_limit + + +# Default retryable exception types. Subclasses of these are also +# retried — :func:`retry_with_backoff` accepts a custom tuple via the +# ``retry_on`` argument. +_DEFAULT_RETRYABLE: tuple[type[BaseException], ...] = ( + RetryableHTTPError, + ConnectionError, + TimeoutError, + OSError, +) + + +def _compute_backoff( + attempt: int, + base_delay: float, + max_delay: float, + jitter: bool, +) -> float: + """Decorrelated full-jitter backoff (AWS architecture blog, 2015). + + ``attempt`` is zero-indexed: the first retry uses ``attempt=0``, so + the unjittered delay is ``base_delay * 2**attempt``. When + ``jitter=True`` we sample uniformly from ``[0, capped_delay]`` — + the *full jitter* variant — which is provably the best at avoiding + synchronised retry storms when many clients fail at once. + + The result is always clamped to ``[0, max_delay]``. + """ + capped: float = min(max_delay, base_delay * (2**attempt)) + if capped < 0.0: + capped = 0.0 + if jitter: + return random.uniform(0.0, capped) + return capped + + +R = TypeVar("R") + + +def retry_with_backoff( + fn: Callable[[], R], + *, + max_retries: int = 5, + base_delay: float = 1.0, + max_delay: float = 30.0, + jitter: bool = True, + retry_on: Optional[tuple[type[BaseException], ...]] = None, + sleep: Callable[[float], None] = time.sleep, +) -> R: + """Run ``fn`` with exponential-backoff retry and full jitter. + + Real retry policy — not a stub. Honors :class:`RateLimitInfo` when + the inner function raises :class:`RetryableHTTPError` carrying one, + sleeping until the explicit ``Retry-After`` / reset deadline rather + than retrying blindly. + + Args: + fn: Zero-arg callable. Re-invoked on each attempt. + max_retries: Maximum retry attempts AFTER the initial call. So + ``max_retries=5`` gives up to 6 total invocations. Must be + ``>= 0``. + base_delay: Backoff base in seconds. The unjittered delay for + attempt ``i`` (zero-indexed) is ``base_delay * 2**i``. + max_delay: Hard upper bound on any single sleep, in seconds. + jitter: When ``True`` (default), apply full-jitter randomisation + to each sleep. Set ``False`` only when reproducible timing + is required by the call site (test fixtures excluded — they + should override ``sleep`` instead). + retry_on: Tuple of exception types to retry. Defaults to + :class:`RetryableHTTPError` plus standard transient + failures (connection / timeout / OS errors). Pass an + explicit tuple to extend or restrict. + sleep: Override for ``time.sleep``. Test fixtures pass a + recording function; production code keeps the default. + + Returns: + The value returned by ``fn`` on its first successful call. + + Raises: + ValueError: When ``max_retries < 0``. + BaseException: The exception raised by ``fn`` on the final + attempt — the original type is preserved (no wrapping). + """ + if max_retries < 0: + raise ValueError(f"max_retries must be >= 0, got {max_retries}") + + retryable: tuple[type[BaseException], ...] = retry_on if retry_on is not None else _DEFAULT_RETRYABLE + + last_exc: Optional[BaseException] = None + for attempt in range(max_retries + 1): + try: + return fn() + except retryable as exc: + last_exc = exc + if attempt >= max_retries: + break + + # Honor explicit rate-limit info when present. + rate_limit_sleep = 0.0 + if isinstance(exc, RetryableHTTPError) and exc.rate_limit is not None: + rate_limit_sleep = exc.rate_limit.sleep_seconds() + + if rate_limit_sleep > 0: + # Cap so a misconfigured Retry-After can't block forever. + delay = min(rate_limit_sleep, max_delay) + logger.debug( + "retry_with_backoff: rate-limit-aware sleep %.2fs (attempt %d/%d)", + delay, + attempt + 1, + max_retries, + ) + else: + delay = _compute_backoff(attempt, base_delay, max_delay, jitter) + logger.debug( + "retry_with_backoff: backoff sleep %.2fs (attempt %d/%d)", + delay, + attempt + 1, + max_retries, + ) + sleep(delay) + + # Exhausted: re-raise the last exception, preserving traceback. + assert last_exc is not None # noqa: S101 — invariant: loop must have set this. + raise last_exc + + +async def aretry_with_backoff( + fn: Callable[[], Awaitable[R]], + *, + max_retries: int = 5, + base_delay: float = 1.0, + max_delay: float = 30.0, + jitter: bool = True, + retry_on: Optional[tuple[type[BaseException], ...]] = None, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, +) -> R: + """Async counterpart of :func:`retry_with_backoff`. + + Identical semantics but ``fn`` is awaited and ``sleep`` defaults to + :func:`asyncio.sleep`. See :func:`retry_with_backoff` for argument + descriptions and behaviour. + """ + if max_retries < 0: + raise ValueError(f"max_retries must be >= 0, got {max_retries}") + + retryable: tuple[type[BaseException], ...] = retry_on if retry_on is not None else _DEFAULT_RETRYABLE + + last_exc: Optional[BaseException] = None + for attempt in range(max_retries + 1): + try: + return await fn() + except retryable as exc: + last_exc = exc + if attempt >= max_retries: + break + + rate_limit_sleep = 0.0 + if isinstance(exc, RetryableHTTPError) and exc.rate_limit is not None: + rate_limit_sleep = exc.rate_limit.sleep_seconds() + + if rate_limit_sleep > 0: + delay = min(rate_limit_sleep, max_delay) + else: + delay = _compute_backoff(attempt, base_delay, max_delay, jitter) + await sleep(delay) + + assert last_exc is not None # noqa: S101 + raise last_exc + + +# --------------------------------------------------------------------------- +# BaseImporter +# --------------------------------------------------------------------------- + + +@dataclass +class ImporterResult: + """Standardised result object for batch importer runs. + + Mirrors the Agentforce ``ImportResult`` shape but is generic across + record types so different importers can share a single result API. + Subclasses MAY add fields; the core counters here are enough for + log-line dashboards (``imported / skipped / failed / quarantined``). + """ + + imported_count: int = 0 + skipped_count: int = 0 + failed_count: int = 0 + quarantined_count: int = 0 + errors: List[str] = field(default_factory=list) + duration_ms: float = 0.0 + dry_run: bool = False + + @property + def total(self) -> int: + """Sum of all per-record counters.""" + return self.imported_count + self.skipped_count + self.failed_count + self.quarantined_count + + @property + def ok(self) -> bool: + """True when no errors were recorded.""" + return not self.errors and self.failed_count == 0 + + +C = TypeVar("C") # client type +S = TypeVar("S") # state type + + +class BaseImporter(abc.ABC, Generic[C, S]): + """Abstract base for trace-importer adapters. + + Concrete subclasses (Langfuse, future ServiceNow, Zendesk, HubSpot, + …) implement :meth:`fetch_records` and :meth:`process_record`; + :meth:`run` orchestrates the loop, applying validation, retry, and + rate-limit handling consistently. + + Type parameters: + C: Concrete upstream client type (e.g. ``LangfuseAPIClient``). + S: Concrete state type (e.g. ``SyncState``). + + Subclass contract: + + 1. Override :meth:`fetch_records` to yield records from the + upstream — typically by calling :func:`paginate` against a + client method. + 2. Override :meth:`process_record` to ingest one record into the + LayerLens pipeline. Return ``True`` for success, ``False`` for + a non-fatal skip; raise to register a failure. + 3. Override :meth:`record_id` to return the upstream-side ID for + quarantine tracking. + + The base :meth:`run` method handles: + + * Per-record try/except so one bad record does not abort the run. + * Rate-limit-aware retry of :meth:`fetch_records` and + :meth:`process_record`. + * Result accumulation into :class:`ImporterResult`. + * Duration measurement. + """ + + #: Subclasses MUST override. + SOURCE: str = "" + + def __init__( + self, + client: C, + state: S, + *, + max_retries: int = 5, + base_delay: float = 1.0, + max_delay: float = 30.0, + ) -> None: + self._client: C = client + self._state: S = state + self._max_retries = max_retries + self._base_delay = base_delay + self._max_delay = max_delay + + # --- Subclass hooks --- + + @abc.abstractmethod + def fetch_records(self) -> Iterator[Dict[str, Any]]: + """Yield records from the upstream, one at a time. + + Implementations typically delegate to :func:`paginate` against + a client list endpoint. The iterator MAY raise + :class:`RetryableHTTPError` mid-iteration; the orchestrator + will catch and retry the entire ``fetch_records`` call. + """ + + @abc.abstractmethod + def process_record(self, record: Dict[str, Any]) -> bool: + """Ingest a single record. Return ``True`` on success. + + Implementations should: + + * Validate the record's ID using :func:`require_valid_id`. + * Map to LayerLens canonical events. + * Push to the configured sink / store. + + Returning ``False`` increments ``skipped_count`` (deliberate + skip — e.g. record was already imported, dry-run mode). + + Raising :class:`RetryableHTTPError` triggers retry. Raising + any other exception increments ``failed_count`` and quarantines + the record after :attr:`quarantine_threshold` consecutive + failures. + """ + + @abc.abstractmethod + def record_id(self, record: Dict[str, Any]) -> str: + """Return the upstream-side identifier for quarantine tracking.""" + + #: Number of consecutive failures before a record is quarantined. + quarantine_threshold: int = 3 + + # Subclass MAY override these if its state object exposes a + # different quarantine API. Default is a no-op (no quarantining). + def is_quarantined(self, record_id: str) -> bool: + """Return True when ``record_id`` is in the quarantine set.""" + state = self._state + if hasattr(state, "is_quarantined"): + return bool(state.is_quarantined(record_id)) + return False + + def record_failure(self, record_id: str) -> bool: + """Track a failure; return True when the record is now quarantined.""" + state = self._state + if hasattr(state, "record_failure"): + return bool(state.record_failure(record_id, self.quarantine_threshold)) + return False + + # --- Orchestration --- + + def run(self, dry_run: bool = False) -> ImporterResult: + """Drive the import loop and return a populated result. + + Args: + dry_run: When ``True``, :meth:`process_record` is NOT + called — the loop only counts what *would* have been + imported. Useful for cost estimates before a real run. + + Returns: + :class:`ImporterResult` with per-record counters populated. + ``duration_ms`` is set from :func:`time.monotonic`. + """ + result = ImporterResult(dry_run=dry_run) + started = time.monotonic() + + try: + records = self._fetch_with_retry() + except Exception as exc: + logger.exception("%s: fetch_records failed terminally", self.SOURCE) + result.errors.append(f"fetch_records: {exc}") + result.failed_count = 1 + result.duration_ms = (time.monotonic() - started) * 1000.0 + return result + + for record in records: + rid = self._safe_record_id(record) + + if rid and self.is_quarantined(rid): + result.quarantined_count += 1 + continue + + if dry_run: + result.imported_count += 1 + continue + + try: + ok = self._process_with_retry(record) + except Exception as exc: + quarantined = self.record_failure(rid) if rid else False + logger.warning( + "%s: process_record failed for %r: %s", self.SOURCE, rid or "", exc + ) + result.failed_count += 1 + if quarantined: + result.quarantined_count += 1 + result.errors.append(f"{rid or ''}: {exc}") + continue + + if ok: + result.imported_count += 1 + else: + result.skipped_count += 1 + + result.duration_ms = (time.monotonic() - started) * 1000.0 + return result + + # --- Internals --- + + def _fetch_with_retry(self) -> List[Dict[str, Any]]: + """Materialise :meth:`fetch_records` under :func:`retry_with_backoff`. + + We materialise to a list because :func:`retry_with_backoff` + operates on callables that return a value — supporting a + partially-consumed generator on retry would silently + double-yield records the first attempt already produced. + """ + + def _do() -> List[Dict[str, Any]]: + return list(self.fetch_records()) + + return retry_with_backoff( + _do, + max_retries=self._max_retries, + base_delay=self._base_delay, + max_delay=self._max_delay, + ) + + def _process_with_retry(self, record: Dict[str, Any]) -> bool: + """Invoke :meth:`process_record` under :func:`retry_with_backoff`.""" + + def _do() -> bool: + return self.process_record(record) + + return retry_with_backoff( + _do, + max_retries=self._max_retries, + base_delay=self._base_delay, + max_delay=self._max_delay, + ) + + def _safe_record_id(self, record: Dict[str, Any]) -> str: + """Wrap :meth:`record_id` so a missing/malformed id can't crash the run.""" + try: + rid = self.record_id(record) + except Exception: + logger.debug("%s: record_id() raised on %r", self.SOURCE, record, exc_info=True) + return "" + return rid or "" + + +__all__ = [ + "ID_PATTERN_DATE", + "ID_PATTERN_INTEGER", + "ID_PATTERN_SALESFORCE", + "ID_PATTERN_SLUG", + "ID_PATTERN_TIMESTAMP", + "ID_PATTERN_UUID", + "BaseImporter", + "ImporterResult", + "RateLimitInfo", + "RetryableHTTPError", + "apaginate", + "aretry_with_backoff", + "batched_in", + "paginate", + "parse_rate_limit_headers", + "require_valid_id", + "retry_with_backoff", + "validate_id", +] diff --git a/src/layerlens/instrument/adapters/frameworks/langfuse/client.py b/src/layerlens/instrument/adapters/frameworks/langfuse/client.py index def3e8ea..f0c7fffe 100644 --- a/src/layerlens/instrument/adapters/frameworks/langfuse/client.py +++ b/src/layerlens/instrument/adapters/frameworks/langfuse/client.py @@ -1,14 +1,24 @@ -""" -Langfuse API Client +"""Langfuse API Client. HTTP client for the Langfuse REST API using stdlib urllib. -Supports Basic auth, pagination, and exponential backoff. + +Hardened in PR ``feat/instrument-importer-hardening`` to use the shared +importer base helpers for: + +* Rate-limit-aware retry with exponential backoff and full jitter + (:func:`layerlens.instrument.adapters._base.importer.retry_with_backoff`). +* Rate-limit header parsing + (:func:`layerlens.instrument.adapters._base.importer.parse_rate_limit_headers`) + — Langfuse Cloud emits ``X-RateLimit-Remaining`` / + ``X-RateLimit-Reset`` / ``Retry-After`` and we honour them rather than + retrying blindly into a 429 storm. +* Cursor-based pagination + (:func:`layerlens.instrument.adapters._base.importer.paginate`). """ from __future__ import annotations import json -import time import base64 import logging import contextlib @@ -19,6 +29,14 @@ from urllib.error import URLError, HTTPError from urllib.request import Request, urlopen +from layerlens.instrument.adapters._base.importer import ( + RateLimitInfo, + RetryableHTTPError, + paginate, + retry_with_backoff, + parse_rate_limit_headers, +) + logger = logging.getLogger(__name__) # Langfuse API rate limit: 429 responses trigger backoff @@ -29,7 +47,7 @@ class LangfuseAPIError(Exception): - """Raised when a Langfuse API call fails.""" + """Raised when a Langfuse API call fails (terminal — not retryable).""" def __init__(self, message: str, status_code: int | None = None, body: str = "") -> None: super().__init__(message) @@ -38,11 +56,11 @@ def __init__(self, message: str, status_code: int | None = None, body: str = "") class LangfuseAPIClient: - """ - HTTP client for the Langfuse REST API. + """HTTP client for the Langfuse REST API. Uses Basic auth with base64(public_key:secret_key). - No external dependencies — built on stdlib urllib.request. + No external dependencies — built on stdlib ``urllib.request`` and + the shared importer-base retry/rate-limit helpers. """ def __init__( @@ -62,8 +80,23 @@ def __init__( encoded = base64.b64encode(credentials.encode()).decode() self._auth_header = f"Basic {encoded}" + # Most recent rate-limit observation (set after every successful + # request and after every retryable HTTPError). Exposed via + # :attr:`last_rate_limit` so the importer can surface warnings + # at the orchestration layer. + self._last_rate_limit: RateLimitInfo | None = None + # --- Public API --- + @property + def last_rate_limit(self) -> RateLimitInfo | None: + """Most recent rate-limit info observed on a response. + + Returns ``None`` when no request has been made or when the + upstream did not include rate-limit headers. + """ + return self._last_rate_limit + def health_check(self) -> dict[str, Any]: """Check Langfuse API health.""" return self._request("GET", "/api/public/health") @@ -79,10 +112,10 @@ def list_traces( from_timestamp: datetime | None = None, to_timestamp: datetime | None = None, ) -> dict[str, Any]: - """ - List traces with pagination and filtering. + """List traces with pagination and filtering. - Returns dict with 'data' (list of trace objects) and 'meta' (pagination info). + Returns dict with 'data' (list of trace objects) and 'meta' + (pagination info, including ``totalPages``). """ params: dict[str, Any] = { "page": page, @@ -183,34 +216,66 @@ def get_all_traces( from_timestamp: datetime | None = None, to_timestamp: datetime | None = None, ) -> list[dict[str, Any]]: - """ - Fetch all traces with automatic pagination. + """Fetch all traces with automatic pagination. - Yields all pages until exhausted. + Iterates the Langfuse paged ``/api/public/traces`` endpoint via + the shared :func:`paginate` helper. The Langfuse API uses + page-number pagination (``page`` query parameter) rather than + opaque cursors; we adapt by treating the next page number as + the "cursor" — :func:`paginate` only requires an opaque token. """ - all_traces: list[dict[str, Any]] = [] - page = 1 - while True: - result = self.list_traces( - page=page, + + def _fetch(cursor: str | None) -> dict[str, Any]: + page_num = int(cursor) if cursor is not None else 1 + response = self.list_traces( + page=page_num, limit=limit, tags=tags, from_timestamp=from_timestamp, to_timestamp=to_timestamp, ) - data = result.get("data", []) - if not data: - break - all_traces.extend(data) - meta = result.get("meta", {}) + meta = response.get("meta") or {} total_pages = meta.get("totalPages", 1) - if page >= total_pages: - break - page += 1 - return all_traces + current_page = meta.get("page", page_num) + # Surface a "next cursor" only when more pages remain. + if current_page < total_pages: + response["__next_page"] = str(current_page + 1) + return response + + return list(paginate(_fetch, cursor_field="__next_page", data_field="data")) # --- Internal --- + def _build_url(self, path: str, params: dict[str, Any] | None) -> str: + """Build the full URL with optional query string.""" + url = f"{self._host}{path}" + if not params: + return url + # Handle list params (e.g., tags) — repeat key=value per item. + query_parts: list[str] = [] + for k, v in params.items(): + if isinstance(v, list): + for item in v: + query_parts.append(f"{k}={item}") + else: + query_parts.append(f"{k}={v}") + return f"{url}?{'&'.join(query_parts)}" + + def _read_response_headers(self, resp: Any) -> dict[str, str]: + """Extract a flat header dict from an ``http.client`` response. + + ``urlopen`` returns an ``http.client.HTTPResponse`` whose + ``.headers`` is an ``http.client.HTTPMessage`` (a list-like). + We surface :func:`parse_rate_limit_headers`-friendly access. + """ + headers = getattr(resp, "headers", None) + if headers is None: + return {} + out: dict[str, str] = {} + for key in headers.keys(): + out[key] = headers.get(key, "") + return out + def _request( self, method: str, @@ -218,76 +283,101 @@ def _request( params: dict[str, Any] | None = None, body: dict[str, Any] | None = None, ) -> dict[str, Any]: - """Make an HTTP request with retry and backoff.""" - url = f"{self._host}{path}" - if params: - # Handle list params (e.g., tags) - query_parts = [] - for k, v in params.items(): - if isinstance(v, list): - for item in v: - query_parts.append(f"{k}={item}") - else: - query_parts.append(f"{k}={v}") - url = f"{url}?{'&'.join(query_parts)}" + """Make an HTTP request with rate-limit-aware retry + backoff. + Translates transient HTTP errors (429, 5xx) into + :class:`RetryableHTTPError` so the shared + :func:`retry_with_backoff` driver can apply exponential backoff + with full jitter, and honours ``Retry-After`` / + ``X-RateLimit-Reset`` to sleep until the explicit deadline + rather than retrying blindly. + + Terminal errors (4xx other than 429) are wrapped in + :class:`LangfuseAPIError` and raised without retry. + """ + url = self._build_url(path, params) headers = { "Authorization": self._auth_header, "Content-Type": "application/json", "Accept": "application/json", } - data = json.dumps(body).encode() if body else None - last_error: Exception | None = None - for attempt in range(self._max_retries + 1): + def _attempt() -> dict[str, Any]: try: req = Request(url, data=data, headers=headers, method=method) with urlopen(req, timeout=self._timeout) as resp: + response_headers = self._read_response_headers(resp) + self._last_rate_limit = parse_rate_limit_headers(response_headers) + self._warn_if_throttle_imminent() + resp_body = resp.read().decode() if not resp_body: return {} - return json.loads(resp_body) # type: ignore[no-any-return] - + parsed = json.loads(resp_body) + if not isinstance(parsed, dict): + return {"data": parsed} + return parsed except HTTPError as e: status = e.code error_body = "" with contextlib.suppress(Exception): error_body = e.read().decode() + # HTTPError exposes headers on .headers — parse them so + # rate-limit info from a 429 drives the next sleep. + try: + self._last_rate_limit = parse_rate_limit_headers(e.headers) + except Exception: # pragma: no cover — defensive only + self._last_rate_limit = None if status == 429 or status >= 500: - last_error = LangfuseAPIError( - f"HTTP {status}: {error_body}", status_code=status, body=error_body - ) - if attempt < self._max_retries: - delay = min(_BACKOFF_BASE_S * (2**attempt), _BACKOFF_MAX_S) - logger.debug( - "Langfuse API %s %s returned %d, retrying in %.1fs (attempt %d/%d)", - method, - path, - status, - delay, - attempt + 1, - self._max_retries, - ) - time.sleep(delay) - continue - raise LangfuseAPIError( # noqa: B904 - f"HTTP {status}: {error_body}", status_code=status, body=error_body - ) - + raise RetryableHTTPError( + f"Langfuse {method} {path} returned HTTP {status}: {error_body}", + status_code=status, + rate_limit=self._last_rate_limit, + ) from e + # 4xx other than 429 → terminal (caller bug, not transient). + raise LangfuseAPIError( + f"HTTP {status}: {error_body}", + status_code=status, + body=error_body, + ) from e except URLError as e: - last_error = LangfuseAPIError(f"Connection error: {e}") - if attempt < self._max_retries: - delay = min(_BACKOFF_BASE_S * (2**attempt), _BACKOFF_MAX_S) - logger.debug( - "Langfuse API connection error, retrying in %.1fs (attempt %d/%d)", - delay, - attempt + 1, - self._max_retries, - ) - time.sleep(delay) - continue - raise last_error # noqa: B904 - - raise last_error or LangfuseAPIError("Max retries exceeded") + # Connection-level errors are transient. + raise RetryableHTTPError( + f"Langfuse {method} {path} connection error: {e}", + ) from e + + try: + return retry_with_backoff( + _attempt, + max_retries=self._max_retries, + base_delay=_BACKOFF_BASE_S, + max_delay=_BACKOFF_MAX_S, + jitter=True, + ) + except RetryableHTTPError as exc: + # Exhausted retries — convert to terminal LangfuseAPIError so + # callers continue to see the same exception type they did + # before this refactor. + raise LangfuseAPIError( + str(exc), + status_code=exc.status_code, + ) from exc + + def _warn_if_throttle_imminent(self) -> None: + """Log a warning when ``usage_ratio`` exceeds 80 %. + + Mirrors Agentforce's ``_check_rate_limit`` (auth.py:204-230) so + operators are alerted BEFORE the upstream actually throttles. + """ + info = self._last_rate_limit + if info is None: + return + if info.usage_ratio is not None and info.usage_ratio >= 0.8: + logger.warning( + "Langfuse API rate-limit warning: %.0f%% of quota consumed (remaining=%s, limit=%s)", + info.usage_ratio * 100.0, + info.remaining, + info.limit, + ) diff --git a/src/layerlens/instrument/adapters/frameworks/langfuse/exporter.py b/src/layerlens/instrument/adapters/frameworks/langfuse/exporter.py index d8df541f..31ec01f7 100644 --- a/src/layerlens/instrument/adapters/frameworks/langfuse/exporter.py +++ b/src/layerlens/instrument/adapters/frameworks/langfuse/exporter.py @@ -1,7 +1,18 @@ -""" -Langfuse Trace Exporter +"""Langfuse Trace Exporter. + +Reverse-maps STRATIX events to Langfuse traces and pushes them via the +API. -Reverse-maps STRATIX events to Langfuse traces and pushes them via the API. +Hardened in PR ``feat/instrument-importer-hardening``: + +* Exported batches go through + :func:`layerlens.instrument.adapters._base.importer.retry_with_backoff` + with rate-limit-aware delays, so a transient 429 / 5xx during push + doesn't lose the trace. +* Trace IDs are validated against + :data:`layerlens.instrument.adapters._base.importer.ID_PATTERN_UUID` + before being sent to Langfuse — keeps the payload well-formed and + prevents accidental relay of corrupt IDs. """ from __future__ import annotations @@ -13,6 +24,12 @@ UTC = timezone.utc # Python 3.11+ has datetime.UTC; alias for 3.9/3.10 compat. +from layerlens.instrument.adapters._base.importer import ( + ID_PATTERN_UUID, + RetryableHTTPError, + validate_id, + retry_with_backoff, +) from layerlens.instrument.adapters.frameworks.langfuse.client import LangfuseAPIError, LangfuseAPIClient from layerlens.instrument.adapters.frameworks.langfuse.config import SyncState, SyncResult, SyncDirection from layerlens.instrument.adapters.frameworks.langfuse.mapper import LayerLensToLangfuseMapper @@ -21,24 +38,34 @@ class TraceExporter: - """ - Export pipeline for STRATIX -> Langfuse. + """Export pipeline for STRATIX → Langfuse. Steps: - 1. Group STRATIX events by trace ID - 2. Reverse-map to Langfuse trace + observations - 3. Create trace and observations via Langfuse API - 4. Tag with 'stratix-exported' to prevent re-import + + 1. Group STRATIX events by trace ID. + 2. Validate each trace ID. + 3. Reverse-map to Langfuse trace + observations. + 4. Push the batch via the Langfuse ingestion API under + :func:`retry_with_backoff` so a transient 429 / 5xx doesn't lose + the trace. + 5. Tag with ``stratix-exported`` to prevent re-import. """ def __init__( self, client: LangfuseAPIClient, state: SyncState, + *, + max_retries: int = 3, + base_delay: float = 1.0, + max_delay: float = 16.0, ) -> None: self._client = client self._state = state self._mapper = LayerLensToLangfuseMapper() + self._max_retries = max_retries + self._base_delay = base_delay + self._max_delay = max_delay def export_traces( self, @@ -46,11 +73,11 @@ def export_traces( trace_ids: list[str] | None = None, dry_run: bool = False, ) -> SyncResult: - """ - Export STRATIX traces to Langfuse. + """Export STRATIX traces to Langfuse. Args: - events_by_trace: Dict mapping trace_id -> list of STRATIX event dicts. + events_by_trace: Dict mapping trace_id → list of STRATIX + event dicts. trace_ids: Optional filter — only export these trace IDs. dry_run: If True, count but don't actually export. @@ -67,6 +94,15 @@ def export_traces( result.skipped_count += 1 continue + # Validate trace id BEFORE sending to upstream. + if not validate_id(trace_id, ID_PATTERN_UUID): + logger.warning( + "Langfuse exporter: skipping trace with invalid id %r", trace_id + ) + result.skipped_count += 1 + result.errors.append(f"Invalid trace id: {trace_id!r}") + continue + # Loop prevention: skip traces that were imported from Langfuse if trace_id in self._state.imported_trace_ids: result.skipped_count += 1 @@ -90,9 +126,9 @@ def export_traces( result.errors.append(f"Trace {trace_id} mapping: {e}") continue - # Push to Langfuse + # Push to Langfuse with retry try: - self._push_to_langfuse(langfuse_data) + self._push_with_retry(langfuse_data) except LangfuseAPIError as e: logger.warning("Failed to export trace %s: %s", trace_id, e) result.failed_count += 1 @@ -105,6 +141,30 @@ def export_traces( return result + def _push_with_retry(self, langfuse_data: dict[str, Any]) -> None: + """Push the Langfuse batch under :func:`retry_with_backoff`.""" + + def _do() -> None: + try: + self._push_to_langfuse(langfuse_data) + except LangfuseAPIError as exc: + status = exc.status_code + if status is None or status == 429 or (500 <= status < 600): + raise RetryableHTTPError( + str(exc), status_code=status, rate_limit=None + ) from exc + raise + + try: + retry_with_backoff( + _do, + max_retries=self._max_retries, + base_delay=self._base_delay, + max_delay=self._max_delay, + ) + except RetryableHTTPError as exc: + raise LangfuseAPIError(str(exc), status_code=exc.status_code) from exc + def _push_to_langfuse(self, langfuse_data: dict[str, Any]) -> None: """Push a mapped trace + observations to Langfuse via batch ingestion.""" trace_body = langfuse_data.get("trace", {}) diff --git a/src/layerlens/instrument/adapters/frameworks/langfuse/importer.py b/src/layerlens/instrument/adapters/frameworks/langfuse/importer.py index e309fda3..9a70e3c8 100644 --- a/src/layerlens/instrument/adapters/frameworks/langfuse/importer.py +++ b/src/layerlens/instrument/adapters/frameworks/langfuse/importer.py @@ -1,8 +1,19 @@ -""" -Langfuse Trace Importer +"""Langfuse Trace Importer. Fetches traces from Langfuse, maps them to STRATIX events, deduplicates, and ingests via the STRATIX pipeline. + +Hardened in PR ``feat/instrument-importer-hardening``: + +* Trace IDs are validated against + :data:`layerlens.instrument.adapters._base.importer.ID_PATTERN_UUID` + before being interpolated into the per-trace fetch URL — Langfuse + trace IDs are UUIDs, and an upstream-supplied non-UUID could + otherwise be relayed verbatim into the URL path. +* Per-trace fetches and per-trace ingest calls are wrapped in + :func:`layerlens.instrument.adapters._base.importer.retry_with_backoff` + so transient HTTP failures recover cleanly without quarantining the + trace prematurely. """ from __future__ import annotations @@ -13,6 +24,12 @@ UTC = timezone.utc # Python 3.11+ has datetime.UTC; alias for 3.9/3.10 compat. +from layerlens.instrument.adapters._base.importer import ( + ID_PATTERN_UUID, + RetryableHTTPError, + validate_id, + retry_with_backoff, +) from layerlens.instrument.adapters.frameworks.langfuse.client import LangfuseAPIError, LangfuseAPIClient from layerlens.instrument.adapters.frameworks.langfuse.config import SyncState, SyncResult, SyncDirection from layerlens.instrument.adapters.frameworks.langfuse.mapper import LangfuseToLayerLensMapper @@ -21,25 +38,37 @@ class TraceImporter: - """ - Import pipeline for Langfuse -> STRATIX. + """Import pipeline for Langfuse → STRATIX. Steps: - 1. List traces from Langfuse (with filters) - 2. Fetch full trace with observations - 3. Map to STRATIX events - 4. Deduplicate against previously imported traces - 5. Ingest via STRATIX emit or pipeline + + 1. List traces from Langfuse (with filters), via the client's + :meth:`get_all_traces` (paginated through the shared + :func:`paginate` helper). + 2. Validate each trace ID against + :data:`ID_PATTERN_UUID`. + 3. Fetch the full trace with observations (retry on transient + failure). + 4. Map to STRATIX events. + 5. Deduplicate against previously-imported traces. + 6. Ingest via STRATIX emit (retry on transient failure). """ def __init__( self, client: LangfuseAPIClient, state: SyncState, + *, + max_retries: int = 3, + base_delay: float = 1.0, + max_delay: float = 16.0, ) -> None: self._client = client self._state = state self._mapper = LangfuseToLayerLensMapper() + self._max_retries = max_retries + self._base_delay = base_delay + self._max_delay = max_delay def import_traces( self, @@ -49,8 +78,7 @@ def import_traces( limit: int | None = None, dry_run: bool = False, ) -> SyncResult: - """ - Import traces from Langfuse. + """Import traces from Langfuse. Args: stratix: STRATIX instance for event emission (or pipeline). @@ -64,7 +92,9 @@ def import_traces( """ result = SyncResult(direction=SyncDirection.IMPORT, dry_run=dry_run) - # Fetch trace list + # Fetch trace list (no retry here — get_all_traces internally + # uses the rate-limit-aware client._request which already + # retries). A LangfuseAPIError surfaced here is terminal. try: traces = self._client.get_all_traces( tags=tags, @@ -81,6 +111,17 @@ def import_traces( for trace_summary in traces: trace_id = trace_summary.get("id", "") + # Reject malformed trace IDs before they reach the upstream + # URL builder. Langfuse uses UUIDs; anything else is either + # corruption or an injection attempt. + if not validate_id(trace_id, ID_PATTERN_UUID): + logger.warning( + "Langfuse importer: skipping trace with invalid id %r", trace_id + ) + result.skipped_count += 1 + result.errors.append(f"Invalid trace id: {trace_id!r}") + continue + # Skip quarantined traces if self._state.is_quarantined(trace_id): result.quarantined_count += 1 @@ -101,9 +142,9 @@ def import_traces( result.imported_count += 1 continue - # Fetch full trace + # Fetch full trace with retry/backoff (handles transient 5xx). try: - full_trace = self._client.get_trace(trace_id) + full_trace = self._fetch_trace_with_retry(trace_id) except LangfuseAPIError as e: logger.warning("Failed to fetch trace %s: %s", trace_id, e) is_quarantined = self._state.record_failure(trace_id) @@ -150,6 +191,41 @@ def import_traces( return result + def _fetch_trace_with_retry(self, trace_id: str) -> dict[str, Any]: + """Fetch a single trace under :func:`retry_with_backoff`. + + The client's ``get_trace`` already retries via the shared + retry helper for HTTP-level failures (429/5xx). We add a + second retry layer here so a brief network blip on a single + trace doesn't quarantine it after one attempt — quarantine + should be reserved for traces that are *truly* unfetchable. + """ + + def _do() -> dict[str, Any]: + try: + return self._client.get_trace(trace_id) + except LangfuseAPIError as exc: + # Translate LangfuseAPIError into RetryableHTTPError ONLY + # when the status is transient (network or 5xx). 4xx is + # terminal — re-raise unchanged. + status = exc.status_code + if status is None or status == 429 or (500 <= status < 600): + raise RetryableHTTPError( + str(exc), status_code=status, rate_limit=None + ) from exc + raise + + try: + return retry_with_backoff( + _do, + max_retries=self._max_retries, + base_delay=self._base_delay, + max_delay=self._max_delay, + ) + except RetryableHTTPError as exc: + # Exhausted — surface as LangfuseAPIError for caller parity. + raise LangfuseAPIError(str(exc), status_code=exc.status_code) from exc + def _ingest_events( self, events: list[dict[str, Any]], diff --git a/tests/instrument/adapters/_base/__init__.py b/tests/instrument/adapters/_base/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/instrument/adapters/_base/test_importer.py b/tests/instrument/adapters/_base/test_importer.py new file mode 100644 index 00000000..37c3dfd4 --- /dev/null +++ b/tests/instrument/adapters/_base/test_importer.py @@ -0,0 +1,676 @@ +"""Unit tests for the shared importer base. + +Covers every public symbol exported from +``layerlens.instrument.adapters._base.importer``: + +* :func:`validate_id` / :func:`require_valid_id` against each + ``ID_PATTERN_*``. +* :func:`parse_rate_limit_headers` for the de-facto-standard headers + plus edge cases (missing headers, malformed values, RFC 7231 + ``Retry-After`` HTTP-dates, absolute-vs-relative reset). +* :class:`RateLimitInfo` semantics — :attr:`is_throttled` and + :meth:`sleep_seconds`. +* :func:`paginate` (cursor termination, empty-data short-circuit, + cursor-loop guard, ``max_pages`` bound). +* :func:`apaginate` async generator parity. +* :func:`batched_in` chunking semantics. +* :func:`retry_with_backoff` — success first-try, success after + retries, exhaustion, rate-limit-aware sleep override, jitter + off vs on, ``max_retries`` validation. +* :class:`BaseImporter` end-to-end orchestration including quarantine, + retry, dry-run, and per-record failure isolation. +""" + +from __future__ import annotations + +import asyncio +from typing import Any, Dict, List, Iterator, Optional +from datetime import datetime, timezone, timedelta + +import pytest + +from layerlens.instrument.adapters._base.importer import ( + ID_PATTERN_DATE, + ID_PATTERN_SLUG, + ID_PATTERN_UUID, + ID_PATTERN_INTEGER, + ID_PATTERN_TIMESTAMP, + ID_PATTERN_SALESFORCE, + BaseImporter, + RateLimitInfo, + ImporterResult, + RetryableHTTPError, + paginate, + apaginate, + batched_in, + validate_id, + require_valid_id, + retry_with_backoff, + parse_rate_limit_headers, +) + +UTC = timezone.utc + + +# --------------------------------------------------------------------------- +# validate_id / require_valid_id +# --------------------------------------------------------------------------- + + +class TestValidateId: + def test_uuid_canonical_form_accepts(self) -> None: + assert validate_id("550e8400-e29b-41d4-a716-446655440000", ID_PATTERN_UUID) + + def test_uuid_no_hyphens_accepts(self) -> None: + assert validate_id("550e8400e29b41d4a716446655440000", ID_PATTERN_UUID) + + def test_uuid_uppercase_accepts(self) -> None: + assert validate_id("550E8400-E29B-41D4-A716-446655440000", ID_PATTERN_UUID) + + def test_uuid_rejects_non_hex(self) -> None: + assert not validate_id("550e8400-e29b-41d4-a716-44665544000z", ID_PATTERN_UUID) + + def test_slug_accepts_typical(self) -> None: + assert validate_id("user-123_alpha", ID_PATTERN_SLUG) + + def test_slug_rejects_space(self) -> None: + assert not validate_id("user 123", ID_PATTERN_SLUG) + + def test_slug_rejects_quote(self) -> None: + # Critical: SQL-injection-style payloads MUST be rejected. + assert not validate_id("user'); DROP TABLE--", ID_PATTERN_SLUG) + + def test_integer_positive(self) -> None: + assert validate_id("12345", ID_PATTERN_INTEGER) + + def test_integer_negative(self) -> None: + assert validate_id("-42", ID_PATTERN_INTEGER) + + def test_integer_rejects_decimal(self) -> None: + assert not validate_id("3.14", ID_PATTERN_INTEGER) + + def test_salesforce_15_char(self) -> None: + assert validate_id("0011x00000abcde", ID_PATTERN_SALESFORCE) + + def test_salesforce_18_char(self) -> None: + assert validate_id("0011x00000abcdeAAA", ID_PATTERN_SALESFORCE) + + def test_salesforce_rejects_short(self) -> None: + assert not validate_id("0011x", ID_PATTERN_SALESFORCE) + + def test_date_iso_8601(self) -> None: + assert validate_id("2026-04-25", ID_PATTERN_DATE) + + def test_date_rejects_us_format(self) -> None: + assert not validate_id("04/25/2026", ID_PATTERN_DATE) + + def test_timestamp_with_z(self) -> None: + assert validate_id("2026-04-25T13:45:00Z", ID_PATTERN_TIMESTAMP) + + def test_timestamp_with_offset(self) -> None: + assert validate_id("2026-04-25T13:45:00+02:00", ID_PATTERN_TIMESTAMP) + + def test_timestamp_with_microseconds(self) -> None: + assert validate_id("2026-04-25T13:45:00.123456Z", ID_PATTERN_TIMESTAMP) + + def test_non_string_returns_false(self) -> None: + # Must not raise — defensive against upstream payload variation. + assert not validate_id(None, ID_PATTERN_UUID) + assert not validate_id(12345, ID_PATTERN_UUID) + assert not validate_id(["uuid"], ID_PATTERN_UUID) + assert not validate_id({"id": "uuid"}, ID_PATTERN_UUID) + + def test_require_valid_id_returns_input_when_valid(self) -> None: + value = "550e8400-e29b-41d4-a716-446655440000" + assert require_valid_id(value, ID_PATTERN_UUID) is value + + def test_require_valid_id_raises_with_label(self) -> None: + with pytest.raises(ValueError, match="Invalid trace_id format"): + require_valid_id("not-a-uuid", ID_PATTERN_UUID, label="trace_id") + + def test_require_valid_id_raises_on_non_string(self) -> None: + with pytest.raises(ValueError): + require_valid_id(None, ID_PATTERN_UUID) + + +# --------------------------------------------------------------------------- +# parse_rate_limit_headers / RateLimitInfo +# --------------------------------------------------------------------------- + + +class TestParseRateLimitHeaders: + def test_empty_headers_yields_all_none(self) -> None: + info = parse_rate_limit_headers({}) + assert info.remaining is None + assert info.limit is None + assert info.reset_at is None + assert info.retry_after_s is None + assert info.usage_ratio is None + assert not info.is_throttled + + def test_none_headers_safe(self) -> None: + info = parse_rate_limit_headers(None) + assert info.remaining is None + + def test_x_ratelimit_remaining_and_limit(self) -> None: + info = parse_rate_limit_headers( + {"X-RateLimit-Remaining": "10", "X-RateLimit-Limit": "100"} + ) + assert info.remaining == 10 + assert info.limit == 100 + # 90% used → ratio 0.9 + assert info.usage_ratio == pytest.approx(0.9, abs=1e-6) + + def test_unprefixed_ratelimit_headers(self) -> None: + # IETF draft uses unprefixed names. + info = parse_rate_limit_headers( + {"RateLimit-Remaining": "5", "RateLimit-Limit": "20"} + ) + assert info.remaining == 5 + assert info.limit == 20 + + def test_case_insensitive_lookup(self) -> None: + info = parse_rate_limit_headers( + {"x-ratelimit-remaining": "1", "X-RATELIMIT-LIMIT": "2"} + ) + assert info.remaining == 1 + assert info.limit == 2 + + def test_retry_after_seconds(self) -> None: + info = parse_rate_limit_headers({"Retry-After": "30"}) + assert info.retry_after_s == 30.0 + assert info.is_throttled + assert info.sleep_seconds() == 30.0 + + def test_retry_after_http_date(self) -> None: + future = datetime.now(UTC) + timedelta(seconds=120) + # RFC 1123 format + http_date = future.strftime("%a, %d %b %Y %H:%M:%S GMT") + info = parse_rate_limit_headers({"Retry-After": http_date}) + assert info.retry_after_s is not None + # Allow ±5s for parse-vs-now slop. + assert 110 <= info.retry_after_s <= 125 + + def test_retry_after_malformed_returns_none(self) -> None: + info = parse_rate_limit_headers({"Retry-After": "not a duration"}) + assert info.retry_after_s is None + + def test_reset_absolute_timestamp(self) -> None: + # Far-future absolute epoch. + future_epoch = int(datetime.now(UTC).timestamp() + 3600) + info = parse_rate_limit_headers({"X-RateLimit-Reset": str(future_epoch)}) + assert info.reset_at is not None + delta = (info.reset_at - datetime.now(UTC)).total_seconds() + assert 3500 <= delta <= 3700 + + def test_reset_relative_seconds(self) -> None: + info = parse_rate_limit_headers({"X-RateLimit-Reset": "60"}) + # 60 < threshold → treated as relative. + assert info.reset_at is not None + delta = (info.reset_at - datetime.now(UTC)).total_seconds() + assert 55 <= delta <= 65 + + def test_remaining_zero_is_throttled(self) -> None: + info = parse_rate_limit_headers({"X-RateLimit-Remaining": "0"}) + assert info.is_throttled + + def test_sleep_seconds_prefers_retry_after_over_reset(self) -> None: + info = parse_rate_limit_headers( + {"Retry-After": "5", "X-RateLimit-Reset": "300"} + ) + # Retry-After (5) wins over computed reset delta (~300). + assert info.sleep_seconds() == 5.0 + + def test_sleep_seconds_falls_back_to_reset(self) -> None: + info = parse_rate_limit_headers({"X-RateLimit-Reset": "30"}) + slp = info.sleep_seconds() + assert 25 <= slp <= 35 + + def test_sleep_seconds_zero_when_no_signal(self) -> None: + assert RateLimitInfo().sleep_seconds() == 0.0 + + def test_handles_request_response_style_headers(self) -> None: + # requests.structures.CaseInsensitiveDict has .items() + class CaseInsensitive(dict): # type: ignore[type-arg] + pass + + h = CaseInsensitive({"X-RateLimit-Remaining": "7"}) + info = parse_rate_limit_headers(h) + assert info.remaining == 7 + + +# --------------------------------------------------------------------------- +# paginate / apaginate / batched_in +# --------------------------------------------------------------------------- + + +class TestPaginate: + def test_single_page(self) -> None: + def fetch(cursor: Optional[str]) -> Dict[str, Any]: + assert cursor is None + return {"data": [1, 2, 3]} + + assert list(paginate(fetch)) == [1, 2, 3] + + def test_multi_page_with_cursor(self) -> None: + pages = [ + {"data": [1, 2], "next_cursor": "p2"}, + {"data": [3, 4], "next_cursor": "p3"}, + {"data": [5], "next_cursor": None}, + ] + idx = {"i": 0} + cursors_seen: List[Optional[str]] = [] + + def fetch(cursor: Optional[str]) -> Dict[str, Any]: + cursors_seen.append(cursor) + page = pages[idx["i"]] + idx["i"] += 1 + return page + + assert list(paginate(fetch)) == [1, 2, 3, 4, 5] + assert cursors_seen == [None, "p2", "p3"] + + def test_empty_data_terminates_even_with_cursor(self) -> None: + # Defensive: provider sent next_cursor on an empty page. + def fetch(cursor: Optional[str]) -> Dict[str, Any]: + return {"data": [], "next_cursor": "p2"} + + assert list(paginate(fetch)) == [] + + def test_cursor_loop_guard(self, caplog: pytest.LogCaptureFixture) -> None: + # Provider buggily returns same cursor twice → loop, must terminate. + def fetch(cursor: Optional[str]) -> Dict[str, Any]: + return {"data": ["x"], "next_cursor": "stuck"} + + # First call: cursor=None, return cursor "stuck". + # Second call: cursor="stuck", returned cursor matches → terminate. + result = list(paginate(fetch)) + assert result == ["x", "x"] + + def test_max_pages_raises(self) -> None: + def fetch(cursor: Optional[str]) -> Dict[str, Any]: + next_cursor = str(int(cursor or "0") + 1) + return {"data": [int(next_cursor)], "next_cursor": next_cursor} + + with pytest.raises(RuntimeError, match="max_pages"): + list(paginate(fetch, max_pages=3)) + + def test_custom_data_field(self) -> None: + def fetch(cursor: Optional[str]) -> Dict[str, Any]: + return {"records": ["a", "b"], "next_cursor": None} + + assert list(paginate(fetch, data_field="records")) == ["a", "b"] + + def test_apaginate_async(self) -> None: + pages = [ + {"data": [1, 2], "next_cursor": "p2"}, + {"data": [3], "next_cursor": None}, + ] + idx = {"i": 0} + + async def fetch(cursor: Optional[str]) -> Dict[str, Any]: + page = pages[idx["i"]] + idx["i"] += 1 + return page + + async def collect() -> List[Any]: + out: List[Any] = [] + async for record in apaginate(fetch): + out.append(record) + return out + + assert asyncio.run(collect()) == [1, 2, 3] + + +class TestBatchedIn: + def test_chunks_of_size(self) -> None: + result = list(batched_in(["a", "b", "c", "d", "e"], batch_size=2)) + assert result == [["a", "b"], ["c", "d"], ["e"]] + + def test_empty_input(self) -> None: + assert list(batched_in([])) == [] + + def test_batch_size_must_be_positive(self) -> None: + with pytest.raises(ValueError): + list(batched_in(["a"], batch_size=0)) + + +# --------------------------------------------------------------------------- +# retry_with_backoff +# --------------------------------------------------------------------------- + + +class TestRetryWithBackoff: + def test_success_on_first_attempt(self) -> None: + calls: List[int] = [] + + def fn() -> str: + calls.append(1) + return "ok" + + sleeps: List[float] = [] + assert retry_with_backoff(fn, sleep=sleeps.append) == "ok" + assert calls == [1] + assert sleeps == [] + + def test_success_after_retries(self) -> None: + attempts = {"n": 0} + + def fn() -> str: + attempts["n"] += 1 + if attempts["n"] < 3: + raise RetryableHTTPError("transient", status_code=500) + return "done" + + sleeps: List[float] = [] + result = retry_with_backoff( + fn, base_delay=0.001, max_delay=0.01, jitter=False, sleep=sleeps.append + ) + assert result == "done" + assert attempts["n"] == 3 + # Two sleeps before success. + assert len(sleeps) == 2 + + def test_exhaustion_re_raises_last_exception(self) -> None: + def fn() -> None: + raise RetryableHTTPError("always", status_code=503) + + sleeps: List[float] = [] + with pytest.raises(RetryableHTTPError, match="always"): + retry_with_backoff( + fn, + max_retries=2, + base_delay=0.001, + max_delay=0.001, + jitter=False, + sleep=sleeps.append, + ) + # 1 initial + 2 retries → 2 sleeps before giving up. + assert len(sleeps) == 2 + + def test_non_retryable_exception_passthrough(self) -> None: + def fn() -> None: + raise ValueError("hard error") + + sleeps: List[float] = [] + with pytest.raises(ValueError, match="hard error"): + retry_with_backoff(fn, sleep=sleeps.append) + # No retry attempted for non-retryable exception. + assert sleeps == [] + + def test_jitter_off_uses_predictable_backoff(self) -> None: + def fn() -> None: + raise ConnectionError("net") + + sleeps: List[float] = [] + with pytest.raises(ConnectionError): + retry_with_backoff( + fn, + max_retries=3, + base_delay=1.0, + max_delay=10.0, + jitter=False, + sleep=sleeps.append, + ) + # Without jitter: 1.0, 2.0, 4.0 — predictable. + assert sleeps == [1.0, 2.0, 4.0] + + def test_jitter_on_caps_within_bound(self) -> None: + def fn() -> None: + raise ConnectionError("net") + + sleeps: List[float] = [] + with pytest.raises(ConnectionError): + retry_with_backoff( + fn, + max_retries=4, + base_delay=1.0, + max_delay=8.0, + jitter=True, + sleep=sleeps.append, + ) + # Each sleep is in [0, min(8.0, 1.0 * 2**i)]. + assert all(0.0 <= s <= 8.0 for s in sleeps) + + def test_rate_limit_aware_sleep_overrides_backoff(self) -> None: + # First call: 429 with Retry-After=2. + # Second call: success. + attempts = {"n": 0} + + rl = parse_rate_limit_headers({"Retry-After": "2"}) + + def fn() -> str: + attempts["n"] += 1 + if attempts["n"] == 1: + raise RetryableHTTPError("throttled", status_code=429, rate_limit=rl) + return "done" + + sleeps: List[float] = [] + retry_with_backoff( + fn, + base_delay=0.001, # tiny — would be picked if rate-limit ignored + max_delay=10.0, + jitter=False, + sleep=sleeps.append, + ) + # The single sleep should be the rate-limit value (~2.0), NOT + # the tiny backoff base — proves rate_limit took precedence. + assert len(sleeps) == 1 + assert sleeps[0] == pytest.approx(2.0, abs=0.1) + + def test_rate_limit_sleep_is_capped_by_max_delay(self) -> None: + rl = parse_rate_limit_headers({"Retry-After": "9999"}) + + def fn() -> str: + raise RetryableHTTPError("forever", status_code=429, rate_limit=rl) + + sleeps: List[float] = [] + with pytest.raises(RetryableHTTPError): + retry_with_backoff( + fn, + max_retries=1, + max_delay=5.0, + jitter=False, + sleep=sleeps.append, + ) + assert sleeps == [5.0] + + def test_max_retries_negative_raises(self) -> None: + with pytest.raises(ValueError, match="max_retries"): + retry_with_backoff(lambda: None, max_retries=-1) + + def test_custom_retry_on(self) -> None: + class MyError(Exception): + pass + + attempts = {"n": 0} + + def fn() -> str: + attempts["n"] += 1 + if attempts["n"] < 2: + raise MyError("custom") + return "ok" + + sleeps: List[float] = [] + result = retry_with_backoff( + fn, + retry_on=(MyError,), + base_delay=0.001, + max_delay=0.001, + jitter=False, + sleep=sleeps.append, + ) + assert result == "ok" + assert attempts["n"] == 2 + + +# --------------------------------------------------------------------------- +# BaseImporter +# --------------------------------------------------------------------------- + + +class _FakeState: + """Minimal state object exposing the BaseImporter quarantine API.""" + + def __init__(self) -> None: + self.failures: Dict[str, int] = {} + self.threshold = 3 + + def is_quarantined(self, record_id: str) -> bool: + return self.failures.get(record_id, 0) >= self.threshold + + def record_failure(self, record_id: str, max_failures: int = 3) -> bool: + self.failures[record_id] = self.failures.get(record_id, 0) + 1 + self.threshold = max_failures + return self.failures[record_id] >= max_failures + + +class _FakeImporter(BaseImporter[None, _FakeState]): + SOURCE = "fake" + + def __init__( + self, + state: _FakeState, + records: List[Dict[str, Any]], + *, + fail_ids: Optional[set[str]] = None, + fetch_fail_count: int = 0, + ) -> None: + super().__init__(client=None, state=state, max_retries=2, base_delay=0.001, max_delay=0.001) + self._records = records + self._fail_ids = fail_ids or set() + self._fetch_fail_count = fetch_fail_count + self._fetch_attempts = 0 + self.processed: List[str] = [] + + def fetch_records(self) -> Iterator[Dict[str, Any]]: + self._fetch_attempts += 1 + if self._fetch_attempts <= self._fetch_fail_count: + raise RetryableHTTPError("transient list error", status_code=500) + for r in self._records: + yield r + + def process_record(self, record: Dict[str, Any]) -> bool: + rid = record.get("id", "") + if rid in self._fail_ids: + raise RuntimeError(f"forced failure for {rid}") + self.processed.append(rid) + return True + + def record_id(self, record: Dict[str, Any]) -> str: + return str(record.get("id", "")) + + +class TestBaseImporter: + def test_run_imports_all_records(self) -> None: + state = _FakeState() + importer = _FakeImporter( + state, + [{"id": "r1"}, {"id": "r2"}, {"id": "r3"}], + ) + result = importer.run() + assert result.imported_count == 3 + assert result.failed_count == 0 + assert result.ok + assert importer.processed == ["r1", "r2", "r3"] + assert result.duration_ms >= 0.0 + + def test_dry_run_does_not_process(self) -> None: + state = _FakeState() + importer = _FakeImporter(state, [{"id": "r1"}, {"id": "r2"}]) + result = importer.run(dry_run=True) + assert result.imported_count == 2 + assert importer.processed == [] + assert result.dry_run + + def test_per_record_failure_isolated(self) -> None: + state = _FakeState() + importer = _FakeImporter( + state, + [{"id": "r1"}, {"id": "r2"}, {"id": "r3"}], + fail_ids={"r2"}, + ) + result = importer.run() + assert result.imported_count == 2 + assert result.failed_count == 1 + assert "r2" in result.errors[0] + assert importer.processed == ["r1", "r3"] + + def test_quarantine_after_repeated_failures(self) -> None: + state = _FakeState() + # Force r1 to fail every time. + importer = _FakeImporter(state, [{"id": "r1"}], fail_ids={"r1"}) + + # Run 1: fails, failure_count=1. + # The retry loop will attempt 3 times (max_retries=2 + 1 initial) + # but each attempt raises a non-retryable RuntimeError, so retry + # gives up after the first attempt. + result1 = importer.run() + assert result1.failed_count == 1 + assert state.failures["r1"] == 1 + + # Run 2 + 3 push the failure count to 3 → quarantined. + importer.run() + result3 = importer.run() + assert state.failures["r1"] >= 3 + assert result3.quarantined_count >= 1 + + # Run 4: r1 is quarantined, never reaches process_record. + importer.processed.clear() + result4 = importer.run() + assert result4.quarantined_count == 1 + assert importer.processed == [] + + def test_fetch_retried_on_transient_error(self) -> None: + state = _FakeState() + importer = _FakeImporter( + state, + [{"id": "r1"}], + fetch_fail_count=2, # First two fetch attempts raise. + ) + result = importer.run() + # Retry should recover and import. + assert result.imported_count == 1 + assert importer.processed == ["r1"] + + def test_fetch_terminal_failure_records_error(self) -> None: + state = _FakeState() + importer = _FakeImporter( + state, + [{"id": "r1"}], + fetch_fail_count=99, # More than max_retries can cover. + ) + result = importer.run() + assert result.imported_count == 0 + assert result.failed_count == 1 + assert any("fetch_records" in e for e in result.errors) + + def test_record_id_exception_safe(self) -> None: + class BadRecordIdImporter(_FakeImporter): + def record_id(self, record: Dict[str, Any]) -> str: + raise RuntimeError("boom") + + state = _FakeState() + importer = BadRecordIdImporter(state, [{"id": "r1"}]) + # Should not crash — record_id failure is silently logged. + result = importer.run() + assert result.imported_count == 1 + + +class TestImporterResult: + def test_total_sums_counters(self) -> None: + r = ImporterResult( + imported_count=2, skipped_count=3, failed_count=1, quarantined_count=4 + ) + assert r.total == 10 + + def test_ok_when_no_errors_or_failures(self) -> None: + assert ImporterResult().ok + + def test_not_ok_with_errors(self) -> None: + r = ImporterResult(errors=["x"]) + assert not r.ok + + def test_not_ok_with_failures(self) -> None: + r = ImporterResult(failed_count=1) + assert not r.ok diff --git a/tests/instrument/adapters/frameworks/test_langfuse_adapter.py b/tests/instrument/adapters/frameworks/test_langfuse_adapter.py index 0fc633a1..7150da75 100644 --- a/tests/instrument/adapters/frameworks/test_langfuse_adapter.py +++ b/tests/instrument/adapters/frameworks/test_langfuse_adapter.py @@ -160,3 +160,231 @@ def test_serialize_for_replay() -> None: assert rt.framework == "langfuse" assert rt.adapter_name == "LangfuseAdapter" assert "sync_state" in rt.metadata + + +# --------------------------------------------------------------------------- +# Importer-hardening (PR ``feat/instrument-importer-hardening``): +# verify the refactored importer/exporter use the shared base helpers. +# --------------------------------------------------------------------------- + + +class TestImporterHardening: + """Verify refactored TraceImporter/TraceExporter use shared helpers. + + These tests intentionally mock at the client level — no real HTTP — + so the unit boundary is the importer/exporter logic rather than the + Langfuse API itself. + """ + + def _make_state(self) -> Any: + from layerlens.instrument.adapters.frameworks.langfuse.config import SyncState + + return SyncState() + + def test_importer_skips_invalid_trace_ids(self) -> None: + """Trace summaries with non-UUID ids are skipped, not fetched.""" + from layerlens.instrument.adapters.frameworks.langfuse.importer import TraceImporter + + class FakeClient: + def __init__(self) -> None: + self.fetched: List[str] = [] + + def get_all_traces(self, **_kw: Any) -> List[Dict[str, Any]]: + # Two malformed ids and one valid UUID. + return [ + {"id": "not-a-uuid"}, + {"id": ""}, + {"id": "550e8400-e29b-41d4-a716-446655440000"}, + ] + + def get_trace(self, trace_id: str) -> Dict[str, Any]: + self.fetched.append(trace_id) + return {"id": trace_id, "timestamp": "2026-04-25T00:00:00Z"} + + client = FakeClient() + state = self._make_state() + importer = TraceImporter(client, state) # type: ignore[arg-type] + result = importer.import_traces(stratix=None) + # Only the valid UUID was fetched (and skipped because no events + # came out of the empty trace). + assert client.fetched == ["550e8400-e29b-41d4-a716-446655440000"] + # Two skips for invalid ids, plus one for the empty-events trace. + assert result.skipped_count >= 2 + # Two errors recorded (one per invalid id). + invalid_errors = [e for e in result.errors if "Invalid trace id" in e] + assert len(invalid_errors) == 2 + + def test_importer_retries_transient_fetch_failures(self) -> None: + """A transient 503 on get_trace is retried, not quarantined.""" + from layerlens.instrument.adapters.frameworks.langfuse.client import LangfuseAPIError + from layerlens.instrument.adapters.frameworks.langfuse.importer import TraceImporter + + class FakeClient: + def __init__(self) -> None: + self.attempts = 0 + + def get_all_traces(self, **_kw: Any) -> List[Dict[str, Any]]: + return [{"id": "550e8400-e29b-41d4-a716-446655440000"}] + + def get_trace(self, trace_id: str) -> Dict[str, Any]: + self.attempts += 1 + if self.attempts < 2: + raise LangfuseAPIError("transient", status_code=503) + return {"id": trace_id, "timestamp": "2026-04-25T00:00:00Z"} + + client = FakeClient() + state = self._make_state() + importer = TraceImporter( + client, # type: ignore[arg-type] + state, + max_retries=3, + base_delay=0.001, + max_delay=0.001, + ) + result = importer.import_traces(stratix=None) + # Even though first attempt failed, retry recovered → no failure. + assert result.failed_count == 0 + assert client.attempts == 2 + + def test_importer_does_not_retry_terminal_4xx(self) -> None: + """A 404 from get_trace is terminal — no retry, marked as failed.""" + from layerlens.instrument.adapters.frameworks.langfuse.client import LangfuseAPIError + from layerlens.instrument.adapters.frameworks.langfuse.importer import TraceImporter + + class FakeClient: + def __init__(self) -> None: + self.attempts = 0 + + def get_all_traces(self, **_kw: Any) -> List[Dict[str, Any]]: + return [{"id": "550e8400-e29b-41d4-a716-446655440000"}] + + def get_trace(self, trace_id: str) -> Dict[str, Any]: + self.attempts += 1 + raise LangfuseAPIError("not found", status_code=404) + + client = FakeClient() + state = self._make_state() + importer = TraceImporter( + client, # type: ignore[arg-type] + state, + max_retries=5, + base_delay=0.001, + max_delay=0.001, + ) + result = importer.import_traces(stratix=None) + # 404 is terminal — exactly one attempt. + assert client.attempts == 1 + assert result.failed_count == 1 + + def test_exporter_skips_invalid_trace_ids(self) -> None: + """Exporter rejects malformed trace ids before push.""" + from layerlens.instrument.adapters.frameworks.langfuse.exporter import TraceExporter + + class FakeClient: + def __init__(self) -> None: + self.pushed: List[List[Dict[str, Any]]] = [] + + def ingestion_batch(self, events: List[Dict[str, Any]]) -> Dict[str, Any]: + self.pushed.append(events) + return {} + + client = FakeClient() + state = self._make_state() + exporter = TraceExporter(client, state) # type: ignore[arg-type] + result = exporter.export_traces( + events_by_trace={ + "not-a-uuid": [{"event_type": "trace.start", "payload": {}}], + "550e8400-e29b-41d4-a716-446655440000": [ + {"event_type": "trace.start", "payload": {}} + ], + } + ) + # Only the valid UUID was pushed. + assert len(client.pushed) == 1 + # The invalid id was reported. + assert any("Invalid trace id" in e for e in result.errors) + + def test_exporter_retries_transient_push_failures(self) -> None: + """Transient 502 during push is retried.""" + from layerlens.instrument.adapters.frameworks.langfuse.client import LangfuseAPIError + from layerlens.instrument.adapters.frameworks.langfuse.exporter import TraceExporter + + class FakeClient: + def __init__(self) -> None: + self.attempts = 0 + + def ingestion_batch(self, events: List[Dict[str, Any]]) -> Dict[str, Any]: + self.attempts += 1 + if self.attempts < 2: + raise LangfuseAPIError("transient", status_code=502) + return {} + + client = FakeClient() + state = self._make_state() + exporter = TraceExporter( + client, # type: ignore[arg-type] + state, + max_retries=3, + base_delay=0.001, + max_delay=0.001, + ) + result = exporter.export_traces( + events_by_trace={ + "550e8400-e29b-41d4-a716-446655440000": [ + {"event_type": "trace.start", "payload": {}} + ], + } + ) + assert client.attempts == 2 + assert result.failed_count == 0 + assert result.exported_count == 1 + + def test_client_pagination_uses_shared_paginate(self) -> None: + """LangfuseAPIClient.get_all_traces walks via the shared paginate helper. + + We verify behaviour rather than internals by stubbing + list_traces and confirming the iteration pattern matches what + :func:`paginate` produces. + """ + from layerlens.instrument.adapters.frameworks.langfuse.client import LangfuseAPIClient + + client = LangfuseAPIClient(public_key="pk", secret_key="sk") + + pages = [ + {"data": [{"id": "a"}, {"id": "b"}], "meta": {"page": 1, "totalPages": 3}}, + {"data": [{"id": "c"}], "meta": {"page": 2, "totalPages": 3}}, + {"data": [{"id": "d"}], "meta": {"page": 3, "totalPages": 3}}, + ] + idx = {"i": 0} + + def fake_list_traces(*_args: Any, **_kwargs: Any) -> Dict[str, Any]: + page = pages[idx["i"]] + idx["i"] += 1 + return page + + client.list_traces = fake_list_traces # type: ignore[method-assign] + + result = client.get_all_traces(limit=2) + assert [t["id"] for t in result] == ["a", "b", "c", "d"] + + def test_client_warns_on_high_rate_limit_usage( + self, caplog: Any + ) -> None: + """Client emits a warning when usage_ratio >= 80%.""" + import logging + + from layerlens.instrument.adapters._base.importer import parse_rate_limit_headers + from layerlens.instrument.adapters.frameworks.langfuse.client import LangfuseAPIClient + + client = LangfuseAPIClient(public_key="pk", secret_key="sk") + # Simulate a response with 95% usage. + client._last_rate_limit = parse_rate_limit_headers( + {"X-RateLimit-Remaining": "5", "X-RateLimit-Limit": "100"} + ) + + with caplog.at_level(logging.WARNING): + client._warn_if_throttle_imminent() + + assert any( + "rate-limit warning" in rec.message for rec in caplog.records + )