world-model-optimizer 0.2.0__py3-none-any.whl

wmo/harness/workspace_patch.py

@@ -0,0 +1,270 @@
+# Copyright (c) 2026 Experiential Labs. All rights reserved.
+
+"""Hash-guarded incremental patch protocol for synchronized workspaces."""
+
+from __future__ import annotations
+
+import hashlib
+import io
+import json
+import tarfile
+from dataclasses import dataclass
+from pathlib import PurePosixPath
+
+MAX_WORKSPACE_PATCH_BYTES = 50 * 1024 * 1024
+MAX_WORKSPACE_PATCH_UNPACKED_BYTES = 512 * 1024 * 1024
+MAX_WORKSPACE_PATCH_ENTRIES = 100_000
+
+_MANIFEST_PATH = "manifest.json"
+_DATA_PREFIX = "data/"
+_VERSION = 1
+
+
+class WorkspacePatchError(ValueError):
+    """An incremental workspace patch is malformed or unsafe."""
+
+
+@dataclass(frozen=True)
+class PatchFileState:
+    """The content and executable mode expected before or after one operation."""
+
+    sha256: str
+    mode: int
+
+
+@dataclass(frozen=True)
+class WorkspacePatchOperation:
+    """One conditional path transition in an incremental patch."""
+
+    path: str
+    before: PatchFileState | None
+    after: PatchFileState | None
+
+
+@dataclass(frozen=True)
+class WorkspacePatch:
+    """Validated operations and replacement bytes keyed by workspace path."""
+
+    operations: tuple[WorkspacePatchOperation, ...]
+    files: dict[str, bytes]
+
+
+def build_workspace_patch(before_archive: bytes, after_archive: bytes) -> bytes | None:
+    """Build a deterministic patch between two full workspace archives."""
+    before = _read_workspace_archive(before_archive)
+    after = _read_workspace_archive(after_archive)
+    operations: list[WorkspacePatchOperation] = []
+    for path in sorted(set(before) | set(after)):
+        before_state = _state(before.get(path))
+        after_state = _state(after.get(path))
+        if before_state != after_state:
+            operations.append(
+                WorkspacePatchOperation(path=path, before=before_state, after=after_state)
+            )
+    if not operations:
+        return None
+
+    manifest = {
+        "version": _VERSION,
+        "operations": [
+            {
+                "path": operation.path,
+                "before": _state_json(operation.before),
+                "after": _state_json(operation.after),
+            }
+            for operation in operations
+        ],
+    }
+    manifest_bytes = json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    buffer = io.BytesIO()
+    with tarfile.open(fileobj=buffer, mode="w:gz") as archive:
+        _add_bytes(archive, _MANIFEST_PATH, manifest_bytes, 0o600)
+        for operation in operations:
+            if operation.after is None:
+                continue
+            body, _mode = after[operation.path]
+            _add_bytes(archive, f"{_DATA_PREFIX}{operation.path}", body, operation.after.mode)
+    content = buffer.getvalue()
+    if len(content) > MAX_WORKSPACE_PATCH_BYTES:
+        raise WorkspacePatchError(f"workspace patch exceeds {MAX_WORKSPACE_PATCH_BYTES} bytes")
+    return content
+
+
+def parse_workspace_patch(content: bytes) -> WorkspacePatch:
+    """Validate and decode one patch without writing any filesystem paths."""
+    if len(content) > MAX_WORKSPACE_PATCH_BYTES:
+        raise WorkspacePatchError(f"workspace patch exceeds {MAX_WORKSPACE_PATCH_BYTES} bytes")
+    members: dict[str, tuple[bytes, int]] = {}
+    total = 0
+    try:
+        with tarfile.open(fileobj=io.BytesIO(content), mode="r:gz") as archive:
+            entries = archive.getmembers()
+            if len(entries) > MAX_WORKSPACE_PATCH_ENTRIES + 1:
+                raise WorkspacePatchError("workspace patch has too many entries")
+            for member in entries:
+                name = _normalized_name(member.name)
+                if name in members:
+                    raise WorkspacePatchError(f"duplicate workspace patch path: {name}")
+                if not member.isfile():
+                    raise WorkspacePatchError(
+                        f"workspace patch entry must be a regular file: {member.name}"
+                    )
+                total += member.size
+                if total > MAX_WORKSPACE_PATCH_UNPACKED_BYTES:
+                    raise WorkspacePatchError("workspace patch expands beyond its limit")
+                source = archive.extractfile(member)
+                if source is None:
+                    raise WorkspacePatchError(
+                        f"workspace patch entry has no content: {member.name}"
+                    )
+                members[name] = (source.read(), member.mode & 0o777)
+    except WorkspacePatchError:
+        raise
+    except (tarfile.TarError, OSError, EOFError) as error:
+        raise WorkspacePatchError("workspace patch must be a gzip tar archive") from error
+
+    manifest_entry = members.pop(_MANIFEST_PATH, None)
+    if manifest_entry is None:
+        raise WorkspacePatchError("workspace patch has no manifest")
+    try:
+        raw = json.loads(manifest_entry[0])
+    except (json.JSONDecodeError, UnicodeDecodeError) as error:
+        raise WorkspacePatchError("workspace patch manifest is not valid JSON") from error
+    if not isinstance(raw, dict) or raw.get("version") != _VERSION:
+        raise WorkspacePatchError("unsupported workspace patch version")
+    raw_operations = raw.get("operations")
+    if not isinstance(raw_operations, list) or not raw_operations:
+        raise WorkspacePatchError("workspace patch has no operations")
+
+    operations: list[WorkspacePatchOperation] = []
+    files: dict[str, bytes] = {}
+    seen: set[str] = set()
+    for raw_operation in raw_operations:
+        operation = _parse_operation(raw_operation)
+        if operation.path in seen:
+            raise WorkspacePatchError(f"duplicate workspace patch operation: {operation.path}")
+        seen.add(operation.path)
+        data_name = f"{_DATA_PREFIX}{operation.path}"
+        entry = members.pop(data_name, None)
+        if operation.after is None:
+            if entry is not None:
+                raise WorkspacePatchError(
+                    f"deleted workspace patch path has replacement data: {operation.path}"
+                )
+        else:
+            if entry is None:
+                raise WorkspacePatchError(
+                    f"workspace patch has no replacement data: {operation.path}"
+                )
+            body, mode = entry
+            if _digest(body) != operation.after.sha256:
+                raise WorkspacePatchError(f"workspace patch digest mismatch: {operation.path}")
+            if mode != operation.after.mode:
+                raise WorkspacePatchError(f"workspace patch mode mismatch: {operation.path}")
+            files[operation.path] = body
+        operations.append(operation)
+    if members:
+        extra = next(iter(sorted(members)))
+        raise WorkspacePatchError(f"unexpected workspace patch entry: {extra}")
+    return WorkspacePatch(operations=tuple(operations), files=files)
+
+
+def _read_workspace_archive(content: bytes) -> dict[str, tuple[bytes, int]]:
+    """Read regular files from a trusted full snapshot used to calculate a patch."""
+    files: dict[str, tuple[bytes, int]] = {}
+    try:
+        with tarfile.open(fileobj=io.BytesIO(content), mode="r:gz") as archive:
+            entries = archive.getmembers()
+            if len(entries) > MAX_WORKSPACE_PATCH_ENTRIES:
+                raise WorkspacePatchError("workspace archive has too many entries")
+            total = 0
+            for member in entries:
+                name = _normalized_name(member.name)
+                if member.isdir():
+                    continue
+                if not member.isfile():
+                    raise WorkspacePatchError(
+                        f"workspace entry must be a regular file or directory: {member.name}"
+                    )
+                if name in files:
+                    raise WorkspacePatchError(f"duplicate workspace path: {name}")
+                total += member.size
+                if total > MAX_WORKSPACE_PATCH_UNPACKED_BYTES:
+                    raise WorkspacePatchError("workspace archive expands beyond its limit")
+                source = archive.extractfile(member)
+                if source is None:
+                    raise WorkspacePatchError(f"workspace file has no content: {name}")
+                files[name] = (source.read(), member.mode & 0o777)
+    except WorkspacePatchError:
+        raise
+    except (tarfile.TarError, OSError, EOFError) as error:
+        raise WorkspacePatchError("workspace must be a gzip tar archive") from error
+    return files
+
+
+def _parse_operation(value: object) -> WorkspacePatchOperation:
+    if not isinstance(value, dict):
+        raise WorkspacePatchError("workspace patch operation must be an object")
+    path_value = value.get("path")
+    if not isinstance(path_value, str):
+        raise WorkspacePatchError("workspace patch operation has no path")
+    path = _normalized_name(path_value)
+    if path == ".":
+        raise WorkspacePatchError("workspace patch operation cannot target the root")
+    before = _parse_state(value.get("before"))
+    after = _parse_state(value.get("after"))
+    if before == after:
+        raise WorkspacePatchError(f"workspace patch operation does not change: {path}")
+    return WorkspacePatchOperation(path=path, before=before, after=after)
+
+
+def _parse_state(value: object) -> PatchFileState | None:
+    if value is None:
+        return None
+    if not isinstance(value, dict):
+        raise WorkspacePatchError("workspace patch file state must be an object")
+    sha256 = value.get("sha256")
+    mode = value.get("mode")
+    if (
+        not isinstance(sha256, str)
+        or len(sha256) != 64
+        or any(character not in "0123456789abcdef" for character in sha256)
+        or not isinstance(mode, int)
+        or isinstance(mode, bool)
+        or not 0 <= mode <= 0o777
+    ):
+        raise WorkspacePatchError("workspace patch file state is invalid")
+    return PatchFileState(sha256=sha256, mode=mode)
+
+
+def _state(entry: tuple[bytes, int] | None) -> PatchFileState | None:
+    if entry is None:
+        return None
+    body, mode = entry
+    return PatchFileState(sha256=_digest(body), mode=mode)
+
+
+def _state_json(state: PatchFileState | None) -> dict[str, object] | None:
+    if state is None:
+        return None
+    return {"sha256": state.sha256, "mode": state.mode}
+
+
+def _add_bytes(archive: tarfile.TarFile, name: str, body: bytes, mode: int) -> None:
+    info = tarfile.TarInfo(name)
+    info.size = len(body)
+    info.mode = mode
+    info.mtime = 0
+    archive.addfile(info, io.BytesIO(body))
+
+
+def _normalized_name(name: str) -> str:
+    path = PurePosixPath(name)
+    if path.is_absolute() or ".." in path.parts:
+        raise WorkspacePatchError(f"unsafe workspace patch path: {name}")
+    parts = tuple(part for part in path.parts if part not in {"", "."})
+    return PurePosixPath(*parts).as_posix() if parts else "."
+
+
+def _digest(body: bytes) -> str:
+    return hashlib.sha256(body, usedforsecurity=False).hexdigest()

