Skip to content

Commit ac14821

Browse files
YAML frontmatter: always-quote scalars for safe round-trip parsing
Refactor frontmatter export through _session_frontmatter_dict and _append_yaml_value so all string fields are double-quoted or block-scalars. _escape_yaml handles tabs, hashes, colons, bool-like literals, and non-printable Unicode. Add round-trip and Hypothesis property tests; pin PyYAML in dev deps.
1 parent 554b2cf commit ac14821

3 files changed

Lines changed: 231 additions & 44 deletions

File tree

requirements-dev.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,5 @@ pytest-cov>=5.0
66
ruff>=0.9.0
77
pip-audit>=2.7.0
88
hypothesis>=6.100.0
9+
PyYAML>=6.0
910
pytest-benchmark==5.2.3

tests/test_md_exporter_yaml.py

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
"""YAML frontmatter escaping and round-trip tests for md_exporter."""
2+
3+
from __future__ import annotations
4+
5+
import os
6+
import sys
7+
8+
import yaml
9+
from hypothesis import given, settings, strategies as st
10+
11+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
12+
13+
from models.session import SessionDict
14+
from utils.md_exporter import (
15+
_append_yaml_value,
16+
_escape_yaml,
17+
_session_frontmatter_dict,
18+
session_to_markdown,
19+
)
20+
21+
FUZZ_SETTINGS = settings(max_examples=100)
22+
23+
24+
def _extract_frontmatter_dict(markdown: str) -> dict:
25+
lines = markdown.splitlines()
26+
if not lines or lines[0].strip() != "---":
27+
raise ValueError("missing opening frontmatter delimiter")
28+
yaml_lines: list[str] = []
29+
for line in lines[1:]:
30+
if line.strip() == "---":
31+
break
32+
yaml_lines.append(line)
33+
else:
34+
raise ValueError("missing closing frontmatter delimiter")
35+
loaded = yaml.safe_load("\n".join(yaml_lines))
36+
return loaded if isinstance(loaded, dict) else {}
37+
38+
39+
def _base_session(**overrides: object) -> SessionDict:
40+
session: SessionDict = {
41+
"session_id": "sess-001",
42+
"title": "Hello",
43+
"messages": [{"role": "user", "text": "hi"}],
44+
"metadata": {
45+
"session_id": "sess-001",
46+
"models_used": ["claude-sonnet-4-20250514"],
47+
"first_timestamp": "2026-01-02T12:00:00Z",
48+
"last_timestamp": "2026-01-02T12:30:00Z",
49+
"total_input_tokens": 120,
50+
"total_output_tokens": 45,
51+
"total_cache_read_tokens": 10,
52+
"total_tool_calls": 2,
53+
"tool_call_counts": {"Read": 2},
54+
"cwd": "/workspace",
55+
"git_branch": "main",
56+
"version": "1.0.0",
57+
"permission_mode": "default",
58+
},
59+
}
60+
if overrides:
61+
for key, value in overrides.items():
62+
if key == "metadata" and isinstance(value, dict):
63+
session["metadata"].update(value) # type: ignore[typeddict-item]
64+
else:
65+
session[key] = value # type: ignore[literal-required]
66+
return session
67+
68+
69+
class TestYamlFrontmatterRoundtrip:
70+
def test_yaml_frontmatter_roundtrip(self):
71+
session = _base_session(
72+
title="Fix: handle edge case #42",
73+
metadata={
74+
"cwd": r"C:\Users\dev\project",
75+
"git_branch": "feat#yaml",
76+
"permission_mode": "true",
77+
"stop_reasons": {"max_tokens": 1, "end_turn": 2},
78+
"tool_call_counts": {"Read": 1, "Fix: tool": 1},
79+
},
80+
)
81+
md = session_to_markdown(session)
82+
assert yaml.safe_load(md.split("---")[1]) == _session_frontmatter_dict(session)
83+
assert _extract_frontmatter_dict(md) == _session_frontmatter_dict(session)
84+
85+
def test_multiline_title_uses_block_scalar(self):
86+
session = _base_session(title="line one\nline two")
87+
md = session_to_markdown(session)
88+
assert "title: |-" in md.split("---")[1]
89+
assert _extract_frontmatter_dict(md)["title"] == "line one\nline two"
90+
91+
def test_tab_and_hash_in_title(self):
92+
session = _base_session(title="tab\there # not a comment")
93+
md = session_to_markdown(session)
94+
assert _extract_frontmatter_dict(md)["title"] == "tab\there # not a comment"
95+
96+
97+
@FUZZ_SETTINGS
98+
@given(st.text())
99+
def test_escape_yaml_roundtrip(s: str) -> None:
100+
"""Double-quoted scalars round-trip for arbitrary single-line text."""
101+
if "\n" in s or "\r" in s:
102+
return
103+
loaded = yaml.safe_load(f"key: {_escape_yaml(s)}")
104+
assert loaded["key"] == s
105+
106+
107+
@FUZZ_SETTINGS
108+
@given(st.text())
109+
def test_yaml_string_field_roundtrip(s: str) -> None:
110+
"""Frontmatter string serializer round-trips arbitrary text."""
111+
lines: list[str] = []
112+
_append_yaml_value(lines, "key", s)
113+
loaded = yaml.safe_load("\n".join(lines))
114+
assert loaded["key"] == s

