inflating compressed payload…
_Default: `40`_ + +#### `--max-turns` `
_Default: `60`_ + +#### `--content-cap` `
_Default: `8000`_ + +
+ ## `aiperf validate` Validate a benchmark artifact. diff --git a/pyproject.toml b/pyproject.toml index f031e8d2f1..89c6b61909 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -264,4 +264,4 @@ filterwarnings = [ [tool.codespell] skip = "*.pyc,*build*,tests/unit/transports/test_aiohttp_sse.py,tests/integration/assets/canary_reference_inputs.json,src/aiperf/server_metrics/units.py,src/aiperf/api/static/dashboard.html,src/aiperf/config/schema/aiperf-config.schema.json" -ignore-words-list = "timeslice,timeslices,optiona,disjointness,concurency" +ignore-words-list = "timeslice,timeslices,optiona,disjointness,concurency,clen" diff --git a/src/aiperf/analysis/fzstd.umd.js b/src/aiperf/analysis/fzstd.umd.js new file mode 100644 index 0000000000..9f3d2380ce --- /dev/null +++ b/src/aiperf/analysis/fzstd.umd.js @@ -0,0 +1,31 @@ +// fzstd v0.1.1 -- pure-JS Zstandard decompressor, vendored from +// https://github.com/101arrowz/fzstd (npm: fzstd). Inlined into the +// `aiperf turn-messages` viewer to decompress the embedded payload +// client-side (no WASM, no network). See ATTRIBUTIONS.md. +// +// MIT License +// +// Copyright (c) 2020 Arjun Barrett +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// +// NOTE (aiperf): this build silently corrupts frames with windowLog >= 27 +// (128 MiB) -- correct length, WRONG bytes. The generator pins --long=26 +// (64 MiB) and round-trips every payload before writing. Do not bump the window. +!function(f){typeof module!='undefined'&&typeof exports=='object'?module.exports=f():typeof define!='undefined'&&define.amd?define(['fzstd',f]):(typeof self!='undefined'?self:this).fzstd=f()}(function(){var _e={};"use strict";var r=ArrayBuffer,t=Uint8Array,e=Uint16Array,n=Int16Array,a=Uint32Array,s=Int32Array,i=function(r,e,n){if(t.prototype.slice)return t.prototype.slice.call(r,e,n);(null==e||e<0)&&(e=0),(null==n||n>r.length)&&(n=r.length);var a=new t(n-e);return a.set(r.subarray(e,n)),a},o=function(r,e,n,a){if(t.prototype.fill)return t.prototype.fill.call(r,e,n,a);for((null==n||n<0)&&(n=0),(null==a||a>r.length)&&(a=r.length);n
`` is first opened, and a turn's messages only when that turn is
+opened, so the live DOM stays tiny regardless of run size.
+"""
+
+from __future__ import annotations
+
+import base64
+from collections import Counter, defaultdict
+from pathlib import Path
+from typing import Any
+
+import orjson
+import zstandard
+
+from aiperf.config.artifacts import OutputDefaults
+
+VIEWER_TEMPLATE = Path(__file__).with_name("turn_messages_viewer.html")
+FZSTD_JS = Path(__file__).with_name("fzstd.umd.js")
+
+# The inlined fzstd 0.1.1 decoder silently corrupts frames with windowLog >= 27
+# (128 MiB) -- correct length, wrong bytes. 26 (64 MiB) is the largest safe
+# window and captures essentially all the long-range redundancy. Do NOT raise
+# without re-verifying fzstd round-trips; write_turn_messages_html self-checks
+# every payload (and caps the decoder window) before writing.
+ZSTD_WINDOW_LOG = 26
+
+# Generator defaults, shared with the CLI so they document once.
+DEFAULT_LIMIT_CONVERSATIONS = 40
+DEFAULT_MAX_TURNS = 60
+DEFAULT_CONTENT_CAP = 8000
+
+
+class TurnMessagesError(Exception):
+ """Raised when a run directory has no usable raw records to render."""
+
+
+def _text_of(content: Any) -> str:
+ """OpenAI message content is a str or a list of typed parts."""
+ if content is None:
+ return ""
+ if isinstance(content, str):
+ return content
+ if isinstance(content, list):
+ out = []
+ for p in content:
+ if isinstance(p, dict):
+ out.append(p.get("text") or f"<{p.get('type', 'part')}>")
+ else:
+ out.append(str(p))
+ return "\n".join(out)
+ return str(content)
+
+
+def _short(s: str | None, n: int = 14) -> str:
+ """Truncate ``s`` to ``n`` chars, mapping a falsy value to ``"?"``."""
+ return (s or "?")[:n]
+
+
+def _disp_id(conv: str) -> str:
+ """Display id that keeps the DISTINGUISHING part of a conversation id visible.
+
+ Conversation ids in a split corpus share one long trace-id prefix (root
+ ````, subagents ``::sa:``, flat chains ``::fa:NNN``)
+ -- naive prefix truncation renders every panel identically. Shorten only the
+ trace-id component and keep the full suffix.
+ """
+ base, sep, suffix = conv.partition("::")
+ if sep:
+ return f"{_short(base, 10)}…::{suffix}"
+ return f"{_short(base, 20)}…"
+
+
+def _find_raw_records(run_dir: Path) -> list[Path]:
+ """Locate raw-record jsonl for ``run_dir``.
+
+ Accepts either a direct jsonl file or a run directory (a ``raw_records/``
+ shard folder, a single ``profile_export_raw.jsonl``, or a recursive
+ fallback for custom export-base names).
+ """
+ if run_dir.is_file():
+ return [run_dir]
+ folder = run_dir / OutputDefaults.RAW_RECORDS_FOLDER
+ if folder.is_dir():
+ shards = sorted(folder.glob("*.jsonl"))
+ if shards:
+ return shards
+ single = run_dir / OutputDefaults.PROFILE_EXPORT_RAW_JSONL_FILE
+ if single.is_file():
+ return [single]
+ return sorted(run_dir.glob("**/raw_records/*.jsonl")) + sorted(
+ run_dir.glob("**/*_raw.jsonl")
+ )
+
+
+def _flatten_record(r: dict) -> dict:
+ """Flatten one parsed raw record into the fields ``build_payload`` consumes.
+
+ Tolerant of malformed records: ``metadata``/``payload`` that are missing or
+ not objects fall back to empty, and ``messages`` that is not a list is
+ dropped (a non-list would otherwise be iterated character-by-character
+ downstream).
+ """
+ md = r.get("metadata")
+ md = md if isinstance(md, dict) else {}
+ pl = r.get("payload")
+ pl = pl if isinstance(pl, dict) else {}
+ msgs = pl.get("messages")
+ return {
+ "conv": md.get("conversation_id") or "(none)",
+ "turn": md.get("turn_index"),
+ "depth": md.get("agent_depth") or 0,
+ "parent_corr": md.get("parent_correlation_id"),
+ "x_corr": md.get("x_correlation_id"),
+ "model": pl.get("model"),
+ "max_tokens": pl.get("max_completion_tokens") or pl.get("max_tokens"),
+ "phase": md.get("benchmark_phase"),
+ "cancelled": md.get("was_cancelled"),
+ "error": r.get("error"),
+ "messages": msgs if isinstance(msgs, list) else [],
+ "start_ns": md.get("request_start_ns") or 0,
+ }
+
+
+def load_records(files: list[Path]) -> list[dict]:
+ """Parse and flatten the per-turn raw records from the given jsonl files."""
+ recs: list[dict] = []
+ for f in files:
+ with open(f, "rb") as fh:
+ for line in fh:
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ r = orjson.loads(line)
+ except orjson.JSONDecodeError:
+ continue
+ if isinstance(r, dict):
+ recs.append(_flatten_record(r))
+ return recs
+
+
+def build_payload(
+ recs: list[dict], n_files: int, cap: int, limit: int, max_turns: int
+) -> dict:
+ """Compute the compact, interned render tree (data only -- no HTML)."""
+ cap, limit, max_turns = max(cap, 0), max(limit, 0), max(max_turns, 0)
+ corr2conv = {r["x_corr"]: r["conv"] for r in recs if r["x_corr"]}
+ # A conversation can be replayed by multiple concurrent lanes in one run
+ # (e.g. --concurrency N recycling a small dataset, trajectory wrap-fill).
+ # x_correlation_id is stable across every turn of ONE session, so it is the
+ # true lane identity; number a conversation's lanes by first-arrival order.
+ # Fall back to repeated-(turn_index) arrival order only when correlation ids
+ # are absent from the export.
+ lanes_by_conv: dict[str, dict[str, int]] = defaultdict(dict)
+ seen_ti: dict[tuple[str, int], int] = defaultdict(int)
+ for r in sorted(recs, key=lambda r: r["start_ns"] or 0):
+ if r.get("x_corr"):
+ lanes = lanes_by_conv[r["conv"]]
+ if r["x_corr"] not in lanes:
+ lanes[r["x_corr"]] = len(lanes)
+ r["replay"] = lanes[r["x_corr"]]
+ else:
+ key = (r["conv"], r["turn"])
+ r["replay"] = seen_ti[key]
+ seen_ti[key] += 1
+ by_conv: dict[str, list[dict]] = defaultdict(list)
+ for r in recs:
+ replay = r.get("replay", 0)
+ conv_key = r["conv"] if replay == 0 else f"{r['conv']} · replay {replay + 1}"
+ by_conv[conv_key].append(r)
+
+ # Order: roots first, then by earliest request time; cap conversation count.
+ conv_order = sorted(
+ by_conv,
+ key=lambda c: (
+ min(r["depth"] for r in by_conv[c]),
+ min(r["start_ns"] or 0 for r in by_conv[c]),
+ c,
+ ),
+ )
+ total_convs = len(conv_order)
+ shown = conv_order[:limit]
+
+ n_turns = sum(len(v) for v in by_conv.values())
+ n_msgs = sum(len(r["messages"]) for r in recs)
+ role_counts = Counter(
+ (m.get("role", "?") if isinstance(m, dict) else "?")
+ for r in recs
+ for m in r["messages"]
+ )
+
+ # Intern unique messages: every (role, content) stored once in `table`; each
+ # turn references them by integer id. Lossless dedup of the repeated history.
+ intern: dict[tuple[str, str], int] = {}
+ table: list[dict] = []
+
+ convs: list[dict] = []
+ for conv in shown:
+ turns = sorted(
+ by_conv[conv],
+ key=lambda r: (r["turn"] if r["turn"] is not None else 0, r["start_ns"]),
+ )
+ depth = min(r["depth"] for r in turns)
+ kind = "root" if depth == 0 else f"subagent (depth {depth})"
+ parent_corr = next((r["parent_corr"] for r in turns if r["parent_corr"]), None)
+ parent_conv = corr2conv.get(parent_corr)
+ parent_lbl = _disp_id(parent_conv) if parent_conv else None
+ models = ",".join(sorted({r["model"] or "?" for r in turns}))
+ n_t = len(turns)
+ more_t = 0
+ if n_t > max_turns:
+ more_t = n_t - max_turns
+ turns = turns[:max_turns]
+ lane = next((r["x_corr"] for r in turns if r["x_corr"]), None)
+ base_id, _, replay_tag = conv.partition(" · ")
+
+ turns_payload: list[dict] = []
+ for r in turns:
+ ids: list[int] = []
+ for m in r["messages"]:
+ if isinstance(m, dict):
+ role = m.get("role", "?")
+ content = _text_of(m.get("content"))
+ else:
+ role = "?"
+ content = _text_of(m)
+ key = (role, content)
+ mid = intern.get(key)
+ if mid is None:
+ mid = len(table)
+ intern[key] = mid
+ clen = len(content)
+ table.append(
+ {
+ "role": role,
+ "len": clen,
+ "body": content[:cap],
+ "trunc": max(0, clen - cap),
+ }
+ )
+ ids.append(mid)
+ turns_payload.append(
+ {
+ "ti": r["turn"],
+ "phase": r["phase"],
+ "maxTokens": r["max_tokens"],
+ "cancelled": bool(r["cancelled"]),
+ "error": bool(r["error"]),
+ "ids": ids,
+ }
+ )
+ convs.append(
+ {
+ "id": base_id,
+ "disp": _disp_id(base_id),
+ "kind": kind,
+ "isRoot": depth == 0,
+ "parent": parent_lbl,
+ "lane": _short(lane, 8) if lane else None,
+ "models": models,
+ "nt": n_t,
+ "moreTurns": more_t,
+ "replayTag": replay_tag or "",
+ "turns": turns_payload,
+ }
+ )
+
+ stat = (
+ f"{total_convs:,} conversations ({len(shown):,} shown) · {n_turns:,} turns · "
+ f"{n_msgs:,} message refs ({len(table):,} unique) · from {n_files} raw shard(s)"
+ )
+ return {
+ "stat": stat,
+ "legend": role_counts.most_common(),
+ "msgs": table,
+ "convs": convs,
+ }
+
+
+def write_turn_messages_html(
+ run_dir: Path,
+ out: Path | None = None,
+ *,
+ limit_conversations: int = DEFAULT_LIMIT_CONVERSATIONS,
+ max_turns: int = DEFAULT_MAX_TURNS,
+ content_cap: int = DEFAULT_CONTENT_CAP,
+) -> Path:
+ """Render the interactive turn-messages viewer for ``run_dir``.
+
+ Reads the run's ``--export-level raw`` records, interns the messages, zstd
+ compresses the payload (pinned to a fzstd-safe ``windowLog=26`` window and
+ round-trip-verified), embeds it with the inlined decoder, and writes a
+ single self-contained HTML file. ``run_dir`` may be a run directory or a
+ direct raw jsonl file; output defaults to ``turn_messages.html`` beside the
+ records. Returns the output path.
+
+ Raises:
+ TurnMessagesError: if ``run_dir`` has no raw records, none parse, the
+ compressed payload fails its round-trip self-check, or the output
+ file cannot be written.
+ """
+ files = _find_raw_records(run_dir)
+ if not files:
+ raise TurnMessagesError(
+ f"no raw records under {run_dir}; run aiperf with --export-level raw first"
+ )
+ recs = load_records(files)
+ if not recs:
+ raise TurnMessagesError(f"no valid raw records in {run_dir}")
+
+ payload = build_payload(
+ recs, len(files), content_cap, limit_conversations, max_turns
+ )
+ raw = orjson.dumps(payload)
+
+ params = zstandard.ZstdCompressionParameters.from_level(
+ 19, enable_ldm=True, window_log=ZSTD_WINDOW_LOG
+ )
+ compressed = zstandard.ZstdCompressor(compression_params=params).compress(raw)
+
+ # Self-check before writing: a windowLog > 26 frame would decode to the
+ # right length with WRONG bytes under the inlined fzstd decoder, silently.
+ # Verify the payload round-trips, capping the decoder at the safe window.
+ dctx = zstandard.ZstdDecompressor(max_window_size=1 << ZSTD_WINDOW_LOG)
+ if dctx.decompress(compressed) != raw:
+ raise TurnMessagesError("zstd round-trip mismatch — refusing to write")
+
+ payload_b64 = base64.b64encode(compressed).decode()
+ html = (
+ VIEWER_TEMPLATE.read_text()
+ .replace("__FZSTD_JS__", FZSTD_JS.read_text())
+ .replace("__PAYLOAD_B64__", payload_b64)
+ )
+ default_dir = run_dir.parent if run_dir.is_file() else run_dir
+ out_path = out or (default_dir / "turn_messages.html")
+ try:
+ out_path.write_text(html)
+ except OSError as e:
+ raise TurnMessagesError(f"could not write {out_path}: {e}") from e
+ return out_path
diff --git a/src/aiperf/analysis/turn_messages_viewer.html b/src/aiperf/analysis/turn_messages_viewer.html
new file mode 100644
index 0000000000..8f81fcb808
--- /dev/null
+++ b/src/aiperf/analysis/turn_messages_viewer.html
@@ -0,0 +1,158 @@
+
+
+
+
+
+turn input messages
+
+
+
+/turn_messages.html``.
+
+ Examples:
+ # Render the newest run (writes /turn_messages.html)
+ aiperf turn-messages ./artifacts/my-run/
+
+ # Multiple runs in one invocation
+ aiperf turn-messages ./artifacts/run_a/ ./artifacts/run_b/
+
+ # Single run, explicit output path, show more conversations, full bodies
+ aiperf turn-messages ./artifacts/my-run/ -o /tmp/msgs.html \
+ -n 1000 --content-cap 1000000
+
+ Args:
+ run_dirs: One or more AIPerf run directories.
+ out: Output HTML path. Only valid when a single run directory is given.
+ limit_conversations: Max conversations to render (roots first, then by
+ earliest request time).
+ max_turns: Max turns rendered per conversation; the rest are summarized
+ as a hidden count.
+ content_cap: Max characters kept per unique message body; longer bodies
+ are truncated with a remaining-chars note. Raise for full fidelity.
+ """
+ from aiperf.analysis.turn_messages import (
+ TurnMessagesError,
+ write_turn_messages_html,
+ )
+
+ if out is not None and len(run_dirs) > 1:
+ print("error: --out only valid with a single run dir", file=sys.stderr)
+ sys.exit(2)
+
+ failures = 0
+ for run_dir in run_dirs:
+ try:
+ saved = write_turn_messages_html(
+ run_dir,
+ out=out,
+ limit_conversations=limit_conversations,
+ max_turns=max_turns,
+ content_cap=content_cap,
+ )
+ print(f"saved {saved} ({saved.stat().st_size // 1024:,} KB)")
+ except TurnMessagesError as e:
+ print(f"skip {run_dir}: {e}", file=sys.stderr)
+ failures += 1
+ if failures and failures == len(run_dirs):
+ sys.exit(1)
diff --git a/tests/unit/analysis/test_turn_messages.py b/tests/unit/analysis/test_turn_messages.py
new file mode 100644
index 0000000000..ceccd0f76c
--- /dev/null
+++ b/tests/unit/analysis/test_turn_messages.py
@@ -0,0 +1,355 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the turn-messages viewer: interning, zstd round-trip, fzstd safety.
+
+Focus: every turn re-sends accumulated history, so unique ``(role, content)``
+messages must be interned into a shared ``msgs`` table referenced by integer
+``ids``; and the embedded payload must decode under a fzstd-safe ``windowLog``
+window (the inlined decoder silently corrupts windows > 26).
+"""
+
+from __future__ import annotations
+
+import base64
+import re
+from pathlib import Path
+
+import orjson
+import pytest
+import zstandard
+
+from aiperf.analysis.turn_messages import (
+ ZSTD_WINDOW_LOG,
+ TurnMessagesError,
+ _flatten_record,
+ build_payload,
+ write_turn_messages_html,
+)
+from aiperf.cli_commands.turn_messages import turn_messages as cli_turn_messages
+
+NS = 1_000_000_000
+
+
+def _rec(
+ conv: str,
+ turn: int,
+ messages: list[dict],
+ *,
+ start_s: float = 0.0,
+ depth: int = 0,
+ x_corr: str | None = None,
+ parent_corr: str | None = None,
+ model: str = "test-model",
+) -> dict:
+ """Build a minimal --export-level raw record (one turn's request)."""
+ return {
+ "metadata": {
+ "conversation_id": conv,
+ "turn_index": turn,
+ "request_start_ns": int(start_s * NS),
+ "agent_depth": depth,
+ "x_correlation_id": x_corr or f"{conv}#0",
+ "parent_correlation_id": parent_corr,
+ "benchmark_phase": "profiling",
+ },
+ "payload": {"model": model, "messages": messages},
+ }
+
+
+def _accumulated(conv: str, n_turns: int) -> list[dict]:
+ """One conversation of ``n_turns`` turns, each re-sending full history."""
+ history: list[dict] = [{"role": "system", "content": "you are a test agent"}]
+ recs = []
+ for t in range(n_turns):
+ history = history + [{"role": "user", "content": f"question {t}"}]
+ recs.append(_rec(conv, t, list(history), start_s=float(t)))
+ history = history + [{"role": "assistant", "content": f"answer {t}"}]
+ return recs
+
+
+def _extract_payload(html_path: Path) -> dict:
+ """Mirror the viewer's client-side decode: base64 -> zstd -> JSON.
+
+ Caps the decoder window at ``ZSTD_WINDOW_LOG`` so this also asserts the
+ frame is within the fzstd-safe window (a larger window would raise here).
+ """
+ match = re.search(r'atob\("([A-Za-z0-9+/=]+)"\)', html_path.read_text())
+ assert match, "embedded payload not found"
+ dctx = zstandard.ZstdDecompressor(max_window_size=1 << ZSTD_WINDOW_LOG)
+ return orjson.loads(dctx.decompress(base64.b64decode(match.group(1))))
+
+
+def _write_raw(run_dir: Path, recs: list[dict]) -> None:
+ (run_dir / "profile_export_raw.jsonl").write_bytes(
+ b"\n".join(orjson.dumps(r) for r in recs)
+ )
+
+
+class TestBuildPayload:
+ def test_interns_repeated_messages(self):
+ # 4 turns re-sending history: 1 system + 4 user + 3 assistant = 8 unique,
+ # but the raw message refs are far more (system+user repeat every turn).
+ recs = [_flatten_record(r) for r in _accumulated("c1", 4)]
+ total_refs = sum(len(r["messages"]) for r in recs)
+ payload = build_payload(recs, 1, cap=8000, limit=40, max_turns=60)
+ assert len(payload["msgs"]) == 8
+ assert total_refs > len(payload["msgs"]) # interning actually dedups
+ # every turn references the shared table by id, never inlines content
+ turns = payload["convs"][0]["turns"]
+ for turn in turns:
+ assert "ids" in turn and "msgs" not in turn
+ assert all(0 <= i < len(payload["msgs"]) for i in turn["ids"])
+
+ def test_content_cap_truncates_and_records_remainder(self):
+ big = "x" * 5000
+ recs = [_flatten_record(_rec("c1", 0, [{"role": "user", "content": big}]))]
+ payload = build_payload(recs, 1, cap=100, limit=40, max_turns=60)
+ (msg,) = payload["msgs"]
+ assert len(msg["body"]) == 100
+ assert msg["len"] == 5000
+ assert msg["trunc"] == 4900
+
+ def test_limit_and_max_turns_summarize_overflow(self):
+ recs = [
+ _flatten_record(r) for r in _accumulated("c1", 5) + _accumulated("c2", 5)
+ ]
+ payload = build_payload(recs, 1, cap=8000, limit=1, max_turns=3)
+ assert "2 conversations (1 shown)" in payload["stat"]
+ conv = payload["convs"][0]
+ assert conv["nt"] == 5
+ assert conv["moreTurns"] == 2
+ assert len(conv["turns"]) == 3
+
+ def test_negative_args_are_clamped_to_zero(self):
+ recs = [_flatten_record(_rec("c1", 0, [{"role": "user", "content": "abcdef"}]))]
+ payload = build_payload(recs, 1, cap=-10, limit=40, max_turns=60)
+ (msg,) = payload["msgs"]
+ assert msg["body"] == "" and msg["trunc"] == 6 # cap clamped to 0
+ # negative limit clamps to 0 -> nothing shown, no crash
+ assert build_payload(recs, 1, cap=8000, limit=-1, max_turns=-1)["convs"] == []
+
+
+class TestWriteHtml:
+ def test_end_to_end_payload_decodes_within_safe_window(self, tmp_path: Path):
+ _write_raw(tmp_path, _accumulated("c1", 6))
+ out = write_turn_messages_html(tmp_path)
+ assert out == tmp_path / "turn_messages.html"
+
+ payload = _extract_payload(out) # also asserts windowLog <= 26
+ assert len(payload["convs"]) == 1
+ assert payload["msgs"] # interned table present
+ ids = payload["convs"][0]["turns"][0]["ids"]
+ assert all(0 <= i < len(payload["msgs"]) for i in ids)
+
+ def test_self_contained_inline_decoder_no_decompressionstream(self, tmp_path: Path):
+ _write_raw(tmp_path, _accumulated("c1", 3))
+ html = write_turn_messages_html(tmp_path).read_text()
+ assert ".fzstd=f()" in html # decoder inlined as a global
+ assert "fzstd.decompress(" in html
+ assert "DecompressionStream" not in html # no native gzip dependency
+ assert "__FZSTD_JS__" not in html and "__PAYLOAD_B64__" not in html
+
+ def test_reads_raw_records_shard_folder(self, tmp_path: Path):
+ shard_dir = tmp_path / "raw_records"
+ shard_dir.mkdir()
+ (shard_dir / "shard0.jsonl").write_bytes(
+ b"\n".join(orjson.dumps(r) for r in _accumulated("c1", 2))
+ )
+ out = write_turn_messages_html(tmp_path)
+ assert _extract_payload(out)["convs"][0]["id"] == "c1"
+
+ def test_missing_raw_records_raises(self, tmp_path: Path):
+ with pytest.raises(TurnMessagesError, match="no raw records"):
+ write_turn_messages_html(tmp_path)
+
+ def test_explicit_out_path_respected(self, tmp_path: Path):
+ _write_raw(tmp_path, _accumulated("c1", 2))
+ dest = tmp_path / "nested" / "viewer.html"
+ dest.parent.mkdir()
+ assert write_turn_messages_html(tmp_path, out=dest) == dest
+ assert dest.is_file()
+
+
+def _write_lines(run_dir: Path, lines: list[bytes]) -> None:
+ """Write raw (possibly malformed) jsonl lines verbatim."""
+ (run_dir / "profile_export_raw.jsonl").write_bytes(b"\n".join(lines))
+
+
+class TestAdversarial:
+ def test_injected_html_in_content_cannot_escape_payload_blob(self, tmp_path: Path):
+ # Message content is attacker-controlled; it must end up inside the
+ # base64 zstd blob and be rendered via textContent, never as live markup.
+ evil = '"\\'
+ _write_raw(tmp_path, [_rec("c1", 0, [{"role": "user", "content": evil}])])
+ out = write_turn_messages_html(tmp_path)
+ html = out.read_text()
+ # the raw injection never appears outside the (base64) payload literal
+ assert evil not in html
+ assert "alert(String.fromCharCode" not in html
+ # ...and it round-trips byte-for-byte through the payload
+ assert _extract_payload(out)["msgs"][0]["body"] == evil
+
+ def test_template_placeholder_strings_in_content_roundtrip(self, tmp_path: Path):
+ # Content equal to the template placeholders must not corrupt the file
+ # (base64 cannot contain '_', so the substitution can't be fooled).
+ for i, marker in enumerate(("__PAYLOAD_B64__", "__FZSTD_JS__")):
+ d = tmp_path / f"run{i}"
+ d.mkdir()
+ _write_raw(d, [_rec("c1", 0, [{"role": "user", "content": marker}])])
+ out = write_turn_messages_html(d)
+ html = out.read_text()
+ assert "__PAYLOAD_B64__" not in html and "__FZSTD_JS__" not in html
+ assert _extract_payload(out)["msgs"][0]["body"] == marker
+
+ def test_injected_conversation_id_cannot_escape(self, tmp_path: Path):
+ cid = "
::sa:7"
+ _write_raw(tmp_path, [_rec(cid, 0, [{"role": "user", "content": "hi"}])])
+ out = write_turn_messages_html(tmp_path)
+ assert cid not in out.read_text()
+ assert _extract_payload(out)["convs"][0]["id"] == cid
+
+ def test_non_object_json_lines_are_skipped(self, tmp_path: Path):
+ # Valid JSON that is not an object (int/str/array/null) must not crash.
+ good = orjson.dumps(_rec("c1", 0, [{"role": "user", "content": "ok"}]))
+ _write_lines(
+ tmp_path,
+ [
+ b"not json{",
+ b"",
+ b" ",
+ b"123",
+ b'"a string"',
+ b"[1,2,3]",
+ b"null",
+ good,
+ ],
+ )
+ payload = _extract_payload(write_turn_messages_html(tmp_path))
+ assert len(payload["convs"]) == 1
+ assert payload["msgs"][0]["body"] == "ok"
+
+ def test_messages_not_a_list_is_tolerated(self, tmp_path: Path):
+ # payload.messages as a non-list (str/dict/int) must not be iterated raw.
+ recs = [
+ _rec("c1", 0, "i am a string not a list"),
+ _rec("c2", 0, {"role": "user", "content": "a dict not a list"}),
+ _rec("c3", 0, 42),
+ _rec("c4", 0, [{"role": "user", "content": "real"}]),
+ ]
+ _write_raw(tmp_path, recs)
+ payload = _extract_payload(write_turn_messages_html(tmp_path))
+ assert {m["body"] for m in payload["msgs"]} == {"real"}
+ assert len(payload["convs"]) == 4
+
+ def test_non_dict_message_items_are_tolerated(self, tmp_path: Path):
+ msgs = ["bare string", 42, None, {"role": "user", "content": "real"}]
+ _write_raw(tmp_path, [_rec("c1", 0, msgs)])
+ payload = _extract_payload(write_turn_messages_html(tmp_path))
+ bodies = {m["body"] for m in payload["msgs"]}
+ assert "real" in bodies and "bare string" in bodies and "42" in bodies
+
+ def test_missing_metadata_and_payload_fields(self, tmp_path: Path):
+ _write_lines(
+ tmp_path,
+ [
+ orjson.dumps({}),
+ orjson.dumps({"metadata": {}, "payload": {}}),
+ orjson.dumps({"metadata": None, "payload": None}),
+ orjson.dumps(
+ {"payload": {"messages": [{"role": "user", "content": "hi"}]}}
+ ),
+ orjson.dumps(_rec("c1", 0, [{"role": "user", "content": "ok"}])),
+ ],
+ )
+ out = write_turn_messages_html(tmp_path)
+ assert out.is_file()
+ assert "ok" in {m["body"] for m in _extract_payload(out)["msgs"]}
+
+ def test_content_type_variants(self, tmp_path: Path):
+ msgs = [
+ {"role": "system", "content": None},
+ {"role": "user", "content": 42},
+ {
+ "role": "user",
+ "content": [{"type": "text", "text": "hello"}, {"type": "image_url"}],
+ },
+ {"role": "user", "content": [{"type": "image"}]},
+ ]
+ _write_raw(tmp_path, [_rec("c1", 0, msgs)])
+ bodies = {
+ m["body"]
+ for m in _extract_payload(write_turn_messages_html(tmp_path))["msgs"]
+ }
+ assert "" in bodies
+ assert "42" in bodies
+ assert "hello\n" in bodies
+ assert "" in bodies
+
+ def test_unicode_and_zero_content_cap(self, tmp_path: Path):
+ body = "héllo 世界 🤖\n\ttabs\x00nul"
+ _write_raw(tmp_path, [_rec("c1", 0, [{"role": "user", "content": body}])])
+ out = write_turn_messages_html(tmp_path, content_cap=0)
+ msg = _extract_payload(out)["msgs"][0]
+ assert msg["body"] == ""
+ assert msg["len"] == len(body)
+ assert msg["trunc"] == len(body)
+
+ def test_all_invalid_records_raises(self, tmp_path: Path):
+ _write_lines(tmp_path, [b"garbage", b"{bad json", b"42", b"null"])
+ with pytest.raises(TurnMessagesError, match="no valid raw records"):
+ write_turn_messages_html(tmp_path)
+
+ def test_replayed_conversation_splits_into_lanes(self, tmp_path: Path):
+ recs = [
+ _rec(
+ "c1", 0, [{"role": "user", "content": "a"}], x_corr="lane-A", start_s=0
+ ),
+ _rec(
+ "c1", 0, [{"role": "user", "content": "b"}], x_corr="lane-B", start_s=1
+ ),
+ ]
+ _write_raw(tmp_path, recs)
+ payload = _extract_payload(write_turn_messages_html(tmp_path))
+ assert len(payload["convs"]) == 2
+ tags = {c["replayTag"] for c in payload["convs"]}
+ assert "" in tags and any("replay" in t for t in tags)
+
+
+class TestCli:
+ def test_out_with_multiple_run_dirs_exits_2(self, tmp_path: Path):
+ a, b = tmp_path / "a", tmp_path / "b"
+ a.mkdir()
+ b.mkdir()
+ _write_raw(a, _accumulated("c1", 1))
+ _write_raw(b, _accumulated("c2", 1))
+ with pytest.raises(SystemExit) as e:
+ cli_turn_messages([a, b], out=tmp_path / "x.html")
+ assert e.value.code == 2
+
+ def test_all_invalid_run_dirs_exit_1(self, tmp_path: Path, capsys):
+ with pytest.raises(SystemExit) as e:
+ cli_turn_messages([tmp_path / "missing1", tmp_path / "missing2"])
+ assert e.value.code == 1
+ assert capsys.readouterr().err.count("skip") == 2
+
+ def test_partial_failure_exits_zero_and_writes_good(self, tmp_path: Path):
+ good = tmp_path / "good"
+ good.mkdir()
+ _write_raw(good, _accumulated("c1", 1))
+ # one good + one missing -> not all failed -> returns normally (no exit)
+ cli_turn_messages([good, tmp_path / "missing"])
+ assert (good / "turn_messages.html").is_file()
+
+ def test_accepts_direct_jsonl_file_writes_beside_it(self, tmp_path: Path):
+ _write_raw(tmp_path, _accumulated("c1", 2))
+ cli_turn_messages([tmp_path / "profile_export_raw.jsonl"])
+ assert (tmp_path / "turn_messages.html").is_file()
+
+ def test_unwritable_out_is_reported_not_traceback(self, tmp_path: Path, capsys):
+ good = tmp_path / "good"
+ good.mkdir()
+ _write_raw(good, _accumulated("c1", 1))
+ with pytest.raises(SystemExit) as e:
+ cli_turn_messages([good], out=tmp_path / "missing_dir" / "x.html")
+ assert e.value.code == 1
+ assert "could not write" in capsys.readouterr().err
diff --git a/tools/ergonomics_baseline.json b/tools/ergonomics_baseline.json
index 3585618a31..77a676c225 100644
--- a/tools/ergonomics_baseline.json
+++ b/tools/ergonomics_baseline.json
@@ -26,6 +26,16 @@
"src/aiperf/analysis/sweepline_stats.py",
"compute_active_weighted_stats"
],
+ [
+ "keyword-only-args",
+ "src/aiperf/analysis/turn_messages.py",
+ "build_payload"
+ ],
+ [
+ "keyword-only-args",
+ "src/aiperf/cli_commands/turn_messages.py",
+ "turn_messages"
+ ],
[
"keyword-only-args",
"src/aiperf/common/mixins/task_manager_mixin.py",
diff --git a/tools/ruff_baseline.json b/tools/ruff_baseline.json
index e0d93179fb..7a9facaaee 100644
--- a/tools/ruff_baseline.json
+++ b/tools/ruff_baseline.json
@@ -234,6 +234,11 @@
"src/aiperf/accuracy/graders/math.py",
"_math_equal"
],
+ [
+ "C901",
+ "src/aiperf/analysis/turn_messages.py",
+ "build_payload"
+ ],
[
"C901",
"src/aiperf/common/bootstrap.py",
Turn input messages — what the run actually sent, by conversation → turn → role
+inflating payload…
+
+
+
+
+
+
+