wmo/ingest/__init__.py

@@ -0,0 +1,47 @@
+"""Trace ingestion: file uploads and vendor pulls -> normalized `Trace` objects.
+
+A `TraceAdapter` turns one source's telemetry into the generic `Trace` schema. Adapters register
+themselves on import and are looked up by name (`get_adapter`) or listed (`list_adapters`). The
+span-based adapters share one normalizer (`wmo.ingest.normalize`) and the `BaseTraceAdapter`
+scaffolding, so adding a new source is transport + attribute-mapping, not a rewrite. See
+`docs/reference/ingest.md`.
+
+Bundled adapters:
+  - `otel-genai`  : OTLP-JSON spans following the OTel GenAI semantic conventions (file or pull).
+  - `chat-json`   : recorded OpenAI-style chat/tool-call conversations (file).
+Provider adapters (Braintrust, Phoenix/Arize, Langfuse, LangSmith) register when their module is
+imported; their heavy SDKs are optional extras, imported lazily inside the adapter.
+"""
+
+# Import for the registration side effect so `get_adapter(...)` works on package import. The
+# provider adapters are SDK-free (they parse exports as JSON and pull over httpx), so importing them
+# here is cheap and brings no heavy dependency — their optional extras only matter if a user drives
+# the provider's own SDK alongside `wmo ingest`.
+from wmo.ingest import braintrust as braintrust  # noqa: F401
+from wmo.ingest import langfuse as langfuse  # noqa: F401
+from wmo.ingest import langsmith as langsmith  # noqa: F401
+from wmo.ingest import mastra as mastra  # noqa: F401
+from wmo.ingest import messages as messages  # noqa: F401
+from wmo.ingest import otel_genai as otel_genai  # noqa: F401
+from wmo.ingest import phoenix as phoenix  # noqa: F401
+from wmo.ingest import postgres as postgres  # noqa: F401
+from wmo.ingest import posthog as posthog  # noqa: F401
+from wmo.ingest.adapter import (
+    TraceAdapter,
+    VendorPull,
+    get_adapter,
+    list_adapters,
+    register_adapter,
+)
+from wmo.ingest.base import BaseTraceAdapter
+from wmo.ingest.quality import drop_degenerate_traces
+
+__all__ = [
+    "BaseTraceAdapter",
+    "TraceAdapter",
+    "VendorPull",
+    "drop_degenerate_traces",
+    "get_adapter",
+    "list_adapters",
+    "register_adapter",
+]

