stepfork 0.1.0a2__py3-none-any.whl

stepfork/trace/storage.py

@@ -0,0 +1,250 @@
+"""Filesystem persistence for `.sftrace` directory bundles."""
+
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+from uuid import uuid4
+
+from pydantic import TypeAdapter, ValidationError
+
+from stepfork.trace.hashing import hash_event_payloads
+from stepfork.trace.manifest import (
+    EnvironmentInfo,
+    FailureInfo,
+    RunStatus,
+    TraceManifest,
+    TraceTotals,
+)
+from stepfork.trace.models import Event, EventType, Trace
+from stepfork.trace.redaction import RedactionManifest, prepare_trace_for_persistence
+
+SCHEMA_VERSION = "0.1"
+MANIFEST_FILE = "manifest.json"
+EVENTS_FILE = "events.jsonl"
+REDACTIONS_FILE = "redactions.json"
+REQUIRED_FILES = (MANIFEST_FILE, EVENTS_FILE, REDACTIONS_FILE)
+
+_event_adapter: TypeAdapter[Event] = TypeAdapter(Event)
+
+
+class TraceStorageError(ValueError):
+    """Raised when a `.sftrace` bundle cannot be saved or loaded."""
+
+
+def build_manifest(
+    trace: Trace,
+    *,
+    status: RunStatus | str | None = None,
+    failure: FailureInfo | None = None,
+    environment: EnvironmentInfo | None = None,
+) -> TraceManifest:
+    """Construct a manifest from an in-memory trace."""
+    manifest_status = RunStatus(status) if status is not None else trace.status
+    manifest_failure = trace.failure if failure is None else failure
+    manifest_environment = trace.environment if environment is None else environment
+    totals = _computed_totals(trace)
+
+    return TraceManifest(
+        run_id=trace.run_id,
+        agent_name=trace.agent_name,
+        created_at=trace.created_at,
+        status=manifest_status,
+        failure=manifest_failure,
+        environment=manifest_environment,
+        totals=totals,
+    )
+
+
+def save_trace(
+    trace: Trace,
+    path: Path,
+    *,
+    status: RunStatus | str | None = None,
+    failure: FailureInfo | None = None,
+    environment: EnvironmentInfo | None = None,
+    overwrite: bool = False,
+) -> Path:
+    """Persist a trace as a `.sftrace` directory and return its path."""
+    path = path.expanduser()
+    _require_sftrace_suffix(path)
+    manifest = build_manifest(
+        trace,
+        status=status,
+        failure=failure,
+        environment=environment,
+    )
+    prepared = prepare_trace_for_persistence(trace, manifest)
+    persisted_events = [hash_event_payloads(event) for event in prepared.trace.events]
+
+    if path.exists() and not path.is_dir():
+        raise TraceStorageError(f"{path} exists and is not a directory")
+    if path.exists() and any(path.iterdir()):
+        if not overwrite:
+            raise TraceStorageError(f"{path} already exists and is not empty")
+        if not _has_required_bundle_files(path):
+            raise TraceStorageError(f"{path} is not a valid .sftrace bundle")
+
+    parent = path.parent
+    parent.mkdir(parents=True, exist_ok=True)
+    temporary = parent / f".{path.name}.tmp-{uuid4().hex}"
+    temporary.mkdir()
+
+    try:
+        _write_manifest(temporary / MANIFEST_FILE, prepared.manifest)
+        _write_events(temporary / EVENTS_FILE, persisted_events)
+        _write_redactions(temporary / REDACTIONS_FILE, prepared.redactions)
+        from stepfork.trace.integrity import write_bundle_integrity
+
+        write_bundle_integrity(temporary)
+
+        if path.exists():
+            if any(path.iterdir()):
+                backup = parent / f".{path.name}.bak-{uuid4().hex}"
+                path.rename(backup)
+                try:
+                    temporary.rename(path)
+                except Exception:
+                    backup.rename(path)
+                    raise
+                shutil.rmtree(backup)
+            else:
+                path.rmdir()
+                temporary.rename(path)
+        else:
+            temporary.rename(path)
+    except Exception:
+        if temporary.exists():
+            shutil.rmtree(temporary)
+        raise
+
+    return path
+
+
+def load_trace(path: Path) -> Trace:
+    """Load a `.sftrace` directory without repairing invalid data."""
+    path = path.expanduser()
+    _require_sftrace_suffix(path)
+    if not path.exists():
+        raise TraceStorageError(f"{path} does not exist")
+    if not path.is_dir():
+        raise TraceStorageError(f"{path} is not a directory")
+
+    for filename in REQUIRED_FILES:
+        file_path = path / filename
+        if not file_path.exists():
+            raise TraceStorageError(f"missing required file: {file_path}")
+        if not file_path.is_file():
+            raise TraceStorageError(f"required path is not a file: {file_path}")
+
+    manifest = _read_manifest(path / MANIFEST_FILE)
+    events = _read_events(path / EVENTS_FILE)
+    _read_redactions(path / REDACTIONS_FILE)
+
+    return Trace(
+        run_id=manifest.run_id,
+        agent_name=manifest.agent_name,
+        created_at=manifest.created_at,
+        status=manifest.status,
+        failure=manifest.failure,
+        environment=manifest.environment,
+        totals=manifest.totals,
+        events=events,
+    )
+
+
+def _computed_totals(trace: Trace) -> TraceTotals:
+    return TraceTotals(
+        events=len(trace.events),
+        llm_calls=sum(event.type is EventType.LLM_REQUEST for event in trace.events),
+        tool_calls=sum(event.type is EventType.TOOL_CALL for event in trace.events),
+        cost_usd=trace.totals.cost_usd,
+        duration_ms=trace.totals.duration_ms,
+    )
+
+
+def _require_sftrace_suffix(path: Path) -> None:
+    if path.suffix != ".sftrace":
+        raise TraceStorageError("trace bundle path must end with .sftrace")
+
+
+def _has_required_bundle_files(path: Path) -> bool:
+    return all((path / filename).is_file() for filename in REQUIRED_FILES)
+
+
+def _write_manifest(path: Path, manifest: TraceManifest) -> None:
+    path.write_text(
+        manifest.model_dump_json(indent=2) + "\n",
+        encoding="utf-8",
+    )
+
+
+def _write_events(path: Path, events: list[Event]) -> None:
+    with path.open("w", encoding="utf-8") as file:
+        for event in events:
+            file.write(_event_adapter.dump_json(event).decode("utf-8"))
+            file.write("\n")
+
+
+def _write_redactions(path: Path, redactions: RedactionManifest) -> None:
+    path.write_text(
+        redactions.model_dump_json(indent=2) + "\n",
+        encoding="utf-8",
+    )
+
+
+def _read_manifest(path: Path) -> TraceManifest:
+    try:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+    except json.JSONDecodeError as exc:
+        raise TraceStorageError(f"{path}: invalid JSON: {exc.msg}") from exc
+
+    if not isinstance(payload, dict):
+        raise TraceStorageError(f"{path}: invalid manifest: expected object")
+    if payload.get("schema_version") != SCHEMA_VERSION:
+        raise TraceStorageError(
+            f"{path}: unsupported schema version {payload.get('schema_version')!r}"
+        )
+
+    try:
+        return TraceManifest.model_validate(payload)
+    except ValidationError as exc:
+        raise TraceStorageError(f"{path}: invalid manifest") from exc
+
+
+def _read_events(path: Path) -> list[Event]:
+    events: list[Event] = []
+    with path.open(encoding="utf-8") as file:
+        for line_number, line in enumerate(file, start=1):
+            if line.strip() == "":
+                raise TraceStorageError(f"{path}:{line_number}: blank JSONL line")
+            try:
+                payload = json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise TraceStorageError(
+                    f"{path}:{line_number}: invalid JSON: {exc.msg}"
+                ) from exc
+            try:
+                events.append(_event_adapter.validate_python(payload))
+            except ValidationError as exc:
+                raise TraceStorageError(f"{path}:{line_number}: invalid event") from exc
+    return events
+
+
+def _read_redactions(path: Path) -> None:
+    try:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+    except json.JSONDecodeError as exc:
+        raise TraceStorageError(f"{path}: invalid JSON: {exc.msg}") from exc
+    if not isinstance(payload, dict):
+        raise TraceStorageError(f"{path}: invalid redactions object")
+    if payload.get("schema_version") != SCHEMA_VERSION:
+        raise TraceStorageError(
+            f"{path}: unsupported redactions schema version "
+            f"{payload.get('schema_version')!r}"
+        )
+    try:
+        RedactionManifest.model_validate(payload)
+    except ValidationError as exc:
+        raise TraceStorageError(f"{path}: invalid redactions metadata") from exc

