diff --git a/docs/concepts/trace-schema.md b/docs/concepts/trace-schema.md new file mode 100644 index 0000000..5005b9d --- /dev/null +++ b/docs/concepts/trace-schema.md @@ -0,0 +1,106 @@ +# Trace/Result Schema & Migration Policy + +`rampart.core.serialization` defines RAMPART's canonical, versioned +`Result`-record format. The current implementation provides the neutral +`ResultRecord.to_dict()` / `ResultRecord.from_dict()` round-trip and the +`serialize_result()` / `deserialize_result()` convenience functions. Existing +xdist and reporting consumers are not yet wired to this module. + +This page defines how the schema may evolve as consumers adopt it. + +## Versioning + +- Every serialized record carries one root `version` field. The current schema + is **`rampart.trace.v1`**. +- The record version is **independent** of transport or projection versions, + including the existing xdist envelope version (`rampart.xdist.v2`). Each + version describes its own layer and may evolve separately. +- There is a **single root version** — nested types (`Turn`, `Payload`, + `EvalResult`, …) do not carry their own versions. + +## What is and is not a breaking change + +- **Additive-optional = no bump.** A new optional field that older readers may + ignore, and whose absence has a defined default, does not change the major. +- **Missing = not recorded (not "false").** An absent optional field means the + producer *did not record it* — never that its value was empty, false, or zero. + Readers supply a default for *shape* only; consumers must not infer a semantic + negative from absence. A v1 record with no `manifest_snapshot` means "the + manifest was not captured," not "there was no manifest." +- **Structural change = major bump.** Removing, renaming, or retyping a field, + or changing its meaning or nesting, bumps `vN → vN+1` with a changelog and a + migration note. + +## Reader posture + +- Readers tolerate unknown fields and **fail closed on an unknown major** — a + record is never best-effort parsed across a major boundary. +- Forward compatibility is **additive-only within a major**. A newer major read + by an older framework fails closed by design. +- Schema descriptions and validators derived from this format must remain open + to unknown properties within a major version. + +## Enum posture + +- The closed enums — `SafetyStatus`, `EvalOutcome`, `ObservabilityLevel`, and + `PayloadFormat` — **fail closed** on an unknown value. A serialized safety + result must never silently misread one; there is no warn-and-degrade path. +- `HarmCategory` is the sole exception: it travels as a **passthrough string** + and is never coerced, so a new harm label from a future producer round-trips + unchanged on an older reader. + +## Value domain + +- Free-form mappings must already contain JSON-safe values. The canonical codec + does not coerce unsupported objects with `str()` or `repr()`. +- Numeric values must be finite. Transport-specific normalization is outside + the canonical schema. +- `rampart.trace.v1` does not define a durable representation for binary or + opaque payload artifacts. Encoding or decoding one fails closed rather than + coercing it to text. +- Transport bookkeeping keys are removed from top-level `Result.metadata`; + nested user mappings are preserved. + +## Migration mechanics + +Only `rampart.trace.v1` exists today. If a later structural change introduces a +new major: + +- writers emit the latest supported major; +- support for an older major uses an explicit adjacent upcaster + (`vN-1 → vN`); +- migrating persisted data is an explicit operation; reading never rewrites an + artifact in place; and +- encountering an unsupported major fails closed. + +## Reserved additive fields (named now, populated later) + +These record-level wire-only collar slots are reserved by name so they can be +added without a major bump: +`manifest_snapshot`, `evaluation_fingerprint`, `replay_provenance`, +`population_ref`, plus `artifacts` / `target` / `provenance`. A field that is +truly *intrinsic to a result* instead lands as an additive-optional field on +`Result`, inside the referenced `result` body. Either way each is +additive-optional; none is populated at v1. + +Other future fields follow the same general rule: optional additions with a +defined absence behavior do not require a major bump; structural changes do. + +## Support window + +Starting with the first release that writes durable trace records by default, +RAMPART supports reading `vN` and `vN-1` for **two subsequent framework +releases** (one deprecation cycle). The window is keyed on releases, not time. +Any major bump includes a changelog entry and migration note. + +```mermaid +flowchart TD + change([proposed schema change]) --> q1{"adds a field only?"} + q1 -- no --> struct["structural:
remove / rename / retype /
change meaning or nesting"] + q1 -- yes --> q2{"optional with a
well-defined default?"} + q2 -- no --> struct + q2 -- yes --> add["additive-optional"] + add --> nobump["NO bump
(new optional fields)
old readers ignore unknown keys"] + struct --> bump["bump major vN → vN+1
+ changelog + migration note"] + bump --> reader["readers: fail closed on
unknown major"] +``` diff --git a/mkdocs.yml b/mkdocs.yml index c74a9f5..606a913 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -151,6 +151,7 @@ nav: - Attacks: concepts/attacks.md - Probes: concepts/probes.md - PyRIT Integration: concepts/pyrit.md + - Trace Schema & Migration: concepts/trace-schema.md - Attacks: - attacks/index.md - XPIA: attacks/xpia.md diff --git a/rampart/core/serialization.py b/rampart/core/serialization.py new file mode 100644 index 0000000..46414ce --- /dev/null +++ b/rampart/core/serialization.py @@ -0,0 +1,954 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Canonical, versioned trace/result serialization for RAMPART. + +This module owns the *single* full-fidelity ``Result`` <-> ``dict`` round-trip +for the whole framework. + +The canonical layer defines the supported *value domain* and nothing else. It +does not apply transport hygiene. When a value falls outside the canonical +domain the codec fails closed with a field path rather than coercing it. + +Every serialized record carries a single root ``version`` field +(:data:`TRACE_SCHEMA_VERSION`). Decoding dispatches on that version and fails +closed on an unknown major. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from enum import Enum +from typing import TYPE_CHECKING, Any, ClassVar, TypeVar + +from rampart.core.result import ( + InjectionRecord, + PopulationRef, + Result, + SafetyStatus, +) +from rampart.core.types import ( + EvalOutcome, + EvalResult, + ObservabilityLevel, + Payload, + PayloadFormat, + Request, + Response, + SideEffect, + ToolCall, + Turn, +) + +if TYPE_CHECKING: + from collections.abc import Callable + +EnumT = TypeVar("EnumT", bound=Enum) + +# Single root schema version stamped on every serialized record. +TRACE_SCHEMA_VERSION = "rampart.trace.v1" + +# Top-level ``Result.metadata`` keys owned by the xdist transport. These +# scheduling and bookkeeping values are stripped from the canonical body; +# nested user maps are never touched. +_RESERVED_METADATA_KEYS: frozenset[str] = frozenset( + { + "_pytest_nodeid", + "_pytest_test_name", + "_rampart_result_index", + "_rampart_transport_truncated", + "_rampart_original_size_bytes", + "_rampart_limit_bytes", + "_rampart_worker_format", + "_rampart_worker_artifact_path", + } +) + + +class SchemaError(Exception): + """Raised when a value falls outside the canonical trace schema domain. + + The message carries the offending field path so producers can locate the + out-of-domain value instead of the codec silently coercing it. + """ + + +class UnsupportedSchemaVersionError(SchemaError): + """Raised when decoding a record whose ``version`` has no decoder. + + Readers fail closed on an unknown major rather than guessing at a shape. + """ + + +@dataclass(frozen=True, kw_only=True) +class ResultRecord: + """The canonical, versioned envelope around a single ``Result``. + + This is the public surface: :meth:`to_dict` / :meth:`from_dict` are the one + round-trip every durable consumer uses. The ``result`` is referenced, not + copied. The ``pytest_nodeid`` and ``result_index`` collar fields are + wire-only attribution stamped once at the producing boundary. + + Args: + result (Result): The single-run verdict being serialized. + pytest_nodeid (str | None): The pytest node id the result came from. + result_index (int | None): Ordinal of this result within its test node. + """ + + VERSION: ClassVar[str] = TRACE_SCHEMA_VERSION + + result: Result + pytest_nodeid: str | None = None + result_index: int | None = None + + def __post_init__(self) -> None: + """Validate record attribution fields. + + Raises: + SchemaError: If ``result_index`` is not an integer or ``None``. + """ + _validate_result_index(value=self.result_index) + + def to_dict(self) -> dict[str, Any]: + """Encode the record into a canonical, JSON-safe dict. + + Returns: + dict[str, Any]: The versioned record with the encoded result body + and wire-only collar. Fails closed via :class:`SchemaError` on + any value outside the canonical domain. + """ + encoded = { + "version": self.VERSION, + "result": _encode_result(result=self.result, path="result"), + } + if self.pytest_nodeid is not None: + encoded["pytest_nodeid"] = self.pytest_nodeid + if self.result_index is not None: + encoded["result_index"] = self.result_index + return encoded + + @classmethod + def from_dict(cls, data: object) -> ResultRecord: + """Decode a canonical dict back into a record, dispatching on version. + + Args: + data (object): A previously encoded record mapping. + + Returns: + ResultRecord: The decoded record. + + Raises: + SchemaError: If ``data`` is not a mapping. + UnsupportedSchemaVersionError: If the record version has no decoder. + """ + if not isinstance(data, Mapping): + msg = f"Expected mapping for record, got {type(data).__name__}." + raise SchemaError(msg) + version = data.get("version") + decoder = _DECODERS.get(version) if isinstance(version, str) else None + if decoder is None: + msg = f"No decoder registered for trace schema version {version!r}." + raise UnsupportedSchemaVersionError(msg) + return decoder(data) + + +def serialize_result( + *, + result: Result, + pytest_nodeid: str | None = None, + result_index: int | None = None, +) -> dict[str, Any]: + """Serialize a result to the canonical, versioned dict. + + Args: + result (Result): The verdict to serialize. + pytest_nodeid (str | None): The pytest node id the result came from. + result_index (int | None): Ordinal of this result within its test node. + + Returns: + dict[str, Any]: The canonical record dict, ready for any durable sink. + """ + record = ResultRecord( + result=result, + pytest_nodeid=pytest_nodeid, + result_index=result_index, + ) + return record.to_dict() + + +def deserialize_result(*, data: object) -> ResultRecord: + """Deserialize a canonical record dict back into a :class:`ResultRecord`. + + Args: + data (object): A previously encoded record mapping. + + Returns: + ResultRecord: The decoded record. + """ + return ResultRecord.from_dict(data) + + +def _encode_result(*, result: Result, path: str) -> dict[str, Any]: + """Encode a ``Result`` body into canonical primitives. + + Returns: + dict[str, Any]: The encoded result with every field represented. + """ + metadata = { + key: value + for key, value in result.metadata.items() + if key not in _RESERVED_METADATA_KEYS + } + return { + "status": _encode_enum(value=result.status, path=f"{path}.status"), + "summary": result.summary, + "observability_level": _encode_enum( + value=result.observability_level, + path=f"{path}.observability_level", + ), + "turns": [ + _encode_turn(turn=turn, path=f"{path}.turns[{index}]") + for index, turn in enumerate(result.turns) + ], + "duration_seconds": _encode_float( + value=result.duration_seconds, + path=f"{path}.duration_seconds", + ), + "harm_category": _encode_harm_category( + value=result.harm_category, + path=f"{path}.harm_category", + ), + "strategy": result.strategy, + "injections": [ + _encode_injection(record=record) for record in result.injections + ], + "population": _encode_population( + value=result.population, + path=f"{path}.population", + ), + "metadata": _encode_json(value=metadata, path=f"{path}.metadata"), + } + + +def _encode_turn(*, turn: Turn, path: str) -> dict[str, Any]: + """Encode a ``Turn``. + + Returns: + dict[str, Any]: The encoded turn. + """ + eval_result = ( + None + if turn.eval_result is None + else _encode_eval_result(value=turn.eval_result, path=f"{path}.eval_result") + ) + return { + "request": _encode_request(request=turn.request, path=f"{path}.request"), + "response": _encode_response(response=turn.response, path=f"{path}.response"), + "eval_result": eval_result, + "turn_number": turn.turn_number, + "timestamp": _encode_datetime(value=turn.timestamp), + "driver_reasoning": turn.driver_reasoning, + } + + +def _encode_request(*, request: Request, path: str) -> dict[str, Any]: + """Encode a ``Request``. + + Returns: + dict[str, Any]: The encoded request. + """ + return { + "prompt": request.prompt, + "attachments": [ + _encode_payload(payload=payload, path=f"{path}.attachments[{index}]") + for index, payload in enumerate(request.attachments) + ], + } + + +def _encode_response(*, response: Response, path: str) -> dict[str, Any]: + """Encode a ``Response``. + + Returns: + dict[str, Any]: The encoded response. + """ + return { + "text": response.text, + "tool_calls": [ + _encode_tool_call(call=call, path=f"{path}.tool_calls[{index}]") + for index, call in enumerate(response.tool_calls) + ], + "side_effects": [ + _encode_side_effect(effect=effect, path=f"{path}.side_effects[{index}]") + for index, effect in enumerate(response.side_effects) + ], + "metadata": _encode_json(value=response.metadata, path=f"{path}.metadata"), + } + + +def _encode_tool_call(*, call: ToolCall, path: str) -> dict[str, Any]: + """Encode a ``ToolCall``. + + Returns: + dict[str, Any]: The encoded tool call. + """ + return { + "name": call.name, + "arguments": _encode_json(value=call.arguments, path=f"{path}.arguments"), + "result": call.result, + "timestamp": _encode_datetime(value=call.timestamp), + } + + +def _encode_side_effect(*, effect: SideEffect, path: str) -> dict[str, Any]: + """Encode a ``SideEffect``. + + Returns: + dict[str, Any]: The encoded side effect. + """ + return { + "kind": effect.kind, + "details": _encode_json(value=effect.details, path=f"{path}.details"), + } + + +def _encode_payload(*, payload: Payload, path: str) -> dict[str, Any]: + """Encode a ``Payload``. + + Returns: + dict[str, Any]: The encoded payload. + + Raises: + SchemaError: If the payload uses a binary format. + """ + if payload.format.is_binary: + msg = ( + f"{path}: binary payload format {payload.format.value!r} is unsupported " + f"in {TRACE_SCHEMA_VERSION}." + ) + raise SchemaError(msg) + return { + "content": payload.content, + "id": payload.id, + "format": _encode_enum(value=payload.format, path=f"{path}.format"), + "artifact": None, + "metadata": _encode_json(value=payload.metadata, path=f"{path}.metadata"), + } + + +def _encode_eval_result(*, value: EvalResult, path: str) -> dict[str, Any]: + """Encode an ``EvalResult``. + + Returns: + dict[str, Any]: The encoded evaluation result. + """ + return { + "outcome": _encode_enum(value=value.outcome, path=f"{path}.outcome"), + "confidence": _encode_float( + value=value.confidence, + path=f"{path}.confidence", + ), + "evidence": list(value.evidence), + "rationale": value.rationale, + "undetermined_operands": list(value.undetermined_operands), + } + + +def _encode_injection(*, record: InjectionRecord) -> dict[str, Any]: + """Encode an ``InjectionRecord``. + + Returns: + dict[str, Any]: The encoded injection record. + """ + return { + "payload_id": record.payload_id, + "surface_name": record.surface_name, + } + + +def _encode_population( + *, value: PopulationRef | None, path: str +) -> dict[str, Any] | None: + """Encode an optional ``PopulationRef``. + + Returns: + dict[str, Any] | None: The encoded reference, or ``None``. + """ + if value is None: + return None + return { + "id": value.id, + "index": value.index, + "size": value.size, + "threshold": _encode_float(value=value.threshold, path=f"{path}.threshold"), + } + + +def _encode_enum(*, value: Enum, path: str) -> str: + """Encode an enum member to its wire value. + + Returns: + str: The enum ``.value``. + + Raises: + SchemaError: If ``value`` is not an enum member. + """ + if not isinstance(value, Enum): + msg = f"{path}: expected enum, got {type(value).__name__}." + raise SchemaError(msg) + return str(value.value) + + +def _encode_harm_category(*, value: object, path: str) -> str | None: + """Encode a harm category as a passthrough string. + + Returns: + str | None: The category string, or ``None`` when unset. + + Raises: + SchemaError: If ``value`` is not a string or ``None``. + """ + if value is None: + return None + if not isinstance(value, str): + msg = f"{path}: expected a string or null, got {type(value).__name__}." + raise SchemaError(msg) + return value + + +def _encode_datetime(*, value: datetime | None) -> str | None: + """Encode a datetime to ISO 8601. + + Returns: + str | None: The ISO timestamp, or ``None``. + """ + if value is None: + return None + return value.isoformat() + + +def _encode_float(*, value: object, path: str) -> float: + """Validate and pass through a float within the canonical domain. + + Returns: + float: The finite float value. + + Raises: + SchemaError: If ``value`` is not a finite real number. Normalizing + non-finite floats is transport hygiene, not a canonical concern. + """ + if isinstance(value, bool) or not isinstance(value, int | float): + msg = f"{path}: expected a real number, got {type(value).__name__}." + raise SchemaError(msg) + if not math.isfinite(value): + msg = f"{path}: expected a finite number, got {value!r}." + raise SchemaError(msg) + return float(value) + + +def _encode_json(*, value: object, path: str) -> object: + """Validate that ``value`` is JSON-safe, failing closed otherwise. + + Recurses through lists and string-keyed maps of primitives. Anything + outside the domain (bytes, ``Path``, arbitrary objects, non-finite floats, + non-string map keys) raises rather than being coerced via ``repr()``. + + Returns: + Any: A JSON-safe copy of ``value``. + + Raises: + SchemaError: If ``value`` contains anything outside the JSON domain. + """ + if value is None or isinstance(value, str | bool): + return value + if isinstance(value, int): + return value + if isinstance(value, float): + return _encode_float(value=value, path=path) + if isinstance(value, Mapping): + return _encode_json_map(value=value, path=path) + if isinstance(value, list | tuple): + return [ + _encode_json(value=item, path=f"{path}[{index}]") + for index, item in enumerate(value) + ] + msg = f"{path}: value of type {type(value).__name__} is outside the JSON domain." + raise SchemaError(msg) + + +def _encode_json_map(*, value: Mapping[Any, Any], path: str) -> dict[str, Any]: + """Validate and copy a JSON-safe string-keyed map. + + Returns: + dict[str, Any]: A JSON-safe copy of the map. + + Raises: + SchemaError: If any key is not a string. + """ + encoded: dict[str, Any] = {} + for key, item in value.items(): + if not isinstance(key, str): + msg = f"{path}: map key {key!r} is not a string." + raise SchemaError(msg) + encoded[key] = _encode_json(value=item, path=f"{path}.{key}") + return encoded + + +def _decode_v1(data: Mapping[str, Any]) -> ResultRecord: + """Decode a ``rampart.trace.v1`` record. + + Returns: + ResultRecord: The decoded record. + + Raises: + SchemaError: If the record body is not a mapping. + """ + body = data.get("result") + if not isinstance(body, Mapping): + msg = f"record 'result' body must be a mapping, got {type(body).__name__}." + raise SchemaError(msg) + result_index = data.get("result_index") + pytest_nodeid = data.get("pytest_nodeid") + if pytest_nodeid is not None and not isinstance(pytest_nodeid, str): + msg = ( + "record 'pytest_nodeid' must be a string or null, " + f"got {type(pytest_nodeid).__name__}." + ) + raise SchemaError(msg) + return ResultRecord( + result=_decode_result(data=body, path="result"), + pytest_nodeid=pytest_nodeid, + result_index=result_index, + ) + + +def _decode_result(*, data: Mapping[str, Any], path: str) -> Result: + """Decode a ``Result`` body. + + Returns: + Result: The reconstructed result. + """ + return Result( + status=_decode_enum( + enum=SafetyStatus, + value=data.get("status"), + path=f"{path}.status", + ), + summary=_decode_str(value=data.get("summary"), path=f"{path}.summary"), + observability_level=_decode_enum( + enum=ObservabilityLevel, + value=data.get("observability_level"), + path=f"{path}.observability_level", + ), + turns=[ + _decode_turn(data=item, path=f"{path}.turns[{index}]") + for index, item in enumerate( + _decode_list(value=data.get("turns"), path=f"{path}.turns") + ) + ], + duration_seconds=_encode_float( + value=data.get("duration_seconds", 0.0), + path=f"{path}.duration_seconds", + ), + harm_category=_decode_harm_category( + value=data.get("harm_category"), + path=f"{path}.harm_category", + ), + strategy=_decode_str(value=data.get("strategy", ""), path=f"{path}.strategy"), + injections=[ + _decode_injection(data=item, path=f"{path}.injections[{index}]") + for index, item in enumerate( + _decode_list(value=data.get("injections"), path=f"{path}.injections") + ) + ], + population=_decode_population( + value=data.get("population"), + path=f"{path}.population", + ), + metadata=dict( + _decode_optional_map(value=data.get("metadata"), path=f"{path}.metadata") + ), + ) + + +def _decode_turn(*, data: object, path: str) -> Turn: + """Decode a ``Turn``. + + Returns: + Turn: The reconstructed turn. + """ + typed = _decode_map(value=data, path=path) + raw_eval = typed.get("eval_result") + eval_result = ( + None + if raw_eval is None + else _decode_eval_result(data=raw_eval, path=f"{path}.eval_result") + ) + return Turn( + request=_decode_request(data=typed.get("request"), path=f"{path}.request"), + response=_decode_response(data=typed.get("response"), path=f"{path}.response"), + eval_result=eval_result, + turn_number=_decode_int( + value=typed.get("turn_number", 0), path=f"{path}.turn_number" + ), + timestamp=_decode_datetime( + value=typed.get("timestamp"), path=f"{path}.timestamp" + ), + driver_reasoning=_decode_str( + value=typed.get("driver_reasoning", ""), + path=f"{path}.driver_reasoning", + ), + ) + + +def _decode_request(*, data: object, path: str) -> Request: + """Decode a ``Request``. + + Returns: + Request: The reconstructed request. + """ + typed = _decode_map(value=data, path=path) + raw_prompt = typed.get("prompt") + prompt = raw_prompt if isinstance(raw_prompt, str) else None + return Request( + prompt=prompt, + attachments=[ + _decode_payload(data=item, path=f"{path}.attachments[{index}]") + for index, item in enumerate( + _decode_list( + value=typed.get("attachments"), + path=f"{path}.attachments", + ) + ) + ], + ) + + +def _decode_response(*, data: object, path: str) -> Response: + """Decode a ``Response``. + + Returns: + Response: The reconstructed response. + """ + typed = _decode_map(value=data, path=path) + return Response( + text=_decode_str(value=typed.get("text", ""), path=f"{path}.text"), + tool_calls=[ + _decode_tool_call(data=item, path=f"{path}.tool_calls[{index}]") + for index, item in enumerate( + _decode_list( + value=typed.get("tool_calls"), + path=f"{path}.tool_calls", + ) + ) + ], + side_effects=[ + _decode_side_effect(data=item, path=f"{path}.side_effects[{index}]") + for index, item in enumerate( + _decode_list( + value=typed.get("side_effects"), + path=f"{path}.side_effects", + ) + ) + ], + metadata=dict( + _decode_optional_map(value=typed.get("metadata"), path=f"{path}.metadata") + ), + ) + + +def _decode_tool_call(*, data: object, path: str) -> ToolCall: + """Decode a ``ToolCall``. + + Returns: + ToolCall: The reconstructed tool call. + """ + typed = _decode_map(value=data, path=path) + raw_result = typed.get("result") + return ToolCall( + name=_decode_str(value=typed.get("name", ""), path=f"{path}.name"), + arguments=dict( + _decode_optional_map(value=typed.get("arguments"), path=f"{path}.arguments") + ), + result=raw_result if isinstance(raw_result, str) else None, + timestamp=_decode_datetime( + value=typed.get("timestamp"), path=f"{path}.timestamp" + ), + ) + + +def _decode_side_effect(*, data: object, path: str) -> SideEffect: + """Decode a ``SideEffect``. + + Returns: + SideEffect: The reconstructed side effect. + """ + typed = _decode_map(value=data, path=path) + return SideEffect( + kind=_decode_str(value=typed.get("kind", ""), path=f"{path}.kind"), + details=dict( + _decode_optional_map(value=typed.get("details"), path=f"{path}.details") + ), + ) + + +def _decode_payload(*, data: object, path: str) -> Payload: + """Decode a ``Payload``. + + Returns: + Payload: The reconstructed payload. + + Raises: + SchemaError: If the payload declares a binary format. + """ + typed = _decode_map(value=data, path=path) + payload_format = _decode_enum( + enum=PayloadFormat, + value=typed.get("format", PayloadFormat.TEXT.value), + path=f"{path}.format", + ) + if payload_format.is_binary: + msg = ( + f"{path}: binary payload format {payload_format.value!r} is unsupported " + f"in {TRACE_SCHEMA_VERSION}." + ) + raise SchemaError(msg) + return Payload( + content=_decode_str(value=typed.get("content", ""), path=f"{path}.content"), + id=_decode_str(value=typed.get("id", ""), path=f"{path}.id"), + format=payload_format, + artifact=None, + metadata=dict( + _decode_optional_map(value=typed.get("metadata"), path=f"{path}.metadata") + ), + ) + + +def _decode_eval_result(*, data: object, path: str) -> EvalResult: + """Decode an ``EvalResult``. + + Returns: + EvalResult: The reconstructed evaluation result. + """ + typed = _decode_map(value=data, path=path) + return EvalResult( + outcome=_decode_enum( + enum=EvalOutcome, + value=typed.get("outcome"), + path=f"{path}.outcome", + ), + confidence=_encode_float( + value=typed.get("confidence", 1.0), + path=f"{path}.confidence", + ), + evidence=_decode_str_list(value=typed.get("evidence"), path=f"{path}.evidence"), + rationale=_decode_str( + value=typed.get("rationale", ""), path=f"{path}.rationale" + ), + undetermined_operands=_decode_str_list( + value=typed.get("undetermined_operands"), + path=f"{path}.undetermined_operands", + ), + ) + + +def _decode_injection(*, data: object, path: str) -> InjectionRecord: + """Decode an ``InjectionRecord``. + + Returns: + InjectionRecord: The reconstructed injection record. + """ + typed = _decode_map(value=data, path=path) + raw_payload_id = typed.get("payload_id") + return InjectionRecord( + payload_id=raw_payload_id if isinstance(raw_payload_id, str) else None, + surface_name=_decode_str( + value=typed.get("surface_name", ""), + path=f"{path}.surface_name", + ), + ) + + +def _decode_population(*, value: object, path: str) -> PopulationRef | None: + """Decode an optional ``PopulationRef``. + + Returns: + PopulationRef | None: The reconstructed reference, or ``None``. + """ + if value is None: + return None + typed = _decode_map(value=value, path=path) + return PopulationRef( + id=_decode_str(value=typed.get("id"), path=f"{path}.id"), + index=_decode_int(value=typed.get("index"), path=f"{path}.index"), + size=_decode_int(value=typed.get("size"), path=f"{path}.size"), + threshold=_encode_float(value=typed.get("threshold"), path=f"{path}.threshold"), + ) + + +def _decode_enum(*, enum: type[EnumT], value: object, path: str) -> EnumT: + """Decode an enum member from its wire value, failing closed on unknown. + + Returns: + EnumT: The enum member. + + Raises: + SchemaError: If ``value`` is not a member of ``enum``. + """ + try: + return enum(value) + except ValueError as exc: + msg = f"{path}: {value!r} is not a valid {enum.__name__}." + raise SchemaError(msg) from exc + + +def _decode_harm_category(*, value: object, path: str) -> str | None: + """Decode a harm category as a passthrough string. + + Returns: + str | None: The category string, or ``None``. + + Raises: + SchemaError: If ``value`` is not a string or ``None``. + """ + if value is None: + return None + if not isinstance(value, str): + msg = f"{path}: expected a string or null, got {type(value).__name__}." + raise SchemaError(msg) + return value + + +def _decode_datetime(*, value: object, path: str) -> datetime | None: + """Decode an ISO 8601 timestamp. + + Returns: + datetime | None: The parsed datetime, or ``None``. + + Raises: + SchemaError: If ``value`` is neither ``None`` nor a valid ISO string. + """ + if value is None: + return None + if not isinstance(value, str): + msg = f"{path}: expected an ISO timestamp string, got {type(value).__name__}." + raise SchemaError(msg) + try: + return datetime.fromisoformat(value) + except ValueError as exc: + msg = f"{path}: {value!r} is not a valid ISO 8601 timestamp." + raise SchemaError(msg) from exc + + +def _decode_str(*, value: object, path: str) -> str: + """Decode a required string field. + + Returns: + str: The string value. + + Raises: + SchemaError: If ``value`` is not a string. + """ + if not isinstance(value, str): + msg = f"{path}: expected a string, got {type(value).__name__}." + raise SchemaError(msg) + return value + + +def _decode_int(*, value: object, path: str) -> int: + """Decode a required integer field. + + Returns: + int: The integer value. + + Raises: + SchemaError: If ``value`` is not an integer. + """ + if isinstance(value, bool) or not isinstance(value, int): + msg = f"{path}: expected an integer, got {type(value).__name__}." + raise SchemaError(msg) + return value + + +def _decode_list(*, value: object, path: str) -> list[Any]: + """Decode an optional wire list. + + Returns: + list[Any]: The list, or an empty list when absent or ``None``. + + Raises: + SchemaError: If ``value`` is present but not a list. + """ + if value is None: + return [] + if not isinstance(value, list): + msg = f"{path}: expected a list or null, got {type(value).__name__}." + raise SchemaError(msg) + return value + + +def _decode_str_list(*, value: object, path: str) -> list[str]: + """Decode a list of strings. + + Returns: + list[str]: The decoded strings. + """ + return [ + _decode_str(value=item, path=f"{path}[{index}]") + for index, item in enumerate(_decode_list(value=value, path=path)) + ] + + +def _validate_result_index(*, value: object) -> int | None: + """Validate an optional result index. + + Returns: + int | None: The validated index. + + Raises: + SchemaError: If ``value`` is not an integer or ``None``. + """ + if value is not None and (isinstance(value, bool) or not isinstance(value, int)): + msg = ( + "record 'result_index' must be an integer or null, " + f"got {type(value).__name__}." + ) + raise SchemaError(msg) + return value + + +def _decode_map(*, value: object, path: str) -> Mapping[str, Any]: + """Decode a required mapping field. + + Returns: + Mapping[str, Any]: The mapping value. + + Raises: + SchemaError: If ``value`` is not a mapping. + """ + if not isinstance(value, Mapping): + msg = f"{path}: expected a mapping, got {type(value).__name__}." + raise SchemaError(msg) + return value + + +def _decode_optional_map(*, value: object, path: str) -> Mapping[str, Any]: + """Decode an optional mapping field, defaulting to empty when absent. + + Returns: + Mapping[str, Any]: The mapping value, or an empty mapping when ``None``. + + Raises: + SchemaError: If ``value`` is present but not a mapping. + """ + if value is None: + return {} + return _decode_map(value=value, path=path) + + +_DECODERS: dict[str, Callable[[Mapping[str, Any]], ResultRecord]] = { + TRACE_SCHEMA_VERSION: _decode_v1, +} diff --git a/tests/unit/core/test_serialization.py b/tests/unit/core/test_serialization.py new file mode 100644 index 0000000..09cad12 --- /dev/null +++ b/tests/unit/core/test_serialization.py @@ -0,0 +1,351 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Unit tests for the canonical trace/result serializer.""" + +from __future__ import annotations + +import math +from dataclasses import fields +from datetime import UTC, datetime + +import pytest + +from rampart.core.result import ( + InjectionRecord, + PopulationRef, + Result, + SafetyStatus, +) +from rampart.core.serialization import ( + TRACE_SCHEMA_VERSION, + ResultRecord, + SchemaError, + UnsupportedSchemaVersionError, + deserialize_result, + serialize_result, +) +from rampart.core.types import ( + EvalOutcome, + EvalResult, + ObservabilityLevel, + Payload, + PayloadFormat, + Request, + Response, + SideEffect, + ToolCall, + Turn, +) + +_TIMESTAMP = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) + + +def _make_eval_result() -> EvalResult: + return EvalResult( + outcome=EvalOutcome.DETECTED, + confidence=0.75, + evidence=["saw the thing", "and another"], + rationale="because reasons", + undetermined_operands=["left operand undetermined"], + ) + + +def _make_turn() -> Turn: + request = Request( + prompt="do the thing", + attachments=[ + Payload( + content="poisoned doc text", + id="payload-1", + format=PayloadFormat.MARKDOWN, + metadata={"persona": "attacker"}, + ), + ], + ) + response = Response( + text="agent said this", + tool_calls=[ + ToolCall( + name="send_email", + arguments={"to": "a@b.com", "nested": {"count": 2}}, + result="ok", + timestamp=_TIMESTAMP, + ), + ], + side_effects=[SideEffect(kind="http_request", details={"url": "http://x"})], + metadata={"latency_ms": 12}, + ) + return Turn( + request=request, + response=response, + eval_result=_make_eval_result(), + turn_number=3, + timestamp=_TIMESTAMP, + driver_reasoning="escalate", + ) + + +def _make_full_result(*, metadata: dict | None = None) -> Result: + return Result( + status=SafetyStatus.UNSAFE, + summary="a violation was detected", + observability_level=ObservabilityLevel.TOOL_AND_SIDE_EFFECTS, + turns=[_make_turn()], + duration_seconds=1.5, + harm_category="prompt_injection", + strategy="xpia", + injections=[InjectionRecord(payload_id="payload-1", surface_name="SharePoint")], + population=PopulationRef(id="pop-1", index=0, size=5, threshold=0.8), + metadata={"note": "user data", "nested": {"k": [1, 2]}} + if metadata is None + else metadata, + ) + + +def _minimal_record_dict() -> dict: + return { + "version": TRACE_SCHEMA_VERSION, + "result": { + "status": "safe", + "summary": "clean", + "observability_level": "response_only", + }, + } + + +class TestRoundTrip: + def test_full_result_round_trips_to_equal_value(self) -> None: + original = _make_full_result() + encoded = ResultRecord(result=original).to_dict() + + decoded = deserialize_result(data=encoded).result + + assert decoded == original + + def test_version_is_stamped_on_the_record(self) -> None: + encoded = ResultRecord(result=_make_full_result()).to_dict() + + assert encoded["version"] == TRACE_SCHEMA_VERSION + assert ResultRecord.VERSION == "rampart.trace.v1" + + def test_serialize_result_builds_attribution_collar(self) -> None: + encoded = serialize_result( + result=_make_full_result(), + pytest_nodeid="tests/test_x.py::test_x", + result_index=2, + ) + + assert encoded["pytest_nodeid"] == "tests/test_x.py::test_x" + assert encoded["result_index"] == 2 + + def test_serialize_result_omits_attribution_when_unset(self) -> None: + encoded = serialize_result(result=_make_full_result()) + + assert "pytest_nodeid" not in encoded + assert "result_index" not in encoded + + def test_attribution_collar_round_trips(self) -> None: + encoded = serialize_result( + result=_make_full_result(), + pytest_nodeid="tests/test_x.py::test_x", + result_index=2, + ) + + decoded = deserialize_result(data=encoded) + + assert decoded.pytest_nodeid == "tests/test_x.py::test_x" + assert decoded.result_index == 2 + + def test_nested_values_survive_the_round_trip(self) -> None: + decoded = deserialize_result( + data=ResultRecord(result=_make_full_result()).to_dict() + ).result + + turn = decoded.turns[0] + assert turn.request.attachments[0].format is PayloadFormat.MARKDOWN + assert turn.response.tool_calls[0].arguments == { + "to": "a@b.com", + "nested": {"count": 2}, + } + assert turn.response.tool_calls[0].timestamp == _TIMESTAMP + assert turn.response.side_effects[0].kind == "http_request" + assert turn.eval_result is not None + assert turn.eval_result.outcome is EvalOutcome.DETECTED + assert decoded.injections[0].surface_name == "SharePoint" + assert decoded.population == PopulationRef( + id="pop-1", index=0, size=5, threshold=0.8 + ) + + +class TestFieldExhaustiveness: + def test_every_field_of_every_type_is_serialized(self) -> None: + body = ResultRecord(result=_make_full_result()).to_dict()["result"] + turn = body["turns"][0] + + cases = [ + (Result, body), + (Turn, turn), + (Request, turn["request"]), + (Payload, turn["request"]["attachments"][0]), + (Response, turn["response"]), + (ToolCall, turn["response"]["tool_calls"][0]), + (SideEffect, turn["response"]["side_effects"][0]), + (EvalResult, turn["eval_result"]), + (InjectionRecord, body["injections"][0]), + (PopulationRef, body["population"]), + ] + + for dataclass_type, encoded in cases: + expected = {field.name for field in fields(dataclass_type)} + assert expected == set(encoded), dataclass_type.__name__ + + +class TestVersionDispatch: + def test_unknown_major_fails_closed(self) -> None: + data = {"version": "rampart.trace.v2", "result": {}} + + with pytest.raises(UnsupportedSchemaVersionError, match="v2"): + deserialize_result(data=data) + + def test_missing_version_fails_closed(self) -> None: + with pytest.raises(UnsupportedSchemaVersionError): + deserialize_result(data={"result": {}}) + + def test_non_mapping_record_fails_closed(self) -> None: + with pytest.raises(SchemaError, match="mapping"): + deserialize_result(data=[1, 2, 3]) + + +class TestMigrationTolerance: + def test_unknown_extra_fields_decode(self) -> None: + encoded = ResultRecord(result=_make_full_result()).to_dict() + encoded["future_collar"] = {"anything": True} + encoded["result"]["future_intrinsic"] = 42 + + decoded = deserialize_result(data=encoded).result + + assert decoded.status is SafetyStatus.UNSAFE + + def test_missing_optional_fields_use_defaults(self) -> None: + decoded = deserialize_result(data=_minimal_record_dict()).result + + assert decoded.status is SafetyStatus.SAFE + assert decoded.turns == [] + assert decoded.duration_seconds == pytest.approx(0.0) + assert decoded.harm_category is None + assert decoded.injections == [] + assert decoded.population is None + assert decoded.metadata == {} + + def test_malformed_present_list_fails_closed(self) -> None: + data = _minimal_record_dict() + data["result"]["turns"] = "not-a-list" + + with pytest.raises(SchemaError, match=r"result\.turns"): + deserialize_result(data=data) + + def test_incomplete_population_reference_fails_closed(self) -> None: + data = _minimal_record_dict() + data["result"]["population"] = {} + + with pytest.raises(SchemaError, match=r"result\.population\.id"): + deserialize_result(data=data) + + +class TestValueDomain: + def test_reserved_metadata_keys_are_stripped(self) -> None: + result = _make_full_result( + metadata={"_pytest_nodeid": "x::y", "note": "keep me"}, + ) + + encoded = ResultRecord(result=result).to_dict() + + assert encoded["result"]["metadata"] == {"note": "keep me"} + + def test_harm_category_is_passed_through_as_string(self) -> None: + result = _make_full_result() + result.harm_category = "custom_product_risk" + + encoded = ResultRecord(result=result).to_dict() + decoded = deserialize_result(data=encoded).result + + assert encoded["result"]["harm_category"] == "custom_product_risk" + assert decoded.harm_category == "custom_product_risk" + + def test_non_finite_float_fails_closed(self) -> None: + result = _make_full_result() + result.duration_seconds = math.inf + + with pytest.raises(SchemaError, match="duration_seconds"): + ResultRecord(result=result).to_dict() + + def test_non_json_metadata_fails_closed(self) -> None: + result = _make_full_result(metadata={"blob": object()}) + + with pytest.raises(SchemaError, match="metadata"): + ResultRecord(result=result).to_dict() + + def test_bad_enum_value_fails_closed_on_decode(self) -> None: + data = _minimal_record_dict() + data["result"]["status"] = "not_a_status" + + with pytest.raises(SchemaError, match="status"): + deserialize_result(data=data) + + def test_non_string_harm_category_fails_closed_on_encode(self) -> None: + result = _make_full_result() + result.__dict__["harm_category"] = 42 + + with pytest.raises(SchemaError, match="harm_category"): + ResultRecord(result=result).to_dict() + + def test_non_string_harm_category_fails_closed_on_decode(self) -> None: + data = _minimal_record_dict() + data["result"]["harm_category"] = {"category": "custom"} + + with pytest.raises(SchemaError, match="harm_category"): + deserialize_result(data=data) + + def test_boolean_result_index_fails_before_encoding(self) -> None: + with pytest.raises(SchemaError, match="result_index"): + serialize_result(result=_make_full_result(), result_index=True) + + +class TestBinaryPayloadFailsClosed: + def test_encoding_a_binary_payload_fails_closed(self, tmp_path) -> None: + artifact = tmp_path / "doc.pdf" + artifact.write_bytes(b"%PDF-1.4 fake") + result = _make_full_result() + result.turns = [ + Turn( + request=Request( + attachments=[ + Payload( + content="binary doc", + format=PayloadFormat.PDF, + artifact=artifact, + ), + ], + ), + response=Response(text="ok"), + ), + ] + + with pytest.raises(SchemaError, match="binary payload"): + ResultRecord(result=result).to_dict() + + def test_decoding_a_binary_payload_fails_closed(self) -> None: + data = _minimal_record_dict() + data["result"]["turns"] = [ + { + "request": { + "prompt": None, + "attachments": [{"content": "x", "id": "p", "format": "pdf"}], + }, + "response": {"text": "ok"}, + }, + ] + + with pytest.raises(SchemaError, match="binary payload"): + deserialize_result(data=data)