wmo/ingest/adapter.py

@@ -0,0 +1,72 @@
+"""TraceAdapter protocol + a small registry.
+
+Sources differ in two ways: *transport* (file vs. vendor SDK) and *schema* (which OTel semantic
+convention the spans follow). An adapter owns both: it pulls/reads raw spans and normalizes them.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from pydantic import BaseModel
+
+from wmo.core.types import Trace
+
+
+class SourceCredentialError(PermissionError):
+    """A source rejected the supplied credentials (API key, DSN password).
+
+    Adapters raise this instead of bare `PermissionError` so the streaming ingest can map it
+    to the `bad_credentials` wire code without misclassifying OS-level permission failures
+    (e.g. an unreadable local file) as credential problems. Subclassing `PermissionError`
+    keeps it a stdlib type: callers never need a driver import to catch it.
+    """
+
+
+class VendorPull(BaseModel):
+    """Parameters for pulling traces from an observability vendor's API or a database."""
+
+    api_key: str | None = None  # falls back to the vendor's env var when None
+    project: str | None = None  # vendor project / workspace to pull from
+    since: str | None = None  # ISO-8601 lower bound on trace start time
+    limit: int | None = None  # max traces to pull
+
+    # Database-source transport params (the `postgres` adapter); API-vendor adapters ignore them.
+    dsn: str | None = None  # connection string; falls back to $WMO_POSTGRES_DSN
+    table: str | None = None  # table holding the trace rows
+    trace_id_column: str | None = None  # override; default "trace_id" (absent -> row-per-trace)
+    payload_column: str | None = None  # override; default "payload"
+    order_column: str | None = None  # override; default "created_at" when the table has it
+
+
+@runtime_checkable
+class TraceAdapter(Protocol):
+    """Turns one source's raw telemetry into normalized `Trace` objects."""
+
+    name: str
+
+    def from_file(self, path: str) -> list[Trace]:
+        """Read traces from an exported file (OTLP-JSON / vendor JSONL)."""
+        ...
+
+    def from_vendor(self, pull: VendorPull) -> list[Trace]:
+        """Pull traces via a vendor SDK/API."""
+        ...
+
+
+_ADAPTERS: dict[str, TraceAdapter] = {}
+
+
+def register_adapter(adapter: TraceAdapter) -> None:
+    _ADAPTERS[adapter.name] = adapter
+
+
+def get_adapter(name: str) -> TraceAdapter:
+    if name not in _ADAPTERS:
+        raise ValueError(f"no trace adapter registered for {name!r}; have {list(_ADAPTERS)}")
+    return _ADAPTERS[name]
+
+
+def list_adapters() -> list[str]:
+    """Names of all registered trace adapters, sorted (what the build source picker shows)."""
+    return sorted(_ADAPTERS)