stepfork/trace/validation.py

@@ -0,0 +1,285 @@
+"""Structured validation for Stepfork traces and `.sftrace` bundles."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Self
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+from stepfork.trace.hashing import (
+    event_payload_hash_fields,
+    verify_event_payload_hashes,
+)
+from stepfork.trace.integrity import IntegrityStatus, verify_bundle_integrity
+from stepfork.trace.manifest import RunStatus
+from stepfork.trace.models import EventType, Trace
+from stepfork.trace.storage import REQUIRED_FILES, TraceStorageError, load_trace
+
+
+class ValidationIssue(BaseModel):
+    """Machine-readable validation issue."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    code: str
+    message: str
+    event_id: str | None = None
+    step: int | None = None
+    location: str | None = None
+
+
+class ValidationResult(BaseModel):
+    """Structured validation result."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    valid: bool = True
+    issues: list[ValidationIssue]
+
+    @model_validator(mode="after")
+    def valid_matches_issues(self) -> Self:
+        """Keep `valid` derived from the issue list."""
+        self.valid = not self.issues
+        return self
+
+
+def validate_trace(
+    trace: Trace,
+    *,
+    strict: bool = True,
+    verify_payload_hashes: bool = False,
+) -> ValidationResult:
+    """Validate cross-event structure for an in-memory trace."""
+    issues: list[ValidationIssue] = []
+
+    if trace.totals.events != len(trace.events):
+        issues.append(
+            ValidationIssue(
+                code="event_count_mismatch",
+                message=(
+                    "Manifest event count "
+                    f"{trace.totals.events} does not match {len(trace.events)}."
+                ),
+                location="manifest.json:totals.events",
+            )
+        )
+
+    seen: dict[str, int] = {}
+    first_positions: dict[str, int] = {}
+    for index, event in enumerate(trace.events):
+        first_positions.setdefault(event.id, index)
+    previous_step: int | None = None
+    run_start_count = 0
+    has_error_event = False
+
+    for index, event in enumerate(trace.events):
+        if event.id in seen:
+            issues.append(
+                ValidationIssue(
+                    code="duplicate_event_id",
+                    message=f"Event {event.id} appears more than once.",
+                    event_id=event.id,
+                    step=event.step,
+                )
+            )
+        else:
+            seen[event.id] = index
+
+        if event.run_id != trace.run_id:
+            issues.append(
+                ValidationIssue(
+                    code="run_id_mismatch",
+                    message=(
+                        f"Event {event.id} run_id {event.run_id!r} does not "
+                        f"match manifest run_id {trace.run_id!r}."
+                    ),
+                    event_id=event.id,
+                    step=event.step,
+                )
+            )
+
+        if previous_step is not None and event.step < previous_step:
+            issues.append(
+                ValidationIssue(
+                    code="step_order",
+                    message=(
+                        f"Event {event.id} step {event.step} is less than "
+                        f"the previous step {previous_step}."
+                    ),
+                    event_id=event.id,
+                    step=event.step,
+                )
+            )
+        previous_step = event.step
+
+        if event.parent_id is not None:
+            parent_index = first_positions.get(event.parent_id)
+            if parent_index is None:
+                issues.append(
+                    ValidationIssue(
+                        code="invalid_parent",
+                        message=(
+                            f"Event {event.id} references missing parent "
+                            f"{event.parent_id}."
+                        ),
+                        event_id=event.id,
+                        step=event.step,
+                    )
+                )
+            elif parent_index >= index:
+                issues.append(
+                    ValidationIssue(
+                        code="parent_not_earlier",
+                        message=(
+                            f"Event {event.id} parent {event.parent_id} does "
+                            "not appear earlier in the trace."
+                        ),
+                        event_id=event.id,
+                        step=event.step,
+                    )
+                )
+
+        if event.type is EventType.RUN_START:
+            run_start_count += 1
+        if event.type is EventType.ERROR:
+            has_error_event = True
+
+        if verify_payload_hashes:
+            expected_hash_fields = event_payload_hash_fields(event)
+            for field in sorted(expected_hash_fields):
+                if getattr(event, field) is None:
+                    issues.append(
+                        ValidationIssue(
+                            code="integrity_unverified",
+                            message=f"Event {event.id} is missing {field}.",
+                            event_id=event.id,
+                            step=event.step,
+                            location=f"events.jsonl:{field}",
+                        )
+                    )
+            for field in verify_event_payload_hashes(event):
+                issues.append(
+                    ValidationIssue(
+                        code="payload_hash_mismatch",
+                        message=f"Event {event.id} has a {field} mismatch.",
+                        event_id=event.id,
+                        step=event.step,
+                        location=f"events.jsonl:{field}",
+                    )
+                )
+
+    if strict and run_start_count != 1:
+        issues.append(
+            ValidationIssue(
+                code="run_start_count",
+                message=(
+                    "Strict validation requires exactly one run_start; "
+                    f"found {run_start_count}."
+                ),
+                location="events.jsonl",
+            )
+        )
+
+    if strict and trace.status is RunStatus.FAILED and not has_error_event:
+        issues.append(
+            ValidationIssue(
+                code="missing_error_event",
+                message=(
+                    "Failed traces require at least one error event in strict mode."
+                ),
+                location="events.jsonl",
+            )
+        )
+
+    return ValidationResult(issues=issues)
+
+
+def validate_bundle(
+    path: str | Path,
+    *,
+    strict: bool = True,
+    verify_integrity: bool = False,
+) -> ValidationResult:
+    """Load and validate a `.sftrace` bundle."""
+    bundle_path = Path(path)
+    missing = _missing_files(bundle_path)
+    if missing:
+        return ValidationResult(
+            issues=[
+                ValidationIssue(
+                    code="missing_file",
+                    message=f"Missing required file: {file_path}.",
+                    location=str(file_path),
+                )
+                for file_path in missing
+            ]
+        )
+
+    try:
+        trace = load_trace(bundle_path)
+    except TraceStorageError as exc:
+        message = str(exc)
+        return ValidationResult(
+            issues=[
+                ValidationIssue(
+                    code=_storage_error_code(message),
+                    message=message,
+                    location=_storage_error_location(message),
+                )
+            ]
+        )
+
+    result = validate_trace(
+        trace,
+        strict=strict,
+        verify_payload_hashes=verify_integrity,
+    )
+    if not verify_integrity:
+        return result
+
+    integrity_result = verify_bundle_integrity(bundle_path)
+    integrity_issues = [
+        ValidationIssue(
+            code=issue.code,
+            message=issue.message,
+            location=issue.file,
+        )
+        for issue in integrity_result.issues
+    ]
+    if integrity_result.status is IntegrityStatus.VERIFIED:
+        return result
+    return ValidationResult(issues=[*result.issues, *integrity_issues])
+
+
+def _missing_files(path: Path) -> list[Path]:
+    if not path.exists() or not path.is_dir():
+        return [path]
+    return [
+        path / filename
+        for filename in REQUIRED_FILES
+        if not (path / filename).is_file()
+    ]
+
+
+def _storage_error_code(message: str) -> str:
+    if "unsupported" in message:
+        return "unsupported_schema"
+    if "invalid JSON" in message or "blank JSONL line" in message:
+        return "invalid_json"
+    if "invalid manifest" in message:
+        return "invalid_manifest"
+    if "invalid event" in message:
+        return "invalid_event"
+    if "missing required file" in message:
+        return "missing_file"
+    return "invalid_manifest"
+
+
+def _storage_error_location(message: str) -> str | None:
+    if ":" not in message:
+        return None
+    first, second, *_ = message.split(":")
+    if second.isdigit():
+        return f"{first}:{second}"
+    return first

stepfork/version.py

@@ -0,0 +1,5 @@
+"""Single source of truth for the Stepfork package version."""
+
+from __future__ import annotations
+
+__version__ = "0.1.0a2"