stepfork 0.1.0a2__py3-none-any.whl

stepfork/trace/models.py

@@ -0,0 +1,271 @@
+"""Typed event models for the Stepfork `.sftrace` v0.1 contract."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from enum import StrEnum
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeAlias, TypeVar, cast
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+from typing_extensions import TypeAliasType
+
+from stepfork.trace.manifest import (
+    EnvironmentInfo,
+    FailureInfo,
+    RunStatus,
+    TraceTotals,
+)
+from stepfork.trace.replay_policy import ReplayPolicy
+
+JsonPrimitive: TypeAlias = str | int | float | bool | None
+if TYPE_CHECKING:
+    JsonValue: TypeAlias = JsonPrimitive | list[Any] | dict[str, Any]
+else:
+    JsonValue = TypeAliasType(
+        "JsonValue",
+        JsonPrimitive | list["JsonValue"] | dict[str, "JsonValue"],
+    )
+E = TypeVar("E", bound="EventBase")
+
+
+def _new_event_id() -> str:
+    """Return an opaque local event ID."""
+    return f"evt_{uuid4().hex}"
+
+
+def _new_run_id() -> str:
+    """Return an opaque local run ID."""
+    return f"run_{uuid4().hex}"
+
+
+def _utc_now() -> datetime:
+    """Return the current timezone-aware UTC timestamp."""
+    return datetime.now(UTC)
+
+
+class EventStatus(StrEnum):
+    """Status for a recorded event."""
+
+    OK = "ok"
+    ERROR = "error"
+    SKIPPED = "skipped"
+
+
+class EventType(StrEnum):
+    """The deliberately small v0.1 event taxonomy."""
+
+    RUN_START = "run_start"
+    LLM_REQUEST = "llm_request"
+    LLM_RESPONSE = "llm_response"
+    TOOL_CALL = "tool_call"
+    TOOL_RESULT = "tool_result"
+    STATE_CHANGE = "state_change"
+    ERROR = "error"
+    RUN_END = "run_end"
+
+
+class EventBase(BaseModel):
+    """Common fields shared by all persisted trace events."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    id: str = Field(default_factory=_new_event_id, min_length=1)
+    run_id: str = Field(default_factory=_new_run_id, min_length=1)
+    parent_id: str | None = Field(default=None, min_length=1)
+
+    step: int = Field(default=0, ge=0)
+
+    type: EventType
+    timestamp: datetime = Field(default_factory=_utc_now)
+    duration_ms: int | None = Field(default=None, ge=0)
+
+    status: EventStatus = EventStatus.OK
+    replay_policy: ReplayPolicy = ReplayPolicy.FROZEN
+
+    input_hash: str | None = None
+    output_hash: str | None = None
+
+    redactions: list[str] = Field(default_factory=list)
+
+    @field_validator("timestamp")
+    @classmethod
+    def timestamp_must_be_timezone_aware(cls, value: datetime) -> datetime:
+        """Reject naive datetimes instead of assuming UTC."""
+        if value.tzinfo is None or value.utcoffset() is None:
+            raise ValueError("timestamp must be timezone-aware")
+        return value
+
+
+class RunStart(EventBase):
+    """Event marking the start of a run."""
+
+    type: Literal[EventType.RUN_START] = EventType.RUN_START
+
+    input: JsonValue | None = None
+
+
+class LLMRequest(EventBase):
+    """Request sent to an LLM provider."""
+
+    type: Literal[EventType.LLM_REQUEST] = EventType.LLM_REQUEST
+
+    provider: str | None = Field(default=None, min_length=1)
+    model: str = Field(min_length=1)
+    input: JsonValue
+
+
+class LLMResponse(EventBase):
+    """Response received from an LLM provider."""
+
+    type: Literal[EventType.LLM_RESPONSE] = EventType.LLM_RESPONSE
+
+    model: str | None = Field(default=None, min_length=1)
+    output: JsonValue
+
+    prompt_tokens: int | None = Field(default=None, ge=0)
+    completion_tokens: int | None = Field(default=None, ge=0)
+    total_tokens: int | None = Field(default=None, ge=0)
+    cost_usd: float | None = Field(default=None, ge=0)
+
+
+class ToolCall(EventBase):
+    """Call made to a tool.
+
+    Tool calls can be constructed directly with only `name` and JSON `input`.
+    Runtime defaults fill in event identity, run identity, 0-based step, and a
+    timezone-aware timestamp; `Trace.add()` normalizes run ID and step.
+    """
+
+    type: Literal[EventType.TOOL_CALL] = EventType.TOOL_CALL
+
+    name: str = Field(min_length=1)
+    input: JsonValue
+
+
+class ToolResult(EventBase):
+    """Result returned by a tool."""
+
+    type: Literal[EventType.TOOL_RESULT] = EventType.TOOL_RESULT
+
+    name: str = Field(min_length=1)
+    output: JsonValue
+
+
+class StateChange(EventBase):
+    """Change to agent state represented as JSON-compatible values."""
+
+    type: Literal[EventType.STATE_CHANGE] = EventType.STATE_CHANGE
+
+    key: str | None = Field(default=None, min_length=1)
+    before: JsonValue | None = None
+    after: JsonValue | None = None
+
+
+class ErrorEvent(EventBase):
+    """Error observed during a run."""
+
+    type: Literal[EventType.ERROR] = EventType.ERROR
+
+    status: EventStatus = EventStatus.ERROR
+    error_type: str = Field(min_length=1)
+    message: str = Field(min_length=1)
+    details: JsonValue | None = None
+
+
+class RunEnd(EventBase):
+    """Event marking the end of a run."""
+
+    type: Literal[EventType.RUN_END] = EventType.RUN_END
+
+    run_status: RunStatus
+    output: JsonValue | None = None
+
+
+Event: TypeAlias = Annotated[
+    RunStart
+    | LLMRequest
+    | LLMResponse
+    | ToolCall
+    | ToolResult
+    | StateChange
+    | ErrorEvent
+    | RunEnd,
+    Field(discriminator="type"),
+]
+
+
+class Trace(BaseModel):
+    """Minimal in-memory trace container.
+
+    `Trace.add()` uses 0-based logical steps: the first added event receives
+    step `0`, the second receives step `1`, and so on. Partial traces can be
+    saved for development and later validated in partial or strict mode.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    run_id: str = Field(default_factory=_new_run_id, min_length=1)
+    agent_name: str = Field(min_length=1)
+    created_at: datetime = Field(default_factory=_utc_now)
+    status: RunStatus = RunStatus.RUNNING
+    failure: FailureInfo | None = None
+    environment: EnvironmentInfo = Field(default_factory=EnvironmentInfo)
+    totals: TraceTotals = Field(default_factory=TraceTotals)
+    events: list[Event] = Field(default_factory=list)
+
+    @field_validator("created_at")
+    @classmethod
+    def created_at_must_be_timezone_aware(cls, value: datetime) -> datetime:
+        """Reject naive datetimes instead of assuming UTC."""
+        if value.tzinfo is None or value.utcoffset() is None:
+            raise ValueError("created_at must be timezone-aware")
+        return value
+
+    def add(self, event: E) -> E:
+        """Append an event after normalizing run ID and 0-based logical step."""
+        normalized = event.model_copy(
+            update={"run_id": self.run_id, "step": len(self.events)}
+        )
+        self.events.append(cast(Event, normalized))
+        self.totals = self._computed_totals()
+        return normalized
+
+    def _computed_totals(self) -> TraceTotals:
+        """Return manifest counters derived from the current event sequence."""
+        return TraceTotals(
+            events=len(self.events),
+            llm_calls=sum(event.type is EventType.LLM_REQUEST for event in self.events),
+            tool_calls=sum(event.type is EventType.TOOL_CALL for event in self.events),
+            cost_usd=self.totals.cost_usd,
+            duration_ms=self.totals.duration_ms,
+        )
+
+    def save(
+        self,
+        path: str | Path,
+        *,
+        status: RunStatus | str | None = None,
+        failure: FailureInfo | None = None,
+        environment: EnvironmentInfo | None = None,
+        overwrite: bool = False,
+    ) -> Path:
+        """Persist this trace as a `.sftrace` directory bundle."""
+        from stepfork.trace.storage import save_trace
+
+        return save_trace(
+            self,
+            Path(path),
+            status=status,
+            failure=failure,
+            environment=environment,
+            overwrite=overwrite,
+        )
+
+    @classmethod
+    def load(cls, path: str | Path) -> Trace:
+        """Load a trace from a `.sftrace` directory bundle."""
+        from stepfork.trace.storage import load_trace
+
+        return load_trace(Path(path))

stepfork/trace/redaction.py

@@ -0,0 +1,294 @@
+"""Best-effort secret redaction for trace persistence."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Any, Literal, cast
+
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
+
+from stepfork.trace.manifest import TraceManifest
+from stepfork.trace.models import Event, JsonValue, Trace
+
+REDACTION_REPLACEMENT = "[REDACTED]"
+MANIFEST_REDACTION_ID = "__manifest__"
+
+SENSITIVE_KEY_NORMALIZED = {
+    "apikey",
+    "accesstoken",
+    "refreshtoken",
+    "authtoken",
+    "authorization",
+    "password",
+    "passwd",
+    "secret",
+    "clientsecret",
+    "privatekey",
+    "sessiontoken",
+    "bearertoken",
+}
+
+PROTECTED_EVENT_FIELDS = {
+    "id",
+    "run_id",
+    "parent_id",
+    "step",
+    "type",
+    "timestamp",
+    "duration_ms",
+    "status",
+    "replay_policy",
+    "input_hash",
+    "output_hash",
+    "redactions",
+    "prompt_tokens",
+    "completion_tokens",
+    "total_tokens",
+    "cost_usd",
+    "run_status",
+}
+
+PROTECTED_MANIFEST_FIELDS = {
+    "format",
+    "schema_version",
+    "run_id",
+    "created_at",
+    "status",
+    "totals",
+}
+
+_event_adapter: TypeAdapter[Event] = TypeAdapter(Event)
+
+_CREDENTIAL_PATTERNS = [
+    re.compile(r"(?i)(Authorization:\s*Bearer\s+)([A-Za-z0-9._~+/=-]{8,})"),
+    re.compile(r"(?i)(Bearer\s+)([A-Za-z0-9._~+/=-]{8,})"),
+    re.compile(r"\b(sk-[A-Za-z0-9_-]{8,})\b"),
+    re.compile(r"\b(ghp_[A-Za-z0-9_]{8,})\b"),
+    re.compile(r"\b(gho_[A-Za-z0-9_]{8,})\b"),
+    re.compile(r"\b(github_pat_[A-Za-z0-9_]{8,})\b"),
+]
+
+
+class RedactionError(ValueError):
+    """Raised when a value cannot be safely redacted for persistence."""
+
+
+class RedactionEntry(BaseModel):
+    """Persisted redaction metadata without raw secret values."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    event_id: str = Field(min_length=1)
+    path: str = Field(min_length=1)
+    reason: str = Field(min_length=1)
+    replacement: Literal["[REDACTED]"] = "[REDACTED]"
+
+
+class RedactionManifest(BaseModel):
+    """Redactions recorded for a `.sftrace` bundle."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    schema_version: Literal["0.1"] = "0.1"
+    entries: list[RedactionEntry] = Field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class RedactionFinding:
+    """Internal redaction finding before an event ID is attached."""
+
+    path: str
+    reason: str
+
+
+@dataclass(frozen=True)
+class RedactionResult:
+    """Sanitized JSON value plus non-secret metadata."""
+
+    value: JsonValue
+    findings: list[RedactionFinding] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class PreparedTrace:
+    """Trace data prepared for persistence."""
+
+    trace: Trace
+    manifest: TraceManifest
+    redactions: RedactionManifest
+
+
+def redact_json(value: JsonValue) -> RedactionResult:
+    """Recursively redact sensitive keys and credential-like string spans."""
+    sanitized, findings = _redact_value(value, path="")
+    return RedactionResult(value=cast(JsonValue, sanitized), findings=findings)
+
+
+def prepare_trace_for_persistence(
+    trace: Trace,
+    manifest: TraceManifest,
+) -> PreparedTrace:
+    """Return sanitized trace data and redaction metadata for writing."""
+    manifest_payload = manifest.model_dump(mode="json")
+    sanitized_manifest_payload, manifest_findings = _redact_mapping(
+        manifest_payload,
+        path="",
+        protected_fields=PROTECTED_MANIFEST_FIELDS,
+    )
+    sanitized_manifest = TraceManifest.model_validate(sanitized_manifest_payload)
+
+    sanitized_events: list[Event] = []
+    entries: list[RedactionEntry] = []
+    for event in trace.events:
+        event_payload = event.model_dump(mode="json")
+        sanitized_event_payload, findings = _redact_mapping(
+            event_payload,
+            path="",
+            protected_fields=PROTECTED_EVENT_FIELDS,
+        )
+        event_entries = [
+            RedactionEntry(
+                event_id=event.id,
+                path=finding.path,
+                reason=finding.reason,
+            )
+            for finding in findings
+        ]
+        if event_entries:
+            existing = sanitized_event_payload.get("redactions", [])
+            sanitized_event_payload["redactions"] = [
+                *existing,
+                *(entry.path for entry in event_entries),
+            ]
+        sanitized_events.append(_event_adapter.validate_python(sanitized_event_payload))
+        entries.extend(event_entries)
+
+    entries.extend(
+        RedactionEntry(
+            event_id=MANIFEST_REDACTION_ID,
+            path=finding.path,
+            reason=finding.reason,
+        )
+        for finding in manifest_findings
+    )
+
+    sanitized_trace = Trace(
+        run_id=sanitized_manifest.run_id,
+        agent_name=sanitized_manifest.agent_name,
+        created_at=sanitized_manifest.created_at,
+        status=sanitized_manifest.status,
+        failure=sanitized_manifest.failure,
+        environment=sanitized_manifest.environment,
+        totals=sanitized_manifest.totals,
+        events=sanitized_events,
+    )
+    return PreparedTrace(
+        trace=sanitized_trace,
+        manifest=sanitized_manifest,
+        redactions=RedactionManifest(entries=_dedupe_entries(entries)),
+    )
+
+
+def _redact_value(value: Any, *, path: str) -> tuple[Any, list[RedactionFinding]]:
+    if isinstance(value, dict):
+        return _redact_mapping(value, path=path, protected_fields=set())
+    if isinstance(value, list):
+        findings: list[RedactionFinding] = []
+        redacted_items: list[Any] = []
+        for index, item in enumerate(value):
+            redacted_item, item_findings = _redact_value(
+                item,
+                path=f"{path}/{index}",
+            )
+            redacted_items.append(redacted_item)
+            findings.extend(item_findings)
+        return redacted_items, findings
+    if isinstance(value, str):
+        return _redact_string(value, path=path)
+    return value, []
+
+
+def _redact_mapping(
+    value: dict[str, Any],
+    *,
+    path: str,
+    protected_fields: set[str],
+) -> tuple[dict[str, Any], list[RedactionFinding]]:
+    findings: list[RedactionFinding] = []
+    redacted: dict[str, Any] = {}
+
+    for key in sorted(value):
+        item = value[key]
+        item_path = f"{path}/{_escape_json_pointer(key)}"
+        if key in protected_fields:
+            _assert_protected_value_is_safe(item)
+            redacted[key] = item
+            continue
+
+        if _is_sensitive_key(key):
+            if item == REDACTION_REPLACEMENT:
+                redacted[key] = item
+            else:
+                redacted[key] = REDACTION_REPLACEMENT
+                findings.append(
+                    RedactionFinding(path=item_path, reason="sensitive_key")
+                )
+            continue
+
+        redacted_item, item_findings = _redact_value(item, path=item_path)
+        redacted[key] = redacted_item
+        findings.extend(item_findings)
+
+    return redacted, findings
+
+
+def _redact_string(value: str, *, path: str) -> tuple[str, list[RedactionFinding]]:
+    if value == REDACTION_REPLACEMENT:
+        return value, []
+
+    redacted = value
+    changed = False
+    for pattern in _CREDENTIAL_PATTERNS:
+        redacted, replacements = pattern.subn(_replace_credential_match, redacted)
+        changed = changed or replacements > 0
+
+    if not changed:
+        return value, []
+    return redacted, [RedactionFinding(path=path, reason="credential_pattern")]
+
+
+def _replace_credential_match(match: re.Match[str]) -> str:
+    if len(match.groups()) == 2:
+        return f"{match.group(1)}{REDACTION_REPLACEMENT}"
+    return REDACTION_REPLACEMENT
+
+
+def _is_sensitive_key(key: str) -> bool:
+    return _normalize_key(key) in SENSITIVE_KEY_NORMALIZED
+
+
+def _normalize_key(key: str) -> str:
+    camel_split = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", key)
+    return re.sub(r"[^a-z0-9]", "", camel_split.lower())
+
+
+def _escape_json_pointer(segment: str) -> str:
+    return segment.replace("~", "~0").replace("/", "~1")
+
+
+def _assert_protected_value_is_safe(value: Any) -> None:
+    if isinstance(value, str) and _redact_string(value, path="")[0] != value:
+        raise RedactionError("credential pattern found in protected structural field")
+
+
+def _dedupe_entries(entries: list[RedactionEntry]) -> list[RedactionEntry]:
+    seen: set[tuple[str, str, str]] = set()
+    unique: list[RedactionEntry] = []
+    for entry in entries:
+        key = (entry.event_id, entry.path, entry.reason)
+        if key in seen:
+            continue
+        seen.add(key)
+        unique.append(entry)
+    return unique