wmo/ingest/base.py

@@ -0,0 +1,114 @@
+"""`BaseTraceAdapter` — shared scaffolding so a new source is transport + mapping, not a rewrite.
+
+A span-based provider adapter only needs to answer two questions:
+  1. how do I get raw bytes/objects? (a file path, or a vendor API/SDK pull)
+  2. how do I turn one raw payload into `SpanRecord`s? (the provider's export shape)
+Everything after that — JSON/JSONL loading, grouping spans into `Trace`s, honoring `wmo.*`
+enrichments — is shared (`wmo.ingest.normalize`). `BaseTraceAdapter` wires (1) and (2) together so
+a concrete adapter is typically ~30 lines: set `name`, implement `spans_from_payload`, and (if it
+supports live pulls) `_pull_payloads`.
+
+Subclasses override:
+  - `name`: the registry key (e.g. "phoenix").
+  - `spans_from_payload(payload) -> list[SpanRecord]`: map ONE decoded JSON payload to spans. The
+    default delegates to `wmo.ingest.normalize.collect_spans` (OTLP/OpenInference-JSON); override it
+    when the provider's export is not OTLP-shaped.
+  - `_pull_payloads(pull) -> list[JsonValue]`: yield raw payloads from the vendor API/SDK. The
+    default raises a friendly "not implemented" — so `from_file` works with no override, and
+    `from_vendor` is opt-in.
+
+`from_file` accepts a single JSON document (object or array) OR JSONL (one payload per line), and
+skips a single corrupt JSONL line rather than aborting the whole ingest.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from pydantic import JsonValue
+
+from wmo.core.types import Trace
+from wmo.ingest.adapter import VendorPull
+from wmo.ingest.normalize import SpanRecord, collect_spans, spans_to_traces
+
+
+def load_payloads(text: str) -> list[JsonValue]:
+    """Parse a whole-document JSON payload, or per-line JSONL on a top-level decode failure.
+
+    Module-level (not only a `BaseTraceAdapter` method) because format auto-detection and the
+    streaming ingest need the payloads BEFORE an adapter has been chosen.
+    """
+    try:
+        return [json.loads(text)]
+    except json.JSONDecodeError:
+        payloads: list[JsonValue] = []
+        for line in text.splitlines():
+            stripped = line.strip()
+            if not stripped:
+                continue
+            try:
+                payloads.append(json.loads(stripped))
+            except json.JSONDecodeError:
+                continue  # tolerate a truncated/corrupt line; keep the rest
+        return payloads
+
+
+class BaseTraceAdapter:
+    """Default file+vendor plumbing around the shared span normalizer."""
+
+    name: str = "base"
+
+    # --- mapping hook (override for non-OTLP export shapes) -----------------------------------
+
+    def spans_from_payload(self, payload: JsonValue) -> list[SpanRecord]:
+        """Map ONE decoded payload to `SpanRecord`s. Default: OTLP/OpenInference-JSON collection."""
+        return collect_spans(payload)
+
+    # --- transport hooks ----------------------------------------------------------------------
+
+    def _pull_payloads(self, pull: VendorPull) -> list[JsonValue]:
+        """Fetch raw payloads from the vendor API/SDK. Override to support live pulls."""
+        raise ValueError(
+            f"{self.name!r} does not support live vendor pulls yet; export traces to a file and "
+            f"use `from_file` (or `wmo ingest --source {self.name} --file <export>`)"
+        )
+
+    # --- public API (rarely overridden) -------------------------------------------------------
+
+    def collect_all(self, payloads: list[JsonValue]) -> list[SpanRecord]:
+        """Map every payload to spans and re-stamp `span_id` globally unique in emission order.
+
+        Adapters assign `span_id`/`start_nano` per payload, so a trace split across payloads (e.g.
+        one row/observation/run per JSONL line) would otherwise emit colliding ids and identical
+        `start_nano`. `spans_to_traces` sorts by `(start_nano, span_id)`, so the collision scrambles
+        the action/observation pairing. Stamping a globally monotonic `span_id` makes equal-time
+        spans (the row-adapter case, all `start_nano=0`) order by emission, while real timestamps
+        (e.g. Phoenix) still dominate the sort. Uniqueness is owned here so adapters can't get it
+        wrong individually.
+        """
+        spans: list[SpanRecord] = []
+        for payload in payloads:
+            for span in self.spans_from_payload(payload):
+                span.span_id = f"{len(spans):012d}-{span.span_id}"
+                spans.append(span)
+        return spans
+
+    def spans_from_file(self, path: str) -> list[SpanRecord]:
+        """All (uniquely re-stamped) spans in a file export: the streaming ingest's seam."""
+        return self.collect_all(load_payloads(Path(path).read_text(encoding="utf-8")))
+
+    def spans_from_pull(self, pull: VendorPull) -> list[SpanRecord]:
+        """All (uniquely re-stamped) spans from a vendor pull: the streaming ingest's seam."""
+        return self.collect_all(self._pull_payloads(pull))
+
+    def from_file(self, path: str) -> list[Trace]:
+        return spans_to_traces(self.spans_from_file(path), source=f"{self.name}:{path}")
+
+    def from_vendor(self, pull: VendorPull) -> list[Trace]:
+        traces = spans_to_traces(
+            self.spans_from_pull(pull), source=f"{self.name}:{pull.project or 'vendor'}"
+        )
+        if pull.limit is not None:
+            traces = traces[: pull.limit]
+        return traces