utils/md_exporter.py

Lines changed: 116 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Markdown export. Produces a .md with YAML frontmatter, a summary section
22
(cost, files touched, commands run), and the full conversation."""
33

4+
import re
45
from datetime import datetime
56
from typing import Any
67

@@ -25,63 +26,72 @@ def session_to_markdown(session: SessionDict, stats: SessionStatsDict | None = N
2526

2627

2728
def _build_frontmatter(session: SessionDict) -> str:
28-
meta = session["metadata"]
2929
lines = ["---"]
30-
lines.append(f'title: "{_escape_yaml(session["title"])}"')
31-
if meta["first_timestamp"]:
32-
lines.append(f"created: {meta['first_timestamp']}")
33-
if meta["last_timestamp"]:
34-
lines.append(f"updated: {meta['last_timestamp']}")
35-
lines.append(f"session_id: {session['session_id']}")
36-
if meta["models_used"]:
37-
lines.append(f"models_used: {', '.join(meta['models_used'])}")
38-
lines.append(f"total_input_tokens: {meta['total_input_tokens']}")
39-
lines.append(f"total_output_tokens: {meta['total_output_tokens']}")
40-
lines.append(f"total_cache_read_tokens: {meta['total_cache_read_tokens']}")
30+
for key, value in _session_frontmatter_dict(session).items():
31+
_append_yaml_value(lines, key, value)
32+
lines.append("---")
33+
return "\n".join(lines)
34+
35+
36+
def _session_frontmatter_dict(session: SessionDict) -> dict[str, Any]:
37+
"""Canonical frontmatter payload; used for export and round-trip tests."""
38+
meta = session["metadata"]
39+
data: dict[str, Any] = {
40+
"title": session["title"],
41+
"session_id": session["session_id"],
42+
"total_input_tokens": meta.get("total_input_tokens", 0),
43+
"total_output_tokens": meta.get("total_output_tokens", 0),
44+
"total_cache_read_tokens": meta.get("total_cache_read_tokens", 0),
45+
"total_tool_calls": meta.get("total_tool_calls", 0),
46+
"message_count": len(session["messages"]),
47+
}
48+
if meta.get("first_timestamp"):
49+
data["created"] = meta["first_timestamp"]
50+
if meta.get("last_timestamp"):
51+
data["updated"] = meta["last_timestamp"]
52+
if meta.get("models_used"):
53+
data["models_used"] = ", ".join(meta["models_used"])
4154
if meta.get("total_cache_creation_tokens", 0) > 0:
42-
lines.append(f"total_cache_creation_tokens: {meta['total_cache_creation_tokens']}")
43-
lines.append(f"total_tool_calls: {meta['total_tool_calls']}")
44-
if meta["tool_call_counts"]:
45-
lines.append("tool_call_breakdown:")
46-
for tool, count in sorted(meta["tool_call_counts"].items(), key=lambda x: -x[1]):
47-
lines.append(f" {tool}: {count}")
55+
data["total_cache_creation_tokens"] = meta["total_cache_creation_tokens"]
56+
if meta.get("tool_call_counts"):
57+
data["tool_call_breakdown"] = dict(
58+
sorted(meta["tool_call_counts"].items(), key=lambda item: -item[1])
59+
)
4860
if meta.get("stop_reasons"):
49-
lines.append("stop_reasons:")
50-
for reason, count in sorted(meta["stop_reasons"].items(), key=lambda x: -x[1]):
51-
lines.append(f" {reason}: {count}")
52-
if meta["cwd"]:
53-
lines.append(f'working_directory: "{_escape_yaml(meta["cwd"])}"')
54-
if meta["git_branch"]:
55-
lines.append(f"git_branch: {meta['git_branch']}")
56-
if meta["version"]:
57-
lines.append(f"claude_code_version: {meta['version']}")
58-
if meta["permission_mode"]:
59-
lines.append(f"permission_mode: {meta['permission_mode']}")
61+
data["stop_reasons"] = dict(
62+
sorted(meta["stop_reasons"].items(), key=lambda item: -item[1])
63+
)
64+
if meta.get("cwd"):
65+
data["working_directory"] = meta["cwd"]
66+
if meta.get("git_branch"):
67+
data["git_branch"] = meta["git_branch"]
68+
if meta.get("version"):
69+
data["claude_code_version"] = meta["version"]
70+
if meta.get("permission_mode"):
71+
data["permission_mode"] = meta["permission_mode"]
6072
if meta.get("service_tiers"):
61-
lines.append(f"service_tiers: {', '.join(meta['service_tiers'])}")
62-
lines.append(f"message_count: {len(session['messages'])}")
63-
if meta["compactions"] > 0:
64-
lines.append(f"compactions: {meta['compactions']}")
73+
data["service_tiers"] = ", ".join(meta["service_tiers"])
74+
if meta.get("compactions", 0) > 0:
75+
data["compactions"] = meta["compactions"]
6576
if meta.get("api_errors", 0) > 0:
66-
lines.append(f"api_errors: {meta['api_errors']}")
77+
data["api_errors"] = meta["api_errors"]
6778
if meta.get("sidechain_messages", 0) > 0:
68-
lines.append(f"sidechain_messages: {meta['sidechain_messages']}")
79+
data["sidechain_messages"] = meta["sidechain_messages"]
6980
wall = meta.get("session_wall_time_seconds")
7081
if wall is not None:
71-
lines.append(f"wall_clock_seconds: {int(wall)}")
82+
data["wall_clock_seconds"] = int(wall)
7283
files_r = meta.get("files_read", [])
7384
files_w = meta.get("files_written", [])
7485
files_c = meta.get("files_created", [])
7586
if files_r or files_w or files_c:
76-
lines.append(f"files_read: {len(files_r)}")
77-
lines.append(f"files_written: {len(files_w)}")
78-
lines.append(f"files_created: {len(files_c)}")
87+
data["files_read"] = len(files_r)
88+
data["files_written"] = len(files_w)
89+
data["files_created"] = len(files_c)
7990
if meta.get("bash_commands"):
80-
lines.append(f"commands_run: {len(meta['bash_commands'])}")
91+
data["commands_run"] = len(meta["bash_commands"])
8192
if meta.get("web_fetches"):
82-
lines.append(f"web_fetches: {len(meta['web_fetches'])}")
83-
lines.append("---")
84-
return "\n".join(lines)
93+
data["web_fetches"] = len(meta["web_fetches"])
94+
return data
8595

8696

8797
def _build_header(session: SessionDict) -> str:
@@ -459,8 +469,70 @@ def _format_ts(ts: str | None) -> str:
459469
return ts
460470

461471

472+
_PLAIN_YAML_KEY = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
473+
474+
475+
def _yaml_mapping_key(key: str) -> str:
476+
"""Format a nested mapping key, quoting when it is not a plain identifier."""
477+
if _PLAIN_YAML_KEY.match(key):
478+
return key
479+
return _escape_yaml(key)
480+
481+
462482
def _escape_yaml(s: str) -> str:
463-
return s.replace('"', '\\"').replace("\n", " ")
483+
"""Return a YAML double-quoted scalar for a single-line string value."""
484+
parts: list[str] = []
485+
for ch in s:
486+
if ch == "\\":
487+
parts.append("\\\\")
488+
elif ch == '"':
489+
parts.append('\\"')
490+
elif ch == "\t":
491+
parts.append("\\t")
492+
elif ch == "\n":
493+
parts.append("\\n")
494+
elif ch == "\r":
495+
parts.append("\\r")
496+
elif not ch.isprintable():
497+
code = ord(ch)
498+
if code <= 0xFF:
499+
parts.append(f"\\x{code:02x}")
500+
elif code <= 0xFFFF:
501+
parts.append(f"\\u{code:04x}")
502+
else:
503+
parts.append(f"\\U{code:08x}")
504+
else:
505+
parts.append(ch)
506+
return f'"{"".join(parts)}"'
507+
508+
509+
def _append_yaml_value(lines: list[str], key: str, value: Any, *, indent: int = 0) -> None:
510+
"""Append one frontmatter field, quoting strings and nesting mappings safely."""
511+
prefix = " " * indent
512+
if isinstance(value, dict):
513+
lines.append(f"{prefix}{key}:")
514+
for nested_key, nested_value in value.items():
515+
_append_yaml_value(
516+
lines,
517+
_yaml_mapping_key(nested_key),
518+
nested_value,
519+
indent=indent + 1,
520+
)
521+
return
522+
if isinstance(value, int):
523+
lines.append(f"{prefix}{key}: {value}")
524+
return
525+
if not isinstance(value, str):
526+
raise TypeError(f"unsupported frontmatter value type: {type(value).__name__}")
527+
if "\n" in value:
528+
if all(ch in "\n\r\t" or ch.isprintable() for ch in value):
529+
lines.append(f"{prefix}{key}: |-")
530+
for line in value.splitlines():
531+
lines.append(f"{prefix} {line}")
532+
return
533+
lines.append(f"{prefix}{key}: {_escape_yaml(value)}")
534+
return
535+
lines.append(f"{prefix}{key}: {_escape_yaml(value)}")
464536

465537

466538
def _truncate(s: str, max_len: int) -> str:

0 commit comments

Comments
 (0)