stepfork/trace/replay_policy.py

@@ -0,0 +1,18 @@
+"""Replay policy values for recorded trace dependencies."""
+
+from enum import StrEnum
+
+
+class ReplayPolicy(StrEnum):
+    """Describe how a dependency should behave during replay.
+
+    `frozen` returns captured output, `live` re-executes the dependency,
+    `manual` requires explicit user approval, `forbidden` prevents automatic
+    execution, and `derived` recomputes locally from deterministic inputs.
+    """
+
+    FROZEN = "frozen"
+    LIVE = "live"
+    MANUAL = "manual"
+    FORBIDDEN = "forbidden"
+    DERIVED = "derived"

stepfork/trace/schema.py

@@ -0,0 +1,25 @@
+"""JSON Schema helpers for the Stepfork trace contract."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import TypeAdapter
+
+from stepfork.trace.manifest import TraceManifest
+from stepfork.trace.models import Event, Trace
+
+
+def event_json_schema() -> dict[str, Any]:
+    """Return the generated JSON Schema for the discriminated Event union."""
+    return TypeAdapter(Event).json_schema()
+
+
+def manifest_json_schema() -> dict[str, Any]:
+    """Return the generated JSON Schema for trace manifests."""
+    return TraceManifest.model_json_schema()
+
+
+def trace_json_schema() -> dict[str, Any]:
+    """Return the generated JSON Schema for the in-memory Trace model."""
+    return Trace.model_json_schema()